Keep attention logit statistics valid across CUDA Graph replay

[ZenAlexa]

Why

CUDA Graph replay executes captured device operations without repeating Python assignments. Replacing current_max_attn_logits after each forward therefore fails to maintain a cross-microbatch accumulator across graph replays and optimizer steps. Clipping and log-only consumers need a stable accumulator that they can reset without replacing the captured storage.

Fixes #7475. This change addresses the reproduced accumulation failure; the separate freed-buffer hypothesis in the issue has not been established.

What changed

• Register one nonpersistent FP32 accumulator per attention module, sized for the local tensor-parallel heads and initialized to negative infinity. Preserve its FP32 dtype through Float16Module using the existing mark_keep_in_fp32 contract.
• Accumulate detached per-forward maxima with torch.maximum(..., out=...), retaining the same storage through capture, replay, consumption, and reset.
• Reset the accumulator in place after standard-attention clipping, MLA clipping, and log-only consumption. Treat an untouched all-negative-infinity buffer as having no collected statistics.
• Clear capture and warmup statistics from the same buffers during TE capture cleanup.
• Extend existing coverage with consecutive optimizer-step comparisons between eager execution and captured forward/backward, including gradients, clipping, parameter updates, buffer reuse, and TE attention replay with logging enabled and disabled.

Runtime context-parallel group selection and restoration retain the current main implementation. The accumulator is runtime state and is omitted from checkpoints.

Verification

The September 25 H100 80GB run incorporated main at 8c487209a2ae6eeb276c8244a6495a28950321af and used PyTorch 2.8.0a0+5228986c39.nv25.06, Transformer Engine 2.19.0, and cuDNN 9.26:

Check	Observed result
Baseline graph replay across three steps	Incomplete first-step accumulation and missing statistics in later steps
Patched replay across three steps	Expected per-head maxima and an unchanged buffer address
BF16 SelfAttention, eager versus graph forward/backward	Matching microbatch gradients, clipping, and parameter updates over two optimizer steps
TE fused replay, maximum-logit output enabled and disabled	Byte-exact replay checks passed; TE selected FusedAttention sub-backend 1
Existing pytest entry points and fixtures	21 passed, 1 unsupported FlashAttention maximum-logit case skipped
Runtime CP context	Original group, runtime CP1, and restored original group exercised

Scoped Black, isort, and Ruff checks passed for that revision. The current head 8114206dacb5ac81a253d419111606c8171c0bc9 merges main at 159b443c3e1a23ff68321973b936fef810bb3c61. Current-source checks covered the buffer dtype-preservation path, capture cleanup, standard-attention and MLA clipping, optimizer logging, and runtime CP-group restoration; git diff --check passed.

The focused GPU entry point is:

uv run python -m torch.distributed.run --standalone --nproc-per-node 1 \
  -m pytest -q \
  tests/unit_tests/transformer/test_attention.py::TestClipQK \
  tests/unit_tests/transformer/test_multi_latent_attention.py \
  tests/unit_tests/determinism/kernels/test_te_wrappers.py::TestTEWrappers::test_te_dot_product_attention_replays \
  -k 'clip_qk or dot_product_attention_replays'

The current merged head has no new GPU execution; the historical measurements cover single-rank execution, and multi-rank execution, full MLA graph training, and the complete TE capture pipeline require validation. The file-level determinism gate's mapping to unrelated compiled-activation and packed-runtime-CP kernels needs a maintainer decision on its coverage requirement.

current mainline validation

merged c4334de9 into 7b887f87 on 2026-10-13.

the merge commit is signed, and the eight changed python files passed black 26.3.0 and isort 5.13.2. current source review followed the statistic buffer through precision conversion, accumulation, optimizer reset and graph-capture cleanup. cuda graph replay and transformer-engine execution were not rerun for this merge. earlier gpu results remain tied to their recorded revisions.

Project CI: Please run current-head multi-rank graph training and the complete Transformer Engine capture/replay checks on the project runners, including MLA, and confirm the applicable determinism selection. The CPU checks and historical single-rank GPU results above retain their stated coverage.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[ZenAlexa]

I merged current main and reran the attention and MLA clipping paths plus TE replay on H100. Runtime-CP group restoration and logit-buffer reuse are covered in the updated replay path. The remaining file-level determinism gate maps attention.py to unchanged activation and packed-CP tests; would a determinism-exempt label fit this change?

[ZenAlexa]

@asolergi-nv, 8114206d merges current main and uses the same FP32 accumulator storage for forward accumulation, clipping, log-only reset, and TE capture cleanup. Could you review the lifecycle boundary identified in #7475 and confirm the required multi-rank and full-TE-capture checks? The single-H100 results remain tied to the September 25 revision.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.