[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 D=256 head-dim backward pass support for cuDNN fused attention on Blackwell (SM10x) GPUs by adding a new version-gated condition in the C++ backend selector and matching test coverage across PyTorch, JAX, and distributed CP test suites. BSHD support is gated at cuDNN 9.23+ and THD at cuDNN 9.25+ (the THD path is forward-looking and acknowledged as not yet stable pending a cuDNN fix).

fused_attn.cpp adds a new head-dim condition: D=256 + SM10x + training + specific constraints (no bias, no dropout, vanilla softmax, full-window or causal-only SWA) routes to the NVTE_F16_arbitrary_seqlen backend.
New tests in test_attention.py, test_attention_with_cp.py, test_fused_attn.py, and test_distributed_fused_attn.py exercise the new path with per-layout cuDNN version skip guards.
test.sh reorders the distributed fused-attn test suite to run before the softmax workaround block.

Confidence Score: 5/5

Safe to merge once the THD cuDNN version question is resolved and the SM10x guard style is harmonised; the core C++ backend logic is correctly constrained and no existing paths are altered.

The C++ change adds a tightly scoped, additive condition that only activates for the exact combination of D=256+SM10x+training+no-bias+no-dropout+vanilla-softmax, leaving every other code path untouched. The THD path is intentionally forward-looking and acknowledged by the author as not yet merged.

transformer_engine/common/fused_attn/fused_attn.cpp — confirm whether 9.25 or 9.30 is the correct cuDNN minimum for THD before merging. test_attention.py and test_attention_with_cp.py — consider widening the SM compute-capability allowlist to the full SM10x range to match the backend.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn.cpp	Adds D=256 BWD condition for SM10x; logic is well-guarded. `is_thd_layout` helper added at top of function. THD path is intentionally forward-looking and the PR is not intended to be merged until cuDNN confirms THD support.
tests/pytorch/attention/test_attention.py	Adds `test_dpa_fused_attn_d256` with four model configs covering no_mask, padding, causal+SWA, and GQA variants; uses explicit tuple-style SM compute capability skip guard `not in ((10, 0), (10, 3))`.
tests/pytorch/attention/test_attention_with_cp.py	Adds `model_configs_fused_attn_d256` and `test_cp_with_fused_attention_d256` covering BSHD/THD x p2p/all_gather x causal/SWA; uses same explicit tuple SM check as test_attention.py.
tests/pytorch/attention/run_attention_with_cp.py	Minor update: imports `model_configs_fused_attn_d256` and adds an elif branch to route D=256 CP model names to the right config dict.
tests/jax/test_fused_attn.py	Adds D=256 skip logic with detailed guards (head_dim, cuDNN version, bias, dropout, softmax, SWA) and two new parametrized test cases for BSHD and THD KV-packed layouts.
tests/jax/test_distributed_fused_attn.py	Adds `test_context_parallel_ring_attn_d256` and `test_context_parallel_allgather_attn_d256` with `skip_if_d256_cp_unsupported`; uses a range-style SM check consistent with the C++ backend.
qa/L1_jax_distributed_unittest/test.sh	Moves `test_distributed_fused_attn.py` to run before the XLA_FLAGS workaround block and updates a TODO comment author name.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["nvte_get_fused_attn_backend()"] --> B{head_dim_qk == 256 and head_dim_v == 256?}
    B -- No --> C[Other head_dim conditions]
    B -- Yes --> D{is_training and sm_arch_ in range 100-109?}
    D -- No --> C
    D -- Yes --> E{is_thd_layout?}
    E -- THD --> F{cuDNN >= 9.25?}
    E -- BSHD --> G{cuDNN >= 9.23?}
    F -- No --> C
    G -- No --> C
    F -- Yes --> H{bias=NO_BIAS, dropout=0.0, softmax=VANILLA, not paged-KV?}
    G -- Yes --> H
    H -- Fail --> C
    H -- Pass --> I{Window: full-window OR causal-mask with right in -1 or 0?}
    I -- No --> C
    I -- Yes --> J[flag_arb = true: NVTE_F16_arbitrary_seqlen D=256 BWD]

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

cyanguwa · 2026-06-01T13:58:11Z

+              attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK ||
+              attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK ||
+              attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) &&
+             (window_size_right == -1 || window_size_right == 0)))) ||


