1. 别急着写kernel先理解“取行”和“取列”到底在说什么1.1 行主序矩阵的物理存储与访存方向我做CUDA优化这些年越来越觉得GEMM是所有并行编程课程里最适合用来练内功的题目。Stanford CS217里有一整套关于矩阵乘法的讨论但我印象最深的不是某个trick而是教授反复强调的一个基础问题你的矩阵到底是怎么放进内存的。行主序下A[i][j]的地址是i * ld j也就是说你访问A[i][0]、A[i][1]、A[i][2]时地址是连续递增的但访问A[0][j]、A[1][j]、A[2][j]时地址每跳一次就跨过一个完整的行宽。这个“取行”和“取列”的差别看起来只是方向问题实际上决定了GEMM后续所有优化手段的走向。全局内存里一个warp的32个线程去读32个连续float只需要一个128字节事务如果这32个线程读的是同一列的不同行地址分散在32个不同的行区间里可能就要产生几十个访存事务带宽利用率直接掉一个数量级。共享内存更敏感它有32个bank同一指令里不同线程访问同一个bank的不同地址就会发生bank conflict轻则多几个周期重则把一次load变成32次串行操作。所以“取行还是取列”从来不是代码风格问题而是算法复杂度级别的取舍。用生活里的例子说取行就像在书架上从左到右拿一排书手顺着扫过去就行取列就像每次从不同层拿第3本每拿一本都要踮脚或弯腰。程序虽然不会抱怨“累”但profiler会诚实地显示带宽利用率上不去只是很多同学看不出来是因为访存方向错了。1.2 向量内积一个输出元素牵涉两个方向矩阵乘法的教科书定义就是向量内积C[i][j]等于A[i][:]和B[:][j]这两个向量做点乘。按这个定义写kernel每个线程负责一个输出元素内层循环对k做累加。这个写法最直观但有两个天生的问题。第一个问题是行主序下A[i][k]是连续取行B[k][j]是跨步取列。虽然你可以通过调整线程映射让一个warp里的线程分别处理相邻的j这样同一时刻warp访问的是B的第k行里一段连续列看起来合并了但你的A访问变成了同一地址的广播。说得直白一点内积公式里总有一个方向是逆着物理存储方向走的你只能选择把惩罚放在哪个矩阵上没法同时让两个矩阵都按照连续方向访问。第二个问题是数据复用率低。A[i][k]会被所有需要第i行输出的线程重复读B[k][j]会被所有需要第j列输出的线程重复读。如果不加共享内存tiling每个输入元素从全局内存被读的次数会随着矩阵规模线性放大访存流量直接到了2 * M * N * K的水平。矩阵稍微大一点这个kernel就完全是带宽瓶颈GPU的浮点算力根本用不上。1.3 向量外积把“点乘”改成“秩1更新”换个角度看矩阵乘法也可以写成C Σ_{k} A[:,k] ⊗ B[k,:]。这就是向量外积视角把A的第k列看成一个向量u把B的第k行看成一个向量v两者做外积得到一个M×N的秩1矩阵把所有k的秩1更新累加得到C。这个写法不会改变数学结果但它把“取行”和“取列”分开了B[k,:]天然是连续的一行适合做合并访问和向量化加载A[:,k]虽然在全局内存里是跨步的列但我们可以先把它搬进共享内存或寄存器再在计算阶段按需要的方向去读。也就是说外积视角让数据搬运和计算可以分开设计而不是像内积那样把跨步访问塞进最内层循环。外积更核心的优势是复用。做秩1更新时A[:,k]这一列可以被这一列对应行上的所有列复用B[k,:]这一行可以被这一行对应列上的所有行复用。如果我们在一个block里同时加载了A的一小块和B的一小块那么每个从全局内存读进来的元素都可能参与很多次乘加。这个复用程度直接决定了你能把全局内存带宽需求压到多低也是GEMM优化能做到多少TFLOPs的关键。2. 内积与外积GEMM并行化结构的分叉点2.1 内积结构线程数与输出元素一一对应最朴素的CUDA kernel就是每个线程算一个C[i][j]Thread与输出元素一一对应。这种结构的最大优点是线程之间完全独立没有跨线程归约不需要__syncthreads()写出来的代码几乎不会有同步错误非常适合做正确性baseline。但它的缺点也非常明显。每个线程只有一个累加器但内层循环每步要从A和B各读一个float然后做一次乘加。你可以把它想象成一条流水线每个线程像一个小工每天只生产一个零件却要跑遍整个仓库去提两种原料。线程数量虽然多但单线程算力利用率很低访存次数和乘加次数的比值高得惊人。另一个问题是向量化困难。A行连续理论上可以一次读4个float但B列跨步一次读4个B[k][j]就是4次独立访存编译器很难自动优化。如果你写内积循环建议不要期待编译器能帮你生成高效的float4指令它通常会被跨步访问拖住最终只能生成一堆标量load。2.2 外积结构粗粒度累加与寄存器复用外积结构在CUDA里的典型形态是寄存器分块也就是register tiling。每个线程不再只算一个C元素而是负责C矩阵里一个TM×TN的子块。每个线程需要维护TM×TN个累加寄存器内层循环每遍历一个k就从共享内存或者寄存器里读入TM个A元素和TN个B元素然后做TM×TN次乘加更新整个子块。这才是我理解的“外积路径”一次读入一小段A列和一小段B行做一次秩1更新。相比内积结构A和B的每个元素会被更多输出重复使用内存访问次数被摊薄了。假设TM4, TN4每个线程一次内层迭代做16次乘加但只需要分别读4个A元素和4个B元素计算访存比立刻变得好看很多。寄存器tiling还能让访存方向更可控。你可以把加载tile和计算tile拆开加载阶段从全局内存按连续方向搬运数据计算阶段再从共享内存按固定方向读取。想优化的取行就连续读想优化的取列就通过padding、转置或者swizzle来把bank conflict压下去。所以在实际高性能GEMM里几乎看不到一个线程只算一个输出元素的版本大家都在做各种尺寸的register tiling。2.3 什么时候该用谁K维度、矩阵规模与硬件限制我把内积和外积结构做一个对比方便你看着选线程输出数内积为1外积为TM×TN线程总数内积需要M×N外积只需要(M×N)/(TM×TN)累加寄存器内积1个外积TM×TN个访存复用内积低外积高向量化友好度内积的B列跨步很难向量化外积连续读行寄存器复用更友好bank conflict风险内积更容易在取列方向踩坑外积可以通过padding/swizzle控制典型场景内积适合小矩阵、正确性验证和动态形状外积适合大矩阵、追求峰值性能的训练/推理GEMM看起来外积全面占优但实际工程里也不是无脑选。外积要占用大量寄存器比如TM4, TN4就需要16个累加器再加上中间变量一个线程轻松用到40~80个寄存器在部分GPU上会压低occupancy。矩阵规模特别小或者K特别短的时候tiling带来的搬运开销可能比计算收益还大这时候简单内积反而更稳。真正的产品级GEMM通常会用TM/TN的各种组合适配不同的矩阵形状再叠加double buffering、ldmatrix、swizzle等高级手段。3. 一步步把kernel从内积改成外积3.1 完整代码骨架与编译环境我用CUDA 12.x编译目标用-archsm_80A100或者sm_86/sm_89Ampere/Ada命令大概是nvcc -O3 -archsm_80 -o gemm gemm.cu矩阵约定为行主序先写一个计时函数后面所有版本都用同一套计时逻辑cudaEvent_t start, stop; cudaEventCreate(start); cudaEventCreate(stop); cudaEventRecord(start); gemm_kernelgrid, block(A, B, C, M, N, K); cudaEventRecord(stop); cudaEventSynchronize(stop); float ms 0.0f; cudaEventElapsedTime(ms, start, stop); double gflops 2.0 * M * N * K / (ms * 1e6);后面每个版本都建议先用小矩阵跟CPU计算做对比确认没算错再谈性能。优化最忌讳一开始就上一堆tiling、double buffering结果错了都不知道是哪一步引入的。3.2 版本一全局内存内积先建立一个正确的baseline第一版就是最简单直接的kernel__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; 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; } }启动配置一般用dim3 block(16, 16)dim3 grid((N 15) / 16, (M 15) / 16)。这个版本性能很差但它是后续所有版本的正确性参照。如果你连这个都写不对后面加了共享内存更没法debug。我每次做新优化都会先把naive版本跑一遍确定输出矩阵没问题再开始动结构。如果你用Nsight Compute看这个kernel会看到全局内存的访问事务特别多。原因就是前面说的每个线程都要从A和B读大量元素而且中间有跨步取列访存流水线一直在等数据。这个阶段的瓶颈是内存带宽不是计算。3.3 版本二共享内存tiling观察取列带来的bank conflict第二版加共享内存tiling。假设块大小是16×16代码大概是#define BLOCK_SIZE 16 __global__ void gemm_tiled(const float* A, const float* B, float* C, int M, int N, int K) { __shared__ float As[BLOCK_SIZE][BLOCK_SIZE]; __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE]; int bx blockIdx.x, by blockIdx.y; int tx threadIdx.x, ty threadIdx.y; int row by * BLOCK_SIZE ty; int col bx * BLOCK_SIZE tx; float sum 0.0f; for (int kk 0; kk K; kk BLOCK_SIZE) { As[ty][tx] A[row * K kk tx]; Bs[ty][tx] B[(kk ty) * N col]; __syncthreads(); for (int k 0; k BLOCK_SIZE; k) { sum As[ty][k] * Bs[k][tx]; } __syncthreads(); } C[row * N col] sum; }这个版本已经比naive快很多因为A和B的tile都只从全局内存读一次复用到了块内所有输出。但如果你仔细看计算阶段的访问As[ty][k]是固定k列、不同ty访问不同行在共享内存里属于“取列”访问Bs[k][tx]是固定k行、不同tx访问不同列属于“取行”访问。取行的方向通常没什么问题取列的方向在不同线程组织下容易触发bank conflict。这里我不给你一个“肯定冲突”的结论因为bank conflict和线程布局强相关。最简单的自查方法是把共享内存的地址线性化然后看同一条load指令里32个线程分别落在哪些bank。比如As[ty][k]如果BLOCK_SIZE32ty从0到31地址偏移是ty * 32 k对32取模之后全部等于k这就会产生32路bank conflict。如果BLOCK_SIZE16一个warp里有两个不同的ty冲突就变成2路。你只需要记住规则不同线程在同一时刻访问同一个bank的不同地址就是冲突。加padding是解决这类问题的通用手段把As[BLOCK_SIZE][BLOCK_SIZE1]让每一行多出一个float占位列的方向就会被错开很多冲突就消失了。3.4 版本三外积风格的寄存器tiling让取行取列各归其位第三版开始上register tiling也就是外积结构。每个线程负责一个TM×TN子块内层循环反复做秩1更新。我常用的一个演示配置是TILE_M64, TILE_N64, TILE_K16, TM4, TN4这样block里线程数是(64/4) * (64/4) 16 * 16 256比较均衡。#define TILE_M 64 #define TILE_N 64 #define TILE_K 16 #define TM 4 #define TN 4 // blockDim.x TILE_N / TN 16 // blockDim.y TILE_M / TM 16 __global__ void gemm_outer(const float* A, const float* B, float* C, int M, int N, int K) { __shared__ float As[TILE_M][TILE_K 1]; __shared__ float Bs[TILE_K][TILE_N 1]; int tx threadIdx.x; int ty threadIdx.y; int rowBase blockIdx.y * TILE_M ty * TM; int colBase blockIdx.x * TILE_N tx * TN; float rC[TM][TN] {0.0f}; for (int kk 0; kk K; kk TILE_K) { for (int idx tx ty * blockDim.x; idx TILE_M * TILE_K; idx blockDim.x * blockDim.y) { int r idx / TILE_K; int c idx % TILE_K; As[r][c] A[(blockIdx.y * TILE_M r) * K kk c]; } for (int idx tx ty * blockDim.x; idx TILE_K * TILE_N; idx blockDim.x * blockDim.y) { int r idx / TILE_N; int c idx % TILE_N; Bs[r][c] B[(kk r) * N blockIdx.x * TILE_N c]; } __syncthreads(); for (int k 0; k TILE_K; k) { float rA[TM]; float rB[TN]; #pragma unroll for (int i 0; i TM; i) rA[i] As[ty * TM i][k]; #pragma unroll for (int j 0; j TN; j) rB[j] Bs[k][tx * TN j]; #pragma unroll for (int i 0; i TM; i) { #pragma unroll for (int j 0; j TN; j) { rC[i][j] rA[i] * rB[j]; } } } __syncthreads(); } for (int i 0; i TM; i) { float* cptr C (rowBase i) * N colBase; for (int j 0; j TN; j) { cptr[j] rC[i][j]; } } }这个kernel的核心变化在计算循环里。每次k迭代我先从共享内存把TM个A元素和TN个B元素读进寄存器然后一次性做TM×TN次乘加。这其实就是外积视角里的秩1更新用rA这一小段列和rB这一小段行去更新rC这个子块。因为数据在寄存器里被反复复用共享内存的访问次数大大减少性能会明显高于前面版本。代码里我给共享内存加了paddingAs第二维是TILE_K1Bs第二维是TILE_N1。这样做的目的是让取列方向的访问在多数布局下尽量分散到不同bank。需要说明的是这个代码是教学级别的register tiling不是CUTLASS那种极致优化版。比如Bs的rB[j]读取在TN4的情况下可能还有少量bank冲突工程上会继续用swizzle或者ldmatrix来彻底解决。但作为CS217课程作业级别的展示它已经足够说明“外积寄存器复用”带来的收益。如果你用的GPU比较新还可以把全局内存访问改成float4向量化前提是矩阵leading dimension和地址都要对齐到16字节。这一步对带宽利用率提升非常明显但会增加代码复杂度尤其处理边界的时候很麻烦。3.5 性能对比与调参要点我把几个版本的性能趋势列一下注意不要当成固定数值因为不同GPU和矩阵尺寸差异很大版本访存模式性能相对naivenaive内积全局内存反复读B列跨步1xbaseline共享内存tiling每个tile只读一次但计算阶段有取列方向bank conflict3~10倍register tiling外积寄存器复用取行取列分开控制再提升一个数量级调参的时候我一般先固定TILE_K从16开始因为在大多数GPU上16×64的共享内存尺寸不会爆容量。然后调TM×TN4×4是比较安全的起点如果寄存器够用可以试8×4或者8×8。TILE_M和TILE_N不要无脑放大因为共享内存和寄存器的压力会同时上升最终occupancy掉下去性能反而下降。调参没有银弹最好的方法是把矩阵规模固定逐组参数跑一遍再用Nsight Compute看瓶颈是内存还是计算。4. 调试实录踩过的坑与排查技巧4.1 共享内存bank conflict的快速定位我自己调试时最常遇到的坑就是bank conflict。你可以用Nsight Compute跑一下重点看这几个指标shared memory的load/store bank conflict次数、全局内存的吞吐量、achieved occupancy。如果怀疑某个load有bank conflict最直接的办法是在纸上手算一次。取一个warp的32个线程写出它们在同一指令下访问的共享内存地址然后对32取模看落到同一个bank的地址是否互不相同。注意多个线程访问同一个地址属于broadcast不算冲突但访问同一个bank的不同地址就算冲突。很多教程把“同一bank”和“同一地址”混为一谈排查时候容易判断失误。解决bank conflict的常用手段有给共享内存行加padding、把矩阵在共享内存里转置、使用swizzle模式重排地址。对初学者来说padding最简单通常一次就能解决大部分冲突。不过加padding会浪费一点共享内存容量所以tile越大越需要仔细计算是否值得。4.2 float4对齐、矩阵leading dimension和边界处理用float4做向量化是GEMM提升带宽利用率的重要一步但它要求地址16字节对齐。cudaMalloc返回的显存地址通常是256字节对齐问题不大坑往往出现在矩阵的leading dimension上。如果你的矩阵宽度不是4的倍数那么每一行的起始地址可能会错开导致float4访问越界。解决方法有两个一是用cudaMallocPitch对齐行宽二是在分配矩阵时手动往每一行末尾多padding几个float让宽度凑成16字节对齐。注意这种padding会改变矩阵在内存里的布局所以传给kernel的leading dimension也要跟着改。我在实际项目里更喜欢用ld参数单独传给kernel而不是直接用N这样代码更灵活边界处理也更清晰。矩阵的M、N、K不是tile整数倍时主流做法是让tile稍微越界读然后计算阶段用mask把越界位置过滤掉。这个操作很容易出错尤其边界块的写回。我的习惯是先在host端把矩阵padding到tile整数倍再做性能测试等基本逻辑都稳定了再去处理复杂的边界分支不要一开始就混合在一起。4.3 寄存器压力、occupancy与launch boundsregister tiling的累加器数量直接决定寄存器占用。TM4, TN4有16个累加器看起来不多但加上指针计算、索引变量、rA/rB一个线程用到40~64个寄存器很常见。如果GPU每SM的寄存器总数有限线程块能跑多少个block就会变少occupancy下降反而可能掩盖掉计算效率的提升。我的经验是在GEMM这种计算密集型kernel里适当地降低occupancy换取更多寄存器的指令级并行通常是划算的。但你得用__launch_bounds__或者maxrregcount控制好上限不能让它肆意占用导致kernel无法启动。比如一个block有256线程每个线程用到64个寄存器那block一共就需要16384个寄存器在每SM 64K寄存器的GPU上刚好可以放4个block这就是一个比较舒服的配置。调寄存器最忌讳凭空猜测直接用--ptxas-options-v看编译报告或者用cudaOccupancyMaxPotentialBlockSize这类API算一下。很多时候性能瓶颈不是register pressure而是你为了压低occupancy强行减少了寄存器数量结果spill到本地内存性能反而更烂。4.4 我每次写GEMM前必做的3个检查第一先在纸上画出A、B、C的tile内存布局把每个访问方向标清楚哪些是取行哪些是取列。这一步看起来麻烦但能避免一半以上的方向性错误。第二模拟一个warp的第一次内存访问手算共享内存bank index确认没有明显的bank conflict。第三跑之前先写naive版本做正确性校验跑完之后立刻用Nsight Compute看访存指标不要只盯着总耗时。说句实在话CUDA GEMM优化这个题目网上教程特别多但很多人抄了tiling、抄了padding性能还是上不去。问题往往出在最基本的方向判断上你的算法结构到底是内积还是外积每一步访问到底是取行还是取列。把这些底层逻辑想清楚再去看那些高级kernel代码你会突然明白为什么人家要这样排线程、这样存共享内存、这样写累加循环。这个“想清楚”的过程比单纯背几个优化技巧值钱得多。