fix: remove redundant self-assignment out_ = out_

[andrewwhitecdw]

This PR fixes two issues in the context-parallel attention helper and adds regression coverage.

Changes

• tests/pytorch/attention/run_attention_with_cp.py:
  • Remove the redundant self-assignment out_ = out_ in the THD forward-only branch.
  • Gate softmax_offset.grad.zero_() on is_training so eval mode no longer crashes when softmax_offset.grad is None.
• tests/pytorch/attention/test_softmax_offset_inference.py:
  • New regression test verifying softmax_offset.grad stays None in eval mode.
• qa/L0_pytorch_unittest/test.sh:
  • Register the new regression test in the L0 PyTorch unittest job.

Details

diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py
@@ -514,7 +514,7 @@ def run_attention_with_cp(...):
         if qkv_format == "thd":
             if is_training:
                 q, kv, out_ = ctx.grad_inputs[:3]
-            else:
+            else:
                 out = out.index_select(0, seq_idx_q).contiguous()
-            out_ = out_

diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py
@@ -550,7 +550,7 @@ def run_attention_with_cp(...):
-            core_attn.softmax_offset.grad.zero_()
+            if is_training:
+                core_attn.softmax_offset.grad.zero_()

Tests

• New regression test: tests/pytorch/attention/test_softmax_offset_inference.py::test_softmax_offset_grad_none_in_eval
• Registered in qa/L0_pytorch_unittest/test.sh
• Local python -m py_compile passed; full pytest was not feasible because Transformer Engine is not installed in this environment.

Contributor guidelines

• Signed-off-by: Andrew White andrewwhitecdw@users.noreply.github.com

[greptile-apps]

Greptile Summary

This PR prevents the context-parallel attention test helper from clearing an absent softmax-offset gradient during inference and removes a redundant assignment.

• Gates softmax-offset gradient clearing on training mode.
• Adds and registers an eval-mode softmax-offset test.
• Removes the THD forward-only self-assignment.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
tests/pytorch/attention/run_attention_with_cp.py	Gates softmax-offset gradient clearing on training mode and removes a redundant THD forward-only assignment.
tests/pytorch/attention/test_softmax_offset_inference.py	Adds an eval-mode check for the learnable softmax-offset parameter.
qa/L0_pytorch_unittest/test.sh	Registers the new attention test in the L0 PyTorch unit-test job.

_{Reviews (7): Last reviewed commit: "Merge branch 'main' into codequality/run..." | Re-trigger Greptile}

[greptile-apps]

Regression test omitted from CI

The QA jobs enumerate test files explicitly, but none includes this new test, so CI silently skips the intended inference regression coverage.

Knowledge Base Used: Tests and QA

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[cyanguwa]

The changes in run_attention_with_cp.py looks good, but the test test_softmax_offset_inference.py‎ is probably unnecessary. The regular CI should test out the changes if any of them breaks anything. Thanks.

[andrewwhitecdw]

@greptile-apps thanks for the review. I have addressed the blocking feedback in this PR:

• Added the required NVIDIA copyright/license header to tests/pytorch/attention/test_softmax_offset_inference.py.
• Switched the regression test to the supported softmax_type="learnable" so softmax_offset is initialized.
• Registered the new test in qa/L0_pytorch_unittest/test.sh so CI runs it instead of skipping it.

The branch has been amended to a single signed-off commit (3b8c96b) and force-pushed. Local pytest was not feasible because transformer-engine is not installed here, but Python/bash syntax checks passed.

[greptile-apps]

Good progress on addressing the feedback. A few things to note on the current state:

test_softmax_offset_inference.py — assertion tests the wrong property

The docstring correctly identifies the bug: softmax_offset.grad is None in eval mode because no backward has run. But the assertion checks requires_grad, which is unrelated to .eval(). Calling .eval() on a model does not change requires_grad on its nn.Parameters — they remain True. So:

assert not core_attn.softmax_offset.requires_grad  # will fail: requires_grad is still True

This test will likely fail as written. The assertion should check what the docstring describes:

    assert core_attn.softmax_offset.grad is None

That directly validates the original crash condition and accurately documents what the fix prevents.

The is_training and guard in run_attention_with_cp.py — the fix itself is correct. Gating the .grad.zero_() call on is_training is the right approach.

Minor scope note: the PR description only mentions removing out_ = out_, but the diff also includes the is_training guard and a new test file. Worth updating the PR description to reflect both fixes so reviewers don't miss the behavioral change at line 517.

Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.

[andrewwhitecdw]

@greptile-apps Thanks for the follow-up. I fixed the assertion in test_softmax_offset_inference.py to check core_attn.softmax_offset.grad is None instead of requires_grad, and corrected the docstring. The branch is now a single signed-off commit rebased onto main.

[greptile-apps]

The fixes look correct. A few quick notes:

1. is_training and config.softmax_type != "vanilla" — guards the .grad.zero_() call properly; in eval mode is_training is False so .grad is never touched.

2. assert core_attn.softmax_offset.grad is None — this is the right assertion. Checking requires_grad would always be True for a learnable parameter regardless of whether a backward pass ran, so it wouldn't catch the original bug.

3. The test_redundant_self_assignment.py approach (reading the source file as text) is a bit fragile — it would pass even if the line were reintroduced in a comment — but for a regression guard on a one-line removal it's acceptable.

Overall the PR is clean. The three changes (dead assignment removal, guard for inference mode, regression tests) are well-scoped and the test accurately reflects the failure mode. LGTM.