On Hopper, DotProductAttention with FusedAttention + qkv_format=thd + softmax_type=learnable (sink attention) can hit a CUDA illegal memory access during backward:
RuntimeError: /root/transformer-engine/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu:1056 in function fused_attn_arbitrary_seqlen_bwd_impl: cuDNN Error: execute(handle, plan->get_raw_desc(), variant_pack_descriptor.get_ptr()) failed with message: err 700 != CUDA_SUCCESS :: shimCuLaunchKernelEx(&config, func, kernelParams, nullptr)
at: err != CUDA_SUCCESS, and code: CUDNN_STATUS_EXECUTION_FAILED_CUDA_DRIVER. For more information, enable cuDNN error logging by setting CUDNN_LOGERR_DBG=1 and CUDNN_LOGDEST_DBG=stderr in the environment.
This is not CP-related. It reproduces with a bare DotProductAttention call (no Megatron / no CP).
It looks related to known THD + sink / Hopper SDPA-bwd issues, but still reproduces on cuDNN 9.21.1 (after the 9.18/9.19 fixes mentioned in #2483 / #2616 / cudnn-frontend#200).
in both fused attn fwd/bwd paths. IMA still happens, so this does not appear to be the same “dirty workspace → wrong grads/NaN” bug as #2186.
#!/usr/bin/env python3
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Minimal repro: cuDNN FusedAttention backward IMA with
# softmax_type=learnable + qkv_format=thd + padding_causal
#
# Trigger (fresh process, observed on Hopper / TE FusedAttention path):
# learnable + THD + FusedAttention bwd, roughly when lengths cross ~598
# Examples that FAIL: [2]+[598]*8, and also equal [600]*9
# Examples that OK: [16]+[600]*8, [2]+[597]*8, vanilla+[2]+[598]*8
#
# Expected (see PRESET_EXPECT):
# FAIL -> CUDA illegal memory access / cuDNN CUDNN_STATUS_EXECUTION_FAILED
# OK -> short>=16 with longs, or long<=597, or vanilla softmax
#
# Usage:
# python repro_thd_learnable_ima.py # run ALL presets (isolated)
# python repro_thd_learnable_ima.py --case fail # single preset
# python repro_thd_learnable_ima.py --seqs 2,598,...
# CUDA_LAUNCH_BLOCKING=1 compute-sanitizer --tool memcheck \
# python repro_thd_learnable_ima.py --case fail
from __future__ import annotations
import argparse
import os
import subprocess
import sys
import traceback
# Force FusedAttention (cuDNN SDPA). Must be set before TE import side-effects matter.
os.environ["NVTE_FLASH_ATTN"] = "0"
os.environ["NVTE_FUSED_ATTN"] = "1"
os.environ["NVTE_UNFUSED_ATTN"] = "0"
os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")
PRESETS = {
# Mixed short+long (classic)
"fail": ([2] + [598] * 8, "learnable"),
# Equal long: also FAILS on H800/cuDNN 9.21.1 (not a safe control)
"equal-600": ([600] * 9, "learnable"),
# Safe controls
"ok-equal597": ([597] * 9, "learnable"),
"ok-short16": ([16] + [600] * 8, "learnable"),
"ok-long597": ([2] + [597] * 8, "learnable"),
"vanilla": ([2] + [598] * 8, "vanilla"),
}
# True => expect OK; False => expect FAIL
PRESET_EXPECT_OK = {
"fail": False,
"equal-600": False,
"ok-equal597": True,
"ok-short16": True,
"ok-long597": True,
"vanilla": True,
}
ALL_CASES = list(PRESETS.keys())
def _print_env() -> None:
import torch
import transformer_engine as te
print("=== env ===", flush=True)
print(f"torch={torch.__version__}", flush=True)
print(f"cuda={torch.version.cuda}", flush=True)
print(f"te={getattr(te, '__version__', '?')}", flush=True)
if torch.cuda.is_available():
print(f"device={torch.cuda.get_device_name(0)}", flush=True)
print(f"capability={torch.cuda.get_device_capability(0)}", flush=True)
try:
import transformer_engine_torch as tex
if hasattr(tex, "get_cudnn_version"):
print(f"cudnn={tex.get_cudnn_version()}", flush=True)
except Exception as e: # noqa: BLE001
print(f"cudnn=unknown ({e})", flush=True)
print(flush=True)
def _rng_tracker():
from transformer_engine.pytorch.distributed import CudaRNGStatesTracker
tracker = CudaRNGStatesTracker()
tracker.add("model-parallel-rng", 1234)
return tracker
def run(seqs: list[int], softmax_type: str = "learnable") -> bool:
import torch
from transformer_engine.pytorch.attention import DotProductAttention
# Config matching the failing SWA / GQA layer shape.
num_heads = 64
num_gqa_groups = 8
head_dim = 192
dtype = torch.bfloat16
qkv_format = "thd"
attn_mask_type = "padding_causal"
window_size = (511, 0)
assert len(seqs) >= 1
batch = len(seqs)
max_seqlen = max(seqs)
total = sum(seqs)
# After IMA the CUDA context is often poisoned; reset if possible.
try:
torch.cuda.empty_cache()
torch.cuda.synchronize()
except Exception: # noqa: BLE001
pass
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
seqlens = torch.tensor(seqs, dtype=torch.int32, device="cuda")
cu_seqlens = torch.zeros(batch + 1, dtype=torch.int32, device="cuda")
cu_seqlens[1:] = torch.cumsum(seqlens, dim=0)
q = (0.1 * torch.randn(total, num_heads, head_dim, dtype=dtype, device="cuda")).requires_grad_(
True
)
k = (
0.1 * torch.randn(total, num_gqa_groups, head_dim, dtype=dtype, device="cuda")
).requires_grad_(True)
v = (
0.1 * torch.randn(total, num_gqa_groups, head_dim, dtype=dtype, device="cuda")
).requires_grad_(True)
dout = 0.001 * torch.randint(
0, 200, (total, num_heads * head_dim), dtype=dtype, device="cuda"
)
attn = DotProductAttention(
num_heads,
(head_dim, head_dim),
num_gqa_groups=num_gqa_groups,
attention_dropout=0.0,
qkv_format=qkv_format,
attn_mask_type=attn_mask_type,
sequence_parallel=False,
tp_size=1,
get_rng_state_tracker=_rng_tracker,
tp_group=None,
layer_number=1,
attention_type="self",
softmax_type=softmax_type,
).to(dtype=dtype, device="cuda")
attn.train()
if softmax_type == "learnable":
attn.softmax_offset.requires_grad_(True)
print(
f"run softmax={softmax_type} nseq={batch} seqs={seqs} "
f"max={max_seqlen} total={total} window={window_size}",
flush=True,
)
try:
out = attn(
q,
k,
v,
window_size=window_size,
attention_mask=None,
qkv_format=qkv_format,
max_seqlen_q=max_seqlen,
max_seqlen_kv=max_seqlen,
cu_seqlens_q=cu_seqlens,
cu_seqlens_kv=cu_seqlens,
cu_seqlens_q_padded=cu_seqlens,
cu_seqlens_kv_padded=cu_seqlens,
attn_mask_type=attn_mask_type,
core_attention_bias_type="no_bias",
fast_zero_fill=True,
pad_between_seqs=False,
)
out.backward(dout)
torch.cuda.synchronize()
print("RESULT: OK", flush=True)
return True
except Exception:
print("RESULT: FAIL", flush=True)
traceback.print_exc()
# Best-effort clear sticky CUDA error so same-process follow-ups can try
# (still prefer --case all with subprocess isolation).
try:
torch.cuda.synchronize()
except Exception: # noqa: BLE001
pass
return False
def run_isolated(case: str | None, seqs_csv: str | None, softmax: str | None) -> bool:
"""Run one case in a fresh process so IMA does not poison later cases."""
env = os.environ.copy()
env["INNER"] = "1"
if case is not None:
env["CASE"] = case
else:
env.pop("CASE", None)
if seqs_csv is not None:
env["SEQS"] = seqs_csv
else:
env.pop("SEQS", None)
if softmax is not None:
env["SOFTMAX"] = softmax
else:
env.pop("SOFTMAX", None)
label = case or f"custom({seqs_csv})"
print(f"\n===== {label} (fresh process) =====", flush=True)
rc = subprocess.call([sys.executable, __file__], env=env)
return rc == 0
def main() -> int:
parser = argparse.ArgumentParser(
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
)
parser.add_argument(
"--case",
choices=(*ALL_CASES, "all"),
default="all",
help="preset: fail / equal-600 (expect FAIL); ok-equal597 / ok-short16 / "
"ok-long597 / vanilla (expect OK). Default 'all' uses isolated subprocesses.",
)
parser.add_argument(
"--seqs",
type=str,
default=None,
help="override as comma-separated lengths, e.g. 2,598,598,598,598,598,598,598,598",
)
parser.add_argument(
"--softmax",
choices=("learnable", "vanilla"),
default=None,
help="override softmax_type (default depends on --case)",
)
parser.add_argument(
"--in-process",
action="store_true",
help="when --case all, run all presets in one process (not recommended after IMA)",
)
args = parser.parse_args()
inner = os.environ.get("INNER", "0") == "1"
env_case = os.environ.get("CASE")
env_seqs = os.environ.get("SEQS")
env_softmax = os.environ.get("SOFTMAX")
# Worker process: run exactly one case and exit.
if inner:
import torch
if not torch.cuda.is_available():
print("CUDA required", file=sys.stderr)
return 2
_print_env()
if env_seqs:
seqs = [int(x) for x in env_seqs.split(",") if x.strip()]
softmax = env_softmax or "learnable"
tag = "custom"
else:
case = env_case or "fail"
seqs, softmax = PRESETS[case]
if env_softmax:
softmax = env_softmax
tag = case
print(f"=== CASE={tag} ===", flush=True)
ok = run(seqs, softmax_type=softmax)
print("===", "OK" if ok else "FAIL", f"({tag}) ===", flush=True)
return 0 if ok else 1
# Parent: custom seqs -> one isolated run
if args.seqs is not None:
ok = run_isolated(None, args.seqs, args.softmax)
return 0 if ok else 1
# Parent: all presets, each in fresh process (default)
if args.case == "all":
if args.in_process:
import torch
if not torch.cuda.is_available():
print("CUDA required", file=sys.stderr)
return 2
_print_env()
results = {}
for name in ALL_CASES:
seqs, softmax = PRESETS[name]
if args.softmax is not None:
softmax = args.softmax
print(f"\n===== {name} (in-process) =====", flush=True)
results[name] = run(seqs, softmax_type=softmax)
else:
results = {}
for name in ALL_CASES:
results[name] = run_isolated(name, None, args.softmax)
print("\n=== SUMMARY ===", flush=True)
unexpected = False
for name, ok in results.items():
expect_ok = PRESET_EXPECT_OK[name]
expect = "OK" if expect_ok else "FAIL"
got = "OK" if ok else "FAIL"
mark = "as-expected" if (ok == expect_ok) else "UNEXPECTED"
if ok != expect_ok:
unexpected = True
print(f" {name:<14} {got:<4} (expect {expect}, {mark})", flush=True)
return 1 if unexpected else 0
# Parent: single named case -> isolated for consistency
ok = run_isolated(args.case, None, args.softmax)
return 0 if ok else 1
if __name__ == "__main__":
raise SystemExit(main())
Minimal core (same config as failing SWA/GQA layer): heads=64, gqa=8, head_dim=192, bf16, attn_mask_type=padding_causal, window_size=(511,0), softmax_type=learnable, THD packed cu_seqlens.
Backward should complete without CUDA IMA / cuDNN CUDNN_STATUS_EXECUTION_FAILED for valid THD + learnable configurations that TE selects FusedAttention for (including cuDNN ≥ 9.18 after PR #2568 enabled this path).
Describe the bug
On Hopper, DotProductAttention with FusedAttention + qkv_format=thd + softmax_type=learnable (sink attention) can hit a CUDA illegal memory access during backward:
error infomation:
This is not CP-related. It reproduces with a bare DotProductAttention call (no Megatron / no CP).
It looks related to known THD + sink / Hopper SDPA-bwd issues, but still reproduces on cuDNN 9.21.1 (after the 9.18/9.19 fixes mentioned in #2483 / #2616 / cudnn-frontend#200).
We also tried the workspace zero-fill workaround from #2186:
allocateSpace(workspace.shape(), workspace.dtype(), true);in both fused attn fwd/bwd paths. IMA still happens, so this does not appear to be the same “dirty workspace → wrong grads/NaN” bug as #2186.
Failure site (TE):
transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu → fused_attn_arbitrary_seqlen_bwd_impl → mha_graph->execute (with set_sink_token / set_dsink_token).Steps/Code to reproduce bug
Force FusedAttention and run a minimal THD + learnable bwd example below:
Minimal core (same config as failing SWA/GQA layer): heads=64, gqa=8, head_dim=192, bf16, attn_mask_type=padding_causal, window_size=(511,0), softmax_type=learnable, THD packed cu_seqlens.
Notes:
Expected behavior
Backward should complete without CUDA IMA / cuDNN CUDNN_STATUS_EXECUTION_FAILED for valid THD + learnable configurations that TE selects FusedAttention for (including cuDNN ≥ 9.18 after PR #2568 enabled this path).
Environment details
Device details