[JAX] Align fused attn backward output gradient sharding and other test strengthening

[KshitijLakhani]

Description

Fix JAX fused-attention bwd sharding and THD metadata handling, while strengthening bwd test coverage with nonuniform output gradients.

The starting point was fixing the JAX fused-attention bwd sharding, but along the way as I decided to strengthen the testing as part of root causing and then just fixed any ripple effects it had in exposing smaller issues in our feature code

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

1. Align bwd doutput sharding with the saved fwd output sharding : JAX can infer a doutput sharding that is incompatible with the cuDNN bwd kernel, particularly when CP reorder operations are differentiated through. TE now explicitly aligns doutput sharding with the saved output sharding.

2. Use nonuniform output gradients in CP tests: Existing tests generated a uniform output gradient from mean(output). Uniform values are unchanged by permutation and therefore masked incorrect CP placement.
   NOTE: This was the reason why TE tests did not catch the doutput sharding bug from above and so I decided to go with a more generic case.

3. Route rotated THD ring metadata through the general path: The THD fast metadata path assumes Q and KV describe matching token metadata. Rotated ring-attention steps may have different local Q and KV boundaries, so those steps now explicitly disable the fast path.
   NOTE: This was found while strengthening the CP test with a nonuniform doutput as mentioned above

4. Make general THD path selection metadata-aware: Non-CP calls retain the fast path when Q and KV metadata match. Equal-length cross-attention with independent metadata uses the general path.

5. Strengthen bwd fused attn tests:

   1. All bwd runner cases now use deterministic nonuniform doutput, including max-logit, score-mod, and distributed self/cross attention. The runner rejects an accidental missing seed. Same motivation as earlier.
   2. The cancellation-sensitive scalar objective comparison was removed; fwd outputs and bwd gradients continue to be compared element wise.
   3. Distributed BF16 dBias uses a reduction-aware default tolerance because its collective accumulation order differs from the unsharded reference.

6. Correct DP/FSDP test data placement (side effect of Bugfix spec check for DP+FSDP configurations #3272) :

   1. The eight-GPU dp2 × fsdp2 × tp2 test previously sharded activations over DP but replicated them over FSDP. TE then correctly reduced dBias across both axes, summing duplicate FSDP contributions.
   2. Batch-dependent test inputs and metadata are now sharded over the compound (dp, fsdp) axis. Related batch-divisibility checks use the combined data-parallel size.

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

[KshitijLakhani]

/te-ci jax L0 L1

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; no outstanding correctness, security, or repository-rule violation was identified.

Summary

The PR corrects JAX fused-attention backward sharding and THD metadata routing while strengthening gradient and distributed-placement tests.

• Aligns backward doutput placement with the saved forward output across fused and context-parallel primitives.
• Routes rotated striped-THD ring steps through the general metadata path and selects the fast path only for matching Q/KV metadata.
• Uses deterministic, nonuniform output gradients to expose permutation and sharding errors.
• Corrects compound DP/FSDP test-data placement and adjusts distributed BF16 dBias tolerances for reduction order.
• Adds explicit self-versus-cross-attention metadata coverage, including equal-length cross-attention.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    F[Saved forward output] --> B[Fused-attention backward]
    D[Incoming doutput] --> S[Align doutput sharding to output]
    S --> B
    M[THD Q/KV metadata] --> E{Metadata equal and fast path allowed?}
    E -->|Yes| FP[Boundary-only fast path]
    E -->|No| GP[General metadata path]
    FP --> K[cuDNN fused-attention kernel]
    GP --> K
    B --> K
    K --> G[Q/K/V and optional bias gradients]

Loading

_{Reviews (2) · Last reviewed commit: "Merge branch 'main' into klakhani/fix/ja..."}

[jberchtold-nvidia]

Was this assert_allclose also part of the debugging code you're removing here? Thanks for removing the debug prints

[jberchtold-nvidia]

What is the performance impact of not picking the fast path?

If it is significant, it might be nice to log a warning with the reasoning the fast path was unable to be taken. This would give users insight if they see a sudden TE perf regression if they take a path where the fast path is not supported

[jberchtold-nvidia]

When would a user want to intentionally set this to False? It seems internally we already automatically switch to a fallback if this allow_fast_causal_path is True but it isn't supported in that configuration.

If there isn't a valid use case for allow_fast_causal_path=False at the user level, maybe we keep this logic internal and don't expose this knob on this public API function.

If there is a valid use case where setting it to False is required/recommended, then let's expand the docstring to document that use case

[jberchtold-nvidia]

This is okay for now and pretty much the best we can do in the current sharding paradigm with automatic sharding propagation. But it could in theory trigger an AG and replication if one of the shardings was replicated, since it would copy that to the other arg's sharding. However, that is correct behavior since it ensures correct outputs and is mitigating the user-provided incorrect sharding

Good news, when we move to explicit sharding, the shardings are propagated at trace time so we can actually raise an error in cases like this and require the user to provide matching sharding. This will avoid cases like this where we have to do something safe but un-optimal (due to incorrect user-provided sharding)

[jberchtold-nvidia]

It's nice that you refactored this commonly re-used code into a helper so when we move to explicit sharding we can handle this error raising in one place