[PyTorch][Fused Attn] Add support for cuDNN to return Softmax Stats always and Max when return_max_logit=True

[sudhakarsingh27]

Description

cuDNN recently made returning any subset of {Stats, SumExp, Max} possible. This PR adapts TE to always get Stats from cuDNN and Max tensor if return_max_logit=True. (Note that Stats = log(SumExp)+Max)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• fused_attn_f16_arbitrary_seqlen.cu
  • Removed references to SumExp tensor as it's not needed since cuDNN returns Stats by default.
  • set generate_stats=True which forces cuDNN to always return Stats tensor (needed in the backward pass)
• transformer_engine/pytorch/cpp_extensions/fused_attn.py
  • Remove code that manually did Stats = log(SumExp) + Max since cuDNN returns Stats directly and TE doesn't need SumExp from cuDNN
• Corresponding documentation

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR adapts TransformerEngine's fused attention implementation to always request the Stats tensor (= log(SumExp) + Max) from cuDNN, and optionally also request the Max tensor when return_max_logit=True. Previously, cuDNN was either asked for Stats (training path) or {Max, SumExp} (return_max_logit path), and TE computed Stats manually from the latter pair. The new cuDNN frontend allows returning any subset, making the manual computation unnecessary.

Key changes:

• generate_stats is now always true; set_generate_stats(true) is called unconditionally on the cuDNN SDPA graph.
• The Aux_CTX_Tensors pack is restructured: Stats is always tensor [0], Max occupies tensor [1] only when return_max_logit=True, followed by rng_state and optional Bias/SoftmaxOffset.
• The Python fused_attn_fwd wrapper no longer manually computes Stats = log(SumExp) + Max; it reads Stats directly from output_tensors[1] and constructs max_logit from output_tensors[2] (Max) when return_max_logit=True.
• FADescriptor_v1::generate_max_sum_exp is renamed to return_max_logit for clarity, correctly remaining part of the graph cache key.
• API documentation in fused_attn.h is updated at both return_max_logit parameter occurrences.
• One minor cleanup opportunity: the generate_stats local variable (line 107) is now always true and could be inlined directly into the .set_generate_stats() call.

Confidence Score: 4/5

• This PR is safe to merge; the tensor ordering changes are internally consistent across all C++ and Python layers.
• The change is well-scoped: the Aux_CTX_Tensors pack layout is updated consistently in both the forward allocation pass (size==0) and the data pass (size>=2), the Python layer correctly routes Stats to aux_ctx_tensors and Max to max_logit, and the backward pass always reads Stats from index 0. No regression risk for backward compatibility since the new tensor ordering matches what both C++ callers and Python callers now expect. Score is 4 rather than 5 only because the submodule bump (cuDNN frontend) is not reviewable here, and no new tests were added.
• No files require special attention; all layers are consistent with the new Stats-first tensor ordering.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu	Core change: generate_stats is now always true, Stats tensor is always returned first in Aux_CTX_Tensors, and Max tensor is included at index 1 only when return_max_logit=True. The tensor pack ordering is consistent between the size==0 allocation pass and the size>=2 data pass. generate_stats variable could be inlined as true, but this is cosmetic.
transformer_engine/common/fused_attn/utils.h	Renames generate_max_sum_exp to return_max_logit in FADescriptor_v1. The field is correctly used in the operator< comparison (cache key), so graph caching correctly differentiates runs with vs. without Max output.
transformer_engine/common/include/transformer_engine/fused_attn.h	Documentation updated at both return_max_logit parameter occurrences (lines 209 and 272) to reflect new semantics: "Whether to produce Max along with Stats."
transformer_engine/pytorch/cpp_extensions/fused_attn.py	Removes the manual Stats = log(SumExp) + Max computation; now constructs aux_ctx_tensors = [Stats, rng_state, ...] and max_logit from output_tensors[2] (Max). Tensor order matches the updated C++ output pack ordering.
transformer_engine/pytorch/csrc/extensions/attention.cpp	Updates the allocation comments and comment for the second auxiliary tensor. Allocation logic is unchanged: first tensor is always Stats (S/M), and Max is allocated as the second tensor when return_max_logit=True.
3rdparty/cudnn-frontend	Submodule bump from 8d19d31 to a5ca04f to pick up the cuDNN frontend support for returning Stats always and Max when requested.

Sequence Diagram

sequenceDiagram
    participant Py as fused_attn_fwd (Python)
    participant Cpp as attention.cpp (C++)
    participant CUDA as fused_attn_f16_arbitrary_seqlen.cu
    participant cuDNN as cuDNN Frontend

    Py->>Cpp: tex.fused_attn_fwd(..., return_max_logit)
    Cpp->>Cpp: Allocate aux_tensor_pack<br/>[0]=Stats, [1]=Max(if rml), [n]=rng_state
    Cpp->>CUDA: nvte_fused_attn_fwd(... Aux_CTX_Tensors)
    CUDA->>CUDA: generate_stats=true (always)
    CUDA->>cuDNN: sdpa with set_generate_stats(true)<br/>+ set_logit_max(Max) if return_max_logit
    cuDNN-->>CUDA: O, Stats, [Max if return_max_logit]
    CUDA-->>Cpp: Aux_CTX_Tensors filled:<br/>[Stats, [Max], rng_state, ...]
    Cpp-->>Py: output_tensors=[O, Stats, [Max], rng_state, ...]

    Note over Py: if return_max_logit:<br/>aux=[Stats, rng_state,...]<br/>max_logit=amax(Max)<br/>else:<br/>aux=output_tensors[1:]

    Py->>Cpp: tex.fused_attn_bwd(..., aux_ctx_tensors=[Stats, rng_state,...])
    Cpp->>CUDA: Aux_CTX_Tensors[0]=Stats, [1]=rng_state
    CUDA->>cuDNN: sdpaBwd(Stats as softmax_stats)
    cuDNN-->>CUDA: dQ, dK, dV

Loading

_{Last reviewed commit: ef0d7ec}

[greptile-apps]

_{5 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{3 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{3 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{3 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

transformer_engine/pytorch/cpp_extensions/fused_attn.py
Stale docstring: wrong formula for softmaxStats

The public docstring still describes softmaxStats as log(sum(e^(x - max(x)))), which is log(SumExp). However, with this PR, the returned tensor is cuDNN's Stats = log(SumExp) + Max, not just log(SumExp). This formula was already incorrect before this PR (the old code computed Max + log(SumExp) and stored it as stats), but the PR is an opportunity to correct it.

                       softmaxStats: torch.Tensor
                           log(sum(e^(x - max(x)))) + max(x), where x=Q*K.T (i.e. Stats = log(SumExp) + Max)
                           shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32

[greptile-apps]

_{3 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[cyanguwa]

Do we need the "there's no typo here" :)

[sudhakarsingh27]

I deliberately added it because I didn't believe it and checked the shapes myself :P

[greptile-apps]

_{3 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{3 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{5 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{5 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

transformer_engine/common/include/transformer_engine/fused_attn.h
Entire file has been reformatted with unintentional 3-space indentation changes. This creates a large diff unrelated to the actual feature changes. Revert the formatting to match the original file structure.

[greptile-apps]

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[sudhakarsingh27]

/te-ci L2