thd + attention_dropout routes to the composite cuDNN SDPA engine and costs 5x on Blackwell

[bzantium]

Describe the bug

With qkv_format="thd", turning on attention_dropout makes the fused (cuDNN) attention backend about 5x slower on B300. The same work in sbhd pays almost nothing for the same dropout, so the cost is not dropout itself — it is dropout combined with the packed layout.

Forward + backward through a single DotProductAttention, bf16, 4096 tokens, 16 heads, head_dim 128, 4 sequences, fused backend forced, median of 50 iterations after 20 warmup:

GPU	TE	layout	p=0.0	p=0.1	dropout cost
B300 (sm103)	2.14.1	thd	0.585 ms	3.060 ms	+2.476 ms (5.24x)
B300 (sm103)	2.14.1	sbhd	0.480 ms	0.525 ms	+0.045 ms (1.09x)
H200 (sm90)	2.15.0	thd, MLA dims	0.726 ms	0.917 ms	+0.191 ms (1.26x)
H200 (sm90)	2.15.0	sbhd	0.594 ms	0.614 ms	+0.020 ms (1.03x)

Same token count, same attention work: on B300 the packed layout pays 55x more for the same
dropout than the batched one. The asymmetry is present on H200 too, at roughly 10x, but the
absolute cost there is +0.19 ms rather than +2.48 ms.

For reference, FlashAttention on the same B300 pays +0.024 ms for the same p=0.1 at thd, and the fused backend at thd with MLA dims (head_dim_qk 192, head_dim_v 128) pays +2.485 ms — the cost tracks the attention matrix, not the head dimension.

This looks like the restriction already flagged in fused_attn_f16_arbitrary_seqlen.cu:

  const bool use_cu_seqlens_directly =
      ...
      // This extra restriction is needed because cuDNN frontend doesn't yet allow
      // the combination of dropout and stats generation for the fprop unified engine,
      // so any such request would always get routed to the old composite SDPA engine
      // (which doesn't support cu_seqlens). Remove this restriction when possible.
      !is_dropout;

sbhd does not need cu_seqlens, so it is unaffected; thd requests with dropout fall to the composite engine. A PyTorch profile of an 8-layer training step on B300 shows where the time goes. Top CUDA kernels over 5 steps:

cudnn::fusion::gen_dropout_mask_4bit_transpose      63.99 ms   25.28%   40 calls
cudnn::fusion::gen_dropout_mask_4bit                45.38 ms   17.92%   40 calls
...
cudnn_generated_fort_native_sdpa_sm100_flash_bprop   4.14 ms    1.64%   40 calls

Mask generation is 43% of all CUDA time in the step, roughly 26x the attention backward it serves. Total CUDA time for that model is 253 ms per 5 steps on the fused backend against 154 ms on FlashAttention.

This is opening an issue rather than a PR because the fix looks like it belongs to the cuDNN frontend, not to Transformer Engine — TE is already requesting the RNG form of dropout via set_dropout(probability, seed, offset) and working around the routing. Filing it to put a number on the existing TODO.

On how much this matters in practice: most current LLM configs set attention_dropout to 0 and never hit this. The exposure is that Megatron-Core's TransformerConfig defaults it to 0.1, so a packed-sequence run that does not explicitly set it inherits the slow path silently. That is how I found this — a step-time comparison came out backwards until a profile showed 43% of CUDA time in dropout mask generation, from a default I had not set.

Steps/Code to reproduce bug

# NVTE_FUSED_ATTN=1 NVTE_FLASH_ATTN=0 NVTE_UNFUSED_ATTN=0 python repro.py thd  0.0
# NVTE_FUSED_ATTN=1 NVTE_FLASH_ATTN=0 NVTE_UNFUSED_ATTN=0 python repro.py thd  0.1
# NVTE_FUSED_ATTN=1 NVTE_FLASH_ATTN=0 NVTE_UNFUSED_ATTN=0 python repro.py sbhd 0.0
# NVTE_FUSED_ATTN=1 NVTE_FLASH_ATTN=0 NVTE_UNFUSED_ATTN=0 python repro.py sbhd 0.1
import math, statistics, sys
import torch
import transformer_engine.pytorch as te
from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends

