From 67154c99a846bae85a598da78c028af295f5968d Mon Sep 17 00:00:00 2001 From: andrewwhitecdw Date: Wed, 12 Aug 2026 15:03:06 -0500 Subject: [PATCH 1/2] fix: remove redundant self-assignment out_ = out_ Squashed to single commit for review. Original PR: https://github.com/andrewwhitecdw/TransformerEngine/pull/11 --- .../pytorch/attention/run_attention_with_cp.py | 3 +-- .../attention/test_softmax_offset_inference.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) create mode 100644 tests/pytorch/attention/test_softmax_offset_inference.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..3d9a83b525 --- /dev/null +++ b/tests/pytorch/attention/test_softmax_offset_inference.py @@ -0,0 +1,18 @@ +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="softmax_offset" + ).cuda().eval() + assert not core_attn.softmax_offset.requires_grad From cb09267d4c4076f948a24f7933f7cf30eabeac56 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 13 Aug 2026 00:15:15 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/attention/test_softmax_offset_inference.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/attention/test_softmax_offset_inference.py b/tests/pytorch/attention/test_softmax_offset_inference.py index 3d9a83b525..103cb29b40 100644 --- a/tests/pytorch/attention/test_softmax_offset_inference.py +++ b/tests/pytorch/attention/test_softmax_offset_inference.py @@ -12,7 +12,9 @@ def test_softmax_offset_grad_none_in_eval(): 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="softmax_offset" - ).cuda().eval() + core_attn = ( + DotProductAttention(8, (64, 64), num_gqa_groups=4, softmax_type="softmax_offset") + .cuda() + .eval() + ) assert not core_attn.softmax_offset.requires_grad