Could these changes be moved to before "\ bias type" just so it's following an increasing order of the cuDNN version?

Does this new feature support BSHD/SBHD and THD? It looks like the tests are focused on BSHD/SBHD only.

RE: THD support
I did test BSHD and BSHD+CP and it did pass on the JAX side and the CI for the PyT side did not fail either so I think that works.
My testing revealed that THD support is not yet available (Bwd plan compialtion issue) so I've filed a bug and shared a reproducer for the same with the cuDNN team: NVIDIA/cudnn-frontend#276

Could these changes be moved to before "\ bias type" just so it's following an increasing order of the cuDNN version?

Fixed in a264de1

cyanguwa · 2026-06-01T14:11:46Z

+                        "D=256 BWD on Blackwell only supports right window -1 or 0"
+                        " for causal masks."
+                    )
+


Aren't these checks duplicate to the checks we added on the C++ side? Would the call FusedAttnHelper().get_fused_attn_backend() give you the same gating effect?

So if we are just interested in the gating effect, you are right. The get_fused_attn_backend() will return NVTE_No_Backend and then there's a catch-all at the end which basically skip the tests as there is no fused attn backend avalable.

However, the reason for this to be here is to give a meaningful reason as to why a test is being skipped as compared to a generic "Unsupported inputs combination or device compute capability." message which does not qualify the reason for the skip. Unfortunately, on the JAX attn side we do not log the reason for disabling fused attn in the feature code like we have on the Pytorch side in d_p_a/utils.py. So there is no way for the user to know why the test was skipped. Hence, we need to rely on test code to log this on the JAX side.

I'd suggest we leave this in here for now. And when your PR for generating log messages in the C++ level when selecting the attn backend is ready, I can plumb it through onto the JAX side and then as part of that clean up, get rid of all the skip messages in check_configs()

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

jepio · 2026-07-08T10:59:54Z

+          // 9.23: d_qk = d_v = 256 + SM10x (cuDNN FE 1.24 / BE 9.23+) + bprop + non-paged.
+          // THD layouts require cuDNN FE 1.26 / BE 9.30+ for execution-plan support.
+          (head_dim_qk == 256 && head_dim_v == 256 && is_training && sm_arch_ >= 100 &&
+           sm_arch_ < 110 && cudnn_runtime_version >= (is_thd_layout ? 93000 : 92300) &&


The PR body says this should be 92500, not 93000. Which one is correct?

Good catch !
I've updated the PR body to reflect cuDN 9.30. I believe the body was a preemptive action from me while I was waiting for cuDNN to add support for THD via the C++ API.

I think cuDNN has renamed 9.30 to 9.25? Please double check!

KshitijLakhani · 2026-07-08T05:04:17Z

+            # TODO(KshitijLakhani): cuDNN FE can model bias input separately from dBias,
+            # but TE does not yet plumb whether dBias is requested into the common backend selector.
+            # Until that distinction is available, the D=256 SM10x gate requires no bias.


I've opened a GH issue for this: #3188

#2964 will plumb bias/dbias shapes/requirements to the C side, so hopefully when that's merged, this wouldn't be an issue.

KshitijLakhani · 2026-07-08T17:29:59Z

+          // 9.23: d_qk = d_v = 256 + SM10x (cuDNN FE 1.24 / BE 9.23+) + bprop + non-paged.
+          // THD layouts require cuDNN FE 1.26 / BE 9.30+ for execution-plan support.
+          (head_dim_qk == 256 && head_dim_v == 256 && is_training && sm_arch_ >= 100 &&
+           sm_arch_ < 110 && cudnn_runtime_version >= (is_thd_layout ? 93000 : 92300) &&


Good catch !
I've updated the PR body to reflect cuDN 9.30. I believe the body was a preemptive action from me while I was waiting for cuDNN to add support for THD via the C++ API.

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>

…ersions 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:24:12Z

  NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout);
+  const bool is_thd_layout = q_format == NVTE_QKV_Format::NVTE_THD ||
+                             kv_format == NVTE_QKV_Format::NVTE_THD ||
+                             qkv_format == NVTE_QKV_Format::NVTE_THD;


is_thd_layout ORs q_format, kv_format, and the combined qkv_format against NVTE_THD — the third is redundant with the first two, but harmless.

That's right. Let me get rid of it and push a commit

cyanguwa · 2026-07-22T16:26:58Z

+    "fused_hd256_padding_causal_gqa": ModelConfig(
+        2, 1024, 16, 256, num_gqa_groups=4, attn_mask_type="padding_causal"
+    ),
+}


