在现代深度学习与高性能计算中,矩阵乘法性能直接决定了训练与推理的吞吐量。在 NVIDIA 的 Blackwell 架构上,借助 Pallas/Mosaic 等工具链,可以通过一系列结构化优化把自定义内核性能推到或超过成熟库的水平。本文从概念到实现,分步介绍在 Blackwell 上编写高性能矩阵乘法内核的关键技巧与实践经验,帮助工程师理解如何平衡算力与内存带宽、如何利用集体 MMA 与线程分工来提高 TensorCore 利用率。\n\n理解问题与架构机会是优化的第一步。Blackwell 引入了高效的 TensorCore 指令与跨块(cluster/CTA)通信能力,允许不同的 SM 通过集体 MMA 协同完成一个更大的乘法子任务。与前代 GPU 相比,这提供了额外的手段来提高算密度(每字节加载对应的浮点运算量),从而摆脱单纯受内存带宽限制的瓶颈。
但要真正发挥这些能力,需要在内核层面做精细设计:合适的分块策略、SMEM 与 TMEM 的布局与双缓冲、同步与信号的使用、以及程序级别的持续化(persistence)与 warp 专用化。\n\n开始实现时,按照由简到繁的策略逐步演进能快速验证每一步优化的收益。最基础的实现是单 CTA、单 warpgroup 的端到端流程:从全局内存(GMEM)发起异步拷贝(TMA)到共享内存(SMEM),由单个 warpgroup 执行一系列 MMA 指令,最后将累加结果从 TMEM 取回写回 GMEM。这个实现便于验证正确性、设置基本 tiling 参数(tile_m、tile_n、tile_k),并确保 SMEM 数据格式与 MMA 指令期望的内存格式对齐。一个常见技巧是在 SMEM 中应用行/列排列变换(swizzle),用 plgpu.find_swizzle 与 TilingTransform/SwizzleTransform 来匹配硬件期望的访问模式。初始实现通常能达到大约 30%~40% 的 TensorCore 利用率,是验证流水线与基本资源使用的良好起点。
\n\n性能进一步提升的第一步是 warp 专用化(warp specialization)。Blackwell 的 TMA 与 tcgen05_mma 指令只需要单个 CUDA lane 来下发。因此可以把一个 warpgroup 分解为若干个 warp,每个 warp 只承担特定角色。利用 pl.core_map 与 plgpu.WarpMesh,可以在进入内核后让每个 warp 独立执行只读或只写任务,例如只发起 GMEM->SMEM 的异步加载,或只执行 MMA 运算。这样做的核心收益是降低控制流与内存发起的干扰,使得 MMA 的发起与资源占用更稳定。需要注意的是,warp 之间通过 barrier 协调数据可用性:为每个加载槽配置 load_barriers,计算端等待加载完成;为避免覆盖被正在使用的 SMEM,加载端等待 consumed_barriers。
另一个关键点是使用 TensorCore 专用的完成信号(mma_done_barrier),由执行 MMA 的 warp 在完成时请求 TensorCore 在该 barrier 上到达,从而保证后续 TMEM 读取或 epilogue 能安全进行。采用 warp 专用化后,通常可以看到显著的提升,TensorCore 利用率大幅上涨。\n\n优化 epilogue(累加结果写回)的效率也能带来明显收益。把一次写回整个 tile 的做法替换为对输出列分块(epilogue_tile_n)的小切片循环,并将 TMEM->SMEM 的加载与 SMEM->GMEM 的拷贝交错执行,可以把计算与写回更好地流水化。通过分配更小的 SMEM 缓冲区并进行双缓冲(例如 acc_smem 的两个槽),可以在一个槽由 TMA 发起写回时同时在另一个槽上装载新数据,这样可以减少 SMEM 占用并增加并行度。实现时注意 commit 与 wait 等调用的顺序:在完成 async_load_tmem 后需要 plgpu.commit_smem,然后发起 copy_smem_to_gmem,再通过 plgpu.wait_smem_to_gmem 确认拷贝完成。
当 Epilogue 分块得当,SMEM 与 TMEM 的重叠使用可以显著降低写回延迟。\n\n尽管以上优化能提升效率,但当内存仍然成为限制时,需要提升算密度。Blackwell 的集体 MMA(collective MMA,也称 2CTA 或跨 CTA 协作)是解决路径之一。其思想是用簇内多个块共同计算一个更大的输出 tile,每个块只加载自身负责的一部分输入数据,但 MMA 指令能够在执行时跨块读取对方 SMEM 的数据,从而把每次加载放大为更多的计算量。实现上,需要在 plgpu.kernel 中使用 cluster 参数并把 TMEM 标记为 collective=True。GMEM->SMEM 的复制也需以 collective_axes 指定集体操作,并用 partitioned_axis 指明哪个轴在集群间划分。
要牢记:在使用 partitioned collective copy 时,只有 leader block(簇内的领导)会看到加载完成的 barrier;非 leader block 不会等待该 barrier。集体 MMA 还需要把 tcgen05_mma、tcgen05_commit_arrive 等调用标注 collective_axis,以确保跨块的完成信号能够传达给集群内所有参与者。通过将 tile 的尺寸放大为 cluster_tile_m、cluster_tile_n,并把网格相应缩放,可以在相同的硬件上翻倍算密度,从而大幅提升利用率。\n\n提升内核效率的另一个关键策略是持久化内核(persistent kernel)。传统方法会为 GPU 上的每个输出 tile 启动一个或多个 block,而持久化是在每个 SM 上仅启动有限数量的簇(例如总 SM 数除以簇大小),然后让这些簇循环处理多个输出 tile。这种方式减少了块初始化开销,并能通过把 epilogue 的 SMEM->GMEM 拷贝与下一轮计算重叠来提升吞吐量。
实现时可以使用 plgpu.nd_loop 来把整个 (m_iters, n_iters) 迭代空间分配给集群网格(collective_axes),并通过 loop_info 提供的本地索引判断当前是第几轮,从而在第一次迭代避免等待旧数据,但在后续迭代中正确等待 consumed_barriers。别忘了在每轮结束前用 plgpu.wait_load_tmem 等待 TMEM 的异步加载完成,以免被后续的 MMA 覆盖。持久化内核在中等与大规模矩阵上通常能显著提升整体吞吐。\n\n为了让计算与结果写回并行进行,采用专门的 epilogue warpgroup(或称专用于写回的 Pallas 线程)能进一步提升并发度。把内核设置为拥有两个 Pallas 线程(num_threads=2),并用第一个线程负责加载与 MMA 发起,第二个线程负责 TMEM->SMEM->GMEM 的写回,可以在保持资源隔离的同时避免互相阻塞。该方案要求把 TMEM 做双缓冲:acc_tmem 的尺寸要包含两个槽,每次迭代切换使用哪一半。
相应地,mma_done_barrier 与 store_done_barrier 等都需要有两个槽位,与 acc_slot 对应,保障不同槽位的独立同步。实现细节包括在计算侧在适当时机等待 store_done_barrier,以确保 TMEM 的目标槽已经被写回并可复用;写回侧在完成写回后要发出 store_done_barrier 的 arrive。注意虽然只有 leader block 发起 MMA,但所有参与集群的块仍然应在必要时等待 store_done_barrier,避免重复 arrive 导致硬件异常。专门的写回线程能进一步提升 TensorCore 的连续性,减少计算单元空闲。\n\n最后一个高效技巧是网格分块(grid tiling,又称 rasterization order 或 program re-ordering)。访问顺序对 L2 缓存重用有巨大影响。
简单地按行或列顺序遍历输出 tile 会使某一侧的输入被频繁反复加载,降低 L2 利用率。通过将遍历顺序改为"平面蛇形"(planar snake)或切片化的小 tile 扫描,并让输出的次级维度在一定 tile 宽度内更快地变化,可以提高对某一操作数的局部性。CUTLASS、Triton 等项目都采用类似策略。Pallas 提供 plgpu.planar_snake 来辅助生成这种访问模式。合理选择 grid_minor_dim 与 grid_tile_width 可以在不改动算子核心的情况下显著提高 L2 命中率,从而进一步提升整体性能。实测表明,在多个优化步骤完成后,启用网格分块往往能把性能从与 cuBLAS 持平推升至超越的水平。
\n\n实践中的调优建议与常见陷阱值得特别强调。首先,选择 tile_k 的经验规则是 tile_k 接近 128 除以输入类型字节宽度(例如 float16 时约为 8),这样能匹配 MMA 的内部分组。max_concurrent_steps(流水线深度)决定了预取与计算重叠的能力,通常通过实验选取 2、3、4 等小整数。epilogue_tile_n 决定了 TMEM->SMEM 与 SMEM->GMEM 的分块粒度,太大减少重叠,太小增加管理开销。SMEM 中采用的 swizzle 与 tiling transforms 对避免 bank 冲突与匹配 MMA 期望布局至关重要。barrier 的 num_arrivals、orders_tensor_core 等参数必须正确配置,否则会出现同步错误或性能不可预测。
集体复制的 partitioned_axis 要与集群布局一致,否则会导致数据错位或性能大幅下降。最后,基准测试时务必说明输入数据的分布;不同数据分布(例如 iid 正态与均匀或稀疏矩阵)会对性能产生显著影响。团队内部测得的基线采用 iid normal float16 输入,某些分布会更快或更慢。\n\n结合上述一系列优化步骤,可以把最初的基础内核性能逐步推向极限。以作者的测试为参考,基础内核的 TensorCore 利用率约为 37.6%,引入 warp 专用化后提升到约 45.5%,加入 tiled epilogue 提升到约 55.8%,采用集体 MMA 后升至约 59.4%,引入持久化与专用 epilogue warpgroup 能分别提升至约 61.5% 和 63.4%,最终通过网格分块将吞吐提升到约 69.4%,超过了同场景下 CUTLASS 的结果。真实工程中,具体百分比会随硬件配置、矩阵大小和数据分布而变化,但整体趋势表明逐步分离责任、提升算密度、并行化写回以及优化遍历顺序是提高 Blackwell 上矩阵乘法性能的核心途径。
\n\n总结而言,在 Blackwell 上为 TensorCore 写出高性能矩阵乘法内核不是单一技巧可以完成的任务,而是多个协同优化的集合。设计合理的 tiling 策略、使用 SMEM/TMEM 双缓冲与合适的 transforms、利用 warp 专用化以减少控制开销、采用集体 MMA 来提高算密度、用持久化减少启动开销并重叠操作、用专门的写回线程提升并发,以及通过网格分块优化 L2 重用,所有这些技术共同作用,才能实现接近或超过成熟库的性能。掌握这些模式并在特定应用场景中进行系统化调优,将帮助工程团队在 Blackwell 平台上获得优异的线性代数性能。若希望快速上手,建议从最小可运行的单 CTA 实现开始,逐步引入 warp 专用化、tiled epilogue 与集体 MMA,并在每一步进行基准测量以量化收益;同时务必利用 Pallas 提供的工具(如 plgpu.find_swizzle、WarpMesh、nd_loop、planar_snake 等)来简化实现并保证与底层硬件约定一致。 。
