[Common] Pass cu_seqlens and token-unit ragged offsets directly to cuDNN SDPA fprop

[egilliam-nv]

Description

cuDNN >= 9.24 SDPA (unified engine) accepts cumulative sequence lengths directly (cu_seq_len_q/kv) and can scale ragged offsets stored in coarser units (e.g. tokens) back to element offsets via a per-tensor ragged offset multiplier. This PR uses both features in the fused-attention forward passes to skip the conversion kernels that previously ran before every varlen call: cu_seqlens_to_actual_seqlens and (f16 THD only) cu_seqlens_padded_to_offsets.

f16/bf16 forward (cuDNN >= 9.24): binds the user's int32 cu_seqlens buffers as CU_SEQ_LEN_Q/KV for the padding mask, and the token-unit cu_seqlens_padded buffers as ragged offsets for Q/K/V/O/Stats with elements-per-token multipliers. Both conversion kernels are skipped and no conversion workspace is allocated.

fp8/mxfp8 forward (cuDNN >= 9.25 and cudnn-frontend >= 1.26): same for cu_seqlens (the fp8 path has no THD support, so only the one conversion kernel existed there). The frontend is header-only, so its version is a compile-time property; the gate is a constant-folded CUDNN_FRONTEND_VERSION check — all referenced symbols exist in 1.25 (SDPA_fp8_attributes is an alias of SDPA_attributes), so no preprocessor guards are needed and the fp8 direct path is simply dead code on 1.25 builds.

Notes for reviewers:

• CU_SEQ_LEN_* inputs pin implementation selection to the unified engine (the composite support surface rejects them), so the gates keep composite-served configs on the legacy path: training with dropout (the frontend rejects dropout together with generated stats on the unified engine, and TE always generates stats), and for f16 the direct path only engages for THD/ragged configs, which backend selection restricts to unified-capable hardware.
• The direct path keeps the true batch size instead of the quantized max_b: cuDNN reads the user's [actual_b+1] cu_seqlens buffers, so a quantized batch would read out of bounds. Token-dim bucketing is unchanged.
• The layout-group -> multiplier mapping is factored into RaggedOffsetMultipliers (utils.h), shared by the graph builder and the legacy conversion kernel so the two cannot drift. The kernel rewrite also removes a cross-thread read (offsets_v[tid] = offsets_k[cu_seqlens_id]) that raced for quantized-batch tail entries with interleaved layouts.
• Backward passes are unchanged (no backend support yet).
• NVTE_FUSED_ATTN_DIRECT_SEQLENS=0 disables the new path (kept as an escape hatch for A/B testing; can be removed on request).

Validation against cuDNN 9.25, no-fallback design so passes prove the unified-engine path:

• f16: test_dpa_softmax_thd 15/15 with the direct path on and off (H100 + Blackwell), matching pre-change baseline; direct-vs-legacy fused outputs/grads match for all THD layouts (thd_thd_thd, t3hd, th3d, thd_t2hd, thd_th2d) x MHA/GQA x padding/padding_causal x pad_between_seqs {false, true} on H100 and Blackwell (14/14 each).
• fp8 (with cudnn-frontend 1.26 built from source): test_dpa_fp8_vs_f16 padding configs — 56 passed on H100 (delayed + current scaling), 168 passed on Blackwell (adds MXFP8), zero failures.

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

• f16/bf16 fwd: pass cu_seqlens / cu_seqlens_padded directly to cuDNN (unified engine) on cuDNN >= 9.24, skipping both conversion kernels. (Also cudnn-frontend must be >= 1.25, but this is now required across the board by TransformerEngine so no need to check for this in the fp16 code.)
• fp8/mxfp8 fwd: pass cu_seqlens directly on cuDNN >= 9.25 + cudnn-frontend >= 1.26, skipping the conversion kernel
• Factor ragged-offset multipliers into a shared RaggedOffsetMultipliers helper used by both the graph builders and the legacy conversion kernel

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

🤖 Generated with Claude Code

[greptile-apps]

Greptile Summary

This PR eliminates 1-2 CUDA kernel launches per varlen forward pass by passing cu_seqlens and token-unit ragged offsets directly to cuDNN SDPA (unified engine) when cuDNN >= 9.24 (f16/bf16) or >= 9.25 + cudnn-frontend >= 1.26 (fp8/mxfp8). The RaggedOffsetMultipliers struct in utils.h centralises all stride-to-element multiplier logic, shared between the graph builders and the legacy conversion kernel so the two cannot drift.

