[JAX] [PyT] [Common] Enable D=256 BWD cuDNN fused attn for Blackwell CC 10.x by KshitijLakhani · Pull Request #3056 · NVIDIA/TransformerEngine

KshitijLakhani · 2026-05-28T19:51:46Z

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.30 + 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 · 2026-05-29T23:00:49Z

Greptile Summary

Adds Blackwell D=256 backward fused-attention support and corresponding framework coverage.

Gates the common cuDNN backend by architecture, layout-specific cuDNN version, bias, dropout, softmax, mask, and window constraints.
Adds JAX and PyTorch single-device and context-parallel D=256 tests.
Aligns FlashAttention 4 backend selection and tensor materialization with D=256 kernel limitations.
Repairs the context-parallel runner to use the existing fused-attention configuration map.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn.cpp	Adds the architecture-, layout-, version-, and feature-gated D=256 backward backend selection.
tests/jax/test_fused_attn.py	Adds D=256 coverage and skip conditions aligned with the common backend gate.
tests/jax/test_distributed_fused_attn.py	Adds ring and all-gather context-parallel coverage for supported BSHD and THD configurations.
tests/pytorch/attention/test_attention.py	Adds backward D=256 attention configurations while preserving forward-only FA4 coverage.
tests/pytorch/attention/test_attention_with_cp.py	Adds D=256 context-parallel configurations with hardware, dtype, communication, and cuDNN guards.
tests/pytorch/attention/run_attention_with_cp.py	Completes the prior import fix by resolving every fused-attention model through the existing configuration map.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Materializes non-contiguous D=256 tensors before invoking FlashAttention 4 on supported architectures.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Disables FlashAttention 4 for unsupported D=256 local-attention windows.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[D=256 backward request] --> B{SM10x server GPU?}
  B -- No --> Z[Use another backend or fallback]
  B -- Yes --> C{Layout-specific cuDNN version met?}
  C -- No --> Z
  C -- Yes --> D{No bias or dropout and vanilla softmax?}
  D -- No --> Z
  D -- Yes --> E{Window and mask supported?}
  E -- No --> Z
  E -- Yes --> F[Select cuDNN F16 arbitrary-seqlen backend]

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

KshitijLakhani · 2026-07-02T18:22:39Z

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 · 2026-07-08T04:57:23Z

/te-ci L2

KshitijLakhani · 2026-07-22T07:15:37Z

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 · 2026-07-22T16:28:31Z

+            ),
+        ),
+    ],
+)


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

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

Done SHA 5ef0d4f

cyanguwa · 2026-07-22T16:30:18Z

+    "cp_hd256_causal": ModelConfig(2, 1024, 16, 256, attn_mask_type="causal"),
+    "cp_hd256_causal_swa": ModelConfig(
+        2, 1024, 16, 256, attn_mask_type="causal", window_size=(128, 0)
+    ),


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

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 · 2026-07-27T16:29:08Z

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 · 2026-07-27T17:46:17Z

+        ),
+    ],
+)
+def test_dpa_fused_attn_d256(dtype, model_configs, model, qkv_layout):


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.

Done SHA eda1a74