GB300 GEMM优化实战：从能跑到逼近cuBLAS的关键技术

这篇本来是想给自己团队的GPU算子新人写的一份内部总结后来发现内容越写越多干脆整理成一篇完整的技术文章发出来。目标是解决一个非常具体的问题在GB300Blackwell Ultra平台上把GEMM kernel的性能从“能跑”优化到接近cuBLAS的水平。GEMM优化是HPC和AI框架工程师绕不过去的一道坎。你在网上搜“GEMM优化”能翻到一堆讲tiling、register blocking的老文章但那些大多还停留在Ampere甚至Volta时代。GB300带来的新硬件特性——第5代Tensor Core、异步矩阵指令、tensor memory——让旧的优化套路部分失效也让新的优化空间被打开。如果你刚接触GPU优化这篇文章可以当一份进阶路线图看如果你已经在写kernel了那里面提到的坑和排查思路应该能帮你省几个通宵。直接说结论想稳定地全面超过cuBLAS说实话很难但把性能调到cuBLAS的85%~95%让业务场景跑到可接受的水平这是一条非常清晰且可复现的路径。下面从平台特性开始一步步拆解。1. 先搞清楚对手是谁GB300的硬件底牌在动手写代码之前我建议你先花半小时把目标平台的硬件规格吃透。优化这件事本质上是在硬件能力边界内做调度。你不清楚边界在哪就不知道kernel离天花板还有多远也不知道瓶颈到底在计算、访存还是同步。1.1 Blackwell Ultra平台特性速览GB300严格来说是NVIDIA Blackwell Ultra架构下的加速系统由Grace CPU和Blackwell Ultra GPU组合而成。我们聊GEMM优化时关注的核心是GPU那一侧。以公开规格来看这代GPU有几个关键数字值得记住。算力层面BF16/FP16的Tensor Core算力已经在PFLOPs量级FP8还能再翻一档。显存容量达到288GB内存带宽在8TB/s级别。对GEMM来说这意味着纯计算能力非常充裕真正的稀缺资源是数据搬运通道——HBM带宽、L2带宽、shared memory带宽每一条都可能是瓶颈。片上资源也有明显变化。SM数量比Hopper世代更多每个SM的shared memory容量提升到228KB。第5代Tensor Core引入了新的异步执行模型包括tcgen05指令族和tensor memorytmem。这些不是小修小补而是让GEMM kernel的编写方式发生了结构性变化。如果你要从这套硬件上榨出性能第一件事就是把“计算密集”这四个字从脑子里划掉。现代GEMM在大多数shape下根本不是计算瓶颈而是数据供给瓶颈。一个kernel如果没跑到90%以上的算力利用率大概率是A/B矩阵的搬运路径堵住了而不是MMA单元在偷懒。1.2 cuBLAS到底强在哪cuBLAS是NVIDIA官方的闭源BLAS库。很多人一提优化就把cuBLAS当成假想敌但真正要打败它你得先理解它为什么强。第一cuBLAS不是一套kernel打天下。它在每个新架构上都会重写核心实现针对不同shape准备了多种kernel配置。你调用一次cublasGemmEx底层launcher会根据M、N、K的具体数值、矩阵布局、数据类型选择一套最合适的实现。同一份API调用跑大矩阵和小矩阵底层可能是完全不同的kernel。第二它有autotuning机制。在首次调用某个shape时它会在内部做benchmark测试几组候选配置挑最快的那个跑。这个机制保证了它在各种shape上都不会太差。你很难用一个固定配置在全部shape上赢过它原因就在这里。第三它对硬件特性的利用非常激进。从Hopper时代的wgmmaTMA到Blackwell的tcgen05tmemcuBLAS始终是第一批用上最新指令的库。它代表的是NVIDIA内部工程师对自家硬件的理解深度所以它是一个非常优秀的基准参照。但这不代表cuBLAS就是神。它要兼顾海量shape必然会在特定shape上做出妥协。理解这一点你就知道“在某些定制场景下超越cuBLAS”并非天方夜谭后面第7章我会展开聊。2. 从朴素GEMM到可用的Tiled Kernel先解决数据复用大部分人学GEMM优化是从一个三重循环开始的。但那个版本在GB300上跑出来的结果会惨不忍睹因为它几乎没有数据复用。我先把朴素版本为什么慢拆讲清楚再给出一个有实用价值的Tiled版本。2.1 朴素实现的性能瓶颈在哪最简单的GEMM实现长这样对C矩阵的每个元素遍历K维度做乘加。问题在于内存里A的一行和B的一列都是按固定地址存放的每次计算一个输出元素都要重新从全局内存里读一遍A的对应行和B的对应列。矩阵乘法本来有极好的数据复用潜力但朴素实现把这种潜力完全浪费了。量化感受一下C矩阵有M乘N个元素每个元素要做K次乘加所以全局内存访问次数是M乘N乘K量级。但A和B的实际数据量A是M乘KB是K乘N。也就是说全局内存访问量是数据量的好几倍算力根本喂不饱。这个版本在GB300上运行时实际性能大概率只有几十GFLOPS连HBM带宽都摸不到上限。解决思路也很明确把数据切块tiling让一个数据块在shared memory里被多次复用。A的一个tile加载到shared memory后可以被多个输出tile使用B的tile同理。这样全局内存的访问次数就能从M乘N乘K降到M乘K加K乘N再乘一个块因子数据复用率直接提升几个数量级。2.2 一个可运行的Tiled Kernel示例写一个最基础的Tiled版本目标是把概念讲清楚代码能编译能跑。这个版本我故意没用任何异步指令也没有做复杂swizzle就是为了让结构清晰。#define TILE_M 64 #define TILE_N 64 #define TILE_K 16 #define THREADS 256 __global__ void gemm_tiled(const float* A, const float* B, float* C, int M, int N, int K) { __shared__ float As[TILE_M * TILE_K]; __shared__ float Bs[TILE_K * TILE_N]; int blockRow blockIdx.y * TILE_M; int blockCol blockIdx.x * TILE_N; int tx threadIdx.x; float acc[4][4] {0.f}; for (int k0 0; k0 K; k0 TILE_K) { // 协作加载A的分块到shared memory for (int idx tx; idx TILE_M * TILE_K; idx THREADS) { int r idx / TILE_K; int c idx % TILE_K; As[idx] A[(blockRow r) * K k0 c]; } // 协作加载B的分块到shared memory for (int idx tx; idx TILE_K * TILE_N; idx THREADS) { int r idx / TILE_N; int c idx % TILE_N; Bs[idx] B[(k0 r) * N blockCol c]; } __syncthreads(); // 每个线程负责一个4x4的输出tile int threadRow tx / 16; int threadCol tx % 16; for (int kk 0; kk TILE_K; kk) { for (int i 0; i 4; i) { for (int j 0; j 4; j) { int aRow threadRow * 4 i; int bCol threadCol * 4 j; acc[i][j] As[aRow * TILE_K kk] * Bs[kk * TILE_N bCol]; } } } __syncthreads(); } // 写回全局内存 for (int i 0; i 4; i) { for (int j 0; j 4; j) { int r blockRow (tx / 16) * 4 i; int c blockCol (tx % 16) * 4 j; C[r * N c] acc[i][j]; } } }跑一下对比朴素版本你会发现性能提升非常明显。这里A和B的每个数据元素在shared memory里被读TILE_M或TILE_N次数据复用从0变成了几十倍。但我也要提醒这个版本放到GB300上看性能离cuBLAS还差着两个数量级。2.3 Register Blocking的真正意义为什么要让每个线程算4x4而不是1x1这是理解GEMM优化很关键的一步。第一个原因是减少shared memory的load次数。如果不做register blocking每个乘加操作都要从shared memory读A的一个数和B的一个数。做了4x4 blocking之后每个As元素可以跟4个Bs元素复用每个Bs元素也可以跟4个As元素复用。同样是16个乘加shared memory的load数量从32次降到了8次。这个倍率随着register tile增大而提升但代价是寄存器占用增加。第二个原因是提升指令级并行ILP。多个累加器之间没有数据依赖编译器可以把循环展开让多个乘加指令在流水线上并行执行把单条FMA的延迟掩盖住。如果你只算一个输出那整个循环就是一条又一条有依赖的FMA延迟完全暴露。我自己在调优时经常算一笔账每个SM有多少寄存器、每个线程能分多少、register tile每增大一档SMEM load能省多少、但occupancy会掉多少。这些都是此消彼长的关系没有什么绝对最优配置只有针对具体shape的最优配置。3. 流水线改造让内存搬运和计算真正重叠Tiled kernel已经能跑了但你会注意到它的整体节奏是搬数据等所有人搬完计算等所有人算完再搬下一块。同步等待的时间占了很大比重。GEMM想跑得快核心就是在“搬下一块数据”和“算当前这块数据”之间做overlap。3.1 为什么要做多级流水线在上一节的kernel里每个K维度tile的流程是先从HBM加载数据到SMEM然后__syncthreads()等待全部线程加载完成接着所有线程做计算再同步进入下一轮。这两个__syncthreads()意味着加载数据期间SM的计算单元是空闲的。HBM的访问延迟大概是几百个cycle而一个MMA指令的延迟也有几十个cycle。如果让“加载第i1块”和“计算第i块”同时进行kernel的总耗时就从“加载总时间计算总时间”变成“加载总时间计算总时间”中较大的那个。这个重叠对GEMM性能的影响是决定性的。具体做法是分配双份甚至多份SMEM buffer。第i轮计算用buffer 0第i1轮的数据预加载到buffer 1。在Hopper和Blackwell平台上这一步主要靠异步拷贝指令来完成。它的本质是把数据搬运从计算指令中解耦出来让SM在等待数据的时候不用干等。流水线级数也不是越多越好。每多一级SMEM消耗就多一份buffer。一个128x128x32的FP16 tile大约16KBGB300的SMEM有228KB放个三四级完全够。但级数越多同步和依赖管理越复杂收益会边际递减。实际调优中我常用的是3级或4级流水线。3.2 cp.async与L2策略怎么配合Ampere时代之后CUDA提供了cp.async指令它能把全局内存的数据直接拷到shared memory不经过寄存器。你的代码只需要发出一系列cp.async指令然后调用cp.async.commit_group和cp.async.wait_group来管理批次。这些指令由异步单元执行计算单元可以继续做上一轮的MM运算。在Blackwell上还有更高层的TMATensor Memory Accelerator可以干这个活它甚至能从global memory直接搬运到shared memory并且支持多维tile描述。TMA的好处是它只需要一个线程发起就能搬运一大块数据基本不占SM的运算资源。另一个容易被忽略的点是L2缓存策略。你可以用CUDA的access policy window给某个访存区域打标签比如cudaAccessPropertyPersisting让频繁复用的数据尽量留在L2里。GEMM里有些数据会在多个block之间共享访问合理利用L2策略能减少HBM访问。但注意L2策略用错了也会帮倒忙比如给一次性数据标成persisting反而把热数据挤出了L2。3.3 访存模式与Bank冲突处理shared memory的带宽不是无限的。它内部被分成32个bank每个bank每周期只能响应一次访问。如果一个warp里有多个线程同时访问同一个bank的不同地址硬件就要把这次访问拆成多次串行执行这就是bank conflict。GEMM里最常见的冲突来源是行主序的SMEM数组。假设As[TILE_M][TILE_K]按行连续存储线程threadRow访问As[threadRow][kk]那么同一warp里的线程如果threadRow恰好间隔某个值它们就会落到同一个bank产生冲突。解决办法是swizzle——在数据写入SMEM时做一次地址变换让原来会冲突的地址错开。我常用的swizzle方法是XOR模式。具体来说对地址的低bit做异或变换让相邻行错开不同的bank偏移。比如把一个64行宽的数据在写入时把row的高位异或到col的低位读取时同一轮循环里不同线程访问的地址就会均匀分布在32个bank上。这个变换看起来只是个位运算但性能差距可能达到20%到30%。提示不要一上来就做复杂的swizzle。先用Nsight Compute看有没有bank conflict告警再决定是不是要动地址映射。没有冲突的kernel加swizzle反而可能因为地址计算复杂而变慢。4. Blackwell时代的核心变量tcgen05与tmem如果你在GB300上直接套用Hopper的优化经验性能大概率达不到预期。Blackwell的Tensor Core执行模型跟以前完全不一样了最大的区别就是tcgen05指令和tmem。4.1 从Ampere到Blackwell的指令演进Ampere时代的mma.sync指令是warp内所有线程同步执行一次矩阵乘加操作数放在寄存器里累加结果也回到寄存器。编程思路很直接准备好寄存器里的A片段和B片段执行mma等结果。但这种同步等待让指令级并行的空间有限。Hopper引入了wgmma和TMA把矩阵运算的单位从warp扩大到warpgroup并且支持异步执行。数据可以提前放到shared memorywgmma启动后不等结果你可以继续做别的计算。这已经是很大的进步但操作数和结果仍然以寄存器或shared memory为主。Blackwell的tcgen05把异步执行推到了一个新高度。它的矩阵乘加指令本身是一种“发起后完全异步”的操作数据源可以来自shared memory累加结果写入tensor memorytmem。tcgen05一旦发出SM的计算管线就会被这个矩阵运算占用但warp不会阻塞等它完成还可以继续做地址计算、数据搬移、或者其他非矩阵类指令。这个变化直接改变了GEMM kernel的编写方式。以前那种“算完一个tilesync拿结果再算下一个”的节奏已经不合适了。你需要把整个kernel编排成把tile铺好、连续发起多个tcgen05、最后统一回收结果。4.2 tmem的正确使用姿势tmem是Blackwell架构里新增的一块片上存储每个SM有256KB。它不是用来替代寄存器或shared memory的而是专门给Tensor Core的结果当累加器用的。它的访问方式和寄存器不同需要通过tcgen05.ld或tcgen05.st这类指令把数据在tmem和寄存器之间搬动。为什么需要这么一块专用存储因为GEMM的累加器太大了。一个128x128的FP32累加结果每个SM上同时需要存放的中间状态非常可观。如果全部放寄存器会严重挤压寄存器的可用数量限制occupancy。把累加器挪到tmem寄存器就被释放出来可以给数据搬移和地址计算用。典型的使用流程长这样示意伪码实际一般用CUTLASS封装1. TMA把一个A tile和一个B tile搬到shared memory 2. tcgen05.mma发起矩阵乘input指向SMEMaccumulator指向tmem 3. 不等结果继续搬运下一组tile 4. K循环结束后用tcgen05.ld把tmem里的结果读回寄存器 5. 在寄存器里做epilogue如加bias、激活函数写回global memory这种模式要求你对tmem的布局有概念。tmem是按行列组织的不同warp只能访问各自对应的分区。写入和读取时的行列映射关系如果不对结果就是乱序的。这块也是Blackwell上最容易写错的地方之一。4.3 编程模型变化对开发者的影响说实话直接手写tcgen05的PTX指令是非常痛苦的。它不像Ampere时代的mma那样有清晰的指令格式和文档案例。目前更现实的路径是使用CUTLASS 3.x。CUTLASS已经把tcgen05封装成了collective层你只需要配置几个模板参数它就能生成对应的PTX。但这不代表你可以完全不懂原理。CUTLASS里那些模板参数比如MmaShape、PipelineStages、SmemSwizzle每一个都对应着你在手写时会遇到的实际调优点。你不理解tmem是什么就不知道为什么CUTLASS的某些collective配置在特定shape上会快换一个shape就慢成狗。这也就是为什么很多团队拿着CUTLASS却调不出性能因为他们把CUTLASS当黑盒用了。5. 实战对标cuBLAS一套可复用的优化流程前面讲了很多原理这里给一条可以直接上手的路径。我自己调GEMM时的流程已经比较固定了花不了太多时间就能定位到关键瓶颈。5.1 从CUTLASS开始而不是从零手写除非你纯粹为了学习否则我不建议在Blackwell上从零手写tcgen05 kernel。CUTLASS 3.x提供了DeviceGemm的完整实现覆盖了FP16/BF16/FP8等多种数据类型还提供了CollectiveBuilder可以让你像组合积木一样配置kernel。实际做法是找到CUTLASS里的examples/47_blackwell_gemm或者类似的示例把自己的M、N、K和数据布局换进去先跑通。然后用不同的tile大小、流水线级数、swizzle模式去试找到一组性能最好的配置。这个阶段通常能把性能做到cuBLAS的80%~90%。CUTLASS是NVIDIA官方维护的库它对Blackwell的支持非常及时很多优化思路直接内置在代码里。你要做的事情就是做配置搜索然后记录不同配置的耗时。配置搜索这一步千万别手动一个个测。写一个简单的脚本遍历tile大小、stage数、swizzle模式这几个维度自动跑benchmark把结果存下来对比。我见过太多人手动试了一两个配置就说CUTLASS不行其实只是没找到合适的参数组合。5.2 用Nsight Compute判断瓶颈跑完性能测试后用Nsight Compute分析kernel。重点看Speed of Light页面里的三个指标ComputeSM忙碌占比、MemoryHBM吞吐、L2/Shared片上带宽利用。GEMM的理想状态是Compute占用最高Memory次之L2和Shared不成为瓶颈。如果看到Stall Long Scoreboard这类原因占比很高说明线程在等着数据从shared memory或global memory回来这是流水线没有填满的典型信号。如果看到bank conflict警告去调swizzle。如果看到Executed Ipc Active很低可能是寄存器占用过高导致occupancy不足。还有一个容易被忽略的点对比cuBLAS时要用Nsight Compute分别分析cuBLAS kernel和自己kernel的SASS指令。看cuBLAS用了什么指令你就能知道自己的实现缺了什么。比如你发现cuBLAS的kernel里有大量TMA相关的指令而你的kernel还在用普通的LDG加载那差距就一目了然了。5.3 性能差距的合理预期给一个现实的预期管理。cuBLAS在FP16大矩阵GEMM上通常能跑到理论峰值的85%~95%。你的kernel如果经过了合理的tiling、流水线化、swizzle调整达到cuBLAS的85%不是难事90%也很常见。能不能超过cuBLAS我的经验是在做足了针对特定shape的定制后小幅超过是可能的。比如某个固定的M、N、K组合你的kernel可以比cuBLAS快5%到10%。原因是cuBLAS为了通用性launcher本身有判断开销而且它不一定为你这个特定shape选到了最合适的配置。但要记住一个原则不要为benchmark而优化。如果你的业务shape本身就不固定追求单一shape的极致性能没有意义。把时间花在配置搜索和流水线优化上比死磕某几条PTX指令更划算。6. 常见问题与避坑速查表写kernel这几年我总结了一些高频问题列在这里给后来人排雷。6.1 五个最典型的性能陷阱第一个是忽略bank conflict。很多人把tiling做完就收工了结果Nsight Compute一开SMEM有30%的bank conflict白白损失了带宽。这个问题最简单改一下swizzle就能解决。第二个是流水线深度不够。如果你还停留在单缓冲、每次循环都要同步等待的阶段性能天花板就已经锁死了。GB300上的GEMM3级流水线几乎是起步配置。第三个是寄存器占用失控。FP32累加器如果开得过大比如每个线程算8x8甚至更大寄存器可能直接爆到200多个occupancy降到很低。算力和occupancy要找平衡不是寄存器块越大越好。第四个是L2策略误用。给不该驻留的数据打了persisting标签反而把热数据挤掉了。L2策略要用在真正会跨block复用的数据上还要配合每次访问的数据量大小来调整。第五个是过早优化。没先验证正确性就开始调性能结果调了半天发现输出是错的浪费时间。6.2 正确性验证与精度问题性能优化的前提是正确性。不要只测一个shape就收工。建议写个测试脚本覆盖小矩阵1x1、2x3这种边界、方阵1024x1024、非对齐shape如M1000,N2000,K500这种跟你确认过的参考实现做对比。精度上要留好余量。FP16/BF16的GEMM累加顺序不同结果会有一些细微差异。你在不同的kernel之间做性能对比时要设置一个合理的误差容忍度比如相对误差1e-2级别不要用绝对相等去判断。还有一个坑是如果你的kernel里做了split-K多个block的部分和做reduction的顺序会影响最终数值误差会比单block大一些这是正常的。注意在对比cuBLAS时cuBLAS可能对不同算法选了不同的累加顺序你和它结果有小差异不代表你错了。判断正确性应该以高精度参考实现为准而不是以cuBLAS为准。6.3 问题速查表下面这个表是我在实际调优过程中最常用的排查清单遇到问题先对照这里能少走很多弯路。现象可能原因排查手段性能远低于预期全局内存访问未合并bank conflict严重打开Nsight Compute的Memory Workload分析检查访存模式SM占用率高但算力利用率低流水线未重叠线程在等待数据查看Stall原因如果是Long Scoreboard加深流水线换了shape性能大幅波动tile大小与shape不匹配做配置搜索针对shape范围找最优配置结果与参考实现不一致SMEM加载顺序或同步错误先跑小shape单步调试检查边界索引输出正确但比cuBLAS慢很多未使用TMA/tcgen05等新指令对比SASS指令差异确认是否用了Blackwell的异步特性寄存器占用过高register tile过大适当缩小寄存器tile或用tmem分担累加器这张表不是万能药但它覆盖了我经历过的大部分坑。遇到没列出来的问题建议抓一个kernel的完整profile逐项看SOL指标通常能找到线索。6.4 关于“热词”里那些不相干优化顺带说一句我在搜资料的时候发现网上关于“优化”的热词特别杂什么慢SQL优化、电脑优化指令、系统优化之类的跟GEMM半毛钱关系都没有。这说明“优化”这个词在不同语境下语义差异极大。在GEMM这个领域“优化”指的是在硬件资源约束下重新组织指令和数据流动把理论算力转化为实际吞吐。它不是调参数更不是清理内存而是对微架构的精确控制。理解这一点你才能分清哪些经验可以直接吸收哪些完全是另一条赛道。7. 想超越cuBLAS可以但要说清场景最后聊聊“超越cuBLAS”这件事。我的态度是纯GEMM在通用场景下想全面超越cuBLAS难度极高性价比很低。但在特定场景下超越不是不可能的而且路径不止一条。7.1 融合算子是更现实的超越路径cuBLAS做的是“纯GEMM”它不会帮你做后续的激活函数、残差连接、量化反量化。但在真实模型里GEMM的结果几乎总是要跟着下一堆操作。如果你把这些操作融合进自己的kernel的epilogue阶段而不是让数据从显存里跑一圈再回来整体延迟能省不少。这种赢法不是靠GEMM本身算得多快而是靠省掉了数据往返的开销。举个例子在做attention推理时QK^T的结果如果直接跟softmax融合能省掉整块中间矩阵的写回和读取。这种融合在cuBLAS的框架里做不到因为它们提供的是最基础的矩阵乘原语不负责业务逻辑。再比如量化场景。你可以在GEMM的epilogue里直接做FP32累加结果到INT8的转换加缩放省掉一个单独的quantize kernel。这种融合对端到端性能的提升非常可观尤其是小shape推理场景。7.2 什么时候不要手写GEMM手写GEMM的成本比很多人想象得高。你需要处理数据布局、同步、指令选择、swizzle、L2策略、TMA descriptor还要花时间做配置搜索和精度验证。如果你的场景是训练大模型直接上cuBLAS或CUTLASS就好框架层面已经集成好了没必要重复造轮子。如果你是在做推理引擎需要针对固定shape和量化格式做极致优化那值得手写一个定制kernel。但我的建议是先基于CUTLASS做修改而不是从零开始。CUTLASS的架构把整个GEMM拆成了很多可替换的部件你只需要替换其中一部分就能实现定制化。如果你是纯粹想学习那我强烈建议你完整地手写一遍。从朴素版本一路优化到接近cuBLAS的水平这个过程对理解GPU微架构的帮助是任何文档都替代不了的。我自己就是写了几版之后才真正理解为什么有时一个看起来应该更快的配置实测却慢得离谱。最后想分享一点个人体会。GEMM优化不是一锤子买卖它需要根据硬件、shape、业务需求反复迭代。今天这个配置最快明天换了batch size可能就慢成渣。所以与其追求一套万能的kernel不如建立起一套能快速测试、快速定位瓶颈的流程。拿CUTLASS做起点用Nsight Compute做眼睛用cuBLAS做尺子这三件套齐了你在GB300上写出接近cuBLAS水平的GEMM只是时间问题。