`current_max_attn_logits` breaks CUDA Graph replay because its Python-level assignment is not re-executed

In megatron/core/extensions/transformer_engine.py (around L2479), the following pattern is used to accumulate attention logits across forward calls:

python self.current_max_attn_logits = torch.max( self.current_max_attn_logits, batch_max_attention_logits, ) This is incompatible with CUDA Graph capture when the enclosing forward() is captured. The assignment is a Python-level state update and is not re-executed on Graph replay, so the tensor address recorded at capture time becomes stale / inconsistent with the intended cross-call accumulation semantics.

Problem CUDA Graph records tensor addresses and kernel launches, not Python control flow or member assignment. During capture:

torch.max(...) allocates a new output tensor at some address C and is recorded into the Graph.

Python executes self.current_max_attn_logits = C, updating the member reference.

During replay:

The recorded torch.max kernel runs again, reading the original inputs and writing to the captured output address C.

The Python assignment is not replayed. self.current_max_attn_logits retains whatever reference it had at capture time, and any other Python-side logic that depends on this member observes a stale state.

The tensor at C may have been freed or reused by the allocator outside the Graph's lifetime. Subsequent kernels that consume self.current_max_attn_logits (e.g. maximum_kernel_cuda) then operate on an address that no longer belongs to the current Graph's live buffers.

The root issue is that current_max_attn_logits is a cross-call persistent state managed by Python, while the tensor it points to is inside the Graph's captured address space. These two lifecycles are not aligned.

This is not a numerical issue. Changing torch.max to torch.maximum does not help, because the problem is the Python assignment not being part of the captured Graph, not the choice of max op.