Wavefront 调度模型详解，理解 AMD GPU 并行计算的核心

从 Warp 到 Wavefront：打破 NVIDIA 思维定式

很多刚从 CUDA 转向 AMD ROCm 的开发者，最容易踩的坑不是环境配置，而是“思维惯性”。在 NVIDIA 的世界里，我们习惯了 Warp（线程束）的概念，默认 32 个线程为一组同步执行。但当你把这套经验直接生搬硬套到 AMD GPU 上时，往往会发现性能不达预期，甚至出现奇怪的延迟抖动。这背后的核心原因，在于 AMD 架构中调度的基本单位是Wavefront，而非 Warp。

虽然两者在概念上都代表一组并行执行的线程，但本质差异巨大。NVIDIA 的 Warp 大小固定为 32，硬件对分支发散（Branch Divergence）有一定的容忍和掩盖机制；而 AMD 的 Wavefront 大小通常是 64（具体取决于架构，如 CDNA 或 RDNA 系列），且其对执行效率的要求更为严苛。在 AMD 架构中，如果一个 Wavefront 内的线程走向了不同的分支路径，硬件不会像某些旧架构那样灵活调度，而是必须串行执行所有分支路径，直到所有线程汇合。这意味着，分支发散在 AMD 上的代价是实打实的性能减半甚至更低，没有任何“魔法”能帮你自动优化。

可视化理解：线程束发散的真实代价

为了更直观地理解这一点，我们可以想象一个典型的条件判断场景。假设在一个并行计算内核中，我们需要根据线程 ID 执行不同的数学运算：

// 伪代码示例：存在潜在的分支发散if(thread_id%2==0){result=do_complex_calc_A(data);}else{result=do_complex_calc_B(data);}

在 NVIDIA 的 Warp 模型下，如果这 32 个线程中有一半走if，一半走else，硬件可能会通过掩码（Masking）技术让两组指令分时复用执行单元，虽然有效率损失，但上下文切换开销相对可控。

但在 AMD 的 Wavefront 模型中，情况则严峻得多。一个 Wavefront 包含 64 个线程。如果这 64 个线程因为上述逻辑被切分成两半，那么整个 Wavefront 必须先完整执行完do_complex_calc_A（此时另一半线程闲置等待），然后再完整执行do_complex_calc_B。这种串行化执行直接导致该 Wavefront 的有效吞吐量下降 50%。如果分支逻辑更复杂，涉及三层嵌套，性能损耗可能呈指数级上升。

这就是为什么直接移植 CUDA 代码往往“能跑但很慢”的原因：原本在 NVIDIA 卡上被硬件部分掩盖的发散问题，在 AMD 卡上被赤裸裸地暴露了出来。要解决这一问题，不能依赖编译器的自动优化，必须从算法设计层面入手，确保同一个 Wavefront 内的线程尽可能走相同的执行路径。

实战优化：SGLang 与 TileLang 的适配策略

理解了 Wavefront 的特性后，我们再来看如何在实际的大模型推理框架中进行优化。以SGLang为例，这是一个专为大语言模型设计的高吞吐推理框架。在将其迁移至 ROCm 平台时，我们发现默认的连续批处理（Continuous Batching）策略在某些场景下会导致严重的 Wavefront 发散。

原因在于，不同长度的请求在动态批处理中可能被分配到同一个 Block 中，导致部分线程提前完成计算进入等待，而另一部分线程仍在处理长序列。在 NVIDIA 上这可能只是轻微的负载不均，但在 AMD 上，这意味着整个 Wavefront 被拖慢到了最慢那个线程的速度。

解决方案是调整 SGLang 的调度策略，使其在分配请求时，不仅考虑显存剩余量，还要尽量保证同一个调度单元内的序列长度相近。通过引入基于序列长度的分组机制，我们可以确保落入同一个 Wavefront 的线程具有相似的计算路径，从而大幅减少分支发散带来的空转。

而对于更底层的算子优化，TileLang则提供了更精细的控制手段。TileLang 允许开发者以高层次的语言描述矩阵分块（Tiling）和数据流动，并针对特定架构生成代码。在优化 Attention 机制时，我们利用 TileLang 重新设计了共享内存（LDS）的访问模式。

具体来说，AMD GPU 的 LDS 带宽极高，但如果访问模式不当（如非对齐访问或 Bank Conflict），会严重阻塞 Wavefront 的执行。通过 TileLang，我们可以显式指定数据在 LDS 中的布局，使其完美匹配 64 线程的 Wavefront 宽度。例如，将矩阵分块大小设置为 64 的倍数，并确保线程索引与内存地址严格对齐。实测表明，经过这种针对性优化的算子，在 MI300X 等卡片上的执行效率相比直接移植的 CUDA 版本提升了近 30%。

拒绝生硬套用：建立 AMD 原生开发直觉

从 Warp 到 Wavefront 的转变，不仅仅是参数的调整，更是开发直觉的重塑。在 NVIDIA 生态中，我们有时可以依赖硬件的“宽容度”写出一些不够严谨的代码；但在 ROCm 生态中，显式的资源控制和严谨的逻辑规划才是高性能的关键。

不要害怕去阅读生成的汇编代码，也不要吝啬使用rocprof进行性能剖析。当你发现性能瓶颈时，首先问自己：这个 Wavefront 内部是否存在分支发散？内存访问是否对齐？线程块大小是否是 64 的倍数？

社区的力量在这一过程中至关重要。在 GitHub 上，许多开发者已经分享了针对 SGLang 和 TileLang 的优化补丁。比如，有人通过提交 PR 修复了特定架构下的分块策略，有人则完善了 LLaMA-Factory 在 ROCm 下的混合精度训练逻辑。参与这些讨论，不仅能解决具体问题，更能让你快速建立起对 AMD 架构的敏感度。

最终，当你不再下意识地将 CUDA 代码逐行翻译，而是开始从 Wavefront 的角度思考并行度与数据流时，你就真正掌握了 AMD GPU 并行计算的核心。这不仅能让你的模型在 AMD 硬件上跑得更快，也能让你在异构计算的广阔天地中，拥有更从容的技术底气。

200小时GPU算力已就位，快来领取：https://marketing.csdn.net/questions/Q2604140858304426315?utm_source=AIpaper