[Bug] FusedAttention (cuDNN) THD backward produces wrong gradients on sm_120 (RTX 5090); forward is correct, Unfused is correct

[qshf]

Description

On an RTX 5090 (sm_120), DotProductAttention with qkv_format="thd" and the cuDNN FusedAttention backend returns incorrect gradients in the backward pass, while the forward output is correct. Switching the same call to the UnfusedDotProductAttention backend (NVTE_FUSED_ATTN=0 NVTE_FLASH_ATTN=0) yields correct gradients with identical inputs.

The error is not related to context parallelism, packing, number of sequences, sequence length being a multiple of 128, or the cu_seqlens_padded arguments — a single sequence (no packing) already reproduces it. This appears distinct from #2186 (which is THD + CP tail-padding specific).

Environment

Component	Version
GPU	NVIDIA GeForce RTX 5090 (sm_120)
Driver	595.84
TransformerEngine	2.17.0
PyTorch	2.11.0+cu130
CUDA	13.0.1
cuDNN	9.25.0.15 (libcudnn9-cuda-13)
Base image	lmsysorg/sglang:latest + system cuDNN 9.25

Reproduction

Minimal, no Megatron dependency. Compares TE's backward gradient against a plain per-sequence SDPA reference (math backend, fp32), using identical q/k/v (fixed seed). We report the cosine similarity of the query gradient and its norm.

import os
import torch
import transformer_engine.pytorch as te

device = "cuda"
H, D = 8, 128

def mk(total):
    torch.manual_seed(42)
    return [torch.randn(total, H, D, device=device, dtype=torch.bfloat16, requires_grad=True) for _ in range(3)]

def ref(q, k, v, lengths):
    outs = []; s = 0
    for l in lengths:
        qi = q[s:s+l].transpose(0, 1); ki = k[s:s+l].transpose(0, 1); vi = v[s:s+l].transpose(0, 1)
        o = torch.nn.functional.scaled_dot_product_attention(qi.float(), ki.float(), vi.float(), is_causal=True)
        outs.append(o.transpose(0, 1)); s += l
    return torch.cat(outs, 0)

def run(lengths):
    total = sum(lengths); offs = [0]
    for l in lengths: offs.append(offs[-1] + l)
    cu = torch.tensor(offs, device=device, dtype=torch.int32)
    qr, kr, vr = mk(total); ref(qr, kr, vr, lengths).float().sum().backward(); gref = qr.grad.float()
    q, k, v = mk(total)
    attn = te.DotProductAttention(num_attention_heads=H, kv_channels=D, attention_dropout=0.0,
        qkv_format="thd", attn_mask_type="padding_causal").to(device).train()
    o = attn(q, k, v, qkv_format="thd", cu_seqlens_q=cu, cu_seqlens_kv=cu,
             max_seqlen_q=max(lengths), max_seqlen_kv=max(lengths)).view(total, H, D)
    o_ref = ref(qr.detach().requires_grad_(), kr.detach().requires_grad_(), vr.detach().requires_grad_(), lengths)
    out_cos = torch.nn.functional.cosine_similarity(o.detach().flatten()[None].float(), o_ref.flatten()[None].float()).item()
    o.float().sum().backward(); g = q.grad.float()
    gcos = torch.nn.functional.cosine_similarity(g.flatten()[None], gref.flatten()[None]).item()
    tag = "OK " if gcos > 0.99 else "BAD"
    print(f"  {tag} lengths={str(lengths):10s} out_cos={out_cos:.6f}  q_grad_cos={gcos:.4f}  q_grad_norm={g.norm():.3e} (ref {gref.norm():.2f})")

for L in [[128], [128, 96], [128, 128], [256, 256]]:
    run(L)

Results

cuDNN FusedAttention (NVTE_FLASH_ATTN=0, default fused on) — forward correct, backward wrong:

[INFO | DotProductAttention]: Running with FusedAttention backend (sub-backend 1)
  BAD lengths=[128]      out_cos=0.999998  q_grad_cos=0.4446  q_grad_norm=2.864e+02 (ref 74.97)
  BAD lengths=[128, 96]  out_cos=0.999998  q_grad_cos=0.0265  q_grad_norm=2.068e+10 (ref 101.93)
  BAD lengths=[128, 128] out_cos=0.999998  q_grad_cos=0.4275  q_grad_norm=4.348e+02 (ref 104.96)
  BAD lengths=[256, 256] out_cos=0.999998  q_grad_cos=0.4303  q_grad_norm=4.262e+02 (ref 126.25)

UnfusedDotProductAttention (NVTE_FUSED_ATTN=0 NVTE_FLASH_ATTN=0) — same inputs, backward correct:

[INFO | DotProductAttention]: Running with UnfusedDotProductAttention backend
  OK  lengths=[128]      out_cos=0.999995  q_grad_cos=1.0000  q_grad_norm=7.494e+01 (ref 74.97)
  OK  lengths=[128, 96]  out_cos=0.999995  q_grad_cos=1.0000  q_grad_norm=1.019e+02 (ref 101.93)
  OK  lengths=[128, 128] out_cos=0.999995  q_grad_cos=1.0000  q_grad_norm=1.049e+02 (ref 104.96)
  OK  lengths=[256, 256] out_cos=0.999994  q_grad_cos=1.0000  q_grad_norm=1.262e+02 (ref 126.25)

The only variable changed between the two runs is the attention backend. Forward out_cos ≈ 1.0 in both cases; only the FusedAttention backward is wrong (gradient direction wrong, q_grad_cos 0.03–0.44, and for [128,96] the norm blows up to 2e10).

What we ruled out (each with a controlled experiment)

• Reference correctness — TE's own UnfusedDotProductAttention agrees with our reference (q_grad_cos=1.0), so the reference is right.
• Packing / multi-sequence / cu_seqlens boundaries — a single sequence [128] (no packing) already fails.
• Sequence length not a multiple of 128 (a known cuDNN varlen caveat) — [128,128] and [256,256] are all-128-multiples and still fail.
• Missing cu_seqlens_q_padded / cu_seqlens_kv_padded — explicitly passing them (equal to cu_seqlens, since there is no padding) does not change the result.
• Context parallelism — not used here (single GPU, no CP), unlike Fused attention with THD format + CP may cause a bug in backward pass, leading to NaN values during training LLM. #2186.

Workaround

NVTE_FUSED_ATTN=0 NVTE_FLASH_ATTN=0 forces UnfusedDotProductAttention, which produces correct gradients (at O(N²) cost).

Possibly related

• Fused attention with THD format + CP may cause a bug in backward pass, leading to NaN values during training LLM. #2186 — FusedAttention THD + CP backward NaN, but that is CP tail-padding specific; ours has no CP and fails with a single sequence.
• Community reports that sm_120 reports arch=120 yet uses SM80-era MMA / has a tighter shared-memory budget, so fused kernels that assume SM90+ features (e.g. TMA output epilogue paths) can misbehave on sm_120. This may be the same class of root cause in the cuDNN fused backward.

Notes

• We could not cross-check on sm_80 (A100) as that instance has been released; happy to re-run additional configs on the RTX 5090 if useful (different head_dim, dtype fp16, attn_mask_type variants, etc.).