fix: guard softmax_offset.grad.zero_() by is_training

[andrewwhitecdw]

This PR addresses the following issue in tests/pytorch/attention/run_attention_with_cp.py: guard softmax_offset.grad.zero_() by is_training.

Changes

• tests/pytorch/attention/run_attention_with_cp.py: guard softmax_offset.grad.zero_() by is_training.

Details

--- a/tests/pytorch/attention/run_attention_with_cp.py
+++ b/tests/pytorch/attention/run_attention_with_cp.py
@@ -1,3 +1,3 @@
-    if config.softmax_type != "vanilla":
-        core_attn.softmax_offset.grad.zero_()
-    if dtype == "fp8":
+    if is_training and config.softmax_type != "vanilla":
+        core_attn.softmax_offset.grad.zero_()
+    if dtype == "fp8":

Tests

Let me know if you want tests added for this fix or not.

[greptile-apps]

Greptile Summary

This PR prevents the context-parallel attention test helper from clearing softmax_offset.grad during inference, where no backward pass populates that gradient.

• Adds is_training to the non-vanilla softmax gradient-reset condition.
• Preserves existing training-mode cleanup and FP8 state reset behavior.

Confidence Score: 5/5

The PR appears safe to merge, with no actionable correctness or security issues identified.

The new training guard matches the lifecycle of gradients while preserving cleanup for training runs and leaving subsequent FP8 state handling unchanged.

Important Files Changed

Filename	Overview
tests/pytorch/attention/run_attention_with_cp.py	Correctly limits softmax-offset gradient clearing to training runs, avoiding access to an absent inference gradient.

_{Reviews (1): Last reviewed commit: "fix: guard softmax_offset.grad.zero_() b..." | Re-trigger Greptile}