1. 从一个“全员同步”的加载-计算循环说起
1.1 最朴素的 Kernel 循环:加载、同步、计算、同步
看标题你可能以为要聊 Linux 内核,或者某个叫 Semantic Kernel 的框架。都不是。这里的 Kernel 是 GPU Kernel,是 CUDA 里每次用<<<grid, block>>>发出去的那个完整执行单元。我最近在调整混合专家模型里的自定义算子,遇到一个典型的性能困惑:显存带宽没吃满,SM 的 FMA 流水线也没吃满,但整个 Kernel 跑下来就是不够快。
把 Nsight Compute 的 warp stall 分布拉出来看,答案其实很老土:大量线程的时间耗在__syncthreads()前后的空等上。很多计算类 Kernel 内部都有类似这样一段循环:
for (int k = 0; k < K; k += BK) { load_tile_global_to_smem(k); // 全部线程一起把下一块数据搬到 shared memory __syncthreads(); // 必须等所有数据到齐 compute_tile(k); // 全部线程一起做计算 __syncthreads(); // 必须等所有线程算完,下一轮才能覆盖 smem }这段代码是正确的,它在数学上每一步都无懈可击,但执行时间上有个很明显的结构性浪费:加载阶段,所有 warp 都被全局内存的长延迟拖住,SM 里的计算单元没有新指令可以发射;计算阶段,访存单元又相对空闲,global memory 和 shared memory 之间的搬运暂停了。换句话说是同一批资源在时间上轮流休息,而不是让不同资源同时满负荷工作。
__syncthreads()这个全块等待是“最粗粒度”的同步。Block 里任何一个 warp 因为某个 global memory 访问多等了几个周期,其他 warp 就得在 barrier 前面陪着等。这种木桶效应在一个内存访问非常不规律的 AI 算子里会被进一步放大。你很难说某一个 LDG 指令真的慢,但统计下来,大量 cycle 都烧在了“等别人”上面。
1.2 跨 Kernel 的 Stream 流水线为什么补不上片上的空泡
面对这种问题,很多人的第一反应是把大 Kernel 拆成多个小 Kernel,然后用多个 CUDA Stream 做流水线。比如把“加载+预处理”放一个 Kernel,“矩阵乘主体”放另一个 Kernel,计算 tile 时让前一个 Stream 已经开始加载下一个 tile 的数据。这种方案确实能产生重叠,也已经有很多框架在用,但它解决不了片上数据复用的问题。
一个 Kernel 结束后,它独占的 shared memory 状态就作废了,数据如果要传给下一个 Kernel,必须经过 global memory 或者 L2,哪怕只是逻辑上的中间结果。对 GEMM、Attention 这种需要反复复用分块数据的算子来说,这意味着本来可以留在片上被计算单元反复读取的 A/B tile,被迫绕了一圈全局内存。再加上 Kernel 启动本身有固定开销,通常几个微秒到十几个微秒,当单个 Kernel 只有几十微秒量级时,流水线还没来得及进入稳态就已经结束了。
还有个更现实的麻烦是 Kernel 之间的依赖粒度太大。Stream 流水线只能做到“第 n 个 Kernel 整体”和“第 n+1 个 Kernel 整体”错开,做不到同一个 Kernel 内部第 3 块和第 4 块数据之间的细粒度同步。真正让性能模型发生变化的是另一条路:在一个 Kernel 内部,在循环处理多个数据块的过程中,让“下一块数据的搬入”和“当前块数据的计算”重叠起来。这个思路就是 Intra-Kernel 流水线。
1.3 Intra-Kernel 流水线的目标:把下一块数据加载藏进当前计算
Intra-Kernel 流水线的目标模型可以写成一个简单的时间关系:原本一个循环要串行完成“加载 tile、等待、计算、等待”,现在要让加载和计算跨越边界重叠。理想情况下,程序执行第 i 块 tile 的计算时,硬件已经或者正在把第 i+1 块 tile 的数据从 global memory 搬到 shared memory。
要做到这件事,不能继续依赖__syncthreads()这种全线程到达型同步,因为你要等的不是“所有线程都执行到了某个点”,而是“我关心的那块 shared memory 缓冲区已经可用了”。这就需要两样东西:能异步完成数据搬移的硬件通道,以及能针对某个缓冲区精确放行/阻塞的同步原语。于是话题自然引到 Hopper 这一代架构给出的两个关键机制,TMA 与 mbarrier。
2. Hopper 的 TMA 与 mbarrier:Kernel 内流水线的两条硬件腿
2.1 TMA:一个类似 DMA 控制器的异步搬运引擎
Hopper 架构引入的 TMA,全称 Tensor Memory Accelerator,本质是一个独立于 SM 计算核心的数据搬运引擎。在没有 TMA 之前,把数据从 global memory 搬到 shared memory,需要每个线程自己发 LDG/STS 指令,哪怕你用上了cp.async,也只是把“load 后不用立刻等数据”这件事变成了异步,但搬运动作本身还是由线程指令驱动的,warp 仍然要为这些指令付出发射开销。
TMA 不一样。它可以在一个线程或一个 warp 发起一次cp.async.bulk或cp.async.bulk.tensor之后,由专门的硬件把整个多维 tile 从 global memory 搬进 shared memory。SM 上的计算 warp 完全不用管这次搬运,期间它们可以继续执行别的计算指令。这个模型和 CPU 世界里用 DMA 控制器做内存拷贝是同一个道理。CPU 发出 memcpy 请求后,DMA 引擎负责搬运,CPU 内核可以去算别的东西,等 DMA 完成后再通过中断或状态位通知 CPU。
对 GEMM 这类算子来说,TMA 还解决了一个很痛的问题:张量形状的边界处理。过去要把一个 M x N 的大矩阵切成 tile,遇到边缘不整块时,每个线程要自己在代码里判断边界、填充零;TMA 配合 Tensor Map,可以在硬件层面直接描述多维张量的形状和 tile 大小,搬运引擎会自动处理越界问题。这不光是代码简洁性的提升,还减少了分支指令,对指令发射效率也有帮助。
2.2 mbarrier:从全块同步变成可编程的状态栅栏
第二块拼图是 mbarrier,可以理解成一个硬件实现的、可编程的异步屏障。__syncthreads()要求整个线程块的所有线程都到达后,所有线程才一起放行,这是匿名广播式的全量同步。mbarrier 则更接近一个计数器:初始化时指定一个预期的到达数,不同的线程或代理可以各自调用 arrive 把它减一,当计数归零时,所有调用 wait 的线程被放行。
这里最有意思的是 mbarrier 可以和 TMA 联动。Hopper 提供了一种机制,把一次 TMA 拷贝的期望字节数登记到 mbarrier 上,叫做 expected transaction count。也就是说,数据搬运这件事不再需要某个线程在搬完之后显式地“到达”一下屏障,而是 TMA 硬件完成拷贝的同时,会自动让 mbarrier 的计数减少。消费者 warp 只需要在这一侧的 mbarrier 上等待,一旦数据真正落到 shared memory,等待队列里的 warp 就被唤醒。
如果只盯着__syncthreads()和 mbarrier 的 API 形式看,你会觉得它们都是“等一个条件满足再继续”,但粒度完全不同。__syncthreads()等待的是全体线程的程序位置,mbarrier 等待的是一个共享状态的变化。前者锁死整个 CTA 里所有 warp 的执行进度,后者允许不同 warp 各自按照数据依赖前进。后者才是 pipeline 能跑起来的关键。
2.3 三层开发视图:PTX、libcu++、CUTLASS
面向这些硬件机制写代码,有三种深度。你完全可以直接写 PTX,比如cp.async.bulk.tensor、mbarrier.arrive.expect_tx这类指令。能看到最底层的语义,但代码可维护性很差,而且不同架构的细节差异都要自己处理,适合做极端性能验证,不适合做产品代码。
更常见的做法是用 libcu++ 提供的cuda::barrier、cuda::memcpy_async、cuda::pipeline这些抽象。它们把 TMA 和 mbarrier 的底层指令封装成了类型安全的 C++ 接口,同时保留足够细的控制粒度。再往上一层是 CUTLASS 或 CUTE 里的 collective 封装,它们已经帮你设计好了 producer/consumer 的角色划分和 stage 管理,普通业务方最合理的姿势是站在这一层做配置和扩展,而不是每次从 memset 开始发明轮子。
我个人的经验是:理解原理时至少看一眼 PTX 一层,知道硬件到底在做什么;真正落产品时尽量往上走,否则光是管理 mbarrier 在不同 phase 下的状态就能让你多调一个星期。
3. Warp 特化:不是分支分歧,而是把角色直接分给不同 Warp
3.1 为什么传统上我们那么怕 warp divergence
初学 CUDA 时我们都被教育过:尽量别让同一个 warp 里的线程走不同分支,因为 warp 是 SIMT 执行的基本单位,一旦同一个 warp 的 32 个线程分别进入 if 和 else,硬件只能先执行 if 路径再执行 else 路径,有效利用率砍半。所以看到“Warp 特化”这个词,很多人第一反应是:这不就是在制造 warp divergence 吗?
不是。warp divergence 说的是同一个 warp 内部不同 lane 执行不同路径,导致串行化。而 warp specialization 说的是线程块内不同的 warp 各自拥有完全不同的指令流,一个 warp 只做搬运,另一个 warp 只做计算,它们在时间上是可以并行执行的。GPU 的 warp scheduler 本来就是多发射、多流水的,一个 SM 里有多个 warp scheduler,不同的 warp 各自维护程序计数器,这些 warp 执行不同的代码路径并不会让彼此串行。真正需要抛弃的是“整个 CTA 里所有 warp 都执行同样的循环体”这个旧直觉。
有点像一个团队去搬仓库:原来 32 个人排成一排,一次一起搬一箱,有一个人路上慢一点,整排人都得停下。现在分成两组,一组持续从外面进货,另一组持续拆箱处理,两组通过传送带衔接。这不会让个人变快,但能让仓库里始终有活干。
3.2 生产者 Warp 与消费者 Warp 如何并行
在典型的 warp-specialized kernel 中,线程块内部通常分成两类角色。一类是 producer warp,职责是驱动 TMA,把 global memory 里的下一块 tile 数据搬到 shared memory 的某个 stage 缓冲区;另一类是 consumer warp,职责是等待数据就绪后,从 shared memory 读入数据做矩阵乘法或别的计算。
Producer warp 不一定很多。很多时候一个 256 线程的 block 里有 1 到 2 个 warp 负责 TMA 就够用,因为 TMA 本身并不需要每个线程去搬一个字节,一条 bulk copy 指令是由硬件引擎完成的,producer warp 的主要工作是提前发起多个 stage 的拷贝请求,以及维护 mbarrier 的到达计数。真正做矩阵乘的 consumer warp 占大多数,这样才能保证计算密度不下降太多。
这种角色分离会改变你对“block 内并行度”的统计方式。表面上看,producer warp 没有参与计算,你损失了几分之一的理论 FLOPS,但如果原来的实现在每次加载时都要让所有计算 warp 停下等数据,那么这点损失换来的往往是整体吞吐的大幅提升。流水线深度足够时,计算单元几乎永远有活可干,那点并行度损失会被隐藏掉。
3.3 从结构上对比:同步粒度彻底变了
传统循环里有两个__syncthreads():一个保证数据加载完,一个保证上一轮计算完。前者是生产者到消费者的同步,后者是消费者到生产者的同步。如果用两个专门化的角色来重新组织,结构会变得很像操作系统里的环形缓冲区。
关键区别在于,consumer 等待“某个 stage 的 full barrier”时,不需要等待其他所有 consumer 也到达;producer 等待“某个 stage 的 empty barrier”时,也不需要等待所有其他 producer。整个 block 的执行流不再是一波一波往前推,而是多个 warp 各自按照自己的依赖关系交错推进。在这种模型下,block 更像一个微型的、生命周期只有几百微秒的协作式操作系统:producer 是 I/O 驱动,consumer 是计算进程,mbarrier 是事件状态,shared memory 中的 stage 缓冲区就是邮箱。
我把这个想法放到实际算子优化中后,对 GPU 执行模型的理解整个翻了一遍。原来的问题一直是“怎么让 100 个 warp 同步地做同一件事”,现在应该问的是“怎么让不同 warp 在各自职责内减少互相等待”。
3.4 一段足够说明问题的伪代码
为了把角色划分讲清楚,我用伪代码示意一个双缓冲的 ring buffer 结构。这不是一段可以直接编译的完整 CUDA 代码,主要目的是展示同步逻辑上发生了什么变化。
// Producer warp 的角色 for (int tile = 0; tile < total_tiles; ++tile) { int slot = tile % stages; empty_bar[slot].arrive_and_wait(); // 等消费者用完这个 slot cuda::memcpy_async( smem_stage[slot], gmem_src + tile_offset(tile), full_bar[slot] // 拷贝完成后自动到达 full_bar ); } // Consumer warp 的角色 for (int tile = 0; tile < total_tiles; ++tile) { int slot = tile % stages; full_bar[slot].arrive_and_wait(); // 等数据真正可见 compute_on_stage(tile, slot); // 做计算,不碰别的 slot empty_bar[slot].arrive_and_wait(); // 通知 producer 这个 slot 可覆盖 }这里每个 stage 需要两个 barrier,一个表示“缓冲区空闲可写”,一个表示“数据已就绪可读”。如果只用一个 barrier,死锁几乎是必然的,因为 producer 可能覆盖掉还在被 consumer 读取的缓冲区。Producer 发起异步拷贝之后立刻返回,并不会阻塞在memcpy_async上,所以它能继续往前走,发起下一轮的拷贝请求。Consumer 则会在full_bar上等待数据真正到达。
这个结构初看和 CPU 上的软件流水线很像,但底层是由 TMA 硬件和 mbarrier 状态机支撑,等待成本低得多,推进粒度也细得多。
4. 从小 GEMM 改造看完整流程:K 循环变成环形缓冲流水线
4.1 选一个带 epilogue 的 SGEMM 当实验台
直接拿完整的 CUTLASS 源码看 warp specialization,容易被一大坨模板参数淹没。我建议第一次做实验时选一个自己能完整掌控的算子,比如一个带 bias 加法和 ReLU 的 SGEMM。算子不算难,但已经包含主循环、shared memory 复用和 epilogue 三个阶段,足够暴露问题。
假设输入是两个 float 矩阵 A[M,K] 和 B[K,N],输出是 C[M,N] = A*B + bias,再过一个 ReLU。朴素实现通常是每个线程负责 C 中的一个或多个元素,外层循环按 K 分块,每一轮把 A 的一块和 B 的一块搬进 shared memory,算完一部分累加。在这个朴素实现里,每个外层 K 迭代都夹着两次全块同步。K 很大时,这个同步等待会累积成非常