ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

AI处理器矩阵乘Tiling深度解析:从Tile切分到数据流优化

AI处理器矩阵乘Tiling深度解析:从Tile切分到数据流优化 无论你手里是一块几十TOPS的AI加速卡还是手机里那颗低功耗NPU矩阵乘法和它的Tile/Tiling设计基本决定了你能从硅片上榨出几分算力。这也是为什么在AI处理器的设计讨论里矩阵乘从不缺席而Tiling往往是最先被拿出来反复推敲的环节。我见过不少团队峰值算力标得很漂亮一跑ResNet或者GPT类模型实际利用率掉到三四成最后排查下来根子大多不在乘法器阵列上而在数据到底怎么切、怎么搬、怎么排。这篇是AI处理器设计系列的第4篇重点聊矩阵乘的Tile/Tiling。我会从“为什么必须做Tiling”这个根本问题出发结合线性代数里“矩阵乘是坐标系变换”的几何直觉再落到三种经典数据流、一次完整的Tile尺寸推导以及DMA突发、Bank冲突这类硬件细节。适合正在做AI芯片架构、编译器映射或者想搞懂加速器底层原理的工程师参考。1. 算力巅峰的假象为什么矩阵乘必须做Tiling1.1 算力与带宽的剪刀差过去十几年AI加速器的MAC阵列规模一路膨胀从早期的64个乘加单元到如今单芯片动辄上千TOPS的算力。但一个扎心的事实是片外DRAM带宽的增长速度远远赶不上算力的增长速度。算力按照面积和工艺红利在涨带宽却受限于引脚数、封装功耗和信号完整性基本是线性甚至更慢地爬升。这就形成了剪刀差。以一枚典型的AI加速器为例假设MAC阵列是64×64跑1GHz峰值算力就是64×64×2×1GHz≈8.2TOPS乘加算两次运算。而HBM这类片外存储能提供的带宽乐观估计在2TB/s左右。看起来都很猛但你把两者放在一起算一笔账就发现问题了8.2TOPS的计算峰值意味着每秒需要消耗约8.2T/24.1T个乘加操作每个乘加操作至少要从存储器读两个操作数也就是每秒至少要喂进来8.2T个数据。在2TB/s、BF16精度2字节下每秒最多从DRAM拿1T个数据。差了一个数量级。结论很直接如果每个数据只用一次那无论阵列有多快实际算力都会被访存死死摁住。要想让算力跑起来唯一的办法是让数据在片上被反复重用。而数据重用靠什么实现靠把矩阵切成一块块适合片上存储的Tile让每一块数据从DRAM被搬上来之后尽量多参与几次计算再被扔掉——这就是Tiling存在的根本理由。1.2 没有Tiling的朴素实现为什么不可行很多人第一次写矩阵乘的硬件映射第一反应是这有什么难的把两个矩阵都读进来逐行逐列算不就行了。但你把规模放大到现实场景就知道这个想法有多天真。一个大模型推理里的典型GEMMM、N、K三个维度动辄4096甚至8192。假设MNK4096BF16精度光A矩阵就是4096×4096×2字节32MBB矩阵同样32MB。片上SRAM一共才几兆到几十兆一次性根本装不下。就算侥幸塞得下如果不做任何分块累加器矩阵C也需要4096×4096×4字节FP32累加64MB这更是任何片上存储都扛不住的。所以必须把输出C分成许多个小块比如每次只算TM×TN这么一小块输出而为了算这一小块输出只需要A的TM行和B的TN列。K维度如果太长还可以再继续切每次只加载K维的一小段BK。这样一来片上同时驻留的数据量就变成TM×BK BK×TN TM×TN不管原矩阵是4096还是十万维只要这个三项之和小于片上SRAM容量理论上就能算。这就是Tiling的字面意思把一个完整的大矩阵计算拆成一组小矩阵块的计算让数据的生命周期完全在片上闭环。2. 坐标系旋转视角Tiling背后的线性代数直觉2.1 从二维旋转矩阵说起很多人刚接触矩阵乘会以为它只是“行乘列”的机械规则。但真正深入硬件设计之后我建议你换个视角——把矩阵乘看成坐标系变换。这也是网上那个热搜问题“矩阵乘为何表示坐标系旋转”背后的直觉。拿最经典的二维旋转矩阵来说[ R(\theta) \begin{bmatrix} \cos\theta -\sin\theta \ \sin\theta \cos\theta \end{bmatrix} ]当你想把一个向量 (v(x,y)) 绕原点旋转 (\theta) 角直接算 (R(\theta) \cdot v) 就能得到旋转后的新坐标。这里的本质是矩阵的每一列表示原坐标系的基向量 ((1,0)) 和 ((0,1)) 在旋转后分别落到了哪里。第一列是 ((\cos\theta, \sin\theta))第二列是 ((-\sin\theta, \cos\theta))。矩阵乘本质上就是在新基底下描述同一个向量。放到更广义的线性变换里一个矩阵A乘上向量x就是把x从原始的坐标系映射到A的列向量张成的新坐标系。矩阵乘矩阵比如 (C A \times B)可以理解成先做B的变换再做A的变换复合后的线性变换就是矩阵乘的结果。2.2 输出块与输入块之间的“只依赖子集”关系坐标系视角对AI处理器设计有一个非常关键的启发线性变换是确定性的、局部的——输出空间的任何一个小区域只需要输入空间中一个对应子区域的信息就能算出不需要全量数据。看矩阵乘的定义输出 (C[i,j]) 等于A的第i行和B的第j列的内积。也就是说如果你想计算输出的一个 (TM \times TN) 子块只需要A的对应TM行和B的对应TN列总共 (TM \times K K \times TN) 个元素或者做了K维切分后更少。这个“只需要”三个字就是Tiling成立的数学基础。它保证了一切切分方案在数学上是无损的无论你怎么切只要行列对应关系正确算出来的结果和一次完整矩阵乘完全一致。从这个意义上说一个Tile其实就是从原坐标系里截取了一个低维子空间上的变换块。A的TM行对应坐标系的若干基方向B的TN列对应目标空间的若干基方向两者在K维求和维上完成内积相当于在子空间内完成了坐标变换的合成。硬件设计者如果要判断“我要算输出哪片区域必须搬入输入哪片区域”这个问题完全可以在编译期离线算清楚因为整个映射关系就是仿射的、可预测的。这也是为什么Tiling的调度逻辑适合用状态机或者简单循环来描述而不是需要运行时动态决策的复杂逻辑。3. 三种数据流分道扬镳Output/Weight/Input Stationary如何支配Tile形状3.1 三种复用模式Tiling的核心是数据复用。但复用谁、以谁为中心会导致完全不同的Tile切分策略。AI加速器设计里通常把数据流分成三类输出固定Output StationaryOS、权重固定Weight StationaryWS、输入固定Input StationaryIS。输出固定意味着累加器C的Tile一经载入片上就一直在寄存器或SRAM里待着直到把对应的所有K维累加完才写回DRAM。这种情况下C的Tile尺寸 (TM \times TN) 就是最重要的设计目标A和B的Tile则围绕它来组织每次加载A的 (TM \times BK) 小块、B的 (BK \times TN) 小块算完一部分累加到C上再换下一段K。输出固定最大的优点是写回DRAM的数据量极小适合输出特征图比较大的卷积层因为C的写带宽往往被忽视但其实很宝贵。权重固定则把权重矩阵B或者A取决于你怎么看作为常驻片上的一方反复用它去处理不同的输入数据。这在推理场景特别自然因为模型训练完权重就冻结了可以提前搬上片一直放在SRAM里被多个Batch、多张输入图重用。权重固定的Tile切分重点在于让权重的生命周期尽可能长输入数据则作为流水线上快速流过的“消耗品”。输入固定面向的是输入特征图被多个不同权重复用的场景比如卷积里同一个输入像素会被多个卷积核使用IS数据流让输入Tile常驻权重不断换入最大化输入数据的重用次数。3.2 量化判断Arithmetic Intensity与访存平衡点选择哪种数据流不能拍脑袋。工程上最实用的量化工具是算术强度Arithmetic IntensityAI定义为一个数据元素从DRAM加载后平均能支撑多少次运算。假设一个Tile配置下每次从DRAM加载A的 (TM \times BK) 和B的 (BK \times TN)在片上完成 (2 \times TM \times TN \times BK) 次FLOPs。那么算术强度大约是[ AI \frac{2 \times TM \times TN \times BK}{TM \times BK BK \times TN} ]分子是计算量分母是搬入的A和B的数据总量C的写回先忽略。如果这个AI值小于硬件平台的平衡点即带宽/算力那说明数据搬一次不够算访存会成为瓶颈如果AI远大于平衡点说明数据复用已经很充分计算端可能变成瓶颈。举个例子TMTN128BK64BF16精度那么AI (2 \times 128 \times 128 \times 64 / (128 \times 64 64 \times 128) 2 \times 128 / 2 128) MAC/B。换算成FLOPs/B要乘2就是256 FLOPs/B。这个数值远超大多数平台的平衡点通常在10~30 FLOPs/B说明只要Tile切得够大访存是可以被喂饱的。但注意AI只衡量了A和B的搬入如果C的写回和读入没有做好比如反复把C写回DRAM再读回来实际的AI会断崖式下跌。这也是为什么不建议把C也当作“流动数据”来对待累加器能留在片上就绝不搬出去。4. 一次完整的Tile尺寸推导从4096矩阵算到片上SRAM容量4.1 目标参数与约束条件讲完原理我带你实际推导一次Tile尺寸。假设我们设计的加速器参数如下MAC阵列64×64主频1GHz片上SRAM总共8MBHBM带宽约2TB/s数据精度BF16累加器用FP32。目标是把 (MNK4096) 的稠密矩阵乘跑出尽可能高的利用率。先列约束条件。MAC阵列64×64意味着理想情况下每个时钟周期能发起64×644096个MAC操作。要让阵列不空转每次从SRAM读取A和B的数据至少要足够阵列连续运行多个周期。此外片上SRAM要同时容纳A的Tile、B的Tile和C的累加Tile三者总大小不能超过8MB。因为数据从DRAM搬上片有延迟实际上还需要给A/B的Tile做双缓冲double buffer也就是一份被MAC阵列读取时另一份正在被DMA填充所以实际可用容量要打个对折考虑。4.2 三段式求解TM、TN、BK的权衡先忽略双缓冲按单缓冲推导。假设C的Tile尺寸是 (TM \times TN)A每次加载 (TM \times BK)B每次加载 (BK \times TN)。那么片上占用A Tile: (TM \times BK \times 2) 字节BF16B Tile: (BK \times TN \times 2) 字节C Tile: (TM \times TN \times 4) 字节FP32累加总占用 (S 2(TM \times BK BK \times TN) 4(TM \times TN))必须小于8MB。考虑双缓冲A和B部分要乘2也就是 (4(TM \times BK BK \times TN) 4(TM \times TN) \le 8MB)。现在开始定参数。MAC阵列是64×64所以TM和TN最好都是64的整数倍不然阵列会有空置的乘法器硬件利用率直接打折。常见的先验选择是TMTN128。因为128×128的输出Tile需要128×128×464KB的累加器为了减少写回C的块在片上待得越久越好。接下来看B和A的缓冲。如果把K整个不切即BK4096那么A和B的Tile就是128×4096×2×2双缓冲×2A和B两块≈4MB加上C的64KB还是勉强放得下的。但这样有个问题K循环没有任何流水块中间累加过程里DMA必须一次性把A和B的整个长条搬上来搬运延时长而且外层循环无法细粒度地交错DMA和计算。工程上通常会把BK切小比如取BK256或者512让DMA可以更细粒度地流水。取BK512。A Tile是128×512×2128KB双缓冲256KBB Tile是512×128×2128KB双缓冲256KBC Tile共64KB。合计约576KB距离8MB还很宽裕。这说明如果单纯看容量TMTN128的配置过于保守可以往大了提。4.3 继续放大64×64阵列下极限Tile与循环组织把TM提高到256TN保持128。A Tile256×512×2256KB双缓冲512KBB Tile512×128×2128KB双缓冲256KBC Tile256×128×4128KB合计约896KB。仍然远小于8MB。但这样做有个问题MAC阵列是64×64也就是每周期最多64×644096个MAC。输出Tile 256×12832768个输出点每个点需要累加K4096次总计算量是32768×4096≈1.34亿MAC。这当然没问题关键是阵列的调度粒度如果输出tile太大内层的累加循环就需要一个复杂的调度状态机来遍历所有输出位置寄存器文件的端口和路由压力会增大。实际工程里我见过很多设计最终的输出Tile落在128×128256×64之类的范围很少直接用512×512。原因是SRAM容量只是瓶颈之一更常见的约束来自片上网络带宽、寄存器文件端口数、以及累加器的读写端口。一个256×128的Tile意味着C累加器必须有能力在单个周期内同时读取/写入足够多的部分和这直接推高了SRAM的端口数和功耗。所以不要只盯着容量算要把MAC阵列的发射宽度、累加器读写带宽、片上网络带宽全部拉进模型做一个联合权衡。假设最终采取TM128、TN128、BK512的方案。遍历矩阵乘的循环组织大致是for mo 0; mo M; mo TM: # 遍历输出行块 for no 0; no N; no TN: # 遍历输出列块 init C_tile[TM][TN] to zero for ko 0; ko K; ko BK: # 遍历K维度块 load A_tile[TM][BK] from DRAM to SRAM load B_tile[BK][TN] from DRAM to SRAM # 内层微内核 for i in 0..TM: for j in 0..TN: for k in 0..BK: C_tile[i][j] A_tile[i][k] * B_tile[k][j] store C_tile[TM][TN] to DRAM这个伪代码其实就是几乎所有GEMM库包括手写汇编级的BLAS内核的骨架。硬件上最内层的三个循环会被展开映射到64×64的MAC阵列上外层循环则由状态机控制DMA的搬运和同步。K维的ko循环是双缓存流水线切换的边界load下一段A/B的同时算当前这一段。再来验证性能。总MAC数是 (M \times N \times K 4096^3 68.7) 亿次MAC。如果阵列每周期跑4096 MAC不考虑停顿需要68.7e9/4096≈16.8M周期主频1GHz下就是约16.8ms。再看访存A矩阵全量32MBB矩阵32MB在不考虑A/B Tile复用的情况下会反复从DRAM读吗并不是。由于外层mo/no把输出块固定后内层ko循环每次加载的新A/B块对于不同的(no, mo)组合是可以重用的。但实际每次切块后A的某个TM×BK块只服务于固定的no方向上的所有TN块所以A Tile会被复用N/TN32次B Tile会被复用M/TM32次。理想情况下DRAM搬运量大约等于A和B各读一遍即64MB加上C写回64MB左右总共约128MB。按2TB/s带宽算搬运耗时约0.064ms。相比16.8ms的计算时间访存远不是瓶颈说明这个Tiling策略在这个平台上是合理的。5. 形状、布局、DMA突发与Bank冲突细节决定Tiling成败5.1 为什么是矩形而不是正方形如果你看很多GEMM内核的Tile参数会发现TM和TN往往不相等比如64×256、128×64这种矩形很常见。有人会问正方形不是更整齐吗答案是Tile的形状必须服务于MAC阵列的形状和内存访问的连续性。如果MAC阵列本身是64×64的二维脉动阵列那么输出Tile在TM、TN两个维度上都会被阵列物理映射。64×64阵列每周期算64×64个MAC对应输出C的一个64×64小块。如果输出Tile取64×256那意味着MAC阵列在一个输出行块上需要沿着TN方向连续走4步这有利于让A的同一行数据在多个TN步中被反复重用数据流更线性。反过来如果TM、TN都取64每个Tile的边界处理次数反而变多——你会在外层循环中频繁切换A和B的TileDMA的启动开销和地址计算开销都会增加。矩形Tile还有助于匹配不同矩阵的形状。比如当M远大于N时你会倾向于TM大、TN小这样外层mo循环迭代次数少A的复用更充分当N远大于M时则反过来。这就是为什么通用GEMM库往往会根据M、N、K动态挑选分块形状而不是所有情况都用同一个正方形。5.2 内存布局与地址对齐Tiling不止是尺寸问题内存布局如果和Tile不匹配性能会差好几倍。这里最典型的例子是A矩阵的行主序row-major与B矩阵的列主序column-major之争。如果A按行主序存储那么A的一个TM×BK Tile在内存中是TM段连续区域每段BK个元素DMA搬运非常友好。但B如果也是行主序就是按K行N列排布。当你要加载B的BK×TN块时实际上是取B的BK行、每行TN个元素。这在内存里是BK段不连续的小块每块TN个元素DMA要发起BK次独立的突发传输效率差很多。解决办法有几个。最直接的是对B提前做转置或重排让B也按Tile友好的布局存储比如按BK×TN分块后连续存放。代价是增加了预处理的开销但推理场景下权重可以离线重排完全值得。另一个办法是让DMA控制器支持二维步进传输2D DMA它能在硬件层面处理跨行步长减少CPU或状态机的干预。我在实际项目里的经验是如果存储面积允许对B做一次Layout Transform明显比依赖2D DMA更稳因为2D DMA的跨步传输在地址非对齐时会浪费不少带宽。对齐细节也容易踩坑。DMA突发传输往往是以32字节、64字节甚至128字节为单位的如果Tile的起始地址没有对齐到突发边界DMA会多读一些用不到的字节白占带宽。更麻烦的是如果每行长度不是对齐粒度的整数倍下一行的起始地址就漂移了。常见的做法是给矩阵的行加paddingpitch让每行的字节数凑成64字节或者128字节的整数倍。比如原本每行是513个BF16元素也就是1026字节就可以把pitch设成1024字节。虽然浪费了约0.2%的存储但换来了DMA效率的大幅提升。5.3 Bank冲突Tiling最隐蔽的性能杀手最后说一个特别容易被初学者忽略的问题SRAM的Bank冲突。片上SRAM为了提高带宽会被划分成多个Bank理想情况下每个周期能同时访问多个无冲突的地址。但如果你加载A Tile时两个不同行的地址落在了同一个Bank的同一个偏移上SRAM就只能串行服务带宽直接减半。举个例子假设SRAM有32个Bank每个Bank宽度是4字节。你给A矩阵设置的pitch是1024字节也就是256个4字节单元。因为256刚好是32的整数倍那么A的第0行的起始地址落在Bank0第1行的起始地址也同样落在Bank0差32个Bank的距离之后又回到了同一个Bank。于是当你试图在同一周期读取两行的数据时就会发生Bank冲突。解决办法是给pitch再额外加一个偏移比如把1024字节改成1028字节让每行的起始Bank位置依次错开。这个“错开”的偏移量怎么选通常取Bank数的奇数倍或者一个与Bank数互质的增量保证连续几行的起始Bank都不重复。这个细节在编译器端可以算出来但需要硬件提供Bank配置的寄存器否则编译器改不了内存布局。所以我一直觉得Tile形状、Pitch、Bank分配这三件事必须放在一起设计硬件和编译器的接口里必须包含这些参数否则后期优化会被锁死。6. Tiling思想延伸卷积、注意力与下一代AI处理器6.1 卷积与矩阵乘的Tiling同构矩阵乘的Tiling思想并不是孤立的。卷积本质上可以转换成矩阵乘——标准做法是im2col把输入特征图上每个卷积窗口拉成一行于是卷积变成了一个更大的GEMM。代价是输入数据被复制了若干倍让原本紧凑的特征图变成膨大的矩阵。但如果你换个角度在硬件实现上不去真的做im2col而是把卷积的遍历本身就理解成一种Tiling输出特征图的每个Tile对应输入特征图的一个感受野区域沿着通道维、空间维分别切块数据复用关系其实和矩阵乘的A、B、C复用完全同构。很多AI加速器在处理3×3卷积时会在片上保留输入特征图的几条行缓冲line buffer这本质上就是沿着H维做的一个特殊Tiling。1×1卷积更是直接退化成GEMM完全套用矩阵乘的Tiling流程。6.2 注意力机制里的Tiling思想Transformer火了之后矩阵乘的Tiling又有了新的用武之地。注意力机制里的QK^T、Softmax、PV本质上都是矩阵运算但和普通GEMM不同的是中间会穿插Softmax的归一化而且序列长度可能远大于芯片能够一次性放下的范围。这就有了以FlashAttention为代表的一类方案把Q、K、V都切成块在片上逐块计算局部注意力分数同时维护一个在线Softmax的统计量最终得到正确的输出。它的核心思想和矩阵乘Tiling一模一样确定输出块的范围反推需要哪些输入块让每个输入块在片上被尽量多地复用。区别只在于中间的计算图上多了一层归一化需要仔细设计状态同步。对AI处理器设计者来说这提醒我们Tiling不只是GEMM专用引擎的问题而是所有张量运算共有的“小数据块调度”问题。6.3 从稠密到稀疏Tile的下一步进化另一个很现实的趋势是稀疏化。大模型训练和推理中权重和激活值都有相当比例的零元素于是出现了各种稀疏加速方案。但稀疏化和Tiling结合时问题会变得更复杂如果A矩阵的某一行恰好全是零传统的稠密Tile切法会白费计算如果稀疏模式是结构化的比如2:4Tile的形状要和稀疏粒度对齐否则剪枝带来的收益会被不规则访存吃掉。我在看一些新架构时发现大家不约而同地在Tile描述符tile descriptor上做文章比如用bitmask标记每个Tile里的有效元素或者按一定粒度做压缩后再搬运。这本质上是把Tiling从“固定形状的矩形块”扩展成“带稀疏掩码的不规则块”。但无论形式怎么变核心思路没有变——尽量让每个从DRAM搬到片上的字节都多参与几次运算让MAC阵列不要在等待数据时空转。我自己做Tiling优化这几年最大的体会是不要一开始就盯着大算力芯片的理论峰值看先把数据流画清楚把每一块数据在片上待多久、被复用几次、什么时候写回算明白再谈MAC阵列效率。Tiling这件事表面上是在切矩阵实际上是在切访存、切调度、切数据生命周期。把这个切明白了AI处理器设计里最难的那一半基本就稳了。
返回列表