[PyTorch] Avoid selecting FA4 for deterministic training on SM120

Start in transformer_engine/pytorch/attention/dot_product_attention/utils.py at get_attention_backend(), then run repro_fa4_sm120.py on an SM120 GPU to observe deterministic FA4 selection and backward failure. Done means deterministic training excludes FA4 on SM120 while evaluation and nondeterministic training remain eligible, with the available fallback or no-backend result reported correctly; the mentioned selector cases should also pass.

Describe the bug

With NVTE_ALLOW_NONDETERMINISTIC_ALGO=0, PyTorch DotProductAttention selects FlashAttention 4 on SM120 even though FA4 does not support deterministic backward on this architecture. Forward completes, then the first backward raises:

  File ".../flash_attn/cute/interface.py", line 2580, in backward
    dq, dk, dv = _flash_attn_bwd(
  File ".../flash_attn/cute/interface.py", line 1401, in _flash_attn_bwd
    assert deterministic is False, "deterministic backward not supported on SM 12.0"
AssertionError: deterministic backward not supported on SM 12.0

Reproduced on one NVIDIA GeForce RTX 5090 on 2026-09-29, using the script below. The same script passed after adding the selector guard, and failed with the same assertion after restoring the original selector. All three runs used the same GPU, installation, inputs, and environment.

TE's get_attention_backend() filters unsupported deterministic configurations for FA2, FA3 and cuDNN FusedAttention, but lacks the FA4/SM120 filter. TE main at 0cc62ecc still lacks the guard, checked 2026-09-30. FA4 main at e9cf2c1 still rejects deterministic backward in its SM12x path.

Steps/Code to reproduce bug

Use one SM120 GPU with the CUDA PyTorch/TE/FA4 environment listed below. Save this as repro_fa4_sm120.py. The executable code is identical to the GPU-tested script; only its descriptive module docstring was updated after verification.

"""GPU-verified on one RTX 5090 (SM120), 2026-09-29."""

import importlib.metadata
import os

# Set before importing Transformer Engine. Keep an unfused fallback available.
for name, value in {
    "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0",
    "NVTE_FLASH_ATTN": "1",
    "NVTE_FLASH_ATTN_V2": "0",
    "NVTE_FLASH_ATTN_V3": "0",
    "NVTE_FLASH_ATTN_V4": "1",
    "NVTE_FUSED_ATTN": "0",
    "NVTE_UNFUSED_ATTN": "1",
    "NVTE_DEBUG": "1",
    "NVTE_DEBUG_LEVEL": "2",
}.items():
    os.environ.setdefault(name, value)

import torch
import transformer_engine
import transformer_engine.pytorch as te

assert torch.cuda.get_device_capability() == (12, 0), "Requires an SM120 GPU"
print("GPU:", torch.cuda.get_device_name())
print("PyTorch:", torch.__version__, "CUDA:", torch.version.cuda)
print("cuDNN:", torch.backends.cudnn.version())
print("TE:", transformer_engine.__version__, transformer_engine.__file__)
print("FA4:", importlib.metadata.version("flash-attn-4"))

torch.manual_seed(0)
shape = (2, 128, 8, 128)  # batch, sequence, heads, head dimension
q, k, v = [
    torch.randn(shape, device="cuda", dtype=torch.bfloat16, requires_grad=True)
    for _ in range(3)
]
dpa = te.DotProductAttention(
    num_attention_heads=shape[2],
    kv_channels=shape[3],
    qkv_format="bshd",
    attn_mask_type="causal",
    attention_dropout=0.0,
).cuda().train()
out = dpa(q, k, v)
out.float().square().mean().backward()
torch.cuda.synchronize()
assert torch.isfinite(out).all()
assert all(x.grad is not None and torch.isfinite(x.grad).all() for x in (q, k, v))
print("Forward and backward completed; inspect the debug log for the selected backend.")

Run in a fresh process with inherited NVTE_* variables unset, allowing the script's defaults to apply:

python -u repro_fa4_sm120.py

The verified inputs are three separate contiguous BF16 tensors in BSHD format: batch 2, sequence length 128, 8 heads, head dimension 128. Attention is causal, dropout is zero, seed is zero, and the module is in training mode. The loss is the mean squared float32 output. FA2, FA3, and fused attention are disabled; unfused attention remains available. This reproduction requires no dataset, FP8, distributed processes, or model-training framework.

Observed before the fix:

Available backends = {FlashAttention=True (4.0.0b24), FusedAttention=False, UnfusedDotProductAttention=True}
Selected backend = FlashAttention (4.0.0b24).
AssertionError: deterministic backward not supported on SM 12.0

Observed with the seven-line selector guard from fix commit 8cec4228 applied to the same installation:

Disabling FlashAttention 4 for deterministic backward on sm120: the CuTe SM 12.0 backward kernel does not support deterministic execution.
Available backends = {FlashAttention=False, FusedAttention=False, UnfusedDotProductAttention=True}
Selected backend = UnfusedDotProductAttention.
Forward and backward completed; inspect the debug log for the selected backend.

The guard was then removed; the original FA4 selection and backward assertion returned.

Additional controls executed in separate processes with the patch applied:

NVTE_UNFUSED_ATTN=0 python -u repro_fa4_sm120.py
NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 python -u repro_fa4_sm120.py

For the executed evaluation control, the script used .eval() instead of .train(), wrapped out = dpa(q, k, v) in torch.no_grad(), omitted backward, and asserted that all three input gradients remained None.

Numerical validation of the patched unfused path used the same BF16 inputs, a random BF16 upstream gradient with seed 1, and an independent float64 causal-attention reference. Output and all input gradients passed torch.testing.assert_close with atol=rtol=0.02, chosen for BF16 output/gradient comparison. Repeating the same forward/backward with the same inputs and upstream gradient produced bitwise-identical output, dQ, dK, and dV.

import json
from pathlib import Path
import runpy

state = runpy.run_path(str(Path(__file__).with_name("repro_fa4_sm120.py")))
torch = state["torch"]
dpa = state["dpa"]
q, k, v = [state[name] for name in ("q", "k", "v")]
shape = state["shape"]
for x in (q, k, v):
    x.grad = None
out = dpa(q, k, v)
torch.manual_seed(1)
upstream = torch.randn_like(out)
out.backward(upstream)
actual = [out.detach().clone()] + [x.grad.detach().clone() for x in (q, k, v)]
qr, kr, vr = [x.detach().double().requires_grad_() for x in (q, k, v)]
scores = torch.einsum("bqhd,bkhd->bhqk", qr, kr) / shape[-1] ** 0.5
mask = torch.ones(shape[1], shape[1], device="cuda", dtype=torch.bool).triu(1)
probs = scores.masked_fill(mask, float("-inf")).softmax(-1)
reference = torch.einsum("bhqk,bkhd->bqhd", probs, vr).reshape_as(out)
reference.backward(upstream.double())
expected = [reference] + [x.grad for x in (qr, kr, vr)]
errors = {}
for name, a, b in zip(("output", "dq", "dk", "dv"), actual, expected):
    errors[name] = {"max_abs": (a.double() - b).abs().max().item(),
                    "rms_abs": (a.double() - b).square().mean().sqrt().item()}
    torch.testing.assert_close(a.double(), b, rtol=0.02, atol=0.02)
for x in (q, k, v):
    x.grad = None
repeated = dpa(q, k, v)
repeated.backward(upstream)
for name, a, b in zip(("output", "dq", "dk", "dv"), actual,
                      [repeated] + [x.grad for x in (q, k, v)]):
    assert torch.equal(a, b), name
print("NUMERICAL_ERRORS", json.dumps(errors))
print("Output and gradients match float64 reference (atol=rtol=0.02); repeated execution is bitwise equal.")

Expected behavior

When deterministic training is requested, TE should exclude FA4 on SM120 and select another enabled compatible backend. If none is available, backend selection should report that fact. The determinism request must be preserved. Evaluation-only FA4 and nondeterministic FA4 training should remain eligible.

Environment overview (please complete the following information)

• Linux container, one host, one GPU; tested 2026-09-29.
• TE was already installed from source at 699ed6ec9e9c600c3d0acd3a488c4ca6f339e5d9, reporting 2.21.0.dev0.
• The reproduction reused that installation; no package installation or native rebuild was performed during this test. The container has no public pull reference. Original installation/container commands are omitted because they are not publicly reproducible.
• The imported package was /usr/local/lib/python3.12/dist-packages/transformer_engine/__init__.py. The unpatched selector's SHA-256 was 2c9a718c5dad6dbe56424e216f69559b2d509e6bd6a25aaea8078ab235ebd3b0, matching utils.py at the source revision above.
• Each control ran in a fresh Python process after clearing inherited NVTE_* settings. Only the Python selector guard and each explicitly listed control setting changed between runs.

Environment details

The PyTorch and CUTLASS development builds are environment information, not publicly installable version pins. OS/Python/CUTLASS versions come from the reused image's inventory; the reproduction printed the framework, attention-library, CUDA, and cuDNN versions, and nvidia-smi recorded the GPU and driver.

Device details

• One NVIDIA GeForce RTX 5090, compute capability (12, 0), 32607 MiB reported by nvidia-smi.
• One process and one GPU; no inter-GPU communication.

Additional context

The prepared fix filters FA4 only when deterministic, is_training, and device_compute_capability == (12, 0). It changes backend selection and does not add a deterministic FA4 kernel. The patch is pure Python.

Implementation: fix branch, commit 8cec4228914cc364dab88f3e1244466852a5fe29.

All 48 new selector cases also passed against the real installed PyTorch/TE with the guard applied. The test function was extracted unchanged from the fix branch and run separately; hardware capability/backend availability are mocked by those unit tests. The GPU integration and numerical results above validate the BF16, head-dimension-128 SM120 case on TE 699ed6ec; they do not represent a GPU build/test of the branch's newer base revision or the full L0 suite. SM121 is outside this patch's scope.

Related upstream work:

• TE #3540 adds an FA4 determinism filter for SM100/SM110 with head dimension 256, explicitly excluding SM120. Its insertion overlaps this fix but does not address this failure.
• TE #3289 and #3125 adjust SM12x/SM121 support checks without adding this FA4 determinism filter.
• FA #2386 concerns separate SM120 bring-up defects.