Inconsistent fused_attn_bwd results of dk, dv under THD layout + GQA on multiple calls

[Big-TRex]

Describe the bug

When using THD layout + GQA, multiple calls to fused_attn_bwd produce unexpected dk and dv results. Observations:

• The second call appears to reuse the workspace from the previous call
• This causes inconsistent and unexpected outputs
• Changing the execution order also affects the results

In MHA mode, the same operations produce correct and consistent outputs.

Steps/Code to reproduce bug

import torch
import transformer_engine as te
import transformer_engine_torch as tex
from transformer_engine.pytorch.constants import TE_DType
from transformer_engine.pytorch.cpp_extensions.fused_attn import (
    fused_attn_bwd,
)

seqlen=12
hq=2
hk=1

torch.manual_seed(42)

for i in range(2):
    cu_seqlens_q_per_step=[
        torch.tensor([0,1,3],dtype=torch.int32,device=torch.device("cuda")),
        torch.tensor([0,0,3],dtype=torch.int32,device=torch.device("cuda")),
    ]
    cu_seqlens_kv_per_step=[
        torch.tensor([0,3,9],dtype=torch.int32,device=torch.device("cuda")),
        torch.tensor([0,3,9],dtype=torch.int32,device=torch.device("cuda")),
    ]
    q=torch.randn((seqlen,hq,8),dtype=torch.float16,device=torch.device("cuda"))
    k=torch.randn((seqlen,hk,8),dtype=torch.float16,device=torch.device("cuda"))
    v=torch.randn((seqlen,hk,8),dtype=torch.float16,device=torch.device("cuda"))
    out=torch.randn((seqlen,hq,8),dtype=torch.float16,device=torch.device("cuda"))
    dout=torch.randn((seqlen,hq,8),dtype=torch.float16,device=torch.device("cuda"))
    lse=torch.randn((seqlen,hq,1),dtype=torch.float16,device=torch.device("cuda"))
    rng=torch.randn(hq,dtype=torch.float16,device=torch.device("cuda"))
    cu_seqlens_q_padded=torch.tensor([0,4,12],dtype=torch.int32,device=torch.device("cuda"))
    cu_seqlens_kv_padded=torch.tensor([0,4,12],dtype=torch.int32,device=torch.device("cuda"))

    fused_attn_meta_args = [
        torch.float16,
        TE_DType[dout.dtype],
        [lse,rng],
        tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen
        ]
    fused_attn_meta_kwargs = {
        "attn_scale": 0.1,
        "dropout": 0.0,
        "qkv_layout": "thd_thd_thd",
        "attn_mask_type": "padding",
        "attn_bias_type": "no_bias",
        "deterministic": True,
    }

    print(cu_seqlens_q_per_step[i])
    print(cu_seqlens_kv_per_step[i])
    dq,dk0,dv0, _, _ = fused_attn_bwd(
        8,
        8,
        cu_seqlens_q_per_step[i],
        cu_seqlens_kv_per_step[i],
        q,
        k,
        v,
        out,
        dout,
        *fused_attn_meta_args,
        cu_seqlens_q_padded=cu_seqlens_q_padded,
        cu_seqlens_kv_padded=cu_seqlens_kv_padded,
        **fused_attn_meta_kwargs,
    )
    print(f"{i=},{dk0=}")

Expected behavior

I simulated two sequences, padded to lengths 4 and 8.

• First call:
  • Query valid lengths: 1 and 2
  • Key/Value valid lengths: 3 and 6
• Second call:
  • Query valid lengths: 0 and 3
  • Key/Value valid lengths: 3 and 6

Expected result for the second call:

• The first 4 tokens of dk and dv should be all zeros.

Observed behavior for the second call:

• In THD layout + GQA, the first 4 tokens of dk and dv are not zero, and appear to reuse the previous dk/dv buffers.
• Changing to MHA mode or swapping the call order produces the correct result, where the first 4 tokens of dk and dv are zero.

