寄存器Tiling：决定GEMM与FlashAttention性能的AI算子优化核心

做AI Infra这几年来我面试别人或者被人追问的时候只要话题绕到算子优化最后几乎都会落在寄存器 tiling 这个点上。它不像共享内存 tiling 那样在教材里占着完整章节也没有太多现成模板可以抄但偏偏是它决定了 GEMM、FlashAttention、MoE 里那些核心 kernel 的性能能到哪一格。寄存器 tiling简单说就是把计算要用的数据切块之后直接塞进寄存器文件里让同一个数据在寄存器里被反复拉出来乘加而不是每一轮都从共享内存甚至全局内存里重新捞。它保证计算单元在每一个时钟周期都能以整块硬件能给的最高带宽拿到操作数也是 Roofline 模型那张图里“ridge point”真正被踩实的地方。这篇文章是“AI Infra 每日一问”系列第 6 天。我想把寄存器 tiling 在不同硬件架构上的形态摊开讲一遍NVIDIA 的 SIMT 线程模型里它怎么组织AMD CDNA 的 wavefront 里它长什么样昇腾这类 AI ASIC 的 Cube 单元里它变成了什么再到 TPU 这种脉动阵列为什么干脆绕开了这条路。适合看的人做算子库或者 kernel 的开发者、推理引擎里做图优化和内存规划的同学、想在 Triton 或者 CUDA 里把 GEMM 跑得更快的人。如果你只是听过 tiling 这个词想彻底搞明白寄存器这一层为什么重要这篇文章也能当一份入门地图来看。1. 先把概念说透寄存器 tiling 到底在 tile 什么1.1 三层 tilingblock、shared memory、register 各切一刀先拿 GEMM 举例。一个 M4096、N4096、K4096 的矩阵乘数据量在几十 MB 量级任何片上存储都放不下。所以大家常说的 tiling 实际上分了三层每一层解决的是不同问题。第一层是线程块级 tiling把输出矩阵切成若干 block tile 分给不同的 SM解决的是并行度和任务分配。第二层是共享内存级 tiling在 block 内部把 A、B 的切片搬进 shared memory解决的是全局内存带宽太慢、数据重复从 HBM 读的问题。第三层才是寄存器级 tiling共享内存里的数据再按每个线程应该负责的输出小块切到线程私有的寄存器里然后做乘加。我把这三层整理成一张表方便对照看层级切分单位核心解决典型粒度block tiling线程块 / SM并行度、任务分配128x128shared memory tilingblock 内所有线程全局内存带宽、跨线程复用128x32register tiling单个线程片上最高带宽、零额外延迟复用8x8 / 16x8寄存器 tiling 关心的就是表里第三行。它在整个优化链路里处在最靠近计算单元的位置往上一层的数据搬运可以被它掩盖往下一层的计算能力能不能吃饱也要靠它。1.2 寄存器凭什么能做 tiling 的主力用个不太严谨但好懂的生活类比寄存器就像工位上手边那排工具共享内存是工位旁边的抽屉全局内存是整个楼层尽头的仓库。你写 report 的时候最称手的思路一定是把最常用的资料放在手边而不是每翻一次就跑到仓库去搬一趟。寄存器 tiling 就是把这个策略用到极致数据进寄存器一次留在那里被多条计算指令反复使用直到这个 tile 算完。硬件层面寄存器的地位更特殊。以 A100 为例每个 SM 的寄存器文件大概 256 KB分成 4 个处理块每个处理块内部每个时钟周期都要能同时喂给成千上万个操作数的读取。这种读取带宽是 shared memory 和 L1 完全够不着的更别提 HBM 了。所以对于计算密度很高的算子有没有把数据正确放进寄存器直接决定运算单元每个周期是吃饱还是在干等。再算一笔账你会更直观地感受到寄存器复用的威力。GEMM 里每算一个输出元素朴素做法要从 K 维度读 A 的一行和 B 的一列总共 2K 个数。如果每个线程持有一个 8x8 的输出 tile那么 A 的行片段可以被 8 个输出列复用B 的列片段可以被 8 个输出行复用需要的访存次数从 2K 降到 K/8 K/8也就是 K/4直接省掉 8 倍。这个复用就发生在寄存器里不占共享内存的带宽也不占 LSU 的发射槽位。1.3 寄存器 tiling 的硬边界寄存器 tiling 不是无限扩张的。寄存器文件总量有限分给每个线程多了能同时驻留的线程数就少占用率就掉。这里面有三条硬约束做 kernel 的人必须心里有数。第一单线程寄存器数量有上限。NVIDIA 上每个线程最多 255 个 32 位寄存器超过之后编译器会把多余变量“溢出”到 local memory其实就是在全局内存里划一块私有区域性能和直接访问全局内存差不多属于断崖式下跌。第二寄存器是线程私有的线程之间不能直接访问对方的寄存器。跨线程的数据复用只能靠 warp shuffle 或者共享内存中转所以寄存器 tiling 只解决线程内部的复用。第三寄存器 tiling 尺寸和占用率永远在打架。A100 一个 SM 最多驻留 2048 个线程如果每个线程用 128 个寄存器最多同时放 1024 线程占用率 50%要是某个 kernel 把每线程寄存器顶到 256 个那就只能放 512 线程占用率掉到 25%。这就引出一个关键判断寄存器 tile 不是越大越好而是要在“单线程计算密度”和“足够多并发的线程”之间找平衡点。后面第 2 章我会用具体例子说明这个平衡怎么找。2. NVIDIA 系寄存器 tiling 的“标准答案”2.1 SIMT 语义下的寄存器 tiling 形态NVIDIA GPU 是 SIMT 模型单指令多线程。一个 warp 有 32 个线程每个线程都有自己独立的寄存器但同一时刻这 32 个线程执行同一条指令。这种模型下寄存器 tiling 的天然单位是“每个线程一小块输出”。经典 GEMM 排布是这样的warp 里 32 个线程组织成 4x8 的线程阵列每个线程负责一个 8x8 的 C tile于是这个 warp 整体覆盖 32x64 的输出块。A 和 B 的片段也按线程各自加载到寄存器里。一个线程的寄存器预算大致是这样8x8 的 C 累加器64 个 float 寄存器A fragment8 个寄存器B fragment8 个寄存器基地址、循环计数、临时变量大约 20 到 40 个合起来每线程大约 100 到 120 个寄存器正好落在 128 以内。如果改成每个线程负责 16x16 的 C tile光累加器就要 256 个寄存器直接爆掉单线程 255 的上限必然 spill性能惨不忍睹。所以 8x8 或者 16x8 这种尺寸在 Ampere 之前几乎是黄金选择不是拍脑袋定出来的。伪码层面这种 tiling 写出来大致是下面这个样子// 每个线程负责 8x8 的 C 子块 float c[8][8] {0}; // 在 K 维上循环每次取 A 的 8 行片段和 B 的 8 列片段 for (int k0 0; k0 K; k0 8) { float a_frag[8]; // 当前线程要用的 A 行片段 float b_frag[8]; // 当前线程要用的 B 列片段 load_a_fragment(a_frag, ...); load_b_fragment(b_frag, ...); // 寄存器里的 8x8x8 乘加 for (int i 0; i 8; i) { for (int j 0; j 8; j) { for (int kk 0; kk 8; kk) { c[i][j] a_frag[kk] * b_frag[kk]; } } } }注意这只是一个概念示意。真实 CUTLASS 这类库里fragment 的布局、循环展开、指令调度都是精心安排的因为编译器默认生成的代码往往达不到最佳状态。但核心思想就是这一段把数据切成小块锁进寄存器拼命乘加。2.2 从手工 fragment 到 WMMA 指令从 Turing 架构开始NVIDIA 引入了硬件矩阵指令也就是 HMMA/WMMA。这类指令把“寄存器 tiling 的乘加”直接硬件化了一条指令完成一个 16x8x16 或者 16x16x16 的矩阵乘加不需要你再写三层循环。用 WMMA API 写的时候A、B、C 都以 fragment 对象出现每个 fragment 内部数据在 warp 内 32 个线程里的分布是硬件规定死的。比如 16x16x16 的 WMMA每个线程持有的 A fragment 是 8 个寄存器B fragment 是 8 个C fragment 是 32 个。这个布局不是程序员能随便改的你只能确保在调用 load_matrix_sync 之前数据在共享内存或者全局内存里的排布满足指令要求。这里有一个新手最容易踩的坑WMMA 的 fragment 布局不是直观的二维数组字节流。你以为是“线程 0 拿矩阵第一行前 8 个数”实际上硬件为了减少 bank conflict 和 shuffle 次数把数据按一种交叉分布放到 32 个线程里。所以千万别手工去推断 fragment 元素的位置老老实实走 API要靠内存布局对齐用不对就等着访问错数据。到了 Hopper 架构WGMMA 更进一步。它不再是“一个 warp 一条指令”而是让 warpgroup4 个 warp128 个线程一起加载更大的寄存器矩阵块累加器可以到 64x64 甚至更大。这种设计本质上是把寄存器 tiling 的规模从线程级提升到 warpgroup 级寄存器文件的整体带宽被更高效地利用了。2.3 算一笔寄存器账再实算一笔账看看为什么 16x16x16 的 WMMA 是 Ampere 时代的甜点尺寸。每个线程的数据类寄存器大概 48 个A 的 8 B 的 8 C 的 32加上地址计算、load 需要的 temporary 寄存器编译器一般会报到 80 到 120 个寄存器。这个数字既能保证一个 SM 里还有足够多的并发 warp 去隐藏访存延迟又不会让单线程太少、计算密度上不去。如果你做的是 double buffering也就是在 K 维循环里同时维护两套 fragment——一套正在被矩阵指令消费另一套正在从共享内存加载——那寄存器数会直接翻倍。这时候就要权衡了流水线掩盖的访存延迟到底值不值牺牲一半占用率。我在 A100 上调 FlashAttention 的时候就反复试过这个临界点最终发现 96 到 160 个寄存器之间通常能找到收益拐点超过 200 寄存器基本就会因为占用率太低反而变慢。这类经验在工程上非常值钱。当年老一点的优化文章喜欢鼓吹“寄存器越多越能提升单线程 ILP”但在现代 GPU 上寄存器文件总量是固定的单线程分得多并发线程就少。最好的策略永远是拿着编译器报告做实验而不是凭空拍一个数字。3. 别家架构寄存器 tiling 的更多姿势3.1 AMD CDNASGPR/VGPR 与 MFMA 指令AMD 的 CDNA 架构从 MI100 到 MI300走的也是类似 SIMT 的路子但有两个明显区别第一执行粒度是 wavefront也就是 64 个线程一组而不是 NVIDIA 的 32 线程 warp第二寄存器文件分成 SGPR标量寄存器和 VGPR向量寄存器两类SGPR 在整个 wavefront 里共享VGPR 才是每个线程私有的。这种划分带来的直接影响是寄存器 tiling 的时候你能把“每个线程都一样”的地址、偏移、循环计数放进 SGPR把真正随线程变化的矩阵元素放进 VGPR。这比 NVIDIA 把所有东西都堆在通用寄存器里要省不少空间也让编译器的寄存器分配更灵活。AMD 的矩阵指令叫 MFMA比如v_mfma_f32_16x16x16_f32。它的参数布局和 NVIDIA 的 WMMA 完全不同因为 64 个线程要分四组分别覆盖矩阵的行或列方向。我在把一些 CUDA kernel 往 ROCm 上移植的时候发现一个非常实际的问题NVIDIA 上最优的 16x8x16 指令形状在 AMD 上并不会自动变成同样高效的 MFMA 形状。因为 wavefront 是 64 线程你往往要把 tile 的边长按 16 的倍数来组织才能让 MFMA 的 fragment 分布正好落满整个 wavefront。这块只能靠对着汇编和 Matrix Core 的文档一点一点试没有捷径。3.2 昇腾达芬奇架构Cube 单元的类寄存器 Buffer昇腾的达芬奇架构走的是另一条路。它的 AICore 里有专门做矩阵乘的 Cube 单元Cube 的输入不是我们熟悉的线程私有寄存器而是 L0 Buffer具体分为 L0A、L0B 和 L0C。L0A 存 A 矩阵的 tileL0B 存 B 矩阵的 tileL0C 既是累加器又是输出缓冲。从功能上说L0 Buffer 扮演的就是“寄存器化存储”的角色数据在进入 Cube 之前必须已经按固定形状摆放在 L0 里。但跟 CUDA 有一个本质区别程序员几乎不手动去控制每个“线程”对应哪几个寄存器。Tiling 策略主要由编译器比如 AKG在编译期生成它会根据算子的 shape 自动推导 Cube 的 tiling 方案、L0 和 L1 之间的搬入搬出节奏。实际调昇腾算子的时候你会感觉到“寄存器 tiling”这个词的存在感被编译器包了一层。你要关注的更多是 L0/L1 容量够不够、双 buffer 有没有铺满、Cube 在切 K 维时是不是一次吞一个整数个 L0 tile。这种体验和 CUDA 那种所有布局都攥在自己手里的感觉差异很大但本质诉求是一致的让 Cube 单元的每个时钟周期都能从最近的存储拿到它要的操作数。3.3 脉动阵列与本地 SRAM 的“非典型”寄存器 tilingTPU 的脉动阵列彻底绕开了“线程 寄存器”这套体系。在脉动阵列里数据一旦从 Unified Buffer 喂进阵列就在 PE 之间一个节拍一个节拍地流动权重驻留在每个 PE 边缘的存储里累加器放在 PE 内部。这里没有传统意义的线程私有寄存器 tiling真正实现数据复用的是“数据流”数据流本身就是复用机制。你不需要在寄存器里保留一整块 A 的行片段因为相邻 PE 会在后续节拍里把这个数据往下传。Graphcore 的 IPU 又是一种极端。它没有硬件缓存每个核心的本地 SRAM 是显式管理的程序员得自己在 SRAM 里规划数据摆放。在这种架构上tiling 的那点“缓存”空间全都在 SRAM 里没有隐藏的寄存器层次可以靠编译器兜底几乎等于把寄存器 tiling 的思想放大成了“SRAM tiling”。各家架构放到一张表里对比形态差异会非常明显架构最小计算粒度寄存器 / buffer 形态tiling 主导者典型机制NVIDIA32 线程 warp线程私有寄存器 fragment 布局程序员 编译器mma.sync / WGMMAAMD CDNA64 线程 wavefrontVGPR SGPR MFMA fragment程序员 编译器v_mfma昇腾Cube 单元L0A / L0B / L0C buffer编译器AKGCube 指令TPU脉动阵列PE 内累加器 权重驻留编译器 / XLA脉动数据流Graphcore IPU单核本地 SRAM 显式管理程序员 / 编译器BSP 同步执行看完这张表你会发现“寄存器 tiling”在不同架构上长得非常不一样但背后的问题始终只有一个计算单元下一拍要的操作数能不能从最近、最快的那层存储直接拿到。4. 实战AI Kernel 里的寄存器 tiling 怎么落地4.1 先画一张寄存器 tiling 草图再写代码我自己的习惯是写任何计算密集 kernel 之前先不碰代码而是在纸上画一张“数据归属图”哪个 warp 负责输出矩阵的哪一块warp 里哪个线程负责哪一小块每个线程手里同时持有哪些 A、B、C 的 fragment。这个习惯是从 CUTLASS 的文档里学来的后来发现极其好用尤其是在跨架构移植的时候。画图时按这个顺序来确定输出 tile 的形状也就是线程块和 warp 的排布。决定单个线程持有多大的寄存器 tile。根据寄存器预算估算每线程大概需要多少寄存器看看是否和硬件上限匹配。确定 K 维循环里一次处理多大片段是否要做多级流水线。最后才是写第一版代码。这个流程里最容易被忽略的是第二步到第三步的衔接。很多人兴致勃勃画了一个 32x32 的大 tile最后编译一看每线程寄存器冲到 300 多spill 成灾性能反而不如十几行朴素写法。记住草图上看着漂亮没用寄存器预算这关过不了所有设计都是空中楼阁。4.2 Triton、CUTLASS、oneDNN 里怎么体现不是所有场景都要手写 CUDA。实际工程里Triton 这类 DSL 已经吸引了很多注意力。你写 Triton 的时候确实不会被要求手动分配寄存器但这不意味着寄存器 tiling 消失了它只是被编译器接管了。你在 Triton 里选的BLOCK_SIZE和num_warps其实就是在给编译器下达不同的寄存器 tiling 策略。比如一个 64x64 的 block 配上 4 个 warp编译器倾向于让每个线程持有 16x8 左右的累加器你把 block 放大到 128x128每个线程要扛的 fragment 就明显变大寄存器压力也随之上升。有一次我把一个 Triton GEMM 的num_warps从 4 改成 8寄存器数从每线程 96 涨到 168性能不升反降就是因为并发线程变少后访存延迟盖不住了。CUTLASS 则是把寄存器 tiling 完全显式化。你通过模板参数指定ThreadblockShape、WarpShape、InstructionShape每个线程持有的 fragment 就是由这几个参数计算出来的。我建议想做 gemm kernel 深入优化的人一定先把 CUTLASS 里 fragment 布局的推导逻辑读一遍那套东西本身就是一部寄存器 tiling 教科书。oneDNN 这边的思路也类似只是它更多面向推理场景会把融合算子直接按寄存器布局来规划尽量避免共享内存里的中转拷贝。4.3 性能分析里的寄存器关键指标代码跑起来之后不能只看端到端时间就完事。我通常这几个指标必看编译日志里的registers per thread这个数字直接反映你的寄存器 tiling 设计得激进不激进。有没有spill loads/stores只要有基本说明寄存器预算爆了。看local memory的大小它侧面反映 spill 的严重程度。用 ncu 看 kernel 的 issue 利用率确认数学指令是不是真的在连续发射。我举个真实例子。之前在 A100 上调一个 attention 类算子第一版每线程寄存器用到了 256编译报告显示有大量 spill。我把寄存器 tiling 从每线程 256 收缩到 168spill 归零同时 occupancy 从 25% 回到 37.5%最终端到端性能反而快了 12%。这个案例很典型你多占的那些寄存器并没有换来同等的计算密度反而因为 spill 频繁跑内存把一切全赔回去了。提示用nvcc --ptxas-options-v编译时就能看到寄存器数和 spill 情况不用等运行完如果在 Triton 里开发也可以用环境变量打开编译诊断信息。5. 踩坑清单与排查技巧实录5.1 RegSpill寄存器用爆了怎么查寄存器溢出是寄存器 tiling 最常踩的坑特征也很明显kernel 运行时间突然比预期慢一个数量级ncu或者编译日志里出现大量local memory访问。排查手法三步走。第一步确认编译器报告里的每线程寄存器数和 spill 读写次数第二步定位是哪个 kernel 爆的很多时候一个项目里不同 kernel 的寄存器压力完全不一样你得逐个看编译日志第三步考虑用__launch_bounds__(maxThreadsPerBlock, minBlocksPerSM)做上限约束强制编译器在寄存器分配上收敛避免它为了激进的指令级并行把寄存器数量推到离谱的地步。有个小经验__launch_bounds__不只是限制并发线程数它其实是在告诉编译器“我要保住多少块 SM 占用率”编译器会按这个目标调整寄存器分配。你如果目标 block 数是 8编译器会把每线程寄存器压到刚好能塞 8 个 block 的水平如果你不设编译器可能为了 ILP 一路放飞最后 spill。5.2 寄存器 tiling 与 Bank Conflict 的真相很多人会把寄存器 tiling 和共享内存的 bank conflict 混在一起聊这里澄清一下寄存器文件本身没有 bank conflict 的概念至少对程序员不可见。真正会踩 bank conflict 的是你从共享内存往寄存器搬 fragment 的那一步。比如你有一个 8x8 的 float tile 要从共享内存加载到寄存器如果共享内存里这 8 个 float 恰好落在同一个 bank 上那么 32 个线程同时访问时会发生冲突反而把加载周期拉长。NVIDIA 的 WMMA 指令在设计 fragment 内存布局时已经考虑过这个问题它会要求你按特定 stride 把数据铺在共享内存里确保访问不冲突。所以最常见的坑是你用共享内存存矩阵时完全按行优先铺然后发现 WMMA 的load_matrix_sync性能很差。这通常不是指令本身的问题而是矩阵在共享内存里的 tread 排列没遵守硬件偏好。5.3 不同架构迁移时的“水土不服”做多架构算子库的同学注意了把 NVIDIA 上最优的 16x8x16 寄存器 tiling 原封不动搬到 AMD 会水土不服因为在 64 线程的 wavefront 下MFMA 的 fragment 分布逻辑变了16x8 这种形状要么填充不满要么多出多余的 shuffle 指令。昇腾上更是直接换了一套 L0 Buffer 编程模型你在 CUDA 里手工安排的 fragment 布局完全没法映射过去。我的建议是在库的抽象层定义自己的“tile 描述”比如用结构体描述每个线程/处理单元应该持有多少 A、B、C 元素然后在每个后端的代码生成阶段各自翻译成对应架构的寄存器布局。这个抽象层初期写着麻烦但只要有第二个后端回报立刻兑现。5.4 一条快速自查清单最后给一张我自己常用的 registry tiling 自查表每次调优半信半疑的时候就拿出来逐项过一遍检查项怎么看健康标准每线程寄存器数ptxas / 编译器日志40 到 168 之间是否发生 spilllocal memory 大小、spill 读写次数0实际 occupancyncu和预期设计一致数学指令连续发射情况issue 利用率越高越好且无明显 stallfragment 加载指令比例加载指令与矩阵指令比值加载明显少于计算按这张表走一遍基本能覆盖 80% 的寄存器 tiling 性能问题。最后再分享一点个人心得。做了这么多架构的矩阵库和算子之后我最大的体会是不同硬件对“寄存器”的定义、容量的理解、以及暴露给程序员的控制方式差得非常多但大家在做的事情本质上是同一件——想办法让计算单元在下一拍开始之前最快拿到它要的操作数。谁能把这件事做到位谁写的 kernel 就离这台机器的极限更近。每次写新算子的第一版我都会先画一张寄存器 tiling 草图哪怕后来编译器帮我改掉一半这张草图也是我理解这台架构的起点。