拓冰网站定制 
先交代一个很现实的场景你花了一周时间照着 CUDA 文档写了一个 GEMM kernelnvcc编译通过结果一跑Nsight Compute 里Achieved Occupancy还行SM Busy也说得过去但算力利用率只有 1%——你的 GPU 号称几十 TFLOPS实际只吐出来零点几。这不是你笨是因为 GEMM 优化这条路上从“能算”到“算得快”中间隔着一整套关于 Tiling、数据复用、指令选择和硬件单元的问题。这一讲就是专门把这个过程拆开从 1% 到 95% 的算力利用率到底是怎么一步步抠出来的每一层优化在解决什么问题Tensor Core 的编程思维和普通 CUDA Core 有什么区别以及你在实际操作中会遇到哪些坑。这讲适合两类人看。第一类是已经写过 CUDA kernel、但性能总是不上不下的开发者你需要理解为什么共享内存分块能带来几十倍提升为什么 Tensor Core 不是简单的“换一条指令”就完事。第二类是读 CUTLASS 源码、看各种 GEMM 优化文章时一头雾水的人——无数文章提到cp.async、swizzle、wgmma但没人告诉你这些到底在解决什么底层痛点。如果你是这两类人之一下面这五千字应该能帮你在脑海里建立起一张完整的 GEMM 优化地图。1. 为什么常规 GEMM 只能跑出 1% 算力先算清这笔账1.1 算子性能的三座大山访存、并行粒度、指令效率先搞清楚一个问题GPU 的核心瓶颈到底在哪里。你可以把一次 GEMM 运算想象成一个仓库配货场景。显存是仓库SM流式多处理器是分拣车间寄存器和共享内存是车间里的操作台。CUDA Core 和 Tensor Core 是操作台上的工人。理论上车间的加工能力算力是固定的但如果你每次加工一个小零件都要跑到仓库里拿原料那时间全花在路上了。朴素 GEMM 就是这个疯狂跑仓库的场景。具体点说一个最简单的 GEMM kernel 长这样__global__ void sgemm_naive(const float* A, const float* B, float* C, int M, int N, int K) { int row blockIdx.y * blockDim.y threadIdx.y; int col blockIdx.x * blockDim.x threadIdx.x; float sum 0.0f; for (int k 0; k K; k) sum A[row * K k] * B[k * N col]; C[row * N col] sum; }这段代码问题很多但最致命的有三条。第一全局内存反复访问。每个线程要计算 C[row][col] 这个点就必须把 A 的第 row 行整行和 B 的第 col 列整列全部读一遍。如果 MNK4096整个 kernel 对全局内存的读取量是 8GB 量级而数据实际只有 64MB。也就是说99% 的访存都是在重复搬运同一个数据。第二并行粒度不是硬件喜欢的形状。每个线程只是算一个点完全没有数据和指令的复用性。第三没有向量化访存。一条 32-bit 的 load 指令一次只搬 4 字节L1/L2 带宽利用率极低。有人可能会说L2 cache 不是能缓解吗没错L2 能帮你挡掉一部分重复访问但 GEMM 的数据访问模式是“行”和“列”交叉的多个线程同时访问同一行、不同列的时候L2 缓存行还勉强能命中可一旦访问模式变成跨行跨列的“跳跃式”缓存行命中率立刻崩掉。所以在 GPU 上靠 cache 兜底不是正路正路是让每个线程自己把数据搬到片上主动制造数据复用。1.2 运算强度判断 GEMM 该不该被算力吃满这里要引出一个关键概念运算强度Arithmetic Intensity也就是一次访存对应的浮点运算次数。GEMM 的理论运算强度很好算——总 FLOPs 是 2×M×N×K数据量是 M×K K×N M×N。当矩阵规模足够大时数据量远小于 FLOPs运算强度可以做到几百 FLOPs/Byte。GPU 的 Roofline 模型告诉我们只有当运算强度超过某个临界点通常是算力除以带宽A100 级别大约是几十 FLOPs/Byte你才能摸到算力天花板。朴素 GEMM 的“实际运算强度”远远达不到这个临界值——因为它每次计算都要重新访存实际的有效运算强度可能只有几个 FLOPs/Byte。所以算力利用率 1% 不是 GPU 的错是你的数据流设计把整个 kernel 压在了访存瓶颈上。反过来这也说明了一个核心结论GEMM 的性能问题本质上是一个数据复用问题。你要做的不是减少计算量计算量固定而是在数据被消费之前尽量让它在片上多待几个时钟周期。这就是 Tiling 存在的全部理由。1.3 “九重天”对应的优化阶段对照下面这张表是我按经验整理的一个 GEMM 优化路线图。不同架构、不同矩阵规模下具体数值会有浮动但大致趋势是稳定的每一行都代表一个独立的优化维度层级关键动作大致算力利用率这一步在解决什么Level 0朴素三重循环约 1%没有任何数据复用访存完全失控Level 1循环重排 寄存器缓存3% - 5%让某些数据在寄存器中被多次使用Level 2向量化访存float48% - 12%降低指令数提高访存带宽利用率Level 3共享内存分块25% - 40%数据从 global 搬到 smem片内复用Level 4消除 Bank Conflict45% - 55%共享内存本身的带宽瓶颈被打开Level 5Double Buffering cp.async60% - 70%用流水线掩盖访存延迟Level 6Tensor CoreFP16/TF3275% - 85%从 CUDA Core 换到专用矩阵单元Level 7Shape 对齐 精度选型88% - 92%让硬件矩阵单元没有一丝浪费Level 8TMA 全局调度优化90% - 95%搬数据这件事交给硬件线程纯算这九级不是必须一步一步走走得顺的话 Level 3 和 Level 4 经常一起做Level 6 和 Level 7 往往绑在一起。但每一级背后的物理约束不同你只有理解了每一层到底在卡什么才能在 Nsight 里看到瓶颈的时候知道该动哪一块代码。2. Tiling 分块把矩阵切成 GPU 能“嚼得动”的大小2.1 为什么需要两级搬运从全局内存到共享内存再到寄存器Tiling 的核心思想用一个做饭的类比就能讲清楚。你要做 100 道菜每道菜需要葱姜蒜。最笨的办法是做一道菜就去冰箱拿一次配料这就是朴素 GEMM。聪明的做法是一次性把一大袋配料全拿到厨房台面上共享内存做菜的时候从台面取再想快一点就把最常用的几样放在手边的小碗里寄存器手一伸就够到。这就是两级分块全局内存 → 共享内存 → 寄存器。为什么不能直接从全局内存到寄存器因为寄存器容量太小一个 SM 的寄存器文件总共才 256KB一个线程最多 255 个寄存器装不下一个能支撑复用的大块。而共享内存每个 SM 有 64KB 到 228KB 不等老的 Volta 是 96KBAmpere 上是 164KBHopper 是 228KB能装下比寄存器大得多的数据块。所以实际最优策略是先用共享内存装一个大 tile再由线程从共享内存里把子块读进寄存器。这个两级结构还对应一个更底层的硬件事实全局内存的延迟是几百个周期共享内存的延迟只有二三十个周期寄存器几乎是零延迟。你的优化目标是把绝大多数访问落在延迟最低的那一级。2.2 Block Tile、Warp Tile、Thread Tile三层名字背后的分工Tiling 不是简单地“把矩阵分成小块”而是分成三个层次对应三种不同的执行单元。Block Tile是整个线程块负责的那一块矩阵比如 128×128。它对应的是共享内存的容量边界。Block Tile 越大数据复用率越高——因为同一个 tile 里的数据会被块内所有线程反复使用。但太大也不行共享内存放不下两张 tile或者放得下却导致占用率暴跌。Warp Tile是线程块内部每个 warp 负责的子块。一个 128×128 的 Block Tile 切成 4 个 warp 来算每个 warp 可能负责一个 64×64 的区域。Warp Tile 是 SIMT 执行的基本单位——一个 warp 里的 32 个线程协同计算如果它们各自算一个独立的小 tile访存会乱但如果它们合作算一个 warp tile 并共享寄存器中的中间结果效率会高很多。Thread Tile是单个线程最终负责的那一小块输出常见的是 4×4、8×8。为什么要让一个线程算 8×8 而不是 1×1因为这样每个从共享内存读进来的 A 和 B 的元素可以被 8 次乘加复用这就是寄存器级的 Tiling。没有这层复用共享内存的带宽又会变成新瓶颈。2.3 经典 128×128/8×8 组合为什么最常见如果你去看 CUTLASS 和各种高性能 GEMM会发现 128×128 的 Block Tile、8×8 的 Thread Tile 非常常见。这是巧合吗不是。算一笔账。Block Tile 选 128×128意味着每个 block 要处理 128×128 的输出。如果用 8×8 的 Thread Tile每个 block 需要 16×16256 个线程来覆盖整个 tile每个线程算 8×8。而 blockDim 是 256 个线程正好一个 block 可以装 8 个 warp。再看共享内存A 的 tile 是 128×BKB 的 tile 是 BK×128。如果 BK8两个 tile 各占 128×8×44KB总共 8KB这对于 164KB 的共享内存来说只占很小一部分。但这只是“能装下”不是“最优”。真正的权衡在于占用率。一个 SM 上到底能跑多少个 block取决于共享内存和寄存器哪个先耗尽。假设一个 block 用 8KB 共享内存、256 个线程每个线程用 128 个寄存器那一个 SM 的 64K 寄存器按 Ampere 算只能容纳 2 个 block共享内存倒是不紧张。这时候 Blocks per SM 就只有 2占用率大概 50%。如果你把 Block Tile 缩到 64×64每个 block 的线程数降到 64SM 上能塞更多 block但共享内存 tile 变小数据复用率下降。所以 128×128 是一个考虑了共享内存容量、寄存器数目和复用率之后的平衡点不是拍脑袋定的。我自己实测的体会是在小矩阵比如 MNK256上用大 Block Tile 反而容易因为边界 padding 浪费计算在 2048 以上的大矩阵上128×128 几乎是稳定好用的万能选择。3. Tensor Core 思维从“线程算一个数”到“Warp 算一块矩阵”3.1 Tensor Core 到底是什么它和 CUDA Core 的区别很多第一次接触 Tensor Core 的人以为它就是把几条 CUDA Core 指令合成一条。这个理解不算错但远远不够。Tensor Core 是一个独立的硬件单元它能在一个时钟周期内完成一个完整的 4×4×4 矩阵乘加不同代际尺寸不同。你不需要让 32 个线程分别去算 32 个点而是让它们协同把一个矩阵块喂给 Tensor Core由硬件一次性算完。这里有个思维上的巨大转变之前你关心的是“每个线程在算什么”用 Tensor Core 之后你要操心的是“一个 warp 如何合作把数据排成硬件要求的形状”。CUDA Core 的编程模型是 SIMT单指令多线程每个线程各自持有一条数据Tensor Core 的编程模型是 warp 级协作一条指令由整个 warp 发出让 Tensor Core 一次性处理一个子矩阵。这就像以前是 32 个工人每个人搬一块砖现在是一个吊车一次吊起 32 块砖。工人线程的职责从“干活”变成了“摆砖”准备数据碎片。3.2 m16n8k8、m16n8k16、m16n16k16这些形状参数怎么读Tensor Core 的 warp-level 矩阵操作有一个三维的形状描述M×N×K。以最经典的m16n8k8为例它表示一个 warp32 个线程合作计算一个 16×8 的输出矩阵块同时沿 K 方向累加 8 个元素。拆解一下输出是 16×8128 个 float或 half。32 个线程怎么分每个线程持有 4 个输出元素作为寄存器里的 fragment。A 矩阵是 16×K16×8B 矩阵是 K×88×8。每个线程负责从 A 和 B 中持有若干片段fragment然后执行mma.sync指令时硬件自动把 32 个线程手里的 fragment 组合成完整的矩阵乘加。到了 Ampere 之后又增加了m16n8k16FP16/BF16 下 K 可以到 16和m16n8k4TF32 下 K 只有 4。Hopper 上的wgmma更进一步直接把 warpgroup4 个 warp作为操作单位K 维可以一次做到 32 甚至 64。为什么 K 维度重要因为 K 是累加的方向K 值越大一次mma指令覆盖的累加量越多循环次数越少效率越高。但是 K 变大意味着每个线程的 fragment 寄存器数变多寄存器压力也上来了。所以在选择 shape 时你要在“一次算得多”和“寄存器够用”之间找平衡。3.3 WMMA 还是 MMA PTX两种编程姿势怎么选NVIDIA 给了两种编程 Tensor Core 的姿势。一种是 C 层的 WMMA API在mma.h里比较友好另一种是 PTX 层的mma.sync.aligned.m16n8k8...内联汇编更底层CUTLASS 内部就是用它。先说 WMMA。它封装了 fragment 的概念你不需要关心线程和数据的映射关系#include mma.h using namespace nvcuda; wmma::fragmentwmma::matrix_a, 16, 16, 16, __half, wmma::row_major a_frag; wmma::fragmentwmma::matrix_b, 16, 16, 16, __half, wmma::col_major b_frag; wmma::fragmentwmma::accumulator, 16, 16, 16, float c_frag; wmma::fill_fragment(c_frag, 0.0f); for (int k 0; k K; k 16) { wmma::load_matrix_sync(a_frag, A ..., lda); wmma::load_matrix_sync(b_frag, B ..., ldb); wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } wmma::store_matrix_sync(C ..., c_frag, ldm, wmma::mem_row_major);WMMA 的优点是开发效率高缺点是可定制性差fragment 的分布你无法干预而且不同架构下同一段代码生成的指令可能不同。等你需要精确控制寄存器布局、想绕过编译器做cp.async和wgmma的流水线时就得用 PTX。PTX 的mma.sync手写起来虽然繁琐但它暴露了真正的硬件接口。实际上 CUTLASS 在 Ampere 之后基本不再走 WMMA而是直接操作 PTX 指令。我的建议是想快速把算法跑通、验证 Tensor Core 加速效果用 WMMA想写生产级 GEMM趁早学 PTX 理解 CUTLASS 的布局逻辑。WMMA 写出来的性能通常只有手写 PTX 的七成左右因为你在 fragment 层面多做的每一次拷贝、每一处对齐都在浪费寄存器带宽。4. 一条可复现的九步进化路线Naive 到 Tensor Core 手把手4.1 起点的朴素实现错在哪第一节已经展示了朴素代码。从性能角度看它的问题可以总结为三个词访存无复用、数据无向量化、寄存器无缓存。对比算力利用率一个 MNK1024 的 FP32 GEMM在 RTX 4090 上用 CUDA Core 算理论峰值大概 80 TFLOPS 左右朴素版本实测往往不到 1 TFLOPS这就是 1% 的地板。要修它第一步不是上 Tensor Core而是先解决数据复用。一个最简单的改进把内层循环换成寄存器缓存——每个线程一次性计算一行输出这样 A 的元素可以缓存在寄存器中复用 N 次。仅仅是这一个改动性能就能提升 3 到 5 倍。但还不够因为 B 的列仍然在反复从全局内存读。4.2 共享内存分块版本把数据留在芯片内进入正题写一个真正的 tiled 版本。这个版本理论上已经达到了“Level 3”的水平__global__ void sgemm_tiled(const float* A, const float* B, float* C, int M, int N, int K) { const int BM 128, BN 128, BK 8; const int TM 8, TN 8; __shared__ float As[BM][BK]; __shared__ float Bs[BK][BN]; int block_row blockIdx.y * BM; int block_col blockIdx.x * BN; float c[TM][TN] {}; for (int k0 0; k0 K; k0 BK) { // 协同加载 A 的 128x8 块到共享内存 for (int idx threadIdx.x; idx BM * BK; idx blockDim.x) As[idx / BK][idx % BK] A[(block_row idx / BK) * K k0 idx % BK]; // 协同加载 B 的 8x128 块到共享内存 for (int idx threadIdx.x; idx BK * BN; idx blockDim.x) Bs[idx / BN][idx % BN] B[(k0 idx / BN) * N block_col idx % BN]; __syncthreads(); for (int i 0; i TM; i) { int row threadIdx.y * TM i; for (int j 0; j TN; j) { int col threadIdx.x * TN j; for (int k 0; k BK; k) c[i][j] As[row][k] * Bs[k][col]; } } __syncthreads(); } int row_base block_row threadIdx.y * TM; int col_base block_col threadIdx.x * TN; for (int i 0; i TM; i) for (int j 0; j TN; j) C[(row_base i) * N col_base j] c[i][j]; }这里的 blockDim.x 相当于 BN/TN 16blockDim.y 相当于 BM/TM 16一共 256 线程。每个线程手里有一个 8×8 的寄存器累加器c[TM][TN]。外层循环每走一步BK8从全局内存读 128×8 的 A 和 8×128 的 B 到共享内存然后所有线程再从共享内存里各取所需做 8 次乘加累加。这个版本和朴素版相比数据复用率提升了多少A 的每个元素从全局内存只读一次被块内 256 个线程复用B 同理。全局访存从 8GB 量级降到了 M×KN×K 量级。实测下来这个版本能到 25% 到 40% 的算力利用率。但这版本有个明显的坑As[row][k]这种访问模式在共享内存上会产生严重的 bank conflict——同一行的线程同时访问同一 bank导致共享内存访问被串行化。这个问题第 5 节再展开。4.3 Tensor Core 版本改变的是数据排布和循环结构现在加入 Tensor Core。这里用 WMMA 接口矩阵类型选__halfFP16shape 用 m16n16k16。注意一个关键点进入 Tensor Core 的矩阵必须在共享内存或全局内存中按硬件要求的布局排好尤其是 B 矩阵的col_major和 A 矩阵的row_major不能弄混。#include mma.h using namespace nvcuda; __global__ void gemm_tensor_core(const __half* A, const __half* B, float* C, int M, int N, int K) { const int WMMA_M 16, WMMA_N 16, WMMA_K 16; int warp_row (blockIdx.y * blockDim.y threadIdx.y) * WMMA_M; int warp_col (blockIdx.x * blockDim.x threadIdx.x) * WMMA_N; wmma::fragmentwmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, __half, wmma::row_major a_frag; wmma::fragmentwmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, __half, wmma::col_major b_frag; wmma::fragmentwmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float c_frag; wmma::fill_fragment(c_frag, 0.0f); for (int k 0; k K; k WMMA_K) { wmma::load_matrix_sync(a_frag, A warp_row * K k, K); wmma::load_matrix_sync(b_frag, B k * N warp_col, N); wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } wmma::store_matrix_sync(C warp_row * N warp_col, c_frag, N, wmma::mem_row_major); }这个版本看起来甚至比共享内存版更简单因为它没有手动 Tiling——你把每个 warp 负责的 16×16 输出块直接映射到矩阵坐标上load_matrix_sync从全局内存加载 fragmentmma_sync执行矩阵乘加。但性能同样不会太好因为缺乏共享内存分块load_matrix_sync本质上还是从全局内存直接取数。真正的生产级写法应该是共享内存 tiling WMMA fragment 从共享内存加载 双缓冲。这里不展开完整代码但思路是把第 4.2 节中的As和Bs保留下来然后在循环内部用load_matrix_sync从As/Bs加载 fragment再mma_sync。这样数据从 global 到 smem 再到寄存器Tensor Core 才会真正喂饱。所以记住一个结论Tensor Core 解决的是“算得快”但“数据到得了”这个问题还是得靠 Tiling 解决。两者不是替代关系是上下游关系。5. 冲上 95% 的关键动作双缓冲、异步拷贝与 Swizzle5.1 Double Buffering用流水线掩盖共享内存加载延迟到了这个阶段你的 kernel 已经从访存瓶颈转向了计算和访存的混合瓶颈。Nsight 里看到的现象是SM 大部分时间在等待数据而非执行 FMA。这时候就要上双缓冲。双缓冲的思路很简单共享内存里准备两份 buffer一份用于当前 K 步的计算另一份预先加载下一步要用的数据。计算第 k 块的时候cp.async已经在往另一个 buffer 里搬运第 k1 块的数据了。这样访存延迟被隐藏在了计算时间里。实现上有两种方式。老式的做法是手动#pragma unroll循环通过奇偶迭代判断用哪一块 buffer__shared__ float As[2][BM][BK]; __shared__ float Bs[2][BK][BN]; for (int k0 0; k0 K; k0 BK) { int cur (k0 / BK) % 2; int next 1 - cur; // 为 next buffer 发起异步拷贝 if (k0 BK K) { cp_async(As[next], A ..., ...); cp_async(Bs[next], B ..., ...); } cp_async_commit(); // 使用当前 buffer 计算 compute(As[cur], Bs[cur]); cp_async_wait(); __syncthreads(); }这里的核心要点是计算和访存必须独立、并行推进。如果你用同步loadstore方式做双缓冲效果会大打折扣因为同步拷贝本身占用线程的指令周期。真正的异步拷贝要用cp.async指令或更高层的cuda::memcpy_async。5.2 cp.async 与 TMA搬数据这件事也别交给线程在 Ampere 之前从全局内存到共享内存的数据搬运必须经过寄存器——线程要先把数据 load 到寄存器再 store 到共享内存。这个“寄存器中转”的代价是每一段搬运数据都要消耗寄存器资源和指令周期。cp.async指令改变了这一点。它允许数据从全局内存直接拷贝到共享内存不经过寄存器由硬件异步完成。这带来的好处不仅是省了几条指令更重要的是把线程从“搬运工”的角色中解放出来让线程可以专心做计算。到了 Hopper这套机制又进化成了 TMATensor Memory Accelerator。TMA 可以在一个 SM、甚至一个 CTA 的维度上描述一个多维张量的拷贝直接把整块数据从全局内存搬到共享内存然后通过 warpgroup 级别的 barrier 通知消费方。TMA 的好处是搬运描述更简洁且能自动处理边界 padding加上wgmma配合使用几乎就是为 GEMM 量身定做的流水线。如果你还在 Ampere 上开发cp.async就足够了如果目标架构是 Hopper 及以上尽量直接上 TMA。这是我和别人讨论时反复强调的一点——很多基于 Ampere 的优化代码直接搬到 Hopper 上跑不出应有的性能就是因为还在用老式的寄存器中转搬运。5.3 Bank Conflict 与 Swizzle同一个坑里别挤两辆车共享内存有一个物理结构它被等分成 32 个 bank每个 bank 每周期只能服务一次访问。当 32 个线程同时访问共享内存时如果它们访问的地址恰好落在同一个 bank 上这些访问就会被硬件串行化——这就是 bank conflict。第 4.2 节的代码里有个典型冲突As[row][k]在一个 warp 内如果 32 个线程的 row 相同而 k 不同那么地址的分布就是row * BK k它们可能全部落在相邻的 bank 上。没有冲突的情况下一次取 32 个 float 只需要一个周期冲突的情况下一次访问可能被拆成多次最极端是 32 次共享内存带宽直接缩水 32 倍。解决 bank conflict 的手段有两类。第一类是 padding把As[BM][BK]改成As[BM][BK 1]人为把每一行的起始地址错开让线程访问落进不同 bank。第二类是 swizzle不再是简单地按行存而是对地址做 XOR 变换把数据的物理排布打乱让任何 warp 访问模式都尽可能均匀地落在所有 bank 上。CUTLASS 里常见的是 XOR swizzle例如把(row, col)映射为(row, col ^ ((row % 8)))。这类操作复杂但效果好实测可以贡献 5% 到 10% 的算力利用率提升。5.4 用 Nsight Compute 去验证每一层优化九重天做不到“闷头写代码一步到位”每一步都需要用 profiling 工具验证。Nsight Compute 里你要盯的核心指标有这么几个指标含义当你看到什么时该做什么Memory Throughput全局内存/共享内存带宽利用率接近 100% 说明访存是瓶颈上双缓冲SM BusySM 算数单元占用率低说明在等数据检查 stall 原因Warp Stall Reasons线程停在哪里Long Scoreboard是访存等待改 cp.asyncShared Memory Conflict共享内存 bank 冲突率不为 0 就做 padding/swizzleAchieved Occupancy实际并发线程数与理论之比低于 50%减小 smem 或寄存器用量Compute Pipe UtilizationFMA/MMA 流水线利用率这就是算力利用率的直接体现优化的逻辑是闭环的先看 SM 为什么闲是数据没到访存瓶颈还是线程不足占用率低然后对症下药。我见过太多人一上来就折腾寄存器分配结果瓶颈根本不在那里。Nsight 的Source View能直接告诉你每行代码的 stall 原因善用这个功能而不是只看一个综合分数。6. 实战中踩过的坑布局、精度与架构差异6.1 矩阵布局错位Tensor Core 的“形状敏感”体质wmma::mma_sync看起来只是三个 fragment 的事但 fragment 里数据在寄存器中的排布完全由加载时的row_major/col_major决定。最常见的问题出现在 B 矩阵上你声明了col_major传入的数据却按 row-major 存结果计算出的 C 每一块都是错的但错得“很局部”导致你 Debug 半天才发现。这个坑的根源在于GEMM 公式里 A 是行主序、B 是列主序最自然C[I][J] 累加 A[I][K]×B[K][J]B 的 K 是行索引、J 是列索引。如果你从 PyTorch 或 NumPy 拿到的 B 是 row-major就必须先转置或者加载时指定wmma::mem_row_major然后配合正确的 stride 计算。我的习惯是写一个最小规模的单元测试MNK16手动算一遍结果对比先把布局验证掉再上大矩阵。6.2 精度TF32、FP16、BF16 怎么选Tensor Core 支持多种精度但它们不是随便选的。FP16 的算力最高但动态范围小容易溢出——K 很大的累加过程中间结果可能在__half范围内溢出为 inf。BF16 的动态范围跟 FP32 一样大但精度只有 8 位适合训练前向、不适合要求精确的累加。TF32 是 FP32 输入的“截断版”精度 10 位指数 10 位尾数是 Ampere 上做 FP32 GEMM 的折中选择。如果你要在 Tensor Core 上算 FP32 精度要求的 GEMMAmpere 和 Ada 上的选择是 TF32如果数据范围可控可以先转成 FP16 再用 split-K 技巧补偿精度损失这种方式在推理场景中很常见。但注意从 FP32 到 TF32 的转换不是免费的——每个数都要过一遍__float_to_tf32如果转换做在循环内部性能会掉。正确做法是提前把 A、B 转成 TF32 格式存到共享内存里一次转换、多次使用。6.3 从 Ampere 到 Hopper指令在变思维不变如果你在不同代的 GPU 上编译同一份 WMMA 代码性能差异可能很大但代码本身通常能跑。问题是当你追求极限性能时必须针对架构调 shapeAmpere 的 FP16mma.sync一次 K 是 16TF32 是 8Hopper 的wgmma一次 K 可以是 32 或 64而且依赖 TMA 的异步数据流根本不是同一套写法。这里有一个实际的兼容性陷阱很多公司代码库为了兼容旧卡只敢用sm_80编译于是在 Hopper 上完全发挥不出性能。遇到这种情况建议按架构拆分代码路径sm_80用 Ampere 优化版sm_90用 TMA wgmma版。编译时用-gencode archcompute_80,codesm_80 -gencode archcompute_90,codesm_90分别生成 cubin运行时按cudaGetDeviceProperties的 major/minor 版本分发。6.4 常见问题速查表症状可能原因排查与解法性能在 10% 原地踏步全局访存未分块检查是否忘了共享内存 TilingNsight 看 L2 命中率提升很小但理论峰值高Bank conflict 严重看 Shared Memory Conflict 指标加 padding 或 swizzle使用 Tensor Core 后反而变慢矩阵太小或 fragment 布局反复拷贝确保数据直接进 fragment避免额外复制结果不正确且有规律的错B 矩阵列/行主序搞错用 MNK16 的最小 case 验证布局算力利用率 70% 上不去缺少双缓冲或占用率不足检查 Stall Reasons上 cp.async换一张卡性能差异巨大编译针对的 SM 版本不对检查 cubin 架构按架构分支编译6.5 一个实操建议从现有高性能库逆向学习最后认真给一个建议不要从零手写终极版 GEMM。CUTLASS 是一个极好的学习对象它的gemm示例从device/到thread/的每个层级都标注了对应的问题域。你可以基于它跑出 90% 以上的性能然后通过修改 tile shape、swizzle 模式、流水线级数来理解每一层的权衡。我自己从 CUTLASS 里学到的最有价值的东西不是哪个函数怎么写而是它如何把“数据排布”和“指令选择”拆成独立的层次让每一层的调优都不牵一发而动全身。总的来说从 1% 到 95% 的九重天跃迁每一步的本质都是在回答同一个问题数据在哪里计算在哪里它们之间的通路是不是最短的。朴素版让数据在全局内存和计算单元之间来回折腾Tensor Core 则要求你把数据喂到硬件最擅长的形状里。等你把 Tiling、双缓冲、Swizzle、精度选型这些环节都打通之后再回头看 CUTLASS 的源码会感觉每一段代码都像老朋友在打招呼——那就说明你把这个领域真正吃透了。这套思路不仅适用于 CUDA你在任何需要手工调度数据的计算平台上写代码都会受益。 从1%到95%：CUDA GEMM优化九重天与Tensor Core实践
先交代一个很现实的场景你花了一周时间照着 CUDA 文档写了一个 GEMM kernelnvcc编译通过结果一跑Nsight Compute 里Achieved Occupancy还行SM Busy也说得过去但算力利用率只有 1%——你的 GPU 号称几十 TFLOPS实际只吐出来零点几。这不是你笨是因为 GEMM 优化这条路上从“能算”到“算得快”中间隔着一整套关于 Tiling、数据复用、指令选择和硬件单元的问题。这一讲就是专门把这个过程拆开从 1% 到 95% 的算力利用率到底是怎么一步步抠出来的每一层优化在解决什么问题Tensor Core 的编程思维和普通 CUDA Core 有什么区别以及你在实际操作中会遇到哪些坑。这讲适合两类人看。第一类是已经写过 CUDA kernel、但性能总是不上不下的开发者你需要理解为什么共享内存分块能带来几十倍提升为什么 Tensor Core 不是简单的“换一条指令”就完事。第二类是读 CUTLASS 源码、看各种 GEMM 优化文章时一头雾水的人——无数文章提到cp.async、swizzle、wgmma但没人告诉你这些到底在解决什么底层痛点。如果你是这两类人之一下面这五千字应该能帮你在脑海里建立起一张完整的 GEMM 优化地图。1. 为什么常规 GEMM 只能跑出 1% 算力先算清这笔账1.1 算子性能的三座大山访存、并行粒度、指令效率先搞清楚一个问题GPU 的核心瓶颈到底在哪里。你可以把一次 GEMM 运算想象成一个仓库配货场景。显存是仓库SM流式多处理器是分拣车间寄存器和共享内存是车间里的操作台。CUDA Core 和 Tensor Core 是操作台上的工人。理论上车间的加工能力算力是固定的但如果你每次加工一个小零件都要跑到仓库里拿原料那时间全花在路上了。朴素 GEMM 就是这个疯狂跑仓库的场景。具体点说一个最简单的 GEMM kernel 长这样__global__ void sgemm_naive(const float* A, const float* B, float* C, int M, int N, int K) { int row blockIdx.y * blockDim.y threadIdx.y; int col blockIdx.x * blockDim.x threadIdx.x; float sum 0.0f; for (int k 0; k K; k) sum A[row * K k] * B[k * N col]; C[row * N col] sum; }这段代码问题很多但最致命的有三条。第一全局内存反复访问。每个线程要计算 C[row][col] 这个点就必须把 A 的第 row 行整行和 B 的第 col 列整列全部读一遍。如果 MNK4096整个 kernel 对全局内存的读取量是 8GB 量级而数据实际只有 64MB。也就是说99% 的访存都是在重复搬运同一个数据。第二并行粒度不是硬件喜欢的形状。每个线程只是算一个点完全没有数据和指令的复用性。第三没有向量化访存。一条 32-bit 的 load 指令一次只搬 4 字节L1/L2 带宽利用率极低。有人可能会说L2 cache 不是能缓解吗没错L2 能帮你挡掉一部分重复访问但 GEMM 的数据访问模式是“行”和“列”交叉的多个线程同时访问同一行、不同列的时候L2 缓存行还勉强能命中可一旦访问模式变成跨行跨列的“跳跃式”缓存行命中率立刻崩掉。所以在 GPU 上靠 cache 兜底不是正路正路是让每个线程自己把数据搬到片上主动制造数据复用。1.2 运算强度判断 GEMM 该不该被算力吃满这里要引出一个关键概念运算强度Arithmetic Intensity也就是一次访存对应的浮点运算次数。GEMM 的理论运算强度很好算——总 FLOPs 是 2×M×N×K数据量是 M×K K×N M×N。当矩阵规模足够大时数据量远小于 FLOPs运算强度可以做到几百 FLOPs/Byte。GPU 的 Roofline 模型告诉我们只有当运算强度超过某个临界点通常是算力除以带宽A100 级别大约是几十 FLOPs/Byte你才能摸到算力天花板。朴素 GEMM 的“实际运算强度”远远达不到这个临界值——因为它每次计算都要重新访存实际的有效运算强度可能只有几个 FLOPs/Byte。所以算力利用率 1% 不是 GPU 的错是你的数据流设计把整个 kernel 压在了访存瓶颈上。反过来这也说明了一个核心结论GEMM 的性能问题本质上是一个数据复用问题。你要做的不是减少计算量计算量固定而是在数据被消费之前尽量让它在片上多待几个时钟周期。这就是 Tiling 存在的全部理由。1.3 “九重天”对应的优化阶段对照下面这张表是我按经验整理的一个 GEMM 优化路线图。不同架构、不同矩阵规模下具体数值会有浮动但大致趋势是稳定的每一行都代表一个独立的优化维度层级关键动作大致算力利用率这一步在解决什么Level 0朴素三重循环约 1%没有任何数据复用访存完全失控Level 1循环重排 寄存器缓存3% - 5%让某些数据在寄存器中被多次使用Level 2向量化访存float48% - 12%降低指令数提高访存带宽利用率Level 3共享内存分块25% - 40%数据从 global 搬到 smem片内复用Level 4消除 Bank Conflict45% - 55%共享内存本身的带宽瓶颈被打开Level 5Double Buffering cp.async60% - 70%用流水线掩盖访存延迟Level 6Tensor CoreFP16/TF3275% - 85%从 CUDA Core 换到专用矩阵单元Level 7Shape 对齐 精度选型88% - 92%让硬件矩阵单元没有一丝浪费Level 8TMA 全局调度优化90% - 95%搬数据这件事交给硬件线程纯算这九级不是必须一步一步走走得顺的话 Level 3 和 Level 4 经常一起做Level 6 和 Level 7 往往绑在一起。但每一级背后的物理约束不同你只有理解了每一层到底在卡什么才能在 Nsight 里看到瓶颈的时候知道该动哪一块代码。2. Tiling 分块把矩阵切成 GPU 能“嚼得动”的大小2.1 为什么需要两级搬运从全局内存到共享内存再到寄存器Tiling 的核心思想用一个做饭的类比就能讲清楚。你要做 100 道菜每道菜需要葱姜蒜。最笨的办法是做一道菜就去冰箱拿一次配料这就是朴素 GEMM。聪明的做法是一次性把一大袋配料全拿到厨房台面上共享内存做菜的时候从台面取再想快一点就把最常用的几样放在手边的小碗里寄存器手一伸就够到。这就是两级分块全局内存 → 共享内存 → 寄存器。为什么不能直接从全局内存到寄存器因为寄存器容量太小一个 SM 的寄存器文件总共才 256KB一个线程最多 255 个寄存器装不下一个能支撑复用的大块。而共享内存每个 SM 有 64KB 到 228KB 不等老的 Volta 是 96KBAmpere 上是 164KBHopper 是 228KB能装下比寄存器大得多的数据块。所以实际最优策略是先用共享内存装一个大 tile再由线程从共享内存里把子块读进寄存器。这个两级结构还对应一个更底层的硬件事实全局内存的延迟是几百个周期共享内存的延迟只有二三十个周期寄存器几乎是零延迟。你的优化目标是把绝大多数访问落在延迟最低的那一级。2.2 Block Tile、Warp Tile、Thread Tile三层名字背后的分工Tiling 不是简单地“把矩阵分成小块”而是分成三个层次对应三种不同的执行单元。Block Tile是整个线程块负责的那一块矩阵比如 128×128。它对应的是共享内存的容量边界。Block Tile 越大数据复用率越高——因为同一个 tile 里的数据会被块内所有线程反复使用。但太大也不行共享内存放不下两张 tile或者放得下却导致占用率暴跌。Warp Tile是线程块内部每个 warp 负责的子块。一个 128×128 的 Block Tile 切成 4 个 warp 来算每个 warp 可能负责一个 64×64 的区域。Warp Tile 是 SIMT 执行的基本单位——一个 warp 里的 32 个线程协同计算如果它们各自算一个独立的小 tile访存会乱但如果它们合作算一个 warp tile 并共享寄存器中的中间结果效率会高很多。Thread Tile是单个线程最终负责的那一小块输出常见的是 4×4、8×8。为什么要让一个线程算 8×8 而不是 1×1因为这样每个从共享内存读进来的 A 和 B 的元素可以被 8 次乘加复用这就是寄存器级的 Tiling。没有这层复用共享内存的带宽又会变成新瓶颈。2.3 经典 128×128/8×8 组合为什么最常见如果你去看 CUTLASS 和各种高性能 GEMM会发现 128×128 的 Block Tile、8×8 的 Thread Tile 非常常见。这是巧合吗不是。算一笔账。Block Tile 选 128×128意味着每个 block 要处理 128×128 的输出。如果用 8×8 的 Thread Tile每个 block 需要 16×16256 个线程来覆盖整个 tile每个线程算 8×8。而 blockDim 是 256 个线程正好一个 block 可以装 8 个 warp。再看共享内存A 的 tile 是 128×BKB 的 tile 是 BK×128。如果 BK8两个 tile 各占 128×8×44KB总共 8KB这对于 164KB 的共享内存来说只占很小一部分。但这只是“能装下”不是“最优”。真正的权衡在于占用率。一个 SM 上到底能跑多少个 block取决于共享内存和寄存器哪个先耗尽。假设一个 block 用 8KB 共享内存、256 个线程每个线程用 128 个寄存器那一个 SM 的 64K 寄存器按 Ampere 算只能容纳 2 个 block共享内存倒是不紧张。这时候 Blocks per SM 就只有 2占用率大概 50%。如果你把 Block Tile 缩到 64×64每个 block 的线程数降到 64SM 上能塞更多 block但共享内存 tile 变小数据复用率下降。所以 128×128 是一个考虑了共享内存容量、寄存器数目和复用率之后的平衡点不是拍脑袋定的。我自己实测的体会是在小矩阵比如 MNK256上用大 Block Tile 反而容易因为边界 padding 浪费计算在 2048 以上的大矩阵上128×128 几乎是稳定好用的万能选择。3. Tensor Core 思维从“线程算一个数”到“Warp 算一块矩阵”3.1 Tensor Core 到底是什么它和 CUDA Core 的区别很多第一次接触 Tensor Core 的人以为它就是把几条 CUDA Core 指令合成一条。这个理解不算错但远远不够。Tensor Core 是一个独立的硬件单元它能在一个时钟周期内完成一个完整的 4×4×4 矩阵乘加不同代际尺寸不同。你不需要让 32 个线程分别去算 32 个点而是让它们协同把一个矩阵块喂给 Tensor Core由硬件一次性算完。这里有个思维上的巨大转变之前你关心的是“每个线程在算什么”用 Tensor Core 之后你要操心的是“一个 warp 如何合作把数据排成硬件要求的形状”。CUDA Core 的编程模型是 SIMT单指令多线程每个线程各自持有一条数据Tensor Core 的编程模型是 warp 级协作一条指令由整个 warp 发出让 Tensor Core 一次性处理一个子矩阵。这就像以前是 32 个工人每个人搬一块砖现在是一个吊车一次吊起 32 块砖。工人线程的职责从“干活”变成了“摆砖”准备数据碎片。3.2 m16n8k8、m16n8k16、m16n16k16这些形状参数怎么读Tensor Core 的 warp-level 矩阵操作有一个三维的形状描述M×N×K。以最经典的m16n8k8为例它表示一个 warp32 个线程合作计算一个 16×8 的输出矩阵块同时沿 K 方向累加 8 个元素。拆解一下输出是 16×8128 个 float或 half。32 个线程怎么分每个线程持有 4 个输出元素作为寄存器里的 fragment。A 矩阵是 16×K16×8B 矩阵是 K×88×8。每个线程负责从 A 和 B 中持有若干片段fragment然后执行mma.sync指令时硬件自动把 32 个线程手里的 fragment 组合成完整的矩阵乘加。到了 Ampere 之后又增加了m16n8k16FP16/BF16 下 K 可以到 16和m16n8k4TF32 下 K 只有 4。Hopper 上的wgmma更进一步直接把 warpgroup4 个 warp作为操作单位K 维可以一次做到 32 甚至 64。为什么 K 维度重要因为 K 是累加的方向K 值越大一次mma指令覆盖的累加量越多循环次数越少效率越高。但是 K 变大意味着每个线程的 fragment 寄存器数变多寄存器压力也上来了。所以在选择 shape 时你要在“一次算得多”和“寄存器够用”之间找平衡。3.3 WMMA 还是 MMA PTX两种编程姿势怎么选NVIDIA 给了两种编程 Tensor Core 的姿势。一种是 C 层的 WMMA API在mma.h里比较友好另一种是 PTX 层的mma.sync.aligned.m16n8k8...内联汇编更底层CUTLASS 内部就是用它。先说 WMMA。它封装了 fragment 的概念你不需要关心线程和数据的映射关系#include mma.h using namespace nvcuda; wmma::fragmentwmma::matrix_a, 16, 16, 16, __half, wmma::row_major a_frag; wmma::fragmentwmma::matrix_b, 16, 16, 16, __half, wmma::col_major b_frag; wmma::fragmentwmma::accumulator, 16, 16, 16, float c_frag; wmma::fill_fragment(c_frag, 0.0f); for (int k 0; k K; k 16) { wmma::load_matrix_sync(a_frag, A ..., lda); wmma::load_matrix_sync(b_frag, B ..., ldb); wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } wmma::store_matrix_sync(C ..., c_frag, ldm, wmma::mem_row_major);WMMA 的优点是开发效率高缺点是可定制性差fragment 的分布你无法干预而且不同架构下同一段代码生成的指令可能不同。等你需要精确控制寄存器布局、想绕过编译器做cp.async和wgmma的流水线时就得用 PTX。PTX 的mma.sync手写起来虽然繁琐但它暴露了真正的硬件接口。实际上 CUTLASS 在 Ampere 之后基本不再走 WMMA而是直接操作 PTX 指令。我的建议是想快速把算法跑通、验证 Tensor Core 加速效果用 WMMA想写生产级 GEMM趁早学 PTX 理解 CUTLASS 的布局逻辑。WMMA 写出来的性能通常只有手写 PTX 的七成左右因为你在 fragment 层面多做的每一次拷贝、每一处对齐都在浪费寄存器带宽。4. 一条可复现的九步进化路线Naive 到 Tensor Core 手把手4.1 起点的朴素实现错在哪第一节已经展示了朴素代码。从性能角度看它的问题可以总结为三个词访存无复用、数据无向量化、寄存器无缓存。对比算力利用率一个 MNK1024 的 FP32 GEMM在 RTX 4090 上用 CUDA Core 算理论峰值大概 80 TFLOPS 左右朴素版本实测往往不到 1 TFLOPS这就是 1% 的地板。要修它第一步不是上 Tensor Core而是先解决数据复用。一个最简单的改进把内层循环换成寄存器缓存——每个线程一次性计算一行输出这样 A 的元素可以缓存在寄存器中复用 N 次。仅仅是这一个改动性能就能提升 3 到 5 倍。但还不够因为 B 的列仍然在反复从全局内存读。4.2 共享内存分块版本把数据留在芯片内进入正题写一个真正的 tiled 版本。这个版本理论上已经达到了“Level 3”的水平__global__ void sgemm_tiled(const float* A, const float* B, float* C, int M, int N, int K) { const int BM 128, BN 128, BK 8; const int TM 8, TN 8; __shared__ float As[BM][BK]; __shared__ float Bs[BK][BN]; int block_row blockIdx.y * BM; int block_col blockIdx.x * BN; float c[TM][TN] {}; for (int k0 0; k0 K; k0 BK) { // 协同加载 A 的 128x8 块到共享内存 for (int idx threadIdx.x; idx BM * BK; idx blockDim.x) As[idx / BK][idx % BK] A[(block_row idx / BK) * K k0 idx % BK]; // 协同加载 B 的 8x128 块到共享内存 for (int idx threadIdx.x; idx BK * BN; idx blockDim.x) Bs[idx / BN][idx % BN] B[(k0 idx / BN) * N block_col idx % BN]; __syncthreads(); for (int i 0; i TM; i) { int row threadIdx.y * TM i; for (int j 0; j TN; j) { int col threadIdx.x * TN j; for (int k 0; k BK; k) c[i][j] As[row][k] * Bs[k][col]; } } __syncthreads(); } int row_base block_row threadIdx.y * TM; int col_base block_col threadIdx.x * TN; for (int i 0; i TM; i) for (int j 0; j TN; j) C[(row_base i) * N col_base j] c[i][j]; }这里的 blockDim.x 相当于 BN/TN 16blockDim.y 相当于 BM/TM 16一共 256 线程。每个线程手里有一个 8×8 的寄存器累加器c[TM][TN]。外层循环每走一步BK8从全局内存读 128×8 的 A 和 8×128 的 B 到共享内存然后所有线程再从共享内存里各取所需做 8 次乘加累加。这个版本和朴素版相比数据复用率提升了多少A 的每个元素从全局内存只读一次被块内 256 个线程复用B 同理。全局访存从 8GB 量级降到了 M×KN×K 量级。实测下来这个版本能到 25% 到 40% 的算力利用率。但这版本有个明显的坑As[row][k]这种访问模式在共享内存上会产生严重的 bank conflict——同一行的线程同时访问同一 bank导致共享内存访问被串行化。这个问题第 5 节再展开。4.3 Tensor Core 版本改变的是数据排布和循环结构现在加入 Tensor Core。这里用 WMMA 接口矩阵类型选__halfFP16shape 用 m16n16k16。注意一个关键点进入 Tensor Core 的矩阵必须在共享内存或全局内存中按硬件要求的布局排好尤其是 B 矩阵的col_major和 A 矩阵的row_major不能弄混。#include mma.h using namespace nvcuda; __global__ void gemm_tensor_core(const __half* A, const __half* B, float* C, int M, int N, int K) { const int WMMA_M 16, WMMA_N 16, WMMA_K 16; int warp_row (blockIdx.y * blockDim.y threadIdx.y) * WMMA_M; int warp_col (blockIdx.x * blockDim.x threadIdx.x) * WMMA_N; wmma::fragmentwmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, __half, wmma::row_major a_frag; wmma::fragmentwmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, __half, wmma::col_major b_frag; wmma::fragmentwmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float c_frag; wmma::fill_fragment(c_frag, 0.0f); for (int k 0; k K; k WMMA_K) { wmma::load_matrix_sync(a_frag, A warp_row * K k, K); wmma::load_matrix_sync(b_frag, B k * N warp_col, N); wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); } wmma::store_matrix_sync(C warp_row * N warp_col, c_frag, N, wmma::mem_row_major); }这个版本看起来甚至比共享内存版更简单因为它没有手动 Tiling——你把每个 warp 负责的 16×16 输出块直接映射到矩阵坐标上load_matrix_sync从全局内存加载 fragmentmma_sync执行矩阵乘加。但性能同样不会太好因为缺乏共享内存分块load_matrix_sync本质上还是从全局内存直接取数。真正的生产级写法应该是共享内存 tiling WMMA fragment 从共享内存加载 双缓冲。这里不展开完整代码但思路是把第 4.2 节中的As和Bs保留下来然后在循环内部用load_matrix_sync从As/Bs加载 fragment再mma_sync。这样数据从 global 到 smem 再到寄存器Tensor Core 才会真正喂饱。所以记住一个结论Tensor Core 解决的是“算得快”但“数据到得了”这个问题还是得靠 Tiling 解决。两者不是替代关系是上下游关系。5. 冲上 95% 的关键动作双缓冲、异步拷贝与 Swizzle5.1 Double Buffering用流水线掩盖共享内存加载延迟到了这个阶段你的 kernel 已经从访存瓶颈转向了计算和访存的混合瓶颈。Nsight 里看到的现象是SM 大部分时间在等待数据而非执行 FMA。这时候就要上双缓冲。双缓冲的思路很简单共享内存里准备两份 buffer一份用于当前 K 步的计算另一份预先加载下一步要用的数据。计算第 k 块的时候cp.async已经在往另一个 buffer 里搬运第 k1 块的数据了。这样访存延迟被隐藏在了计算时间里。实现上有两种方式。老式的做法是手动#pragma unroll循环通过奇偶迭代判断用哪一块 buffer__shared__ float As[2][BM][BK]; __shared__ float Bs[2][BK][BN]; for (int k0 0; k0 K; k0 BK) { int cur (k0 / BK) % 2; int next 1 - cur; // 为 next buffer 发起异步拷贝 if (k0 BK K) { cp_async(As[next], A ..., ...); cp_async(Bs[next], B ..., ...); } cp_async_commit(); // 使用当前 buffer 计算 compute(As[cur], Bs[cur]); cp_async_wait(); __syncthreads(); }这里的核心要点是计算和访存必须独立、并行推进。如果你用同步loadstore方式做双缓冲效果会大打折扣因为同步拷贝本身占用线程的指令周期。真正的异步拷贝要用cp.async指令或更高层的cuda::memcpy_async。5.2 cp.async 与 TMA搬数据这件事也别交给线程在 Ampere 之前从全局内存到共享内存的数据搬运必须经过寄存器——线程要先把数据 load 到寄存器再 store 到共享内存。这个“寄存器中转”的代价是每一段搬运数据都要消耗寄存器资源和指令周期。cp.async指令改变了这一点。它允许数据从全局内存直接拷贝到共享内存不经过寄存器由硬件异步完成。这带来的好处不仅是省了几条指令更重要的是把线程从“搬运工”的角色中解放出来让线程可以专心做计算。到了 Hopper这套机制又进化成了 TMATensor Memory Accelerator。TMA 可以在一个 SM、甚至一个 CTA 的维度上描述一个多维张量的拷贝直接把整块数据从全局内存搬到共享内存然后通过 warpgroup 级别的 barrier 通知消费方。TMA 的好处是搬运描述更简洁且能自动处理边界 padding加上wgmma配合使用几乎就是为 GEMM 量身定做的流水线。如果你还在 Ampere 上开发cp.async就足够了如果目标架构是 Hopper 及以上尽量直接上 TMA。这是我和别人讨论时反复强调的一点——很多基于 Ampere 的优化代码直接搬到 Hopper 上跑不出应有的性能就是因为还在用老式的寄存器中转搬运。5.3 Bank Conflict 与 Swizzle同一个坑里别挤两辆车共享内存有一个物理结构它被等分成 32 个 bank每个 bank 每周期只能服务一次访问。当 32 个线程同时访问共享内存时如果它们访问的地址恰好落在同一个 bank 上这些访问就会被硬件串行化——这就是 bank conflict。第 4.2 节的代码里有个典型冲突As[row][k]在一个 warp 内如果 32 个线程的 row 相同而 k 不同那么地址的分布就是row * BK k它们可能全部落在相邻的 bank 上。没有冲突的情况下一次取 32 个 float 只需要一个周期冲突的情况下一次访问可能被拆成多次最极端是 32 次共享内存带宽直接缩水 32 倍。解决 bank conflict 的手段有两类。第一类是 padding把As[BM][BK]改成As[BM][BK 1]人为把每一行的起始地址错开让线程访问落进不同 bank。第二类是 swizzle不再是简单地按行存而是对地址做 XOR 变换把数据的物理排布打乱让任何 warp 访问模式都尽可能均匀地落在所有 bank 上。CUTLASS 里常见的是 XOR swizzle例如把(row, col)映射为(row, col ^ ((row % 8)))。这类操作复杂但效果好实测可以贡献 5% 到 10% 的算力利用率提升。5.4 用 Nsight Compute 去验证每一层优化九重天做不到“闷头写代码一步到位”每一步都需要用 profiling 工具验证。Nsight Compute 里你要盯的核心指标有这么几个指标含义当你看到什么时该做什么Memory Throughput全局内存/共享内存带宽利用率接近 100% 说明访存是瓶颈上双缓冲SM BusySM 算数单元占用率低说明在等数据检查 stall 原因Warp Stall Reasons线程停在哪里Long Scoreboard是访存等待改 cp.asyncShared Memory Conflict共享内存 bank 冲突率不为 0 就做 padding/swizzleAchieved Occupancy实际并发线程数与理论之比低于 50%减小 smem 或寄存器用量Compute Pipe UtilizationFMA/MMA 流水线利用率这就是算力利用率的直接体现优化的逻辑是闭环的先看 SM 为什么闲是数据没到访存瓶颈还是线程不足占用率低然后对症下药。我见过太多人一上来就折腾寄存器分配结果瓶颈根本不在那里。Nsight 的Source View能直接告诉你每行代码的 stall 原因善用这个功能而不是只看一个综合分数。6. 实战中踩过的坑布局、精度与架构差异6.1 矩阵布局错位Tensor Core 的“形状敏感”体质wmma::mma_sync看起来只是三个 fragment 的事但 fragment 里数据在寄存器中的排布完全由加载时的row_major/col_major决定。最常见的问题出现在 B 矩阵上你声明了col_major传入的数据却按 row-major 存结果计算出的 C 每一块都是错的但错得“很局部”导致你 Debug 半天才发现。这个坑的根源在于GEMM 公式里 A 是行主序、B 是列主序最自然C[I][J] 累加 A[I][K]×B[K][J]B 的 K 是行索引、J 是列索引。如果你从 PyTorch 或 NumPy 拿到的 B 是 row-major就必须先转置或者加载时指定wmma::mem_row_major然后配合正确的 stride 计算。我的习惯是写一个最小规模的单元测试MNK16手动算一遍结果对比先把布局验证掉再上大矩阵。6.2 精度TF32、FP16、BF16 怎么选Tensor Core 支持多种精度但它们不是随便选的。FP16 的算力最高但动态范围小容易溢出——K 很大的累加过程中间结果可能在__half范围内溢出为 inf。BF16 的动态范围跟 FP32 一样大但精度只有 8 位适合训练前向、不适合要求精确的累加。TF32 是 FP32 输入的“截断版”精度 10 位指数 10 位尾数是 Ampere 上做 FP32 GEMM 的折中选择。如果你要在 Tensor Core 上算 FP32 精度要求的 GEMMAmpere 和 Ada 上的选择是 TF32如果数据范围可控可以先转成 FP16 再用 split-K 技巧补偿精度损失这种方式在推理场景中很常见。但注意从 FP32 到 TF32 的转换不是免费的——每个数都要过一遍__float_to_tf32如果转换做在循环内部性能会掉。正确做法是提前把 A、B 转成 TF32 格式存到共享内存里一次转换、多次使用。6.3 从 Ampere 到 Hopper指令在变思维不变如果你在不同代的 GPU 上编译同一份 WMMA 代码性能差异可能很大但代码本身通常能跑。问题是当你追求极限性能时必须针对架构调 shapeAmpere 的 FP16mma.sync一次 K 是 16TF32 是 8Hopper 的wgmma一次 K 可以是 32 或 64而且依赖 TMA 的异步数据流根本不是同一套写法。这里有一个实际的兼容性陷阱很多公司代码库为了兼容旧卡只敢用sm_80编译于是在 Hopper 上完全发挥不出性能。遇到这种情况建议按架构拆分代码路径sm_80用 Ampere 优化版sm_90用 TMA wgmma版。编译时用-gencode archcompute_80,codesm_80 -gencode archcompute_90,codesm_90分别生成 cubin运行时按cudaGetDeviceProperties的 major/minor 版本分发。6.4 常见问题速查表症状可能原因排查与解法性能在 10% 原地踏步全局访存未分块检查是否忘了共享内存 TilingNsight 看 L2 命中率提升很小但理论峰值高Bank conflict 严重看 Shared Memory Conflict 指标加 padding 或 swizzle使用 Tensor Core 后反而变慢矩阵太小或 fragment 布局反复拷贝确保数据直接进 fragment避免额外复制结果不正确且有规律的错B 矩阵列/行主序搞错用 MNK16 的最小 case 验证布局算力利用率 70% 上不去缺少双缓冲或占用率不足检查 Stall Reasons上 cp.async换一张卡性能差异巨大编译针对的 SM 版本不对检查 cubin 架构按架构分支编译6.5 一个实操建议从现有高性能库逆向学习最后认真给一个建议不要从零手写终极版 GEMM。CUTLASS 是一个极好的学习对象它的gemm示例从device/到thread/的每个层级都标注了对应的问题域。你可以基于它跑出 90% 以上的性能然后通过修改 tile shape、swizzle 模式、流水线级数来理解每一层的权衡。我自己从 CUTLASS 里学到的最有价值的东西不是哪个函数怎么写而是它如何把“数据排布”和“指令选择”拆成独立的层次让每一层的调优都不牵一发而动全身。总的来说从 1% 到 95% 的九重天跃迁每一步的本质都是在回答同一个问题数据在哪里计算在哪里它们之间的通路是不是最短的。朴素版让数据在全局内存和计算单元之间来回折腾Tensor Core 则要求你把数据喂到硬件最擅长的形状里。等你把 Tiling、双缓冲、Swizzle、精度选型这些环节都打通之后再回头看 CUTLASS 的源码会感觉每一段代码都像老朋友在打招呼——那就说明你把这个领域真正吃透了。这套思路不仅适用于 CUDA你在任何需要手工调度数据的计算平台上写代码都会受益。