Can we reuse some of the configs in other tests (I think there are some h256 ones as well), or can we keep this dict and remove those other h256 ones to deduplicate?

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.

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?

cyanguwa · 2026-07-22T16:36:33Z

+)
+@pytest.mark.parametrize("dtype", ["bf16", "fp16"])
+@pytest.mark.parametrize("model,qkv_format,cp_comm_type", DPA_CP_FUSED_HD256_CASES)
+def test_cp_with_fused_attention_hdim256(cp_pool, dtype, model, qkv_format, cp_comm_type):


Could we fold this into the regular test_cp_with_fused_attention() tests instead of creating this extra test code and extra tests? You can modify some of the existing tests to cover the d256 cases as well, or force some skips if you think it's adding too many tests (I feel your pytest params above might have enumerated all the combinations anyway). Thanks.

Also, please use d256 as the reference, same as in the title, instead of hdim256. (nit)

Also, please use d256 as the reference, same as in the title, instead of hdim256. (nit)

Done. Do you want me to change it for the Flas Attn tests too ?

Yes, please!

cyanguwa · 2026-07-22T16:42:45Z

+        # cuDNN FE 1.26 / BE 9.30. The kernel rejects dBias, dropout, and ALiBi, supports vanilla
+        # softmax only, and allows SWA together with a causal mask only.
+        is_sm10x = 100 <= compute_capability < 110
+        if self.is_training and is_sm10x and (self.head_dim_qk == 256 or self.head_dim_v == 256):


Shouldn't these checks already be covered on the C side? Are these by any chance duplicate?

IIUC, is your concern the same as you had in your earlier review: #3056 (comment). ? I think my comment goes into more details. Happy to answer any follow ups on that.

I will get rid of these extraneous checks in the tests (if #2964 does not already) once #2964 is merged in. For now, I think we should let them be in. The clean up for these test messages can follow up once we've confirmed that a more descriptive string is plumbed through to the tests

Sounds good - are there any CP specific skips/conditions we should add?

#2964 will surface the reject reason to Jax FusedAttnHelper.get_fused_attn_backend() so it'd be nice to use that as a single source of rejection/pass reasons. The d256 feature reject reason does come from cudnn-frontend so we should be able to remove some of the code (and reduce some maintenance cost).

cyanguwa · 2026-07-22T16:44:01Z

+            # TODO(KshitijLakhani): cuDNN FE can model bias input separately from dBias,
+            # but TE does not yet plumb whether dBias is requested into the common backend selector.
+            # Until that distinction is available, the D=256 SM10x gate requires no bias.


#2964 will plumb bias/dbias shapes/requirements to the C side, so hopefully when that's merged, this wouldn't be an issue.

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

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

KshitijLakhani commented Jul 8, 2026

View reviewed changes

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 11 commits July 21, 2026 10:29

Switching test check logic around

021c450

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

Add THD xfail marked tests for testing

b597eda

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

Fix the jax test bias condition fo skipping for D=256

4882cda

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

nit: Fix comment

4d5ad94

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

Move the D=256 changes so as to follow an order of increasing cuDNN v…

6184e85

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

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

424b2c4

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

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

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: 5/5

Important Files Changed

Flowchart

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!

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

Uh oh!

Choose a reason for hiding this comment

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!

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!

KshitijLakhani commented Jul 22, 2026

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!

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!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

KshitijLakhani commented May 28, 2026 •

edited

Loading

greptile-apps Bot commented May 29, 2026 •

edited

Loading

KshitijLakhani Jun 5, 2026 •

edited

Loading

KshitijLakhani commented Jul 2, 2026 •

edited

Loading