你的显卡能跑多少算子？用 55 个检查项，给 PyTorch GPU 环境做一次冒烟测试

背景

之前的文章: ROCm on Windows 性能排查：RX 6650 XT 跑 PyTorch，为什么加速不明显？

在 Windows 上用 RX 6650 XT 跑自编译 ROCm + PyTorch 的时候，我遇到一个问题：LLM 推理确实跑在 GPU 上了，但加速比只有 1.7-2.0x，任务管理器里 GPU 利用率也很低。

第一反应是——是不是很多操作没有 GPU 路径，PyTorch 悄悄把计算挪回了 CPU？

为了搞清楚这件事，我写了一个简单的 Python 脚本。它做的事情用一句话就能说清楚：对一堆常见的 PyTorch 操作，检查输入是 CUDA/HIP tensor 的时候，输出是不是还在 CUDA/HIP 上。

如果某个操作直接报错，或者输出跑到了 CPU，那就值得深挖一下。

脚本在做什么

核心逻辑非常朴素：

x = torch.randn(128, 512, device="cuda")
out = some_operation(x)
print(out.device)  # 还在 cuda:0 上吗？

再配合 warmup + 计时（前后都 torch.cuda.synchronize()），以及 Python warning 捕获和错误分类，就是一个完整的冒烟测试了。

但这里有个容易误解的地方：输出 tensor 仍在 GPU 上，不等于底层一定调用了高性能 GPU kernel。 PyTorch 的设备语义本身就倾向于让 CUDA tensor 的结果留在同一设备上；缺少对应 accelerator 实现时，常见行为是直接报错，而不是悄悄生成 CPU tensor。所以"输出还在 cuda 上"这个结论的含金量，其实比"证明该算子有高性能 GPU kernel"低得多。

覆盖了哪些操作

脚本测了 9 类常见操作，大概 55 个检查项：

选这些操作的原因很简单：它们基本覆盖了主流深度学习模型（包括 Transformer）里最常见的计算模式。够用，但不代表完整 workload 的性能质量。

计时怎么做

GPU 操作是异步提交的，所以计时前后必须同步：

for _ in range(3):       # warmup
    out = fn()

torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(20):
    out = fn()
torch.cuda.synchronize()
elapsed = (time.perf_counter() - t0) / 20 * 1000

这些时间只能粗略参考。部分检查项会在 lambda 里创建随机输入，计时会混入 tensor 分配和随机数生成的成本。要做严肃性能分析，应该单独固定输入、用更长的 repeat、配合 profiler。

实际跑出来什么样

在 RX 6650 XT (gfx1032) + 自编译 ROCm/PyTorch 环境下，结果类似：

Summary: 54/55 returned CUDA/HIP tensors | 1 errors | 1 warnings

怎么理解这个结果：

• 54 个检查项成功返回了 CUDA/HIP tensor。这是好消息——说明这些操作在当前环境下不会直接崩掉。
• batch_norm 报错了。MIOpen 跑 HIPRTC 编译的时候找不到 <type_traits>，这是 C++ 头文件缺失的问题，装上 VS Build Tools 的 C++ 工具链大概率能修。
• sdpa 能跑，但 PyTorch 没启用 memory efficient attention，走的是数学 fallback。功能没问题，性能不是最优。

但这个结果不能写成"算子覆盖率 98%"——这 54 个只是"输出还在 GPU 上"，不是"证明有高性能 GPU kernel"。如果你看到有人拿类似的数字说"ROCm 在 Windows 上已经完美支持了"，那个结论是过头的。

错误诊断

脚本会对报错做自动分类，直接在终端里给出建议。目前能识别的几类：

BatchNorm MIOpen/HIPRTC error — 最常见的问题。MIOpen 用 HIPRTC 运行时编译 kernel，但 HIPRTC 在 Windows 上经常找不到 C++ 标准库头文件。stderr 里会看到 fatal error: 'type_traits' file not found。修复方法是装 VS Build Tools 并确认 HIPRTC 的 include path 配置正确。注意：这不是 CPU fallback，是编译直接失败。

SDPA: no flash attn — PyTorch 编译时没开启 memory efficient attention / FlashAttention。SDPA 会走数学 fallback，仍在 GPU tensor 上运行，但对长序列、大 batch 的 LLM workload 来说性能差距明显。

GPU OOM — 显存不够，减小输入尺寸或关掉其他 GPU 进程。

MIOpen runtime error — 其他 MIOpen 错误。可以试试清 kernel 缓存（rmdir /s /q %LOCALAPPDATA%\miopen），或者开 MIOpen_ENABLE_LOGGING=2 看详细日志。

这个脚本适合干什么

• 新搭好 ROCm/PyTorch 环境后，快速确认常见操作不会直接报错。
• ROCm 或 PyTorch 升级后，做回归检查。
• 作为 profiler 定位的第一层筛查——先知道哪些操作有问题，再深入分析。

不适合拿来做"算子覆盖率统计"或者"性能瓶颈最终结论"。

如果真的想定位性能问题

光看输出 device 是不够的。要搞清楚 LLM 推理为什么慢，建议继续做这些：

1. 用 torch.profiler 记录 CPU 和 CUDA/HIP activity 的时间分布。
2. 把 prefill 和 decode 分开计时——自回归 decode 的瓶颈和 prefill 完全不同。
3. 用真实 LLM 的 tensor shape 单独测 linear、RMSNorm、RoPE、SDPA、KV cache 路径。
4. 开 MIOpen/rocBLAS 日志，确认关键操作到底走了哪个库。
5. 试不同 batch size、prompt length、max_new_tokens，看性能曲线怎么变。

这样得到的结论，比"看输出还在不在 GPU 上"要靠谱得多。