[Bug已解决] CK SDPA flash attention kernel 在 gfx1200/gfx1201 (RDNA4) 上断言失败解决方案一、现象长什么样你在一块AMD RDNA4 架构的显卡上跑 PyTorch比如基于gfx1200/gfx1201的 GPU这是 AMD 最新的 RDNA4 家族例如 RX 9000 系列。当你调用import torch import torch.nn.functional as F out F.scaled_dot_product_attention(q, k, v)或者模型里依赖 SDPAscaled dot product attention的注意力层如torch.nn.MultiheadAttention、HuggingFace Transformers 的eager/flash_attention_2后端、vLLM、SDPA backend时程序不是报普通的「不支持」或优雅回退而是直接触发内核层的断言assert崩溃典型报错形如CK SDPA flash attention kernel asserts on gfx1200/gfx1201 (RDNA4) ... hipErrorAssert / GPU kernel assert triggered ... RuntimeError: CK (Composable Kernel) SDPA flash attention kernel assertion failed或者更底层一些来自 Composable KernelCK库的内核CK_ASSERT failed: ... in flash attention kernel for gfx1200/gfx1201也就是说PyTorch 的 SDPA 调度器把注意力计算派发到了 CK 的 flash attention 内核但这个内核在 RDNA4 的 gfx1200/gfx1201 上存在一个断言没过导致直接 GPU 端断言失败、进程崩溃而不是退回 math 或 mem-efficient 后端。 本文聚焦这个 CK SDPA 断言是什么、为什么会触发、怎么在 RDNA4 上绕开并继续跑注意力计算。二、背景PyTorch SDPA 与 AMD 的 CK 后端要理解这个 bug先要明白 PyTorch 的scaled_dot_product_attentionSDPA是怎么选内核的。F.scaled_dot_product_attention(q, k, v)本身只是一个调度入口。真正的计算由若干「后端backend」之一完成PyTorch 会按优先级自动选择torch.backends.cuda下的FlashAttention/MemEfficient是 NVIDIA CUDA 的在AMD ROCm上PyTorch 接入的是Composable KernelCK库它提供了 flash attention / mem-efficient attention 的高性能 HIP 内核还有最稳但最慢的math后端纯数学计算无特殊优化。 PyTorch 提供了torch.backends.cuda.sdp_kernel注意 ROCm 下同样走这套 API 调度来强制开/关某个后端import torch import torch.nn.functional as F # 强制只用 math 后端最稳但慢 with torch.backends.cuda.sdp_kernel( enable_flashFalse, enable_mem_efficientFalse, enable_mathTrue, ): out F.scaled_dot_product_attention(q, k, v)关键点是SDPA 调度器默认会尝试 flash attention在 ROCm 上就是 CK 的 flash attention 内核。当你的 GPU 是 RDNA4gfx1200/gfx1201时如果 PyTorch / CK 的预编译内核还没有正确适配这个架构或者内部某个CK_ASSERT对 RDNA4 的某些张量形状/数据类型/对齐条件判断失败就会崩在内核里。 这不是「我没装对驱动」那么简单而是CK 内核在 gfx1200/gfx1201 上的代码路径存在缺陷——属于 PyTorch Composable Kernel 的已知 issuepytorch/pytorch#188113。三、为什么 RDNA4 会触发 CK 断言RDNA4gfx1200 / gfx1201是 AMD 较新的架构它的 ISA、波前wavefront宽度、共享内存LDS布局、向量寄存器文件都与上一代 RDNA3gfx1100/gfx1101/gfx1102不同。CK 的 flash attention 内核里有很多针对具体 gfx 架构特化的模板与CK_ASSERT检查例如对 tile size / block size 是否与当前架构的 wavefront 对齐的断言对 LDS 占用是否超过上限的断言对数据类型fp16 / bf16向量长度是否合法的检查。 在 RDNA4 上这些特化可能还没被充分验证或编译进发布版于是要么调度器错误地认为「RDNA4 支持 CK flash attention」并派发过去但内核里的架构判断分支没覆盖 gfx1200/gfx1201走到一个CK_ASSERT(false)要么内核编译产物针对 RDNA4 有未修复的数值/对齐 bug在运行时命中断言。 无论哪种结果都是内核断言失败 → 进程直接崩而不是优雅回退。这就是它最恼人的地方你指望 PyTorch「不支持就退回 math」但它偏偏在 flash 内核里 assert 了。四、最小可运行复现带 ROCm / RDNA4 守卫下面这段代码在没有 RDNA4 的环境里不会崩溃它只是演示「你会在什么代码路径上触发这个 SDPA 调用」在真 RDNA4 有缺陷的 PyTorch 上它会触发 CK 断言。我们用torch.cuda.is_available()这类守卫来避免在非 AMD GPU 环境误报同时给出真实可用的调用方式。import torch import torch.nn.functional as F def make_qkv(batch2, heads4, seq512, dim64, devicecuda, dtypetorch.float16): # 标准注意力输入 (B, H, N, D) q torch.randn(batch, heads, seq, dim, devicedevice, dtypedtype) k torch.randn(batch, heads, seq, dim, devicedevice, dtypedtype) v torch.randn(batch, heads, seq, dim, devicedevice, dtypedtype) return q, k, v def run_default_sdpa(): 默认后端选择在 RDNA4 上有缺陷的 PyTorch 上可能触发 CK flash assert。 if not torch.cuda.is_available(): print([skip] 当前环境无 GPU仅演示调用路径) return q, k, v make_qkv() # 这一行在 gfx1200/gfx1201 有缺陷的 PyTorch 上会崩在内核断言 out F.scaled_dot_product_attention(q, k, v) print(SDPA 默认后端输出形状:, out.shape) def current_gfx(): 打印当前 AMD GPU 的 gfx 架构仅 ROCm 有意义。 if not torch.cuda.is_available(): return N/A # ROCm 下可通过设备名粗略判断更精确可用 rocminfo return torch.cuda.get_device_name(0) if __name__ __main__: print(设备:, current_gfx()) run_default_sdpa()要点在 RDNA4 上run_default_sdpa()里那一行F.scaled_dot_product_attention就是崩点。要稳住就得在调用前关掉 flash 与 mem_efficient 后端强制走 math。五、解决方案一强制只用 math 后端最稳立即可用这是最推荐的第一步绕过方案。它不改变你的模型结构只是让 SDPA 别去碰有问题的 CK flash 内核import torch import torch.nn.functional as F def run_sdpa_math_only(q, k, v): 强制只用 math 后端绕开 RDNA4 上的 CK flash 断言。 with torch.backends.cuda.sdp_kernel( enable_flashFalse, enable_mem_efficientFalse, enable_mathTrue, ): return F.scaled_dot_product_attention(q, k, v) if torch.cuda.is_available(): q, k, v make_qkv() out run_sdpa_math_only(q, k, v) print(math 后端输出形状:, out.shape)优点100% 稳定不依赖 CK 内核数值正确就是标准注意力公式改动极小一行with即可。 代价速度慢。math 后端不走 flash attention 的分块/在线 softmax 优化显存占用也更高需要物化QK^T的N×N矩阵。对小序列无所谓长序列会明显变慢、吃显存。注意在 ROCm 上torch.backends.cuda.sdp_kernel同样有效因为它控制的是 SDPA 的 backend 调度逻辑而 CK flash 正是被enable_flash控制的那个后端。把enable_flashFalse就绕开了 CK 的 flash 内核。六、解决方案二封装一个「RDNA4 安全」的注意力函数更工程化的做法写一个自动判断、对 RDNA4 自动降级的函数让其它架构仍享受 flash 性能import torch import torch.nn.functional as F def is_rdna4(device_name: str) - bool: 粗略判断是否为 RDNA4。生产环境建议用 rocminfo 取 gfx1200/1201。 name (device_name or ).lower() # 仅示例RX 9000 系列多为 RDNA4不同厂商命名有差异 return any(tag in name for tag in (rx 90, gfx1200, gfx1201, rdna4)) def safe_scaled_dot_product_attention(q, k, v, attn_maskNone, dropout_p0.0): RDNA4 上强制 math其它架构走默认含 flash。 use_math_only False if torch.cuda.is_available(): name torch.cuda.get_device_name(q.device if isinstance(q, torch.Tensor) else 0) use_math_only is_rdna4(name) if use_math_only: with torch.backends.cuda.sdp_kernel( enable_flashFalse, enable_mem_efficientFalse, enable_mathTrue, ): return F.scaled_dot_product_attention(q, k, v, attn_mask, dropout_pdropout_p) return F.scaled_dot_product_attention(q, k, v, attn_mask, dropout_pdropout_p)把模型里所有F.scaled_dot_product_attention替换成safe_scaled_dot_product_attention就能在不改动算法的前提下规避 RDNA4 的 CK 断言。七、解决方案三升级 PyTorch / ROCm / CK 到已修复版本pytorch/pytorch#188113是已知 issue官方与 CK 社区会随版本修复。因此长期正确做法是升级# 查看当前版本 python -c import torch, sys; print(torch.__version__, torch.version.hip if hasattr(torch.version,hip) else n/a) # 升级到修复了该 RDNA4 断言的 nightly / 新正式版ROCm wheel pip install --upgrade torch torchvision --index-url https://download.pytorch.org/whl/rocm6.3判断「是否已修复」的标准升级后在 gfx1200/gfx1201 上直接调用F.scaled_dot_product_attention不再触发 CK 断言用TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL视版本而定或查看 release note 是否提及 RDNA4 / gfx1200 / gfx1201 flash attention 支持跑下面的自检脚本确认后端可用。 自检脚本import torch import torch.nn.functional as F if torch.cuda.is_available(): q, k, v (torch.randn(1, 4, 256, 64, devicecuda, dtypetorch.float16) for _ in range(3)) try: out F.scaled_dot_product_attention(q, k, v) # 默认后端 print([OK] 默认 SDPA 在 RDNA4 上可用CK 断言应已修复) except RuntimeError as e: if assert in str(e).lower() or CK in str(e): print([FAIL] 仍触发 CK 断言请用 math 后端绕过并继续升级) else: print([other error], e)八、解决方案四在 Transformers / vLLM 等上层框架里关掉 flash_attention_2很多上层库会直接请求attn_implementationflash_attention_2这同样会落到 CK 的 flash 内核。需要在配置层关掉from transformers import AutoModelForCausalLM # 不要这样RDNA4 上会崩 # model AutoModelForCausalLM.from_pretrained(xxx, attn_implementationflash_attention_2) # 改成 eager 或 sdpa配合上面 math-only 的 backend 关闭 model AutoModelForCausalLM.from_pretrained( xxx, attn_implementationeager, # 或 sdpa再配合 sdp_kernel 关 flash )对 vLLM 等推理引擎若其默认走 flash attention需要查对应版本的--attention-backend之类参数切到非 flash 的实现如FLASHINFER/XFORMERS/TORCH_SDPA并进一步用sdp_kernel关 flash。九、排查清单RDNA4 CK SDPA 断言按以下顺序排查能快速定位与绕开确认架构rocminfo | grep gfx看是不是gfx1200/gfx1201。非 RDNA4 却崩多半是别的问题驱动 / 显存。确认崩在内核断言报错含CK_ASSERT/hipErrorAssert/flash attention kernel asserts即为本 bug。第一步绕开用torch.backends.cuda.sdp_kernel(enable_flashFalse, enable_mathTrue)强制 math确认能跑通。升级升级 PyTorch ROCm重跑第八章自检脚本验证是否修复。上层框架检查 Transformers / vLLM 是否强制flash_attention_2改为eager/sdpa并关 flash。显存math 后端N×N中间矩阵大长序列可能 OOM必要时减小 batch / 序列长度或等官方修复后用 flash。十、小结CK SDPA flash attention kernel asserts on gfx1200/gfx1201 (RDNA4)的本质是PyTorch 的 SDPA 调度器在 AMD RDNA4 上把注意力派发到了 Composable Kernel 的 flash attention 内核而该内核在 RDNA4 架构上有未修复的断言缺陷导致 GPU 端断言失败、进程直接崩溃而不是优雅回退到 math 后端。 应对三步走立即可用用torch.backends.cuda.sdp_kernel(enable_flashFalse, enable_mathTrue)强制 math 后端100% 稳定代价是慢、吃显存工程化封装safe_scaled_dot_product_attention对 RDNA4 自动降级、其它架构仍用 flash根本解决升级到修复该 RDNA4 断言的 PyTorch / ROCm / CK 版本并用自检脚本验证。 上层框架Transformers、vLLM若强制flash_attention_2需在配置层切回eager/sdpa并配合关闭 flash 后端。只要绕开 CK 的 flash 内核RDNA4 上跑注意力模型完全没问题等官方完善 RDNA4 特化后就能恢复完整性能。