**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: ``` fused_attn_arbitrary_seqlen_bwd_impl → mha_graph->execute(...) → err 700 != CUDA_SUCCESS :: shimCuLaunchKernelEx(...) → CUDNN_STATUS_EXECUTION_FAILED_CUDA_DRIVER ``` error infomation: ``` 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). We also tried the workspace zero-fill workaround from [#2186](https://github.com/NVIDIA/TransformerEngine/issues/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: ``` #!/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. Notes: 1. Short sequence is not the only trigger: equal-length [600]*9 also fails. 2. Crossing ~598 matters: [2]+[597]*8 OK, [2]+[598]*8 FAIL. 3. Padding the short seq to ≥16 avoids the mixed short+long case, but does not fix all-equal [600]*9. 4. vanilla with the same failing lengths is OK → sink / learnable path is required. 5. Plan-cache masking: in the same process, an OK nearby shape (e.g. L=595) then a failing shape (L=599) can make the failing shape appear OK; always use a fresh process per case. **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** - OS version: Linux (DGX) - PyTorch version: 2.11.0+cu130 - Python version: 3.12 - Transformer Engine version: 2.14.1+cc0ca688 - CUDA version: 13.0 - CUDNN version: 92101 (9.21.1) **Device details** - GPU model: NVIDIA H800 (compute capability sm90)