From f3f7e6de2ce761bdcd9bc4f2fcc246de485c029e Mon Sep 17 00:00:00 2001 From: andrewwhitecdw Date: Wed, 12 Aug 2026 15:03:06 -0500 Subject: [PATCH] fix: remove redundant self-assignment out_ = out_ - Remove the no-op out_ = out_ from the forward-only THD path in tests/pytorch/attention/run_attention_with_cp.py. - Guard core_attn.softmax_offset.grad.zero_() with is_training so inference no longer crashes on an uninitialized gradient. - Add the NVIDIA copyright/license header to the new regression test. - Use the supported softmax_type="learnable" in the regression test so softmax_offset is initialized. - Register tests/pytorch/attention/test_softmax_offset_inference.py in qa/L0_pytorch_unittest/test.sh so CI runs it. Signed-off-by: Andrew White --- qa/L0_pytorch_unittest/test.sh | 1 + .../attention/run_attention_with_cp.py | 3 +-- .../test_softmax_offset_inference.py | 22 +++++++++++++++++++ 3 files changed, 24 insertions(+), 2 deletions(-) create mode 100644 tests/pytorch/attention/test_softmax_offset_inference.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 759432857a..8530439285 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -57,6 +57,7 @@ NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hybrid_quantization.xml $TE_PATH/tests/pytorch/test_hybrid_quantization.py || test_fail "test_hybrid_quantization.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_identity_quantizer.xml $TE_PATH/tests/pytorch/test_identity_quantizer.py || test_fail "test_identity_quantizer.py" NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_softmax_offset_inference.xml $TE_PATH/tests/pytorch/attention/test_softmax_offset_inference.py || test_fail "test_softmax_offset_inference.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_flex_attention.xml $TE_PATH/tests/pytorch/attention/test_flex_attention.py || test_fail "test_flex_attention.py" NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_linear_mxfp8_attention.xml $TE_PATH/tests/pytorch/attention/test_linear_mxfp8_attention.py || test_fail "test_linear_mxfp8_attention.py" diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 7c6cdefd15..fa2ac1d299 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -517,7 +517,7 @@ def run_dpa_with_cp( torch.cuda.Stream(), cp_comm_type, ) - if config.softmax_type != "vanilla": + if is_training and config.softmax_type != "vanilla": core_attn.softmax_offset.grad.zero_() if dtype == "fp8": core_attn.fp8_initialized = False @@ -690,7 +690,6 @@ def run_dpa_with_cp( ) else: out = out.index_select(0, seq_idx_q).contiguous() - out_ = out_ atol, rtol, rmse_tol = get_tols(config, dtype) tensors_cp = [out_, dq_, dk_, dv_, dbias_, d_softmax_offset_, max_logit_] diff --git a/tests/pytorch/attention/test_softmax_offset_inference.py b/tests/pytorch/attention/test_softmax_offset_inference.py new file mode 100644 index 0000000000..904fdeb594 --- /dev/null +++ b/tests/pytorch/attention/test_softmax_offset_inference.py @@ -0,0 +1,22 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import pytest +import torch +from transformer_engine.pytorch import DotProductAttention + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_softmax_offset_grad_none_in_eval(): + """Regression test: eval mode leaves softmax_offset.grad as None. + + The context-parallel test helper previously crashed here by calling + core_attn.softmax_offset.grad.zero_() unconditionally for non-vanilla + softmax. In eval mode requires_grad is False and no backward has run, + so .grad must stay None. + """ + core_attn = ( + DotProductAttention(8, (64, 64), num_gqa_groups=4, softmax_type="learnable").cuda().eval() + ) + assert not core_attn.softmax_offset.requires_grad