Environment overview

sudo docker run -it -d --privileged --gpus all --network host --ipc host --name {name} nvcr.io/nvidia/pytorch:25.10-py3

Environment details

None.

Device details

• Nvidia H100 GPU.

[ptrendx]

@cyanguwa Could you take a look?

[dengoswei]

Independent corroboration of this on a different setup, plus a minimal self-contained repro and what looks like the underlying mechanism.

Env: TransformerEngine 2.13.0, cuDNN 9.11.0, torch 2.9.1, H100 (sm90). THD + GQA (n_q=4, n_kv=1), head_dim=256, bf16, backend = NVTE_F16_arbitrary_seqlen, CP=1, no padding (cu_seqlens == cu_seqlens_padded).

Key finding: the wrong dk/dv is tied to the CUDA caching allocator reusing the fused-attention workspace block, which matches the "second call appears to reuse the workspace from the previous call" observation here. Concretely, a single fp32→bf16 free/realloc between two attention backward passes corrupts the next one, and torch.cuda.empty_cache() recovers it. Since empty_cache() only touches allocator free-blocks (not plans/handles/state), this points at workspace memory lifetime, not a math or plan-selection bug.

Minimal repro (pure torch + TE, random inputs, no Megatron):

import torch, torch.nn as nn
import transformer_engine.pytorch as te

DEV, DT = "cuda", torch.bfloat16
N_Q, N_KV, HD = 4, 1, 256
cu = torch.tensor([0, 64, 128], device=DEV, dtype=torch.int32)
T, MS = 128, 64
torch.manual_seed(0)
dpa = te.DotProductAttention(num_attention_heads=N_Q, kv_channels=HD, num_gqa_groups=N_KV,
                             qkv_format="thd", attn_mask_type="padding_causal",
                             softmax_scale=HD**-0.5).to(DEV)

def dk(tag):
    q = torch.randn(T, N_Q,  HD, device=DEV, dtype=DT, requires_grad=True)
    k = torch.randn(T, N_KV, HD, device=DEV, dtype=DT, requires_grad=True)
    v = torch.randn(T, N_KV, HD, device=DEV, dtype=DT, requires_grad=True)
    dpa(q, k, v, cu_seqlens_q=cu, cu_seqlens_kv=cu, max_seqlen_q=MS, max_seqlen_kv=MS).sum().backward()
    d = k.grad.abs().max().item()
    print(f"  {tag:42s} dK.max={d:.3e} ({'NaN' if d != d else 'ok'})")

dk("baseline attention backward")
nn.Linear(4096, 1024).to(DEV).bfloat16()        # free 16MB fp32 / alloc 8MB bf16
dk("after nn.Linear(4096,1024).bfloat16()")
torch.cuda.empty_cache()
dk("after torch.cuda.empty_cache()")

Output:

  baseline attention backward                dK.max=9.000e+00 (ok)
  after nn.Linear(4096,1024).bfloat16()      dK.max=nan          (NaN)
  after torch.cuda.empty_cache()             dK.max=6.938e+00 (ok)

Notes:

• Size-specific: nn.Linear(4096,1024) / torch.randn(4096,1024).bfloat16() (16 MB fp32 block) trigger it; an 8192×8192 (256 MB) churn does not — consistent with the freed block being reused for the workspace only when sizes line up.
• MHA (n_q==n_kv) is unaffected, matching the report here; only GQA breaks.
• dQ stays correct; only dK/dV are wrong (NaN or ~70× too large with real activations).

Real-world impact: this is silently triggered in normal training. Megatron/mbridge wraps the model in Float16Module, whose __init__ does module.bfloat16() during build — exactly the free/realloc churn above. After that, every training-step fused_attn_bwd returns NaN/inflated dK/dV, producing a first-step grad_norm=inf. Switching to the flash-attention backend (which doesn't use this cuDNN workspace) avoids it; so does torch.cuda.empty_cache() after model build.

Happy to provide more detail if useful.