ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

CUDA GEMM九重优化:从1%到95%算力利用率实战指南

CUDA GEMM九重优化:从1%到95%算力利用率实战指南 第一次用 CUDA 把 GEMM 跑起来的时候我一度怀疑手上的显卡是假的。Kernel 跑完打开 NCU 一看SM 吞吐率只有 1%。后来我才明白这不是硬件不行是我把 GPU 用成了单核 CPU——每个线程只算一个输出元素数据在全局内存里进进出出Tensor Core 在旁边干瞪眼。这九重优化就是我从 1% 爬到 95% 算力利用率的完整路径。核心手段无非三样分块Tiling、流水线、Tensor Core但真正值钱的是每层优化背后的数学直觉和一排踩过坑才记得住的细节。这篇内容适合所有被 GEMM 性能困扰的 CUDA 开发者尤其建议先跑过 Naive Kernel、又觉得 Profiler 数字很刺眼的人读。1. 先给“1%”算笔账Naive Kernel 离硬件峰值差在哪1.1 GEMM 的本质与最朴素写法GEMM 就是通用矩阵乘法C A × B其中 A 是 M×KB 是 K×NC 是 M×N。深度学习里的全连接层、卷积的隐式 GEMM、Transformer 里的 Attention 投影最后几乎都会被框架层转成一次 GEMM。所以 GEMM 快不快直接决定了你在 GPU 上等一分钟还是等一小时。Naive Kernel 的写法非常直觉化每个线程负责 C 里的一个元素先累加 K 次乘加最后写回。我最早写出来的就是这样一个版本__global__ void gemm_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; }代码没错结果也对。但你要是把矩阵规模放到 4096×4096然后看 Profiler会发现sm__throughput.avg.pct_of_peak_sustained_elapsed只有 1% 左右。注意这里说的不是 10%不是 5%是 1%。很多新手看到这个数字以为是驱动问题实际是算法结构问题。1.2 从算术强度看懂第一层瓶颈要想说清楚为什么只有 1%得先算清楚这张卡能干什么。以我手头一块 Ampere 架构显卡为例FP32 的 CUDA Core 峰值大概在 35 TFLOPS 左右HBM 显存带宽约 936 GB/s。要让算力单元跑满每读一个字节的数据进来至少要伴随 35T / 936G ≈ 38 次浮点运算。这个比值叫算术强度Arithmetic Intensity单位是 FLOP/Byte。现在回头看 Naive Kernel 的循环体每算一次乘加需要读两个 float一共 8 字节产生 2 次浮点运算乘一次、加一次。算术强度是 2 / 8 0.25 FLOP/Byte。和需要的 38 FLOP/Byte 差了 150 倍以上。换句话说这个 Kernel 的性能天花板物理上就被钉死在峰值算力的 1% 到 2%不管你编译器开多高的优化都救不回来。这是 GEMM 优化里最核心的一个认知转变优化 GEMM 不是优化“怎么算”而是优化“怎么让同一份数据被反复算”。数据从显存搬进计算单元是有带宽成本的只有把单次搬运的数据复用次数提上去算力才有机会跑满。2. 第一重跃迁用分块把算术强度拉过饱和线2.1 分块的数学逻辑一份数据反复用分块Tiling解决的就是复用问题。思路很直接不让每个线程去全局内存里零散地取数而是让一个 Block 先把一块 A 子矩阵和一块 B 子矩阵搬到共享内存里然后这块数据在计算中被反复读取。假设一个 Block 负责 64×64 的 C 子矩阵每次从全局内存搬一块 64×16 的 A 分块和一块 16×64 的 B 分块到共享内存。搬进来的数据量是 (64×16 16×64) × 4 8192 字节。这段时间内这个 Block 能做的计算是 64×64×16 次乘加也就是 131072 次浮点运算。算下来算术强度是 131072 / 8192 16 FLOP/Byte比 Naive 高了 64 倍。如果把 tile 放大到 128×128算术强度还能进一步上涨逼近 32 FLOP/Byte 左右。这就是为什么 GEMM 教程里所有人的第一步都是分块分块的本质是增加数据复用让每一字节从显存搬到芯片上的代价都有足够的计算量来摊薄。不过这里有个现实约束共享内存不是无上限的。以我用的卡为例单个 SM 的共享内存上限大约 100KB 上下。你要是把 A 分块和 B 分块都设成 128×128单缓冲就要 128×128×4×2 128KB一个 Block 都塞不下。所以实际 Tiling 参数要在“算术强度”和“共享内存容量”之间做取舍。我常用的起步配置是 BMBN64、BK16先把正确性和性能趋势验证出来再慢慢往大调。2.2 共享内存 Bank Conflict分块方案头号暗坑分块上到共享内存之后你以为利用率能直接飞到 30%结果一测发现只涨到 8% 左右。这时候十有八九是踩了 Bank Conflict。共享内存的硬件结构是 32 个 Bank每个 Bank 每周期只能提供一个 4 字节的数据。如果同一 warp 里 32 个线程同时访问了同一个 Bank 下的不同地址硬件就必须把这些访问串行化也就是产生冲突。一个最简单的灾难场景是这样你声明了一个__shared__ float As[64][32];然后线程 t 读As[t][kk]。同一个 warp 的线程 t 从 0 到 31地址偏移是t * 32 kk。因为32 % 32 0所有线程访问的都是同一个 Bank硬件会把这次访问拆成 32 个周期执行。明明一次能读完硬生生慢了 32 倍。这种问题在 Profiler 里看不太直观你得打开 NCU 的shared__st_bank_conflicts和shared__ld_bank_conflicts这两个计数器。一旦发现指令数高得离谱、共享内存吞吐率明显低于峰值先怀疑 Bank Conflict 大概率没错。2.3 加个 Padding一行代码治好三十二路冲突解决 Bank Conflict 的办法非常朴素给共享内存数组的 pitch 方向多加一个元素。把上面那个声明改成__shared__ float As[64][32 1];地址偏移变成t * 33 kk。因为33 % 32 132 个线程访问时会依次错开到 32 个不同的 Bank冲突瞬间变成零。这个操作在 CUDA 优化里叫 Padding一个字值。我把这个改动写进 Kernel 后利用率从 8% 直接翻到 15% 左右。没有改任何数学逻辑只是给数组多空了一列。所以后面每次写共享内存数组我都习惯性地在最后一个维度加 1然后看一眼 Bank 冲突计数器再决定要不要留着。注意Padding 加在行 pitch 上不是在数据维度上千万别在逻辑索引里把它当真实数据用。3. 隐藏延迟双缓冲与寄存器分块让 SM 不再干等3.1 计算单元快但你喂不饱分块加 Padding 之后利用率到了 15%但距离 95% 还很远。下一个瓶颈藏在“等待”里。GPU 的计算单元执行一条 FMA 只要几个周期但从全局内存发起一次访问到数据真正回到寄存器可能要几百个周期。如果 Kernel 的逻辑是“先加载再计算等下一次加载”那绝大多数时间 SM 都在干等。你得让访存和计算重叠起来——等数据的时候先干别的活。这就是流水线的出发点。共享内存分块之后一个自然的做法是双缓冲准备两个共享内存 Buffer一个 Buffer 的数据在被计算时另一个 Buffer 同时在从全局内存加载下一批数据。两个 Buffer 轮流切换访存延迟就被藏进了计算时间里。3.2 cp.async 把全局访存变成后台任务在 Ampere 和更新的架构上双缓冲已经不需要手工“加载到寄存器再写共享内存”这种老套路了。硬件提供了cp.async指令可以直接让数据从全局内存飞到共享内存不占用寄存器也不阻塞当前线程的后续指令。配合 CUDA 的__pipeline_内建函数主循环可以写成类似这样的骨架__shared__ float As[2][BM][BK]; __shared__ float Bs[2][BK][BN]; for (int k 0; k K; k BK) { int stage (k / BK) 1; __pipeline_memcpy_async(As[stage][0][0], A[blockIdx.y * BM * K k], BM * BK * sizeof(float)); __pipeline_memcpy_async(Bs[stage][0][0], B[k * N blockIdx.x * BN], BK * BN * sizeof(float)); __pipeline_commit(); __pipeline_wait_prior(1); // 使用 As[stage ^ 1] 和 Bs[stage ^ 1] 做计算 }这里每一轮循环都会发起异步拷贝然后 commit 一组再等待前一组完成。__pipeline_wait_prior(1)的意思是“等到最多只剩 1 组还没完成”这样当前这一组的预取就不会影响上一组正在使用的数据。等计算做完新数据也基本到位了访存延迟要么被藏掉要么被压缩到几乎看不出来。刚上手的人很容易把 wait 放到拷贝之前那就等于没做双缓冲。判断方法也很简单用 NCU 看sm__cycles_elapsed和dram__cycles_elapsed的比值如果计算周期里还夹杂着明显的空闲尾巴说明流水线还没排满。3.3 每个线程多算几个元素指令级并行度拉满双缓冲解决的是“等待显存”但还有一个“等待指令”的问题。Naive Kernel 里每个线程只算一个 C 元素K 循环内部是串行依赖上一次乘加的结果是下一次乘加的输入。就算数据都在寄存器里这一串依赖也会让 FMA 单元每隔几个周期就停一下。解决方法是 Register Blocking让每个线程一次算多个 C 元素。比如每个线程维护一个 4×4 的 C 子矩阵也就是 16 个独立的累加器。读一次 A 分块和 B 分块的元素就能同时喂给 16 次乘加。因为累加器之间互不依赖编译器可以把 16 条 FMA 排队发射指令级并行度一下子拉满。这层做完利用率大约能从 15% 到 30% 以上。你会看到 NCU 里sm__inst_executed_pipe_fma的活跃度明显上升而共享内存的读取次数反而下降了——因为一份 A/B 数据在寄存器里被复用了多次不再需要反复从共享内存里取。Register Blocking 的分块尺寸通常用 4×4 或 8×4取决于目标寄存器数量和 tile 大小Ampere 的寄存器文件比较宽裕可以稍微贪一点。4. Tensor Core 思维跃迁你在搬矩阵不是在算元素4.1 一条 HMMA.16816 顶 4096 条 FMA到这里如果你走的还是 FP32 CUDA Core 路径利用率大概会卡在 70% 到 80% 之间。想再往上走必须把 Tensor Core 请出来。Tensor Core 是 GPU 里专门做矩阵乘法的硬件单元。Ampere 架构上它执行一条mma.sync.aligned.m16n8k16指令能一次性完成 16×16×16 的矩阵乘加。16×16×16 是什么意思就是 16 行乘 16 列乘 16 累加深度一共 4096 次乘加运算但只发射一条指令、只读一次相应的寄存器片段。对比一下如果用 CUDA Core 做同样的事需要发射 4096 条 FMA还要为每条 FMA 准备操作数。Tensor Core 不仅指令数少了三个数量级数据搬运的开销也大幅下降。这就是为什么同一个人写的 Kernel从 CUDA Core 切到 Tensor Core 之后算力利用率可以再跳一大截。但 Tensor Core 有它的脾气。它不接受任意形式的数据它要求你把矩阵切成 16×16 或 16×8 这样的块并且这些块的布局必须严格符合硬件约定。所以 Tensor Core 思维的核心变化是你不再想“这个线程算哪个元素”而是想“这个 Warp 负责搬运哪块矩阵”。4.2 wmma Fragment 的布局逻辑不想一上来就手写 PTX 的话可以用 CUDA 的 wmma APIWarp-level Matrix Multiply-Accumulate。#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::load_matrix_sync(a_frag, A_shared, lda); wmma::load_matrix_sync(b_frag, B_shared, ldb); wmma::fill_fragment(c_frag, 0.0f); wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); wmma::store_matrix_sync(C_shared, c_frag, ldc, wmma::mem_row_major);这段代码的逻辑非常清晰把共享内存里的 A 分块和 B 分块加载到 fragment执行一次矩阵乘加把结果写回共享内存。整个 Warp 协同工作你完全不需要关心具体哪个线程持有哪个元素。但这里有个新手必踩的坑不要试图手工索引 fragment 内部的元素不要假设它按行或按列排列。fragment 的元素分布在 warp 里 32 个 lane 的寄存器上布局由硬件决定不同架构可能不一样。wmma 手册里写得很明白但我们实际调代码时还是会有人写a_frag.x[0]就当它是第一个元素然后开始做各种手工变换。正确的做法是只在需要逐元素操作比如 ReLU、乘 scale时遍历num_elements只做“无差别的”逐元素运算绝不做跨 lane 的数据交换。跨 lane 交换要用wmma::permute这类专门函数或者退回共享内存做中转。4.3 B 矩阵转置与对齐两个必须提前处理的细节wmma 对布局的要求有一个隐藏的账要算。在 GEMM 里 B 矩阵的维度是 K×N行主序存储时连续内存方向是 N。但 Tensor Core 的matrix_bfragment 如果按col_major声明它要求内存连续方向符合另一种语义。我在第一次把 wmma 接进主循环时怎么跑怎么错最后发现是 B 的布局不对。解决方案有两个一个是让 B 分块在加载进共享内存之前先做一次全局层面的预转置。也就是说把原来的 K×N 行主序矩阵转成 N×K 行主序矩阵存好主循环里的 B fragment 统一声明成col_major。这样硬件读数据的时候内存流是顺着连续的不需要额外 shuffle。提前转置 B 的成本是 O(KN) 级的显存读写相对于 GEMM 本身的 O(MNK) 计算量基本可以忽略不计。另一个是在共享内存层面搞一个“转置加载 Kernel”或者用ldmatrix指令实现转置。这个更高级但调试难度也更大。还有个更隐蔽的问题是对齐。wmma 和手工 mma 都要求共享内存里的分块地址按至少 16 字节对齐实际使用中通常要求 128 位对齐。如果你在共享内存数组上做了 Paddingtile pitch 一般会自动满足但如果你为了省空间用了紧凑布局就可能碰到misaligned address的运行时错误。报错的时候别慌90% 是 pitch 忘了对齐回看一眼数组定义就明白了。5. 逼近 95% 的调度细节LDMatrix、多 Stage 与 Wave Tail5.1 从 LDMatrix 到手写 mma.syncwmma API 能帮你把逻辑跑通但它不是性能终点。wmma 在编译器生成指令时会做一些保守的寄存器分配和数据搬运安排。如果你追求最后 10% 左右的性能就得把手伸进更底层用ldmatrix加载数据用手写mma.sync发射矩阵乘法。ldmatrix是一条专门为 Tensor Core 设计的加载指令。它可以从共享内存一次加载多个 8×8 的矩阵块到 warp 的寄存器里并且支持加载时做转置。这意味着你给 A 和 B 设计共享内存布局的时候不用再纠结“要不要先手工转置”加载的同时就把硬件需要的 fragment 形态整理好了。手写 mma 的代码大致长这样asm volatile( mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n : f(c[0]), f(c[1]), f(c[2]), f(c[3]) : r(a0), r(a1), r(a2), r(a3), r(b0), r(b1));写成这样之后编程模型就完全变了所有线程要自己维护好哪些寄存器是 A fragment、哪些是 B fragment、哪些是累加器。调试起来非常折磨但收益也是实打实的。我的经验是先用 wmma 把主循环调通并验证正确性再做一个手写 mma 的版本做 A/B 对比别一上来就手写否则遇到性能问题你分不清是布局问题还是流水问题。5.2 多 Stage 缓冲池用共享内存换延迟双缓冲能藏住一部分延迟但要逼近 95% 峰值双缓冲往往不够。原因很简单双缓冲的意思是“我在算第 k 块时预取第 k1 块”。可如果访存延迟本身大于计算一块的时间计算单元还是会等。多 Stage 流水线把 Buffer 数量从 2 扩到 4、8、甚至 16。每一轮主循环可以同时维护多个正在飞行的cp.async操作用__pipeline_wait_prior(N)控制进度。比如一个 8 Stage 流水线当前计算使用的是 stage i背后可能有 i1、i2、i3 三个 stage 的数据正在从显存往共享内存搬运。只要预取深度够大显存往返延迟就被彻底稀释掉了。注意多 Stage 不是免费的。每个 Stage 都要占一份共享内存8 Stage 的共享内存开销是单缓冲的 8 倍。如果你用 128×128 的 tile 去配 8 Stage共享内存直接爆掉。此时需要缩小 BM/BN/BK。这是一道空间换时间的博弈题得根据具体卡上的共享内存上限来调。跑 NCU 的时候重点看sm__warp_issue_stalled_long_scoreboard这个计数器如果这个值还没降到接近 0说明预取深度还不够继续加 Stage 或增大 tile。5.3 Wave Tail矩阵尺寸和 SM 数量的整除游戏主循环已经接近完美以后有一个容易被忽略的外部因素会吃掉最后几个百分点Wave Tail也叫“尾波效应”。GPU 把线程块调度到 SM 上执行一批能同时塞下多少个 Block取决于 SM 数量和每个 SM 的 Block 容量。假设你显卡有 82 个 SM每 SM 一次能跑 2 个 Block那你一个 Wave 能跑 164 个 Block。如果你启动的 Block 总数是 165那最后一个 Block 会单独成为一个 Wave只占用 1 个 SM剩下 163 个 SM 全部空转。等到这个 Block 跑完整个 Kernel 才结束。对于几百微秒的 Kernel这个尾巴可能占掉 10% 到 20% 的时间。解决办法有两类。一类是把 grid 大小配置成“SM 数量 × 每 SM 最大并发 Block 数”的整数倍。这类调整可以在启动 Kernel 前用cudaOccupancyMaxActiveBlocksPerMultiprocessor查询到设备数据然后去凑启动尺寸。另一类是给主循环加一个 grid-stride 的 Shell让每个 Block 循环处理多个 tile也就是把“Block 总数”从矩阵尺寸里解耦出来彻底消灭最后一波的小尾巴。我在调一个小规模 GEMM 时主循环已经优化得不错但利用率始终卡在 84% 左右。后来查启动参数发现 grid 是 83而设备只有 82 个 SM最后一个 Block 单独跑了一波。把 tile 尺寸调小一点、让 grid 变成 164 的整数倍之后利用率直接到了 95%。这个细节特别便宜但最容易被忽略。6. 精度兜底与性能验收让每个百分点都有据可依6.1 FP16/TF32 不是白给的精度Tensor Core 最常见的输入是 FP16累加器是 FP32。FP16 只有 10 位尾数动态范围也不大。在 K 很大的 GEMM 里直接拿 FP16 做输入精度损失是肉眼可见的。尤其当你训练大模型或者做科学计算时误差累积到一定程度真的会爆。应对精度问题有几个常用手段用 TF32。TF32 是 Ampere 及以后架构专门为“接近 FP32 语义”准备的 Tensor Core 输入格式尾数位比 FP16 多精度更接近 FP32。代价是 Tensor Core 处理 TF32 的峰值通常只有 FP16 路径的一半左右。如果业务允许TF32 是个很省事的折中。用 Split 技术。将 FP32 的 A 拆成 A_hi A_lo 两个 FP16B 同理拆成 B_hi B_lo然后做两次或三次 Tensor Core 乘法组合近似恢复 FP32 精度。这个技术实现起来有一定复杂度但可以同时吃到 Tensor Core 的高吞吐和接近 FP32 的精度典型的 3xTF32 方案在不少计算库里有现成实现可以参考。校验误差。无论用哪种方案都要拿 cuBLAS 的结果当参考算绝对误差、相对误差再和你的业务精度要求对齐。FP16 路径的误差在 1e-2 到 1e-3 量级都算常见别等到模型跑崩了才回头查。6.2 九重优化对照表我把自己从 Naive 到最终 Kernel 的性能变化整理成了一张表。这里的利用率是按每一行采用的指令路径峰值归一的也就是说 CUDA Core 路径按 FP32 CUDA Core 峰值算Tensor Core 路径按 Tensor Core 峰值算。这样看优化趋势最直观。层级关键动作路径实测利用率约起点Naive 三层循环CUDA Core FP321%一重64×64 Tiling 入共享内存CUDA Core FP328%二重Padding 修复 Bank ConflictCUDA Core FP3215%三重Register Blocking 每线程 4×4CUDA Core FP3232%四重双缓冲 cp.asyncCUDA Core FP3248%五重多 Stage 流水线CUDA Core FP3266%六重Tensor Core wmmaTensor Core62%七重手写 mma.sync ldmatrix 调布局Tensor Core82%八重多 Stage Wave Quantization 调启动参数Tensor Core88%九重结合精度方案与编译/启动调优后的最终打磨Tensor Core95%注意第五重到第六重的利用率不一定是单调上升的因为换了指令路径之后峰值标尺变了。但整体趋势很清楚每一步的收益都来自“减少等待”和“增加复用”而不是把循环代码改得更花哨。6.3 我建议的调优顺序如果你也是从 1% 开始调我的建议是不要一上来就奔着 Tensor Core 去。先把 CUDA Core 路径下的分块、Padding、Register Blocking、双缓冲跑熟因为这些基本功里的很多坑在 Tensor Core 路径下会以更复杂的形式重新出现一遍。分块都没分明白直接上 wmma你会被布局对齐、fragment 分布、ldmatrix 转置这些问题淹没。每一步优化做完停一下打开 NCU 截图存档。尤其是sm__throughput、dram__throughput、shared__ld_bank_conflicts、sm__warp_issue_stalled_long_scoreboard这几个指标它们能告诉你当前瓶颈在算力、带宽、共享内存冲突还是访存延迟。我后来回看自己调优的截图时发现很多当时靠直觉瞎猜的地方其实 Profiler 早就把答案写在数字里了。最后一层优化往往是无关紧要的小细节__launch_bounds__限制一下寄存器数量、把循环用#pragma unroll展开、把 grid 尺寸凑成 Wave 的整数倍。但就是这些“小细节”叠加起来才能把利用率从 88% 推到 95%。GEMM 优化拼的不是一天灵感是精确到每个 bank、每个 cycle 的记账。
返回列表