fix: remove redundant self-assignment out_ = out_

[andrewwhitecdw]

This PR addresses the following issue in tests/pytorch/attention/run_attention_with_cp.py: remove redundant self-assignment out_ = out_.

Changes

• tests/pytorch/attention/run_attention_with_cp.py: remove redundant self-assignment out_ = out_.

Details

--- a/tests/pytorch/attention/run_attention_with_cp.py
+++ b/tests/pytorch/attention/run_attention_with_cp.py
@@ -1,3 +1,2 @@
-        else:
-            out = out.index_select(0, seq_idx_q).contiguous()
-            out_ = out_
+        else:
+            out = out.index_select(0, seq_idx_q).contiguous()

Tests

• tests/pytorch/attention/test_redundant_self_assignment.py

--- /dev/null
+++ b/tests/pytorch/attention/test_redundant_self_assignment.py
@@ -0,0 +1,16 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Regression test for removal of redundant self-assignments."""
+
+import os
+
+
+def test_no_redundant_out_self_assignment():
+    """Ensure the forward-only THD branch no longer contains `out_ = out_`."""
+    file_path = os.path.join(os.path.dirname(__file__), "run_attention_with_cp.py")
+    with open(file_path, "r", encoding="utf-8") as f:
+        source = f.read()
+    assert "out_ = out_" not in source, (
+        "redundant self-assignment 'out_ = out_' must be removed"
+    )

---

Greptile feedback addressed

• Added the required NVIDIA copyright/license header to tests/pytorch/attention/test_softmax_offset_inference.py.
• Changed softmax_type from the unsupported "softmax_offset" to "learnable" so softmax_offset is initialized before the assertion.
• Registered the new test in qa/L0_pytorch_unittest/test.sh so CI runs it.

Local verification: python3 -m py_compile on the test file and bash -n on the QA script passed. Full pytest execution was not feasible because transformer-engine is not installed in this environment.

[greptile-apps]

Greptile Summary

The PR removes a redundant context-parallel output self-assignment and attempts to fix and register the softmax-offset inference regression test. The new test currently fails on CUDA because eval mode does not disable gradient tracking on registered parameters.

• Gates softmax-offset gradient clearing to training runs.
• Removes the no-op THD forward assignment.
• Adds the inference regression test to the L0 PyTorch QA script.

Confidence Score: 4/5

The PR is not safe to merge until the CUDA regression test checks softmax_offset.grad rather than expecting eval() to disable parameter gradient tracking.

The newly registered test deterministically fails on CUDA because softmax_offset remains a learnable parameter after eval(), preventing the L0 job from validating the intended inference behavior.

Files Needing Attention: tests/pytorch/attention/test_softmax_offset_inference.py

Important Files Changed

Filename	Overview
tests/pytorch/attention/test_softmax_offset_inference.py	Uses a supported learnable softmax type and includes the required header, but incorrectly expects eval() to set the parameter's requires_grad flag to false.
qa/L0_pytorch_unittest/test.sh	Correctly registers the regression test, which currently causes the CUDA L0 job to fail because of the test's invalid assertion.
tests/pytorch/attention/run_attention_with_cp.py	Safely limits gradient clearing to training and removes a behaviorally inert self-assignment.

_{Reviews (5): Last reviewed commit: "fix: remove redundant self-assignment ou..." | 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.

[greptile-apps]

Eval assertion fails on CUDA

When this test runs on a CUDA-enabled worker, eval() leaves the registered softmax_offset parameter with requires_grad=True, so this assertion fails before the test can validate that inference leaves its gradient unset.

Suggested change

	assert not core_attn.softmax_offset.requires_grad
	assert core_attn.softmax_offset.grad is None

Knowledge Base Used: Tests and QA