TOKENS, HEADS, SEQS, DIM = 4096, 16, 4, 128
layout, dropout = sys.argv[1], float(sys.argv[2])
max_seqlen = TOKENS // SEQS

def rand(*shape):
    return torch.randn(*shape, device="cuda", dtype=torch.bfloat16).requires_grad_(True)

if layout == "thd":
    q, k, v = (rand(TOKENS, HEADS, DIM) for _ in range(3))
    out_shape, mask = (TOKENS, HEADS, DIM), "padding_causal"
    cu = torch.linspace(0, TOKENS, SEQS + 1, dtype=torch.int32, device="cuda")
    extra = dict(cu_seqlens_q=cu, cu_seqlens_kv=cu)
else:
    q, k, v = (rand(max_seqlen, SEQS, HEADS, DIM) for _ in range(3))
    out_shape, mask, extra = (max_seqlen, SEQS, HEADS, DIM), "causal", {}

grad_out = torch.randn(*out_shape, device="cuda", dtype=torch.bfloat16)
attn = te.DotProductAttention(
    num_attention_heads=HEADS, kv_channels=(DIM, DIM), attention_dropout=dropout,
    qkv_format=layout, attn_mask_type=mask, softmax_scale=1.0 / math.sqrt(DIM),
).cuda()
attn.train()

def once():
    out = attn(q, k, v, max_seqlen_q=max_seqlen, max_seqlen_kv=max_seqlen, **extra).view(*out_shape)
    torch.autograd.grad(out, (q, k, v), grad_out)

for _ in range(20):
    once()
torch.cuda.synchronize()

samples = []
for _ in range(50):
    s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
    s.record(); once(); e.record(); torch.cuda.synchronize()
    samples.append(s.elapsed_time(e))
used = [key for key, value in _attention_backends.items() if key.startswith("use_") and value]
print(f"layout={layout} dropout={dropout} used={used} median_ms={statistics.median(samples):.4f}")

NVTE_UNFUSED_ATTN=0 matters: without it, a shape the selected backend cannot serve falls through to the unfused path silently and the timing lands under the wrong label. The used= field is there to catch that.

Expected behavior

Dropout on thd should cost roughly what it costs on sbhd (+0.045 ms) rather than +2.48 ms, i.e. thd requests with dropout should be able to reach the same engine sbhd does.

Environment overview

• Environment location: Cloud (AWS), Slurm with pyxis/enroot
• Method of Transformer Engine install: pre-installed in a NeMo 26.04-based container image
• Also measured on an H200 with a different image (TE 2.15.0) for comparison

Environment details

	B300	H200
GPU	NVIDIA B300 SXM6, sm103	NVIDIA H200, sm90
TE	2.14.1+366798ef	2.15.0+42b84005
cuDNN	9.20	9.20
FlashAttention	2.7.4.post1	2.8.3 (FA2), 3.0.0 (FA3)
Python	3.12	3.13.13

The thd/sbhd asymmetry appears on both stacks, which fits the routing explanation above being
version-independent. What differs is how much the composite engine costs once a request lands on
it: +2.48 ms on B300 against +0.19 ms on H200. The two stacks differ in both compute capability
and TE version, so this data does not establish which of the two accounts for that gap.

One aside

On the H200 image, head_dim 128/128 at thd with dropout=0.0 fails to run on the fused backend, reproducibly across runs, while dropout=0.1 with the same dims works and sbhd works at both:

cuDNN Error: CUDNN_BACKEND_TENSOR_DESCRIPTOR cudnnFinalize failed ptrDesc->finalize()
cudnn_status: CUDNN_STATUS_SUBLIBRARY_LOADING_FAILED
  at transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu:419

Likely a packaging issue in that image rather than a TE bug, noted only because it is why one H200 cell is missing above.