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

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 protects non-FP8 packed-THD learnable-softmax training from a known pre-cuDNN-9.26 backward defect.

• Disables FusedAttention for affected runtime versions and non-specialized value head dimensions.
• Logs an upgrade recommendation and allows existing fallback or no-backend handling to proceed.
• The new selection boundary lacks automated regression coverage.

Confidence Score: 4/5

The PR appears safe to merge, with the non-blocking concern that its narrow backend-selection boundary is not protected by automated regression coverage.

The guard correctly removes the known-bad fused path while preserving supported configurations and existing fallback behavior; the remaining concern is future regression detection.

Files Needing Attention: transformer_engine/pytorch/attention/dot_product_attention/utils.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Adds a targeted cuDNN capability guard with appropriate fallback behavior, but no regression test covers the newly blocked branch or version boundary.

_{Reviews (1): Last reviewed commit: "Guard THD learnable dSink on older cuDNN" | 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!