1. 从“算力焦虑”说起为什么你的 GPU 总是喂不饱先抛一个很多做 AI 训练和推理的同学都遇到过的现象明明买了一堆顶级加速卡跑大模型训练的时候监控面板上的 GPU 利用率却忽高忽低甚至经常只有百分之三四十。很多人第一反应是“GPU 不够用”于是继续加卡、加机器结果钱花了不少吞吐量却没什么起色。我早几年也踩过这个坑。后来一步步追下去发现问题不在显存容量也不在显存带宽而是出在更底层的计算单元上——Tensor Core。简单说Tensor Core 是 NVIDIA 从 Volta 架构开始引入的专用矩阵运算单元专门用来加速深度学习里最常见的通用矩阵乘法GEMM和卷积运算。它的峰值算力比普通 CUDA Core 高出好几倍理论数据非常好看但实际跑起来绝大多数 Kernel 根本没法让 Tensor Core 吃饱。这篇文章想聊的就是“怎么让 Tensor Core 吃饱”这件事。核心概念有两个一个是 Kernel 效率也就是你写的算子到底有没有把硬件资源用到位另一个是算术强度Arithmetic Intensity也就是每个字节的数据搬运对应多少次浮点运算。这两个概念看起来抽象实际上决定了你在同样一块 GPU 上能把性能做到理论峰值的百分之几。这篇文章适合正在做 AI 系统性能优化、算子开发、模型部署加速的同学也适合那些只是想在训练脚本里调几个参数但想知道“为什么调这些参数有用”的工程师。文章不会堆一堆晦涩的底层手册而是从原理讲到实操最后再给你一份可以直接拿去用的排查清单。2. 先搞懂两件事Kernel 效率和算术强度到底是什么2.1 算术强度计算和搬运的“性价比”算术强度这个概念简单理解就是“运算量”除以“数据搬运量”单位是 FLOP/Byte。它的核心逻辑是计算机做一次浮点运算很快但从内存里把数据搬到计算单元却慢得多。如果一个 Kernel 每搬一个字节的数据只做一两次浮点运算那它大概率是被“数据搬运”卡住了也就是所谓的内存带宽瓶颈。反过来如果每搬一个字节就能做几十甚至上百次浮点运算那计算单元才有机会满负荷工作。举个例子假设你有一个矩阵乘法输入是两个 1024x1024 的矩阵。这个运算的浮点运算量大约是 2 × 1024³也就是约 21 亿次 FLOP。而如果你用传统的 CPU 那样把数据从主存搬到片上再搬回去搬运量至少是几个矩阵的大小假设是 12MB。那算术强度大约是 21亿 / 12MB换算一下大概是 175 FLOP/Byte。这个数值在 GPU 上算是中等偏上理论上是可以把 Tensor Core 喂饱的。但问题在于实际 Kernel 并不会一次把整个矩阵都搬进片上。如果你写的是最简单的三层循环矩阵乘法每次只取矩阵中的一个元素过来算那搬运量和计算量就是同一个量级算术强度会降到个位数性能必然惨不忍睹。注意算术强度不是硬件指标而是算法和 Kernel 实现的属性。同样的矩阵乘法不同写法算术强度可以差一个数量级。2.2 Kernel 效率不只是“GPU 利用率”那个数字很多人看 GPU 利用率只看 nvidia-smi 里那个百分比以为到了 90% 就代表优化到位了。这是个很常见的误区。那个百分比代表的是“GPU 上有没有活干”而不是“计算单元忙不忙”。真正衡量 Kernel 效率的是看它的实际吞吐量throughput达到了理论峰值的百分之多少。比如一张 A100 的数据FP16 Tensor Core 的理论峰值大概是 312 TFLOPS如果你的 Kernel 实际只跑出了 100 TFLOPS那即使 GPU 利用率显示 100%Kernel 效率也只有约 32%。Kernel 效率低的原因通常有三类一是数据搬运来不及即算术强度太低二是并行度不足即没有足够的线程和线程块把 GPU 里的所有 SM流式多处理器占满三是访存模式不友好即数据在显存里的分布和你访问它的顺序不匹配导致缓存命中率极低甚至发生大量的 bank conflict。所以“让 Tensor Core 吃饱”这件事本质上就是同时解决三个问题第一让算术强度足够高第二让并行度足够大第三让访存模式足够友好。三者缺一不可。3. 为什么矩阵乘法是 Tensor Core 的“主场”3.1 Tensor Core 的硬件结构决定了它喜欢什么活Tensor Core 本质上是一个专用的矩阵乘累加单元它在一拍之内可以完成一个 4x4 矩阵与 4x4 矩阵的乘累加操作。以 A100 为例每个 SM 上有 4 个 Tensor Core每个时钟周期可以处理 1024 次 FP16 乘加运算。这和普通的 CUDA Core 相比差距是数量级的CUDA Core 处理矩阵乘法要一条指令一个元素地算而 Tensor Core 是一条指令算一个矩阵块。Tensor Core 之所以快是因为它在硬件上直接把“乘加树”做成了脉动阵列systolic array的形式。数据从阵列的一端流入在内部和权重做乘累加再从另一端流出中间不需要频繁读写寄存器文件因此可以做到极高的数据复用率和吞吐量。但它的前提是你喂给它的数据必须是规整的矩阵分块而且这些分块要能一次性地放在寄存器或共享内存里用一条 wmma 指令或者 mma 指令完成计算。这就解释了为什么所有高性能 GEMM Kernel 都在做同一件事把大矩阵切分成小块让每个线程块负责一小块输出矩阵并且让这一小块数据尽可能长时间地停留在片上。3.2 算术强度如何决定 Tensor Core 的“饱腹感”我们继续用矩阵乘法来推演。假设你要计算 C A × B其中 A 是 M×KB 是 K×NC 是 M×N。如果你采用最朴素的分块方案每个线程块负责计算一个 BM×BN 的输出块。为了算这个块你需要把 A 的 BM×K 那一块和 B 的 K×BN 那一块搬到共享内存里。计算量是 2 × BM × BN × K。搬运量大约是 BM × K K × BN。算术强度就是 2 × BM × BN × K / (BM × K K × BN)。当 BM 和 BN 比较大的时候算术强度近似为 2 × BM × BN / (BM BN)。你会看到算术强度随着分块尺寸的增大而增大。比如 BMBN64 的时候算术强度是 2×64×64/(6464)64 FLOP/Byte。而 BMBN128 的时候算术强度就翻倍到了 128 FLOP/Byte。所以“让 Tensor Core 吃饱”的一个核心手段就是尽量加大每个线程块负责的输出块尺寸。但也不能无限增大因为共享内存和寄存器文件的大小是有限的。你必须在“块大一点”和“别爆显存/寄存器”之间找到一个平衡点。这个平衡点通常要通过搜索和实测来确定NVIDIA 的 CUTLASS 库和 Triton 编译器里都内置了这种自动调优autotuning的机制。4. 实操用 CUDA 和 Triton 把 Tensor Core 喂饱4.1 CUDA 层面从朴素 GEMM 到分块 GEMM先看一个最朴素的 CUDA 矩阵乘法 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; if (row M col N) { 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; } }这个 Kernel 每个线程只算一个输出元素每个输出元素都要从全局内存里读 K 次 A 和 K 次 B。算术强度极低几乎完全被内存带宽限制Tensor Core 根本帮不上忙。换成 CUDA 里用 Tensor Core 的标准姿势要靠 wmma API。核心思路是先在共享内存里准备好 A 和 B 的子块然后调用wmma::load_matrix_sync把数据加载到 fragmentwmma 的寄存器数据结构里再用wmma::mma_sync做矩阵乘累加最后用wmma::store_matrix_sync把结果写回。关键代码框架大概是这样的#include mma.h using namespace nvcuda; // 每个 warp 计算一个 64x64 的输出块使用 16x16x16 的 wmma 指令 __global__ void gemm_wmma(const half* A, const half* B, half* C, int M, int N, int K) { 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; // 初始化累加器为 0 wmma::fill_fragment(c_frag, 0.0f); // 确定当前 warp 负责的输出块位置 int warp_id threadIdx.x / 32; int warp_row (blockIdx.y * blockDim.y warp_id / 4) * 64; int warp_col (blockIdx.x * blockDim.x warp_id % 4) * 64; for (int k 0; k K; k 16) { // 从全局内存加载到 fragment这一步其实会自动走共享内存或直接访存 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); }这段代码里wmma::load_matrix_sync会把数据从全局内存搬到 fragment。为了让搬数据的过程更快通常你还会在更外层加一层“从全局内存到共享内存”的分块拷贝然后再从共享内存加载到 fragment。这样可以让共享内存作为缓存反复复用数据减少全局内存访问次数提高算术强度。我实测过同样的矩阵规模比如 4096×4096从朴素 Kernel 换到带共享内存分块 wmma 的版本性能可以提升 8~15 倍当然具体要看矩阵规模和硬件。核心就是算术强度上去了Tensor Core 才有机会发力。4.2 Triton把“喂饱 Tensor Core”的活交给编译器如果你不想手写 CUDA 的 wmma API 和共享内存分配可以考虑 Triton。Triton 是 OpenAI 开源的 GPU 编程语言/编译器它的关键思想是你不需要管理每个线程和共享内存只需要描述“每个块做什么”编译器会自动帮你生成高性能的 CUDA 代码。一个简单的 Triton GEMM Kernel 长这样import triton import triton.language as tl triton.jit def matmul_kernel( a_ptr, b_ptr, c_ptr, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, ): pid tl.program_id(axis0) num_pid_m tl.cdiv(M, BLOCK_SIZE_M) num_pid_n tl.cdiv(N, BLOCK_SIZE_N) pid_m pid // num_pid_n pid_n pid % num_pid_n offs_m pid_m * BLOCK_SIZE_M tl.arange(0, BLOCK_SIZE_M) offs_n pid_n * BLOCK_SIZE_N tl.arange(0, BLOCK_SIZE_N) offs_k tl.arange(0, BLOCK_SIZE_K) a_ptrs a_ptr (offs_m[:, None] * stride_am offs_k[None, :] * stride_ak) b_ptrs b_ptr (offs_k[:, None] * stride_bk offs_n[None, :] * stride_bn) accumulator tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtypetl.float32) for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): a tl.load(a_ptrs) b tl.load(b_ptrs) accumulator tl.dot(a, b, accumulator) a_ptrs BLOCK_SIZE_K * stride_ak b_ptrs BLOCK_SIZE_K * stride_bk c_ptrs c_ptr (offs_m[:, None] * stride_cm offs_n[None, :] * stride_cn) tl.store(c_ptrs, accumulator)这段代码非常简洁但你注意到没有你只需要控制三个常量BLOCK_SIZE_M、BLOCK_SIZE_N、BLOCK_SIZE_K。这三个常量直接影响算术强度。当你把BLOCK_SIZE_M和BLOCK_SIZE_N从 32 提到 128 时每个线程块负责的输出块变大了数据复用率提高算术强度翻了几倍性能自然就上去了。Triton 内部会自动把tl.dot编译成 Tensor Core 指令你不需要关心它到底用的是 wmma 还是 mma 指令也不需要手动管理共享内存。对于大多数不想深入研究底层 CUDA 的同学Triton 是性价比极高的选择。但 Triton 也不是没有代价。它目前对某些特殊的矩阵运算支持还不够灵活而且自动调优有时候不会收敛到全局最优解。我的经验是先让 Triton 跑通再用它的 autotune 功能搜索一遍分块大小一般能拿到和手写 CUDA 接近的性能差距在 10% 以内。4.3 实操案例如何通过调整 BLOCK_SIZE 让 Kernel 提速 3 倍我曾经在一个推荐系统场景里遇到过一个问题一个 embedding 聚合 矩阵乘法的算子在 A100 上只跑到 40 TFLOPS远低于 FP16 Tensor Core 峰值。我一开始以为是数据布局问题查了半天没头绪。后来用 Triton 改写并且跑了一组不同BLOCK_SIZE的对比实验BLOCK_SIZE_MBLOCK_SIZE_NBLOCK_SIZE_K实测 TFLOPS算术强度估算32323258约 32 FLOP/Byte64643286约 64 FLOP/Byte12812832121约 128 FLOP/Byte1281286498约 96 FLOP/Byte从这组数据能看到几个规律第一加大BLOCK_SIZE_M和BLOCK_SIZE_N带来的收益非常明显因为算术强度翻倍了第二BLOCK_SIZE_K不是越大越好因为 K 方向太大意味着每次循环搬运的数据量也变大如果 K 方向的数据复用不够充分反而会拖慢速度。最终我选定了 128×128×32 的组合速度比最初的 32×32×32 提升了约 3 倍。这个案例最想说明的就是优化 Tensor Core Kernel 的时候先别急着换算法先把分块尺寸走一遍 grid search很多时候性能就上来了。5. 影响 Kernel 效率的隐藏因素数据布局与访存模式5.1 layout 选错Tensor Core 再强也得白干在矩阵乘法里A 和 B 矩阵在内存中的布局直接影响 Tensor Core 的取数效率。最常见的两种布局是 row-major行优先和 col-major列优先。平时用 PyTorch矩阵默认都是 row-major 的也就是tensor.stride()返回的是(K, 1)这样的形式。问题出在当你需要转置一个矩阵的时候。如果你直接把一个 row-major 的矩阵当作 col-major 来读取比如在torch.matmul里对 B 矩阵做转置通常需要先做一次显式的contiguous()拷贝否则 Kernel 在访存时会产生大量的非连续读取缓存命中率暴跌。注意不要小看contiguous()这个调用。它相当于做了一次完整的数据重排内存开销和耗时都不可忽略。如果你在训练循环里每次 forward 都触发隐式拷贝性能损耗可能高达 20%。我在实际项目中见过不少因为 layout 不一致导致 Kernel 效率骤降的案例。排查方法很简单在 PyTorch 里打印x.stride()和x.is_contiguous()确认输入是否连续。如果发现不连续优先在数据处理阶段就统一 layout不要每次 forward 都做转换。5.2 浅谈 bank conflict 和共享内存 padding在 CUDA / Triton 的优化过程中还有一个经常咬人的细节叫 bank conflict。共享内存被划分成 32 个 bank每个 bank 在同一时钟周期只能服务一次访问。如果一个 warp 里的 32 个线程在同一时钟周期访问同一个 bank 的不同地址硬件就会把这些访问串行化导致吞吐量下降好几倍。最典型的触发场景是在从共享内存按列方向读取数据的时候。假设你以行优先存了一个 64×64 的浮点矩阵共享内存的地址布局就是每行 64 个 float连续存放。如果 32 个线程同时读取某一列的连续 32 个元素它们访问的地址恰好每隔 64 个 float 一个都落在了同几个 bank 上就会产生严重的 bank conflict。解决办法也很经典给共享内存加 padding。比如把共享内存声明成 64 列 1 个 float 的 padding也就是每行变成 65 个 float。这样按列访问时地址分布被错开了bank conflict 就会大幅减少。这个技巧在旧版 CUDA 代码里很常见Triton 编译器通常会自动处理但如果你手写 CUDA Kernel最好主动加上。5.3 向量化访存128 位对齐是基本要求除了布局和 bank conflict向量化访存也是一个容易被忽略的细节。现代 GPU 从全局内存加载数据时支持一次加载 128 位16 字节的数据。如果你的数据是 FP32那就意味着一次可以加载 4 个 float如果是 FP16一次可以加载 8 个 half。要让编译器生成向量化加载指令数据地址必须按 16 字节对齐。这也是为什么很多高性能 Kernel 会要求输入矩阵的行大小是 8 的倍数FP16或 4 的倍数FP32。如果你的矩阵宽度不满足对齐要求就只能退回到标量加载访存效率可能下降一半以上。我在做推理优化时经常遇到这种情况模型权重尺寸不是对齐的导致 Kernel 性能上不去。我的处理方法是在权重转换阶段就把它 pad 到对齐大小而不是在 Kernel 里处理边界。这样虽然多占了一点显存但性能和效率都划算得多。6. 工具链与排查怎么快速定位 Kernel 效率低的真正原因6.1 用 Nsight Compute 看关键指标说到排查 Kernel 性能问题NVIDIA Nsight Compute 是绕不开的工具。它可以直接分析你跑出来的每个 Kernel给出许多关键指标。我最常看的有这么几个指标含义红线参考Memory Throughput内存子系统利用率超过 80% 说明接近带宽瓶颈SM BusySM 忙碌百分比低于 60% 说明并行度不足或启动开销大Achieved Occupancy实际达到的占用率低于理论值 50% 需检查寄存器或共享内存限制Arithmetic Intensity实际算术强度低于设备拐点值说明是带宽瓶颈这里提一下“设备拐点”。每张卡的算术强度拐点不一样A100 大约是 200 FLOP/Byte 附近。如果实测算术强度远低于这个值那基本可以断定是内存带宽瓶颈反过来如果远高于这个值但性能仍然上不去那可能是计算单元没喂饱或者并行度不够。我遇到过一个经典案例某个 Kernel 在 Nsight Compute 里显示 Memory Throughput 只有 30%算术强度也不算低但 SM Busy 只有 40%。仔细一看原来是线程块数量太少每个 SM 只分配了两个线程块导致延迟无法被隐藏。后来我把线程块尺寸调小增加了线程块总数SM Busy 一下就上去了性能提高了一倍。6.2 PyTorch/Triton 场景下的快速自检清单如果你暂时不想深入 Nsight Compute也可以从几个简单的角度做快速自检。以下是我每次遇到 Kernel 性能不佳时的排查顺序确认数据类型是否使用了 Tensor Core 支持的类型FP16/BF16/INT8如果一直在用 FP32Tensor Core 根本不参与工作。检查输入张量是否连续特别是转置、切片、拼接等操作后是否触发了非连续存储。用 Triton 或 CUTLASS 的示例 Kernel 跑一遍同样规模的矩阵乘法对比你自己的 Kernel看差距有多大。跑一个torch.matmul做基准PyTorch 自带的 cuBLAS 通常已经优化得不错如果差距巨大说明你的 Kernel 还有很大优化空间。用torch.cuda.synchronize()计时别被异步执行的假象骗了。这一步排查做完一般能定位到 80% 的问题。剩下的 20% 才需要深入到 Nsight Compute 的详细报告和汇编层面。6.3 几个容易误判的“伪性能瓶颈”在优化过程中我还遇到过一些乍看是性能瓶颈、实际上另有他因的情况这里集中说一下。第一种是 GPU 利用率很高但吞吐量很低。这种情况往往是因为 Kernel 在忙等内存或者做了大量无用的计算。GPU 利用率只是“有没有活干”不代表“干的活有价值”。第二种是显存占用高导致并行度下降。比如你把共享内存设置得太大每个 SM 能同时驻留的线程块数量减少占用率上不去。很多人以为共享内存越大越好但实际上要综合考虑它带来的算术强度提升和占用率下降的权衡。第三种是 CPU 侧数据预处理拖了后腿。数据从 CPU 拷贝到 GPU 的过程中如果 pipeline 没做好GPU 会在每个 step 之间空等。这时候 GPU 利用率看起来也不行但问题不在 Kernel 内部而在数据加载流水线。7. 我踩过的几个坑希望你不用再踩7.1 迷信“大分块”反而变慢前面我把 BLOCK_SIZE 说得越大越好但实际应用中有个反例。在某次优化中我把 BLOCK_SIZE_M 和 BLOCK_SIZE_N 从 128 调到 256结果性能不升反降。后来看 Nsight Compute发现是因为每个线程块占用的寄存器太多导致占用率骤降延迟隐藏变差。经验是分块尺寸不是越大越好要考虑寄存器和共享内存的预算。每块多算一点如果代价是多等几个周期反而不划算。动手调参时先固定一个维度另一个维度做 grid search别一上来就搞极端值。7.2 数据类型不匹配导致 Tensor Core “隐形失业”有一次我确定自己用了 FP16 的输入但 Tensor Core 就是没启用。查了好久才发现问题出在累加器类型。Tensor Core 的 mma 指令虽然输入可以是 FP16但累加器必须是 FP32。我在代码里把累加器的 fragment 类型写成了 half导致编译器放弃了 Tensor Core 指令降级成普通 CUDA Core 运算。这个问题在 wmma API 里特别容易犯。Triton 里则是通过tl.dot的输出类型来控制的如果你指定输出为tl.float16编译器也不会用 Tensor Core 做 FP32 累加。所以务必要确认你的累加器是 FP32或者 INT32对应 INT8 输入。7.3 忽略非对齐问题导致访存效率崩盘前面提过向量化访存需要 16 字节对齐。这里我再补一个真实案例我在处理一个宽度为 513 的矩阵时因为没有 padKernel 只能走非向量化路径访存效率直接掉了一半。后来在数据加载阶段把宽度 pad 到 520性能立刻回升。提示PyTorch 中可以用torch.nn.functional.pad或tensor.contiguous()配合切片操作把数据 pad 到对齐尺寸。CUTLASS 和 Triton 都有内置的 mask 机制但最稳的还是数据预处理时就处理好。7.4 混合精度不是单纯换 dtype最后提一个和 Tensor Core 相关但经常被误解的点AMP自动混合精度不是简单地把所有运算切成 FP16。AMP 会把某些对精度敏感的运算如 loss、softmax、部分 normalization保持在 FP32而把矩阵乘法和卷积切成 FP16/BF16。如果你发现开了 AMP 后性能没有提升先检查一下是不是所有 GEMM 都被切到了 FP16还是只有少数几层切了。有时候即使切了 FP16如果 Kernel 的算术强度太低带宽瓶颈依然存在性能提升也不会明显。所以 AMP 只是必要条件不是充分条件最终还是要回到算术强度和 Kernel 效率的分析上来。8. 给不同读者的一句话建议如果你只是在训练脚本里用 PyTorch不做算子开发那么“让 Tensor Core 吃饱”这件事你要做的就是确保输入是连续的、确保数据是混合精度、尽量用大的 batch size 来提高计算复用。你不需要手写 CUDA但理解算术强度能帮你判断到底该增大 batch 还是该优化数据管线。如果你是做推理优化、算子开发或者性能工程的那么建议把 Nsight Compute 的用法吃透然后至少研究一个高性能 GEMM 的实现比如 CUTLASS 里的示例或者 Triton 的官方 tutorial。这些代码读一遍比看十篇性能优化文章都管用。最后再分享一个我个人的实操习惯每次优化一个 Kernel 之前先写一个“理想状态”的假设比如“我预期它能到 80% 的 Tensor Core 峰值”然后反推它需要的算术强度、并行度和访存带宽再针对性地去调。这样不会像无头苍蝇一样乱试参数而是每一步都有据可依。Tensor Core 能不能吃饱很多时候不是玄学而是你用没用对方法的问题。
如何让Tensor Core吃饱？从算术强度到Kernel效率的GPU优化实践
1. 从“算力焦虑”说起为什么你的 GPU 总是喂不饱先抛一个很多做 AI 训练和推理的同学都遇到过的现象明明买了一堆顶级加速卡跑大模型训练的时候监控面板上的 GPU 利用率却忽高忽低甚至经常只有百分之三四十。很多人第一反应是“GPU 不够用”于是继续加卡、加机器结果钱花了不少吞吐量却没什么起色。我早几年也踩过这个坑。后来一步步追下去发现问题不在显存容量也不在显存带宽而是出在更底层的计算单元上——Tensor Core。简单说Tensor Core 是 NVIDIA 从 Volta 架构开始引入的专用矩阵运算单元专门用来加速深度学习里最常见的通用矩阵乘法GEMM和卷积运算。它的峰值算力比普通 CUDA Core 高出好几倍理论数据非常好看但实际跑起来绝大多数 Kernel 根本没法让 Tensor Core 吃饱。这篇文章想聊的就是“怎么让 Tensor Core 吃饱”这件事。核心概念有两个一个是 Kernel 效率也就是你写的算子到底有没有把硬件资源用到位另一个是算术强度Arithmetic Intensity也就是每个字节的数据搬运对应多少次浮点运算。这两个概念看起来抽象实际上决定了你在同样一块 GPU 上能把性能做到理论峰值的百分之几。这篇文章适合正在做 AI 系统性能优化、算子开发、模型部署加速的同学也适合那些只是想在训练脚本里调几个参数但想知道“为什么调这些参数有用”的工程师。文章不会堆一堆晦涩的底层手册而是从原理讲到实操最后再给你一份可以直接拿去用的排查清单。2. 先搞懂两件事Kernel 效率和算术强度到底是什么2.1 算术强度计算和搬运的“性价比”算术强度这个概念简单理解就是“运算量”除以“数据搬运量”单位是 FLOP/Byte。它的核心逻辑是计算机做一次浮点运算很快但从内存里把数据搬到计算单元却慢得多。如果一个 Kernel 每搬一个字节的数据只做一两次浮点运算那它大概率是被“数据搬运”卡住了也就是所谓的内存带宽瓶颈。反过来如果每搬一个字节就能做几十甚至上百次浮点运算那计算单元才有机会满负荷工作。举个例子假设你有一个矩阵乘法输入是两个 1024x1024 的矩阵。这个运算的浮点运算量大约是 2 × 1024³也就是约 21 亿次 FLOP。而如果你用传统的 CPU 那样把数据从主存搬到片上再搬回去搬运量至少是几个矩阵的大小假设是 12MB。那算术强度大约是 21亿 / 12MB换算一下大概是 175 FLOP/Byte。这个数值在 GPU 上算是中等偏上理论上是可以把 Tensor Core 喂饱的。但问题在于实际 Kernel 并不会一次把整个矩阵都搬进片上。如果你写的是最简单的三层循环矩阵乘法每次只取矩阵中的一个元素过来算那搬运量和计算量就是同一个量级算术强度会降到个位数性能必然惨不忍睹。注意算术强度不是硬件指标而是算法和 Kernel 实现的属性。同样的矩阵乘法不同写法算术强度可以差一个数量级。2.2 Kernel 效率不只是“GPU 利用率”那个数字很多人看 GPU 利用率只看 nvidia-smi 里那个百分比以为到了 90% 就代表优化到位了。这是个很常见的误区。那个百分比代表的是“GPU 上有没有活干”而不是“计算单元忙不忙”。真正衡量 Kernel 效率的是看它的实际吞吐量throughput达到了理论峰值的百分之多少。比如一张 A100 的数据FP16 Tensor Core 的理论峰值大概是 312 TFLOPS如果你的 Kernel 实际只跑出了 100 TFLOPS那即使 GPU 利用率显示 100%Kernel 效率也只有约 32%。Kernel 效率低的原因通常有三类一是数据搬运来不及即算术强度太低二是并行度不足即没有足够的线程和线程块把 GPU 里的所有 SM流式多处理器占满三是访存模式不友好即数据在显存里的分布和你访问它的顺序不匹配导致缓存命中率极低甚至发生大量的 bank conflict。所以“让 Tensor Core 吃饱”这件事本质上就是同时解决三个问题第一让算术强度足够高第二让并行度足够大第三让访存模式足够友好。三者缺一不可。3. 为什么矩阵乘法是 Tensor Core 的“主场”3.1 Tensor Core 的硬件结构决定了它喜欢什么活Tensor Core 本质上是一个专用的矩阵乘累加单元它在一拍之内可以完成一个 4x4 矩阵与 4x4 矩阵的乘累加操作。以 A100 为例每个 SM 上有 4 个 Tensor Core每个时钟周期可以处理 1024 次 FP16 乘加运算。这和普通的 CUDA Core 相比差距是数量级的CUDA Core 处理矩阵乘法要一条指令一个元素地算而 Tensor Core 是一条指令算一个矩阵块。Tensor Core 之所以快是因为它在硬件上直接把“乘加树”做成了脉动阵列systolic array的形式。数据从阵列的一端流入在内部和权重做乘累加再从另一端流出中间不需要频繁读写寄存器文件因此可以做到极高的数据复用率和吞吐量。但它的前提是你喂给它的数据必须是规整的矩阵分块而且这些分块要能一次性地放在寄存器或共享内存里用一条 wmma 指令或者 mma 指令完成计算。这就解释了为什么所有高性能 GEMM Kernel 都在做同一件事把大矩阵切分成小块让每个线程块负责一小块输出矩阵并且让这一小块数据尽可能长时间地停留在片上。3.2 算术强度如何决定 Tensor Core 的“饱腹感”我们继续用矩阵乘法来推演。假设你要计算 C A × B其中 A 是 M×KB 是 K×NC 是 M×N。如果你采用最朴素的分块方案每个线程块负责计算一个 BM×BN 的输出块。为了算这个块你需要把 A 的 BM×K 那一块和 B 的 K×BN 那一块搬到共享内存里。计算量是 2 × BM × BN × K。搬运量大约是 BM × K K × BN。算术强度就是 2 × BM × BN × K / (BM × K K × BN)。当 BM 和 BN 比较大的时候算术强度近似为 2 × BM × BN / (BM BN)。你会看到算术强度随着分块尺寸的增大而增大。比如 BMBN64 的时候算术强度是 2×64×64/(6464)64 FLOP/Byte。而 BMBN128 的时候算术强度就翻倍到了 128 FLOP/Byte。所以“让 Tensor Core 吃饱”的一个核心手段就是尽量加大每个线程块负责的输出块尺寸。但也不能无限增大因为共享内存和寄存器文件的大小是有限的。你必须在“块大一点”和“别爆显存/寄存器”之间找到一个平衡点。这个平衡点通常要通过搜索和实测来确定NVIDIA 的 CUTLASS 库和 Triton 编译器里都内置了这种自动调优autotuning的机制。4. 实操用 CUDA 和 Triton 把 Tensor Core 喂饱4.1 CUDA 层面从朴素 GEMM 到分块 GEMM先看一个最朴素的 CUDA 矩阵乘法 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; if (row M col N) { 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; } }这个 Kernel 每个线程只算一个输出元素每个输出元素都要从全局内存里读 K 次 A 和 K 次 B。算术强度极低几乎完全被内存带宽限制Tensor Core 根本帮不上忙。换成 CUDA 里用 Tensor Core 的标准姿势要靠 wmma API。核心思路是先在共享内存里准备好 A 和 B 的子块然后调用wmma::load_matrix_sync把数据加载到 fragmentwmma 的寄存器数据结构里再用wmma::mma_sync做矩阵乘累加最后用wmma::store_matrix_sync把结果写回。关键代码框架大概是这样的#include mma.h using namespace nvcuda; // 每个 warp 计算一个 64x64 的输出块使用 16x16x16 的 wmma 指令 __global__ void gemm_wmma(const half* A, const half* B, half* C, int M, int N, int K) { 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; // 初始化累加器为 0 wmma::fill_fragment(c_frag, 0.0f); // 确定当前 warp 负责的输出块位置 int warp_id threadIdx.x / 32; int warp_row (blockIdx.y * blockDim.y warp_id / 4) * 64; int warp_col (blockIdx.x * blockDim.x warp_id % 4) * 64; for (int k 0; k K; k 16) { // 从全局内存加载到 fragment这一步其实会自动走共享内存或直接访存 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); }这段代码里wmma::load_matrix_sync会把数据从全局内存搬到 fragment。为了让搬数据的过程更快通常你还会在更外层加一层“从全局内存到共享内存”的分块拷贝然后再从共享内存加载到 fragment。这样可以让共享内存作为缓存反复复用数据减少全局内存访问次数提高算术强度。我实测过同样的矩阵规模比如 4096×4096从朴素 Kernel 换到带共享内存分块 wmma 的版本性能可以提升 8~15 倍当然具体要看矩阵规模和硬件。核心就是算术强度上去了Tensor Core 才有机会发力。4.2 Triton把“喂饱 Tensor Core”的活交给编译器如果你不想手写 CUDA 的 wmma API 和共享内存分配可以考虑 Triton。Triton 是 OpenAI 开源的 GPU 编程语言/编译器它的关键思想是你不需要管理每个线程和共享内存只需要描述“每个块做什么”编译器会自动帮你生成高性能的 CUDA 代码。一个简单的 Triton GEMM Kernel 长这样import triton import triton.language as tl triton.jit def matmul_kernel( a_ptr, b_ptr, c_ptr, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, ): pid tl.program_id(axis0) num_pid_m tl.cdiv(M, BLOCK_SIZE_M) num_pid_n tl.cdiv(N, BLOCK_SIZE_N) pid_m pid // num_pid_n pid_n pid % num_pid_n offs_m pid_m * BLOCK_SIZE_M tl.arange(0, BLOCK_SIZE_M) offs_n pid_n * BLOCK_SIZE_N tl.arange(0, BLOCK_SIZE_N) offs_k tl.arange(0, BLOCK_SIZE_K) a_ptrs a_ptr (offs_m[:, None] * stride_am offs_k[None, :] * stride_ak) b_ptrs b_ptr (offs_k[:, None] * stride_bk offs_n[None, :] * stride_bn) accumulator tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtypetl.float32) for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): a tl.load(a_ptrs) b tl.load(b_ptrs) accumulator tl.dot(a, b, accumulator) a_ptrs BLOCK_SIZE_K * stride_ak b_ptrs BLOCK_SIZE_K * stride_bk c_ptrs c_ptr (offs_m[:, None] * stride_cm offs_n[None, :] * stride_cn) tl.store(c_ptrs, accumulator)这段代码非常简洁但你注意到没有你只需要控制三个常量BLOCK_SIZE_M、BLOCK_SIZE_N、BLOCK_SIZE_K。这三个常量直接影响算术强度。当你把BLOCK_SIZE_M和BLOCK_SIZE_N从 32 提到 128 时每个线程块负责的输出块变大了数据复用率提高算术强度翻了几倍性能自然就上去了。Triton 内部会自动把tl.dot编译成 Tensor Core 指令你不需要关心它到底用的是 wmma 还是 mma 指令也不需要手动管理共享内存。对于大多数不想深入研究底层 CUDA 的同学Triton 是性价比极高的选择。但 Triton 也不是没有代价。它目前对某些特殊的矩阵运算支持还不够灵活而且自动调优有时候不会收敛到全局最优解。我的经验是先让 Triton 跑通再用它的 autotune 功能搜索一遍分块大小一般能拿到和手写 CUDA 接近的性能差距在 10% 以内。4.3 实操案例如何通过调整 BLOCK_SIZE 让 Kernel 提速 3 倍我曾经在一个推荐系统场景里遇到过一个问题一个 embedding 聚合 矩阵乘法的算子在 A100 上只跑到 40 TFLOPS远低于 FP16 Tensor Core 峰值。我一开始以为是数据布局问题查了半天没头绪。后来用 Triton 改写并且跑了一组不同BLOCK_SIZE的对比实验BLOCK_SIZE_MBLOCK_SIZE_NBLOCK_SIZE_K实测 TFLOPS算术强度估算32323258约 32 FLOP/Byte64643286约 64 FLOP/Byte12812832121约 128 FLOP/Byte1281286498约 96 FLOP/Byte从这组数据能看到几个规律第一加大BLOCK_SIZE_M和BLOCK_SIZE_N带来的收益非常明显因为算术强度翻倍了第二BLOCK_SIZE_K不是越大越好因为 K 方向太大意味着每次循环搬运的数据量也变大如果 K 方向的数据复用不够充分反而会拖慢速度。最终我选定了 128×128×32 的组合速度比最初的 32×32×32 提升了约 3 倍。这个案例最想说明的就是优化 Tensor Core Kernel 的时候先别急着换算法先把分块尺寸走一遍 grid search很多时候性能就上来了。5. 影响 Kernel 效率的隐藏因素数据布局与访存模式5.1 layout 选错Tensor Core 再强也得白干在矩阵乘法里A 和 B 矩阵在内存中的布局直接影响 Tensor Core 的取数效率。最常见的两种布局是 row-major行优先和 col-major列优先。平时用 PyTorch矩阵默认都是 row-major 的也就是tensor.stride()返回的是(K, 1)这样的形式。问题出在当你需要转置一个矩阵的时候。如果你直接把一个 row-major 的矩阵当作 col-major 来读取比如在torch.matmul里对 B 矩阵做转置通常需要先做一次显式的contiguous()拷贝否则 Kernel 在访存时会产生大量的非连续读取缓存命中率暴跌。注意不要小看contiguous()这个调用。它相当于做了一次完整的数据重排内存开销和耗时都不可忽略。如果你在训练循环里每次 forward 都触发隐式拷贝性能损耗可能高达 20%。我在实际项目中见过不少因为 layout 不一致导致 Kernel 效率骤降的案例。排查方法很简单在 PyTorch 里打印x.stride()和x.is_contiguous()确认输入是否连续。如果发现不连续优先在数据处理阶段就统一 layout不要每次 forward 都做转换。5.2 浅谈 bank conflict 和共享内存 padding在 CUDA / Triton 的优化过程中还有一个经常咬人的细节叫 bank conflict。共享内存被划分成 32 个 bank每个 bank 在同一时钟周期只能服务一次访问。如果一个 warp 里的 32 个线程在同一时钟周期访问同一个 bank 的不同地址硬件就会把这些访问串行化导致吞吐量下降好几倍。最典型的触发场景是在从共享内存按列方向读取数据的时候。假设你以行优先存了一个 64×64 的浮点矩阵共享内存的地址布局就是每行 64 个 float连续存放。如果 32 个线程同时读取某一列的连续 32 个元素它们访问的地址恰好每隔 64 个 float 一个都落在了同几个 bank 上就会产生严重的 bank conflict。解决办法也很经典给共享内存加 padding。比如把共享内存声明成 64 列 1 个 float 的 padding也就是每行变成 65 个 float。这样按列访问时地址分布被错开了bank conflict 就会大幅减少。这个技巧在旧版 CUDA 代码里很常见Triton 编译器通常会自动处理但如果你手写 CUDA Kernel最好主动加上。5.3 向量化访存128 位对齐是基本要求除了布局和 bank conflict向量化访存也是一个容易被忽略的细节。现代 GPU 从全局内存加载数据时支持一次加载 128 位16 字节的数据。如果你的数据是 FP32那就意味着一次可以加载 4 个 float如果是 FP16一次可以加载 8 个 half。要让编译器生成向量化加载指令数据地址必须按 16 字节对齐。这也是为什么很多高性能 Kernel 会要求输入矩阵的行大小是 8 的倍数FP16或 4 的倍数FP32。如果你的矩阵宽度不满足对齐要求就只能退回到标量加载访存效率可能下降一半以上。我在做推理优化时经常遇到这种情况模型权重尺寸不是对齐的导致 Kernel 性能上不去。我的处理方法是在权重转换阶段就把它 pad 到对齐大小而不是在 Kernel 里处理边界。这样虽然多占了一点显存但性能和效率都划算得多。6. 工具链与排查怎么快速定位 Kernel 效率低的真正原因6.1 用 Nsight Compute 看关键指标说到排查 Kernel 性能问题NVIDIA Nsight Compute 是绕不开的工具。它可以直接分析你跑出来的每个 Kernel给出许多关键指标。我最常看的有这么几个指标含义红线参考Memory Throughput内存子系统利用率超过 80% 说明接近带宽瓶颈SM BusySM 忙碌百分比低于 60% 说明并行度不足或启动开销大Achieved Occupancy实际达到的占用率低于理论值 50% 需检查寄存器或共享内存限制Arithmetic Intensity实际算术强度低于设备拐点值说明是带宽瓶颈这里提一下“设备拐点”。每张卡的算术强度拐点不一样A100 大约是 200 FLOP/Byte 附近。如果实测算术强度远低于这个值那基本可以断定是内存带宽瓶颈反过来如果远高于这个值但性能仍然上不去那可能是计算单元没喂饱或者并行度不够。我遇到过一个经典案例某个 Kernel 在 Nsight Compute 里显示 Memory Throughput 只有 30%算术强度也不算低但 SM Busy 只有 40%。仔细一看原来是线程块数量太少每个 SM 只分配了两个线程块导致延迟无法被隐藏。后来我把线程块尺寸调小增加了线程块总数SM Busy 一下就上去了性能提高了一倍。6.2 PyTorch/Triton 场景下的快速自检清单如果你暂时不想深入 Nsight Compute也可以从几个简单的角度做快速自检。以下是我每次遇到 Kernel 性能不佳时的排查顺序确认数据类型是否使用了 Tensor Core 支持的类型FP16/BF16/INT8如果一直在用 FP32Tensor Core 根本不参与工作。检查输入张量是否连续特别是转置、切片、拼接等操作后是否触发了非连续存储。用 Triton 或 CUTLASS 的示例 Kernel 跑一遍同样规模的矩阵乘法对比你自己的 Kernel看差距有多大。跑一个torch.matmul做基准PyTorch 自带的 cuBLAS 通常已经优化得不错如果差距巨大说明你的 Kernel 还有很大优化空间。用torch.cuda.synchronize()计时别被异步执行的假象骗了。这一步排查做完一般能定位到 80% 的问题。剩下的 20% 才需要深入到 Nsight Compute 的详细报告和汇编层面。6.3 几个容易误判的“伪性能瓶颈”在优化过程中我还遇到过一些乍看是性能瓶颈、实际上另有他因的情况这里集中说一下。第一种是 GPU 利用率很高但吞吐量很低。这种情况往往是因为 Kernel 在忙等内存或者做了大量无用的计算。GPU 利用率只是“有没有活干”不代表“干的活有价值”。第二种是显存占用高导致并行度下降。比如你把共享内存设置得太大每个 SM 能同时驻留的线程块数量减少占用率上不去。很多人以为共享内存越大越好但实际上要综合考虑它带来的算术强度提升和占用率下降的权衡。第三种是 CPU 侧数据预处理拖了后腿。数据从 CPU 拷贝到 GPU 的过程中如果 pipeline 没做好GPU 会在每个 step 之间空等。这时候 GPU 利用率看起来也不行但问题不在 Kernel 内部而在数据加载流水线。7. 我踩过的几个坑希望你不用再踩7.1 迷信“大分块”反而变慢前面我把 BLOCK_SIZE 说得越大越好但实际应用中有个反例。在某次优化中我把 BLOCK_SIZE_M 和 BLOCK_SIZE_N 从 128 调到 256结果性能不升反降。后来看 Nsight Compute发现是因为每个线程块占用的寄存器太多导致占用率骤降延迟隐藏变差。经验是分块尺寸不是越大越好要考虑寄存器和共享内存的预算。每块多算一点如果代价是多等几个周期反而不划算。动手调参时先固定一个维度另一个维度做 grid search别一上来就搞极端值。7.2 数据类型不匹配导致 Tensor Core “隐形失业”有一次我确定自己用了 FP16 的输入但 Tensor Core 就是没启用。查了好久才发现问题出在累加器类型。Tensor Core 的 mma 指令虽然输入可以是 FP16但累加器必须是 FP32。我在代码里把累加器的 fragment 类型写成了 half导致编译器放弃了 Tensor Core 指令降级成普通 CUDA Core 运算。这个问题在 wmma API 里特别容易犯。Triton 里则是通过tl.dot的输出类型来控制的如果你指定输出为tl.float16编译器也不会用 Tensor Core 做 FP32 累加。所以务必要确认你的累加器是 FP32或者 INT32对应 INT8 输入。7.3 忽略非对齐问题导致访存效率崩盘前面提过向量化访存需要 16 字节对齐。这里我再补一个真实案例我在处理一个宽度为 513 的矩阵时因为没有 padKernel 只能走非向量化路径访存效率直接掉了一半。后来在数据加载阶段把宽度 pad 到 520性能立刻回升。提示PyTorch 中可以用torch.nn.functional.pad或tensor.contiguous()配合切片操作把数据 pad 到对齐尺寸。CUTLASS 和 Triton 都有内置的 mask 机制但最稳的还是数据预处理时就处理好。7.4 混合精度不是单纯换 dtype最后提一个和 Tensor Core 相关但经常被误解的点AMP自动混合精度不是简单地把所有运算切成 FP16。AMP 会把某些对精度敏感的运算如 loss、softmax、部分 normalization保持在 FP32而把矩阵乘法和卷积切成 FP16/BF16。如果你发现开了 AMP 后性能没有提升先检查一下是不是所有 GEMM 都被切到了 FP16还是只有少数几层切了。有时候即使切了 FP16如果 Kernel 的算术强度太低带宽瓶颈依然存在性能提升也不会明显。所以 AMP 只是必要条件不是充分条件最终还是要回到算术强度和 Kernel 效率的分析上来。8. 给不同读者的一句话建议如果你只是在训练脚本里用 PyTorch不做算子开发那么“让 Tensor Core 吃饱”这件事你要做的就是确保输入是连续的、确保数据是混合精度、尽量用大的 batch size 来提高计算复用。你不需要手写 CUDA但理解算术强度能帮你判断到底该增大 batch 还是该优化数据管线。如果你是做推理优化、算子开发或者性能工程的那么建议把 Nsight Compute 的用法吃透然后至少研究一个高性能 GEMM 的实现比如 CUTLASS 里的示例或者 Triton 的官方 tutorial。这些代码读一遍比看十篇性能优化文章都管用。最后再分享一个我个人的实操习惯每次优化一个 Kernel 之前先写一个“理想状态”的假设比如“我预期它能到 80% 的 Tensor Core 峰值”然后反推它需要的算术强度、并行度和访存带宽再针对性地去调。这样不会像无头苍蝇一样乱试参数而是每一步都有据可依。Tensor Core 能不能吃饱很多时候不是玄学而是你用没用对方法的问题。