[All] Guard THD learnable dSink on older cuDNN

[KshitijLakhani]

Description

cuDNN versions before 9.26 have a known backward-pass issue for learnable softmax with packed THD inputs. The generic kernel can incorrectly index ragged softmax statistics, producing an incorrect dSink or a potential out-of-bounds read.

Fixes #3249 (potentially)

Validation

• Locally validated backend selection across affected and unaffected configurations.
• Confirmed that the issue reproducer passes with cuDNN 9.26.0.37.
• Confirmed the existing THD softmax regression test passes.
• Did not add CI tests, even though I validated locally

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

The change in this PR disables FusedAttention when all the following apply:

• Training with packed THD inputs
• Learnable softmax
• cuDNN older than 9.26
• head_dim_v is not 64, 128, or 256

When the guard is triggered, TE warns the user to upgrade to cuDNN 9.26 or later. TE falls back to UnfusedDotProductAttention when available; otherwise, the existing “no backend available” error is raised.
Note: Inference, non-THD layouts, other softmax types, and cuDNN 9.26+ are unchanged.

• Common C++ change: implements the actual guard. It makes nvte_get_fused_attn_backend() return NVTE_No_Backend for the affected configuration. This protects PyTorch, JAX, and direct common callers.

• PyTorch/Jax change: only recognizes why common returned No_Backend and prints the targeted cuDNN 9.26 warning. It does not independently guard the fused backend.

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 pytorch L0 L1

[greptile-apps]

Greptile Summary

The PR prevents affected FP16/BF16 packed-THD learnable-softmax training configurations from selecting fused attention on cuDNN versions older than 9.26.

• Adds the native backend-selection guard while retaining specialized value-head dimensions.
• Adds targeted PyTorch and JAX diagnostics when the native selector rejects the affected configuration.
• Preserves inference and unaffected attention configurations.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn.cpp	Adds the authoritative native capability guard for the affected pre-9.26 THD learnable-softmax training configuration.
transformer_engine/jax/flax/transformer.py	Emits a targeted JAX warning when native backend selection rejects the affected configuration.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Emits a targeted PyTorch warning when the fused backend is unavailable for the affected configuration.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Attention request] --> B{Training + THD + learnable softmax}
  B -->|No| F[Fused backend capability checks continue]
  B -->|Yes| C{cuDNN before 9.26}
  C -->|No| F
  C -->|Yes| D{head_dim_v is 64, 128, or 256}
  D -->|Yes| F
  D -->|No| E[Return No Backend]
  E --> G[Framework warning and fallback]

Loading

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into klakhani/fix/ds..." | Re-trigger Greptile}

[greptile-apps]

Backend guard lacks regression coverage

The new cuDNN-version and head-dimension guard has no automated coverage for either the affected fallback branch or the 9.26 boundary. A regression in condition ordering, version handling, or the specialized-dimension allowlist can therefore re-enable the incorrect-gradient path or unnecessarily disable supported fused execution without CI detecting it.

Knowledge Base Used:

• PyTorch attention execution
• Gate FP8 THD sink-token backward on cuDNN 9.26

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[cyanguwa]

Is this for Hopper only, or all arches?

[KshitijLakhani]

I confirmed with @vedaanta and it is not Hopper specific

[KshitijLakhani]

/te-ci L0 L1