• f16/bf16 fwd: replaces cu_seqlens_to_actual_seqlens + cu_seqlens_padded_to_offsets kernel launches with direct buffer bindings via set_cu_seq_len_q/kv and set_ragged_offset_multiplier; the dropout guard keeps composite-served configs on the legacy path.
• fp8/mxfp8 fwd: same for the single cu_seqlens_to_actual_seqlens kernel; the higher version floor (cuDNN 9.25 + frontend 1.26) is needed for fp8 support.
• SM_120 (GB200) stats fix: the outer allocation code now correctly allocates the stats tensor as dense [batch, h, max_seqlen_q, 1] on GB200, matching the graph's dense-stride path (use_ragged_stats is false there).

Confidence Score: 5/5

The core logic is well-guarded and tested on H100 and Blackwell; the SM_120 stats-shape fix is a genuine correctness improvement with no backward-pass impact.

The RaggedOffsetMultipliers refactor faithfully reproduces every layout-group case from the old kernel and eliminates the pre-existing cross-thread read race on tail entries. Stats and Max multiplier calls are correctly gated inside use_ragged_stats. The two open questions are the missing runtime escape hatch documented in the PR description and whether set_cu_seq_len_q/kv is safe for non-THD padding configs - both are non-blocking.

The use_cu_seqlens_directly activation condition in fused_attn_f16_arbitrary_seqlen.cu lines 100-109 warrants a second look to confirm that non-THD, non-ragged padding configs are handled correctly by the unified engine on cuDNN >= 9.24.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/utils.h	Adds RaggedOffsetMultipliers struct with correct per-layout stride multipliers; updated cu_seqlens_padded_to_offsets declaration matches new kernel signature.
transformer_engine/common/fused_attn/utils.cu	Kernel refactored to accept RaggedOffsetMultipliers; the cross-thread read race on interleaved-layout tail entries is correctly replaced with independent multiplier computation.
transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu	Direct cu_seqlens path implemented correctly; SM_120 stats allocation fixed; PR description references an NVTE_FUSED_ATTN_DIRECT_SEQLENS escape hatch that is absent from the code.
transformer_engine/common/fused_attn/fused_attn_fp8.cu	Direct cu_seqlens path mirrors f16 approach correctly; no ragged-offset changes needed since fp8 path has no THD support.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[fused_attn_fwd_impl called] --> B{use_cu_seqlens_directly}
    B -- YES --> C{is_padding?}
    C -- YES --> D[seq_q/seq_kv shape b+1, set_cu_seq_len_q/kv]
    B -- YES --> F{is_ragged_q?}
    F -- YES --> G[Q and O set_ragged_offset_multiplier]
    B -- YES --> H{is_ragged_kv?}
    H -- YES --> I[K and V set_ragged_offset_multiplier]
    B -- YES --> J{use_ragged_stats?}
    J -- YES --> K[Stats and Max set_ragged_offset_multiplier]
    B -- NO --> L{is_padding?}
    L -- YES --> M[Launch cu_seqlens_to_actual_seqlens kernel]
    B -- NO --> N{is_ragged_q or is_ragged_kv?}
    N -- YES --> O[Launch cu_seqlens_padded_to_offsets kernel]
    D --> P[mha_graph execute]
    G --> P
    I --> P
    K --> P
    M --> P
    O --> P

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A[fused_attn_fwd_impl called] --> B{use_cu_seqlens_directly}
    B -- YES --> C{is_padding?}
    C -- YES --> D[seq_q/seq_kv shape b+1, set_cu_seq_len_q/kv]
    B -- YES --> F{is_ragged_q?}
    F -- YES --> G[Q and O set_ragged_offset_multiplier]
    B -- YES --> H{is_ragged_kv?}
    H -- YES --> I[K and V set_ragged_offset_multiplier]
    B -- YES --> J{use_ragged_stats?}
    J -- YES --> K[Stats and Max set_ragged_offset_multiplier]
    B -- NO --> L{is_padding?}
    L -- YES --> M[Launch cu_seqlens_to_actual_seqlens kernel]
    B -- NO --> N{is_ragged_q or is_ragged_kv?}
    N -- YES --> O[Launch cu_seqlens_padded_to_offsets kernel]
    D --> P[mha_graph execute]
    G --> P
    I --> P
    K --> P
    M --> P
    O --> P

Loading

_{Reviews (7): Last reviewed commit: "[Common] Suppress fn_size lint on fused_..." | Re-trigger Greptile}

[sudhakarsingh27]

Just found some grammar nits, otherwise looks good to me