[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.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]

Greptile Summary

This PR extends Blackwell D=256 fused-attention backward support and its validation coverage.

• Adds architecture, cuDNN-version, layout, bias, dropout, softmax, paging, and sliding-window constraints to common backend selection.
• Adds JAX and PyTorch single-device and context-parallel D=256 tests.
• Updates JAX distributed CI ordering and PyTorch context-parallel test configuration handling.

Confidence Score: 3/5

The PR does not yet appear safe to merge because THD D=256 backward may still be enabled before cuDNN can construct a supported execution plan.

The current selector routes qualifying THD D=256 training calls on cuDNN 9.25 and newer into the cuDNN backend, but the prior thread and PR support statement leave unresolved whether this path is valid before cuDNN 9.30; a failed plan build would surface as a runtime exception rather than 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 the D=256 SM10x backward backend gate, but the THD minimum remains unresolved relative to the previously reported execution-plan failure.
tests/jax/test_fused_attn.py	Adds D=256 configuration filtering and deterministic backward coverage aligned with the common selector.
tests/jax/test_distributed_fused_attn.py	Adds D=256 ring and all-gather context-parallel coverage with hardware and cuDNN guards.
tests/pytorch/attention/test_attention.py	Adds Blackwell D=256 fused-attention backward configurations and version-gated BSHD/THD coverage.
tests/pytorch/attention/test_attention_with_cp.py	Adds D=256 context-parallel configurations and restricts their test matrix to supported combinations.
tests/pytorch/attention/run_attention_with_cp.py	Keeps CP fused-attention configurations in the existing exported mapping and validates unknown model names.
qa/L1_jax_distributed_unittest/test.sh	Reorders the fused-attention distributed test ahead of the NCCL-specific softmax invocation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A["PyTorch / JAX D=256 training"] --> B["Common fused-attention backend selector"]
  B --> C{"SM10x, cuDNN version, layout, and feature guards pass?"}
  C -- No --> D["No fused backend / framework fallback"]
  C -- Yes --> E["cuDNN F16 arbitrary-seqlen backend"]
  E --> F["D=256 backward execution plan"]

Loading

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

[cyanguwa]

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

[cyanguwa]

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

[KshitijLakhani]

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

[KshitijLakhani]

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]

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?

[KshitijLakhani]

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]

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

[jepio]

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

[KshitijLakhani]

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.

[cyanguwa]

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

[KshitijLakhani]

You are right indeed. Thanks for the note. I've made changes for the same

[KshitijLakhani]

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

[cyanguwa]

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

[KshitijLakhani]

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.

[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]

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?

[KshitijLakhani]

I considered reusing the existing D=256 configs, but they are tied to other coverage: decode-style generic tests, mask/layout sweeps (using post scale bias which would be rejected by cuDNN D=256 bprop path), FA4, or MLA where d_v != 256 (hence not the same D=256 cuDNN kernel path either).

Changing them would either remove existing coverage or mix cuDNN D=256-specific gates into generic tests. I kept this as a small dedicated dict so the cuDNN FusedAttention D=256 bprop surface is explicit and isolated.

If you have a strong opinion on this, I can arm twist the code to do so

[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.

[cyanguwa]

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)

[KshitijLakhani]

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 ?

[cyanguwa]

Yes, please!

[KshitijLakhani]

Regarding folding this into the regular test_cp_with_fused_attention() tests, that was my first approach actually but I dropped it in favor of the current approach.

The current test_cp_with_fused_attention() is a pretty broad sweep which would add quite a few test cases (while still missing on quite a few as essential mode does not cover BSHD and FP16, both of which we want for D=256). Most of the collected test cases would just be skips adding unnecessary noise in the logs (e.g. fp8, diff mask combos, etc.) so I decided to take a more funneled approach.

The focused test is clearer: it directly covers BSHD and THD, fp16 and bf16, P2P and all-gather, plus THD SWA, without polluting the broad existing CP matrix.

Are you okay with it being as is @cyanguwa ?

[cyanguwa]

"essential mode does not cover BSHD and FP16, both of which we want for D=256"

Do we have SBHD tests in the essential mode instead? They should be equivalent to testing BSHD though.

I feel the D256 feature shouldn't be that special to have its own test function, even though the kernels are written in cutedsl, different from the regular SDPA. But we seem to be always expanding the range of head dims we support so do we want to add a new test every time, and would the test code become a bit repetitive.

[KshitijLakhani]

Do we have SBHD tests in the essential mode instead? They should be equivalent to testing BSHD though.

You are right about that

[KshitijLakhani]

I folded in the D=256 tests into the test_cp_with_fused_attention and added a bunch of skips in there for D=256 in commit 9921b99
This should cover most of the earlier D=256 tests I had added.
The only differences are:

• Older D=256 tests ran BF16 and FP16 in essential mode but the Newer D=256 tests only run BF16 (FP8 is not supported) and FP16 is ran as non-essential, which is consistent with other fused attn tests
• Older D=256 tests ran BSHD in essential mode but the Newer D=256 tests run SBHD instead, which is consistent with the other fused attn tests

So with this commit I think we are consistent throughout

[cyanguwa]

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

[KshitijLakhani]

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

[cyanguwa]

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]

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