[bvolpato]

Description

Fix silently wrong results from Flex Attention score_mod with nvidia-cudnn-frontend 1.31.0 and newer.

Only the cuDNN backend implements score_mod. The frontend decides whether a graph has a score modifier by reading the record key fn, which is the argument name of sdpa_fp8(), while sdpa() names it score_mod. So the frontend does not see the modifier on these graphs. From 1.31.0, its Python engine sdpa_fwd_prefill_sm100 also accepts 64-dimension heads, so it can serve a Transformer Engine graph as plain attention and drop the modifier without an error.

NVIDIA/cudnn-frontend#1960 fixes the frontend, and the v1.32.1 tag contains that fix. The nvidia-cudnn-frontend==1.32.1 wheel on PyPI does not, so >=1.29.0 still installs the affected code. This PR therefore protects Transformer Engine on all of these versions, without a version pin.

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

• Exclude the Python SDPA engines from Flex Attention score_mod graphs, so that the cuDNN backend serves them.
• Raise an error if a Python engine is still selected, so that a later frontend cannot drop the modifier silently.

Before and after

test_dot_product_attention_score_mod ran on one NVIDIA B200 with each PyPI frontend release.

Frontend	Before	After
1.29.1, 1.30.0	Backend serves every graph, and the test finds no failures	Same
1.31.0	The Python engine serves the softcap and post_scale_bias graphs. Output differs from the reference: up to 38.8% of elements are outside tolerance	Backend serves every graph, and the test finds no failures
1.32.0, 1.32.1	The same engine is selected. SBHD output differs from the reference, and BSHD raises ValueError: cudnn.sdpa: q: this plan was built for 4 heads of dim 64; got (2, 16, 4, 64)	Backend serves every graph, and the test finds no failures

Validation

test_dot_product_attention_score_mod is the regression test: it fails on the affected frontends without this change. The plan names on the run above confirm that backend engines (for example eng7_k24=2_...) do not match the excluded sdpa_fwd and sdpa_bwd names.

The L0 PyTorch commands ran against an untouched tree of b0219803 on B200, with CUDA 13.2, PyTorch 2.15.0+cu132, and frontend 1.32.1. No test failed only with this change. That run also included #3795 and #3797.

JAX is not changed. It builds the same kind of graph, but in a JAX-only environment with frontend 1.32.1 the graph at this shape was served by the backend, with or without the exclusion. A JAX environment that also has PyTorch installed was not tested. A100, H100, and SM120 were not tested. #3693 rewrites this path, so this check needs to move if that PR lands first.

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]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable issues were found.

Summary

Keeps PyTorch Flex Attention score_mod graphs on the cuDNN backend.

• Score-modified Flex Attention graphs stay on the cuDNN backend.

Acknowledgment: bvolpato explicitly leaves JAX unchanged. The PR description reports backend selection in a JAX-only environment and leaves mixed JAX/PyTorch environments untested.

_{Reviews (1) · Last reviewed commit: "[PyTorch] Keep Flex Attention score_mod ..." · Reviewed by Greptile}