[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

This PR enables SM10x D=256 fused-attention backward dispatch and expands corresponding JAX and PyTorch coverage.

Adds architecture, cuDNN-version, layout, bias, dropout, softmax, and window guards to the common backend selector.
Adds single-device and context-parallel D=256 tests for JAX and PyTorch.
Repairs the PyTorch CP runner’s fused-attention configuration import and lookup.

Confidence Score: 3/5

The PR is not yet safe to merge until THD D=256 execution-plan support at the new cuDNN 9.25 threshold is verified or the selector excludes THD.

The selector now routes THD D=256 backward to the cuDNN backend at version 9.25, but the available evidence does not establish that the previously reported execution-plan failure is resolved; if it remains, plan construction raises a hard runtime error without fallback.

Files Needing Attention: transformer_engine/common/fused_attn/fused_attn.cpp

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn.cpp	Adds constrained SM10x D=256 backward backend selection, including a THD-specific cuDNN 9.25 threshold whose runtime support remains unverified.
tests/jax/test_fused_attn.py	Aligns D=256 test skips with the common selector and adds BSHD and THD coverage.
tests/jax/test_distributed_fused_attn.py	Adds version-gated D=256 ring and all-gather context-parallel test coverage.
tests/pytorch/attention/test_attention.py	Adds isolated D=256 attention configurations and copies shared configurations before mutation.
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	Fixes the undefined configuration import and resolves D=256 models through the existing fused-attention dictionary.
qa/L1_jax_distributed_unittest/test.sh	Reorders the distributed fused-attention suite ahead of the softmax-specific XLA flag invocation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A["JAX / PyTorch attention"] --> B["Common fused-attention selector"]
  B --> C{"SM10x D=256 constraints satisfied?"}
  C -- "No" --> D["No fused backend / framework fallback"]
  C -- "Yes" --> E["NVTE_F16_arbitrary_seqlen"]
  E --> F["cuDNN execution-plan validation and build"]
  F --> G["D=256 backward"]

_{Reviews (18): 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

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

for more information, see https://pre-commit.ci

…n fused attn Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

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.

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

for more information, see https://pre-commit.ci

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

cyanguwa

Other than the two nits (test name change and 'id' removal), looks good!

KshitijLakhani · 2026-07-27T21:01:52Z

+           // Blackwell rejects dBias, dropout, and ALiBi (and supports vanilla softmax only).
+           bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && dropout == 0.0 &&
+           softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX &&
+           // Non-causal D=256 supports only full-window attention; SWA is allowed only for causal masks.


@cyanguwa Here's the cuDNN FE restriction in this class for non-causal SWA : https://github.com/NVIDIA/cudnn-frontend/blob/c5012bd3011d28d8ffdf0307b8c213d1b97c68e0/python/cudnn/sdpa/fwd/api.py#L22 and the check is at https://github.com/NVIDIA/cudnn-frontend/blob/c5012bd3011d28d8ffdf0307b8c213d1b97c68e0/python/cudnn/sdpa/fwd/api.py#L201

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

for more information, see https://pre-commit.ci

KshitijLakhani force-pushed the klakhani/feat/enable-headdim256-bwd-sm100 branch from 51ad582 to d177ecf Compare May 28, 2026 23:05

KshitijLakhani changed the title ~~[JAX] [Common] Enable D=256 BWD cuDNN fused attn for Blackwell CC 10.x~~ [JAX] [PyT] [Common] Enable D=256 BWD cuDNN fused attn for Blackwell CC 10.x May 28, 2026

KshitijLakhani marked this pull request as ready for review May 29, 2026 22:56

KshitijLakhani requested a review from cyanguwa as a code owner May 29, 2026 22:56

KshitijLakhani self-assigned this May 29, 2026

greptile-apps Bot reviewed May 29, 2026

View reviewed changes

Comment thread tests/pytorch/attention/test_attention.py Outdated

Comment thread tests/jax/test_fused_attn.py Outdated

cyanguwa reviewed Jun 1, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/fused_attn.cpp Outdated

cyanguwa reviewed Jun 1, 2026

View reviewed changes

Comment thread tests/jax/test_fused_attn.py

cyanguwa requested changes Jun 1, 2026

View reviewed changes

KshitijLakhani force-pushed the klakhani/feat/enable-headdim256-bwd-sm100 branch from e317f99 to e08e9e8 Compare June 3, 2026 21:48

greptile-apps Bot reviewed Jun 3, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/fused_attn.cpp Outdated

KshitijLakhani mentioned this pull request Jun 3, 2026

Missing THD support for cuDNN 9.23 + cuDNN FE 1.24 for D=256 bwd on sm10.x NVIDIA/cudnn-frontend#276

Open

KshitijLakhani force-pushed the klakhani/feat/enable-headdim256-bwd-sm100 branch from 5063b39 to 55f3bb4 Compare July 1, 2026 16:36

KshitijLakhani added the 2.18 label Jul 1, 2026

KshitijLakhani force-pushed the klakhani/feat/enable-headdim256-bwd-sm100 branch from a678b84 to ba0cb94 Compare July 8, 2026 04:55

KshitijLakhani requested a review from cyanguwa July 8, 2026 05:08

jepio reviewed Jul 8, 2026

View reviewed changes

Comment thread transformer_engine/common/fused_attn/fused_attn.cpp Outdated

KshitijLakhani commented Jul 8, 2026

View reviewed changes

Comment thread tests/jax/test_fused_attn.py

Comment thread transformer_engine/common/fused_attn/fused_attn.cpp Outdated

KshitijLakhani and others added 9 commits July 21, 2026 10:29

Refine the support for D=256 on Blackwell server type GPUs

f108426

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Add deterministic tests for D=256 for sm10.x

8ca68e2

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Test clean up

a093127

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Add PyT side tests for D=256 cuDNN fused attn on SM100-110

b7dc18c

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Refine the PyT D=256 tests

28fe7ba

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

dd6cbf3

for more information, see https://pre-commit.ci

Fix the filtering condition for bias type for D=256 on sm10x for cudn…

9920cbe

…n fused attn Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Code clean up

3f06289

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

38de70a

for more information, see https://pre-commit.ci

KshitijLakhani and others added 5 commits July 21, 2026 10:29

Add support for D=256 THD in TE

fe7c6b2

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Add TE JAX CP test support to verify D=256 for AG and Ring

cfed53f

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Add PyTorch D256 CP tests

5767db0

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

3035301

for more information, see https://pre-commit.ci

nit: clean up comments

06eee35

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

KshitijLakhani force-pushed the klakhani/feat/enable-headdim256-bwd-sm100 branch from 88217fd to 06eee35 Compare July 21, 2026 17:30

Correct the atgs passed to test dpa

ee79662

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

cyanguwa requested changes Jul 22, 2026

View reviewed changes

KshitijLakhani added 4 commits July 22, 2026 16:24

Simplify THD fused attention layout check

cbab884

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Rename PyTorch D256 fused attention tests

00d14cd

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Run JAX fused attention before softmax

ed4b745

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Relabel THD D256 cuDNN guard to 9.25

e3032d3

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

KshitijLakhani force-pushed the klakhani/feat/enable-headdim256-bwd-sm100 branch from da25221 to e3032d3 Compare July 23, 2026 00:22

[pre-commit.ci] auto fixes from pre-commit.com hooks

6f2b748

for more information, see https://pre-commit.ci

KshitijLakhani requested a review from cyanguwa July 24, 2026 04:03

greptile-apps Bot reviewed Jul 25, 2026

View reviewed changes

Comment thread tests/pytorch/attention/run_attention_with_cp.py Outdated

Fold D256 PyTorch CP tests

9921b99

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

KshitijLakhani force-pushed the klakhani/feat/enable-headdim256-bwd-sm100 branch from 078c80f to 9921b99 Compare July 27, 2026 00:37

[pre-commit.ci] auto fixes from pre-commit.com hooks

995cea4

for more information, see https://pre-commit.ci

cyanguwa reviewed Jul 27, 2026

View reviewed changes

cyanguwa previously approved these changes Jul 27, 2026

View reviewed changes

KshitijLakhani commented Jul 27, 2026

View reviewed changes

KshitijLakhani added 2 commits July 27, 2026 14:37

Rename D256 PyTorch attention tests

eda1a74

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Remove D256 PyTorch attention param ids

5ef0d4f

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

KshitijLakhani dismissed cyanguwa’s stale review via 5ef0d4f July 27, 2026 21:53

KshitijLakhani requested a review from cyanguwa July 27, 2026 21:54

[pre-commit.ci] auto fixes from pre-commit.com hooks

d2539ac

for more information, see https://pre-commit.ci

cyanguwa approved these changes Jul 28, 2026

View reviewed changes

Uh oh!

Conversation

KshitijLakhani commented May 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented May 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 3/5

Important Files Changed

Flowchart

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

KshitijLakhani commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

KshitijLakhani commented Jul 8, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

KshitijLakhani commented Jul 22, 2026

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

KshitijLakhani commented Jul 27, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cyanguwa left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

KshitijLakhani commented May 28, 2026 •

edited

Loading

greptile-apps Bot commented May 29, 2026 •

edited

Loading

KshitijLakhani commented Jul 2, 2026 •

edited

Loading