[JAX] [PyT] [Common] Enable D=256 BWD cuDNN fused attn for Blackwell CC 10.x

[KshitijLakhani]

Description

Support for D=256 BWD for Blackwell CC 10x via the C++ API (which TE uses) was added in cuDNN 9.23 + cuDNN FE 1.24. Enabling this support in TE attention

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

Add guard when picking the backend (sub backend) in TE common.
Add tests for D=256 case in TE PyT and TE JAX

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 enables cuDNN 9.23 / FE 1.24's dedicated deterministic SDPA backward kernel for head-dim 256 on Blackwell (SM10.x / CC 10.x) GPUs. The C++ backend selector, JAX test skip guards, and PyTorch test cases are all updated in concert.

• fused_attn.cpp adds a new compound condition that selects the new BWD path only when d_qk == d_v == 256, is_training, sm_arch ∈ [100, 110), cuDNN ≥ 9.23, non-paged layout, no bias, no dropout, vanilla softmax, and the window-size rule (full window OR causal mask with right window −1 or 0).
• Both the JAX and PyTorch tests add matching guards and new parametrized cases; the PyTorch comment block contains a duplicated trailing sentence that should be removed, and the JAX skip logic for non-1HSS bias configs diverges slightly from the C++ gate (those configs won't use the new kernel path but will not fail).

Confidence Score: 4/5

The C++ backend gate is additive and well-guarded; existing paths are unchanged and the new path only activates under a very specific combination of hardware, cuDNN version, and kernel parameters.

The core logic in fused_attn.cpp looks correct and conservative. The two test-side issues (a duplicate comment line and a JAX skip-guard that allows some bias configs to pass without actually exercising the new kernel) are minor and do not affect production correctness. No existing behavior is altered.

The JAX test's _check_configs bias skip logic (tests/jax/test_fused_attn.py, lines 465-475) deserves a second look to ensure the skip conditions exactly match the C++ gate.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn.cpp	Adds a new head-dim guard condition enabling the cuDNN 9.23+ deterministic D=256 BWD path on SM10.x. The condition correctly restricts to d_qk == d_v == 256, is_training, sm_arch [100, 110), cuDNN >= 92300, non-paged layout, no bias, no dropout, vanilla softmax, and the appropriate window-size rule (full window OR causal mask with right window -1/0).
tests/pytorch/attention/test_attention.py	Adds test_dpa_fused_attn_hdim256 gated on cuDNN >= (9,23,0) and SM100/SM103. Tests no-mask, padding, causal+SWA, and padding-causal+GQA variants. Contains a duplicated trailing comment fragment.
tests/jax/test_fused_attn.py	Adds skip guards in FusedAttnRunner._check_configs for D=256 BWD on SM10.x (cuDNN version, head dim symmetry, dropout, softmax type, window size). The skip logic for non-1HSS bias diverges from the C++ gate, meaning some configs will silently use a different backend.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[nvte_get_fused_attn_backend called] --> B{dtype FP16/BF16?}
    B -- No --> Z[Other backend]
    B -- Yes --> C{head_dim_qk/v <= 128?}
    C -- Yes --> ARB[NVTE_F16_Arbitrary backend]
    C -- No --> D{d_qk==d_v==256 AND is_training\nAND sm_arch in 100-109\nAND cuDNN >= 9.23?}
    D -- Yes --> E{bias==NO_BIAS\ndropout==0\nsoftmax==VANILLA?}
    E -- No --> SKIP[Skip: Fall through to next condition]
    E -- Yes --> F{window_size == -1,-1\nOR causal mask + right_win in -1,0?}
    F -- No --> SKIP
    F -- Yes --> ARB
    D -- No --> G{d_qk/v <= 256 AND Hopper?\nOR Blackwell fprop?\nOR other existing rules}
    G -- Yes --> ARB
    G -- No --> SKIP

Loading

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[greptile-apps]

Duplicated comment fragment

The comment block ends with a repeated phrase: line 383 (# (for non-causal masks) full-window attention.) is a verbatim fragment of line 382, left over from editing. It should be removed.

Suggested change

	# vanilla type of softmax only, no dropout, no ALiBi, and (for non-causal masks) full-window attention only.
	# (for non-causal masks) full-window attention.
	# vanilla type of softmax only, no dropout, no ALiBi, and (for non-causal masks) full-window attention only.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

JAX skip logic diverges from C++ backend gate for non-1HSS bias

The comment says "frozen/non-learnable bias inputs (i.e. non-1HSS bias shapes) are supported" and the skip block deliberately allows those configs to proceed. However, the C++ gate in fused_attn.cpp requires bias_type == NVTE_NO_BIAS for the new D=256 BWD path, meaning any config with attn_bias_type != NO_BIAS && bias_shape != _1HSS will silently fall back to a different backend rather than exercising the new kernel. The test will not fail, but it also will not validate the D=256 BWD path for those configs, and the inline comment creates a misleading expectation that such configs are actually routed through it.