[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 (BSHD) and cuDNN 9.25 + cuDNN FE 1.26 (THD)
Enabling this support in TE attention and adding conditional tests for the same

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 the cuDNN D=256 backward kernel on Blackwell (SM10x) GPUs: BSHD layouts with cuDNN ≥ 9.23 and THD layouts with cuDNN ≥ 9.25. The C++ backend gate in fused_attn.cpp is carefully guarded (arch range, cuDNN version, no-bias, no-dropout, vanilla softmax, SWA restricted to causal masks), and corresponding skip logic in both PyTorch and JAX test helpers mirrors those constraints faithfully.

• fused_attn.cpp: Adds the D=256 BWD sub-condition with is_thd_layout-based cuDNN version branching, layout/bias/dropout/softmax/window-size guards, and an SM10x arch ceiling (sm_arch_ < 110).
• backends.py / utils.py: Adds .contiguous() materialization for non-contiguous QKV views under FA4 D=256 on SM10x, and disables FA4 for D=256 + SWA on SM10x (deferring to cuDNN fused attn).
• Test files (PyTorch + JAX): New test_dpa_d256, test_dpa_fa4_hdim256 forward-only fix, D=256 CP model configs, D=256 distributed CP ring/all-gather tests, and per-test runtime skip guards consistent with the C++ conditions.

Confidence Score: 5/5

Safe to merge once the developer is satisfied with THD readiness — the C++ gate already requires cuDNN ≥ 9.25 for THD, so production callers with older stacks will not reach the new kernel.

The functional changes are narrowly scoped to the backend-selection gate and two forward-pass tensor preparation fixes; all have correct version and architecture guards. The test additions are conditional and fall back gracefully. The two observations flagged are test coverage gaps that do not affect runtime correctness.

Files Needing Attention: The test coverage gap in tests/pytorch/attention/test_attention.py is worth a second look: test_dpa_d256 with thd_t2hd + non-padding-mask models runs but does not reach the cuDNN D=256 BWD kernel it is nominally guarding against cuDNN 9.25.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn.cpp	Adds D=256 BWD backend condition with well-structured guards: SM10x arch range, cuDNN version thresholds (9.23 BSHD / 9.25 THD), is_thd_layout flag, bias/dropout/softmax constraints, and SWA restriction to causal masks only.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Adds .contiguous() materialization for non-contiguous QKV tensors when FA4 D=256 runs on SM10x. Safe: contiguous() is a no-op for already-contiguous tensors.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Disables FA4 for D=256 on SM10x when sliding-window attention is active (non-causal window), correctly deferring those cases to cuDNN fused attention.
tests/pytorch/attention/test_attention.py	Adds new test_dpa_d256 covering Blackwell D=256 BWD via cuDNN fused attn; makes is_training a caller-supplied parameter. The thd_t2hd layout + non-padding-mask model combos fall back to unfused rather than exercising the guarded cuDNN path.
tests/pytorch/attention/test_attention_with_cp.py	Adds D=256 configs (cp_5_0, cp_5_1) to model_configs_fused_attn and runtime skip guards for non-Blackwell hardware, FP8, unsupported cp_comm_types, and insufficient cuDNN version.
tests/pytorch/attention/run_attention_with_cp.py	Removes the previously-broken model_configs_fused_attn_d256 import; replaces the unconditional dict lookup with an explicit assertion fallback.
tests/jax/test_fused_attn.py	Adds detailed D=256 BWD skip logic in _check_configs() mirroring C++ constraints (bias, dropout, softmax, SWA rules), plus two new D=256 test param entries (BSHD and THD layouts).
tests/jax/test_distributed_fused_attn.py	Adds D=256 distributed CP ring and all-gather tests. skip_if_d256_cp_unsupported uses integer range 100–110 (any SM10x) unlike the explicit (10,0)/(10,3) tuple used in PyTorch tests — a minor scope inconsistency.
qa/L1_jax_distributed_unittest/test.sh	Moves test_distributed_fused_attn.py before the NCCL-workaround block; updates a TODO attribution. No functional change to the test suite.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["nvte_get_fused_attn_backend\nD=256 bprop path"] --> B{SM arch\n100 to 109?}
    B -- No --> Z["NVTE_No_Backend\nfalls back to\nunfused attn"]
    B -- Yes --> C{THD layout?}
    C -- Yes --> D{cuDNN ge 9.25?}
    C -- No --> E{cuDNN ge 9.23?}
    D -- No --> Z
    E -- No --> Z
    D -- Yes --> F
    E -- Yes --> F{NO_BIAS and\nno dropout and\nVANILLA_SOFTMAX?}
    F -- No --> Z
    F -- Yes --> G{full-window or\ncausal SWA only?}
    G -- No --> Z
    G -- Yes --> H["NVTE_F16_arbitrary_seqlen\ncuDNN D=256 BWD kernel"]

    style H fill:#22c55e,color:#fff
    style Z fill:#ef4444,color:#fff

Loading

_{Reviews (23): Last reviewed commit: "Use is_training in DPA test helper" | Re-trigger Greptile}

[KshitijLakhani]

Latest testing: Pipeline [56553160]
I manually set the cuDNN BE in the CI to be 9.30+ and the cuDNN FE to point at latest develop as of 07/01/26 to test the THD support for D=256
I confirmed that the newly added tests for BSHD and THD (PyT and JAX, both) pass

[KshitijLakhani]

/te-ci L2

[KshitijLakhani]

Latest testing: Pipeline [58978691]
I manually set the cuDNN BE in the CI to be 9.30+ and the cuDNN FE to be 1.26 to test the THD support for D=256
I confirmed that the newly added tests for BSHD and THD (PyT and JAX, both) pass for non-CP and CP

[cyanguwa]

No problem with the "id", but it seems inconsistent with other PyTorch tests. Also, please check the 9.30 vs 9.25 version guard.

[KshitijLakhani]

Will do

[cyanguwa]

Could you remove the 'id' fields actually - they just look really odd sitting next to other tests with no 'id's. :D

[KshitijLakhani]

Done SHA 5ef0d4f

[cyanguwa]

Should we test no_mask, padding_causal as well, or not really?

[KshitijLakhani]

I think it's okay to not test those two, if you are fine with it ?
The non-CP tests cover no_mask and I believe THD CP anyways translates to padding_causal under the hood, so I decided to keep this lightweight.

[KshitijLakhani]

Latest testing: Pipeline 59713381
I manually set the cuDNN BE in the CI to be 9.25+ (cuDNN FE defaults to 1.26) to test the full support for D=256. This pipeline uses base image ghcr.io/nvidia/jax:jax for JAX and defaults for PyT and Core.
I confirmed that the newly added tests for BSHD and THD (PyT and JAX, both) pass for non-CP and CP

[cyanguwa]

I think we can remove the "fused_xx" keyword from model_configs_fused_d256, fused_d256_no_mask, and test_dpa_fused_attn_d256. Those tests will run both unfused/flash and fused for comparison. Maybe there's no need to restrict it to "fused". Essentially what we want to test here is d256.

[KshitijLakhani]

Done SHA eda1a74

[KshitijLakhani]

TL; DR: New pipeline passes all tests fully after latest pushed commits

Details:
I added two commits (top two commits in the branch): f61ef61 and 72958ee after your approval @cyanguwa as I noticed a couple failures in the L3 FA tests. These commits were pushed to not fix an issue in the functionality added in this PR (i.e. support for the cuDNN backend for D=256) but rather to fix an issue in flash attention 4 support configuration in TE.
The tests and feature support added in this PR were inadvertently triggering Flash attention 4 code paths that aren't supported (or need some tweaking before being supported) by FA4

_Added two focused FA4 fixes :

• disable FA4 for SM10x D=256 local/SWA, which FA4 does not support yet, and
• materialize non-contiguous packed Q/K/V views before FA4 for SM10x D=256 so packed BSHD/THD layouts compare correctly._

I also added a change that keeps the pre-existing FA4 hdim256 test forward-only, matching its prior effective coverage (in main) before cuDNN fused D=256 backward became available in this PR

Latest PyT only pipeline launched (as changes were local to PyT only): #60370958
Tests look clean ✅