[Common] Gate FP8 THD sink backward on cuDNN 9.26

[sudhakarsingh27]

Description

cuDNN versions before 9.26 can misindex ragged FP8 Stats during sink-token backward. FP8 THD backend selection currently permits this training path with cuDNN 9.23 and newer.

Require cuDNN 9.26 only when selecting FP8 THD training with a non-vanilla softmax mode. This preserves FP8 THD inference and vanilla-softmax training on older supported cuDNN versions.

Testing

• pre-commit run clang-format --files transformer_engine/common/fused_attn/fused_attn.cpp
• git diff --check

GPU runtime validation is left to CI.

[greptile-apps]

Greptile Summary

This PR prevents selection of the affected cuDNN fused-attention backend for FP8 THD sink-softmax training before cuDNN 9.26 while preserving inference and vanilla-softmax training.

• Adds a cuDNN 9.26 minimum for FP8 THD training with off-by-one or learnable sink softmax.
• Leaves FP8 THD inference and vanilla-softmax training available on older supported cuDNN versions.
• Returns no fused backend for known-bad combinations so they fail cleanly or use a framework fallback where available.

Confidence Score: 5/5

The PR appears safe to merge because the new predicate narrowly excludes the documented unsafe cuDNN training combination while retaining unaffected inference and vanilla-softmax paths.

The selector now rejects FP8 THD sink-softmax training before cuDNN 9.26, and the surrounding enum, training, layout, and fallback behavior align with the intended mitigation.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_attn/fused_attn.cpp	Narrows FP8 THD backend selection only for sink-softmax training on cuDNN releases affected by the ragged-statistics backward defect.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[FP8 THD request] --> B{Training?}
  B -- No --> C[Allow supported cuDNN backend]
  B -- Yes --> D{Vanilla softmax?}
  D -- Yes --> C
  D -- No: sink softmax --> E{cuDNN >= 9.26?}
  E -- Yes --> C
  E -- No --> F[Return No Backend]

Loading

_{Reviews (1): Last reviewed commit: "[Common] Gate FP8 THD sink backward on c..." | Re-trigger Greptile}

[sudhakarsingh27]

/te-ci pytorch L0

[KshitijLakhani]

LGTM ! Thanks
Good to merge once CI passes fully for both pipelines