☰
算子访存优化实战:基于L1/L2缓存的分块策略与GEMM性能提升
2026/10/18 7:22:04 网站建设 项目流程

在算子开发这个方向上,我调试过不少矩阵乘和卷积算子,大部分性能问题归根结底不是“算不出来”,而是“数据来得太慢”。前阵子我拿到一个运行在带L1/L2两级缓存架构的NPU平台上的GEMM算子,理论上限有20 TFLOPS,实测只有1.2 TFLOPS,计算单元大部分时间都在空转,等数据从主存搬过来。后来通过一轮一轮调整Tile策略,把分块从L2层级做到L1层级,性能一路提到了17 TFLOPS以上。这篇文章就聊聊我在这个过程中对基于L1/L2缓存的访存优化的一些理解,适合正在做算子开发、AI框架底层优化或者干脆被慢吞吞的GEMM/卷积折磨过的同学参考。

我不打算写成某个平台的具体使用教程,而是把这一类L1/L2缓存架构下的通用思路讲清楚。因为不管换哪款芯片,Tile的切法、循环的顺序、缓存容量的约束,背后的逻辑都是同一套。

1. 访存受限是算子优化躲不开的第一道坎

1.1 算力增长和带宽增长的“剪刀差”

芯片的计算单元这几年一直在膨胀,哪怕不追求极致频率,靠增加乘加单元数量和并行度,单芯片的峰值算力也能轻松做到几十甚至几百TFLOPS。但存储带宽的增速远远赶不上算力。主存带宽受物理封装引脚、HBM堆叠层数和功耗的限制,天花板很明显;片上SRAM的带宽虽然高,但容量有限,放不下大算子的全部中间数据。

这意味着什么?如果一个算子的数据访问方式没规划好,它就会从“计算密集”退化成“访存密集”。ALU管道的利用率再高,也架不住指令周期里有一大半在等待数据返回。做算子优化,第一件事就是判断当前瓶颈到底在算还是在访。

1.2 用算术强度判断一个算子是不是访存受限

判断方法不复杂,算一下算术强度就行了。算术强度的定义是:

算术强度 = 总计算量(FLOPs) / 主存访存量(Byte)

再对比硬件本身的计算强度平衡点:

平衡点 = 峰值算力(FLOPS) / 主存带宽(Byte/s)

如果算子的算术强度远低于平衡点,就处于访存受限状态;反之才可能逼近算力上限。

举个实际例子。一个朴素的三层循环矩阵乘,M=N=K=1024,FP16,计算量是2×10^9 FLOPs。如果循环内层每次都用一次乘加、读两个输入、写一个输出,不利用任何局部性,主存访问量会达到几十GB级别,算术强度可能连1 FLOP/Byte都不到。而现代AI芯片的平衡点往往在60-100 FLOP/Byte以上,这种代码跑起来必然是访存瓶颈。

理解了这条线,后面所有Tile策略的出发点就很清楚了:想办法把主存访存量压下来,把算术强度提上去,提到接近甚至超过硬件平衡点。Tile就是为了达到这个目的而存在的。

2. L1/L2缓存架构的容量悖论:数据放不下,放得下又用不够

2.1 某NPU平台的存储层次与典型参数

为了后面算账方便,我拿一个典型的双缓存架构说明参数。它分四层:最靠近计算单元的是每核私有的L1缓存(片上SRAM),然后是多核共享的L2缓存,再往下是主存。典型容量和带宽大致如下(不同芯片会有差异,但量级上很有代表性):

层级容量带宽访问延迟作用
L1缓存每核256KB-1MB500GB/s-2TB/s几个周期给计算单元直接供数
L2缓存整片32MB-64MB100-300GB/s几十个周期多核数据中转与复用
主存32GB-64GB80-200GB/s数百周期存放权重、激活、中间结果

L1的带宽几乎贴着计算单元的需求走,但容量小。L2的容量大一些,带宽比主存高一档,但也没高到能跟计算单元直接对接的地步。

2.2 缓存带宽与命中率的实际意义

很多人一提到缓存优化就只盯着“命中率”,其实在NPU场景里,更值得关注的是“数据在哪一层被读了多少次”。一个数据如果只在主存里读一次,然后在L2里被反复读几百次,那么沉重的访存开销就被摊销掉了。

给你算一笔账:主存带宽100GB/s,L2带宽300GB/s。假设某个数据从主存读一次耗时1个单位,之后300次访问都命中L2,那么均摊到每次访问的成本大约是(1+300/300)/301,接近于零。这就是Tile为什么能带来数量级提升的底层原因。

但这里有个矛盾:L2放得下的数据块,L1可能放不下;L1放得下的数据块,L2又显得太大,导致复用率不够。于是就有了经典的“二级分块”套路:先在L2层面做一次粗粒度分块,保证大块数据能滞留在片上;再在L1层面做一次细粒度分块,保证每个小循环体内数据完全驻留。

2.3 Tile策略的本质:把数据流掰成缓存装得下的形状

Tile翻译过来是“瓦片”,但不要只理解成把大矩阵切成小方块。它的本质是把算子的数据流拆成一连串满足缓存容量约束的计算片段,让每一份数据在搬进某级缓存之后,尽可能多地被复用几轮之后才被替换出去。

我用一个“食堂打饭”的类比:你手里有个托盘,L2就是那个托盘,L1是你面前的小碟子。如果你每次都只拿一碟菜回座位吃,来回跑几十趟食堂(主存),累死;如果你用大托盘一次多拿几份菜(L2级分块),来回次数就少了;但托盘太大也麻烦,L2装不下,菜会洒一地。所以托盘大小、碟子大小、每趟拿几份菜,就是Tile设计要解决的问题。

3. 从数据复用的角度拆解Tile切分,而不是简单切方块

3.1 三种基础复用模式:输入复用、权重复用、输出部分和复用

做矩阵乘或卷积优化时,一个算子的数据流里通常存在三种复用模式,它们对Tile维度的要求各不相同。

  • 输入复用:输出矩阵里相邻的多个结果会共享同一块输入数据。卷积里尤其明显,一个输入patch会被多个输出位置反复取用。对应到Tile上,输入Tile应该足够大,让同一块输入数据在滑动窗口内被完整消费掉。
  • 权重复用:在卷积和全连接层中,同一份权重会作用到所有输出位置。权重Tile只要放进缓存,理论上可以支撑当前L2块内所有输出点的计算。所以往往把权重设计成常驻缓存的那部分。
  • 输出部分和复用:在归约维度(比如K维)上做累加时,输出Tile一般很小,一直攒在L1里做累加,直到整个归约维度算完,才把结果写回主存。这个“暂不写回”的操作本身就是在复用输出Tile。

三种复用模式叠加在一起,才是选择Tile维度组合的真实依据。单独按输入尺寸切块,或者单独按输出尺寸切块,都不完整。

3.2 循环变换:分块、重排、展开怎么配合

确定了Tile维度后,还要决定循环顺序。对于矩阵乘,常见的顺序选择是:

L2层按M、N方向切大块,K方向作为最外层归约方向先扫一遍,让大块的A和B从主存只加载一次,之后就一直在L2里滚动;L1层再在K方向内做更细的循环块,每次取一小段A、一小段B做乘加。

内层循环顺序通常保持K在最里层,因为K维是归约维度,把它放进去可以让累加操作一直命中寄存器或L1。如果K在最外层,每次做一次乘加就要重新取一遍A和B的行列,访存模型直接崩塌。

循环展开也需要跟K方向配合。假设向量单元一次能处理16个FP16元素,那么K方向按16或32做展开就很自然。展开的本质是让每个数据加载在寄存器或L1里被多条指令复用,而不是加载一次用一次。这等于在指令层面又做了一层更小的“Tile”。

3.3 容量约束方程:决定Tile大小的数学边界

Tile大小不是拍脑袋定的,它需要满足缓存容量约束。对于矩阵乘的一个输出块Br×Bc,涉及的数据是:

A块:Br×Bk,B块:Bk×Bc,输出块:Br×Bc

总数据量 = (Br×Bk + Bk×Bc + Br×Bc) × sizeof(元素)

约束条件是:

层级约束条件
L1级分块A子块+B子块+输出子块 ≤ L1容量×0.8
L2级分块A大块+B大块+输出大块 ≤ L2容量×0.7

为什么留余量?因为缓存里不只有这些数据块,还有地址转换开销、双缓冲预留空间、系统自身占用的临时数据。我见过有人把L1塞到99%,结果双缓冲一开,整个Tile溢出到L2,性能反而崩了。

举个例子。FP16,假设L1容量256KB,选Br=32、Bc=64、Bk=64,那么:

  • A块=32×64×2=4KB
  • B块=64×64×2=8KB
  • 输出块=32×64×2=4KB
  • 总量=16KB,远小于256KB,看起来很宽松,但后面加双缓冲和向量展开后,实际占用会涨到五六倍。

所以在容量边界上,我习惯留20%-30%的安全余量,再用profiling验证。

4. Tile形状不是越方正越好,计算密度才是取舍标准

4.1 为什么同样的面积,不同形状的计算密度不同

拿一张白纸做类比。一张纸面积固定,你可以把它裁成10×10的小方块,也可以裁成2×50的细长条。矩阵乘里,同样面积(缓存容量)的Tile,不同长宽比会导致不同的“计算量/访存量”比例。

假设缓存容量C固定,允许输出的Tile规模为Br×Bc,归约维度的分块为Bk。计算量正比于Br×Bc×Bk,访存量正比于Br×Bk + Bc×Bk。当C固定时,Br和Bc有一个上界。如果只追求Br×Bc最大,可以做到很大的方块,但Bk的空间就被挤占;而Bk太小,每个输入数据被复用的次数就少,算术强度上不去。

实际优化中,我一般把Bk设得尽可能接近L1容量的1/4到1/2,而不是盲目追求Br和Bc的乘积最大。因为Bk直接影响数据复用次数,而Br和Bc决定了并行度。一个比较顺手的初始选择是Br取32、Bc取64、Bk取64,FP16下一个输出块只占4-8KB,L1能装下好几个这样的块,同时Bk又足够长,适合向量化。

4.2 常见算子的Tile形状选择经验

不同算子的最佳Tile形状差异很大,简单总结我常用的选择逻辑:

算子类型推荐Tile形状说明
矩阵乘Br≈Bc≈Bk≈32~64三者权衡,优先保证Bk足够长
卷积空间维按输出像素取8~16,通道维按对齐宽度取16的倍数通道维必须对齐到向量宽度
向量逐元素一维Tile,长度取向量单元宽度的整数倍太短则循环开销高,太长则溢出L1
Attention/Softmax按序列长度切分,保留Batch维整体避免跨Batch维度做归约

卷积比较特殊:如果把通道维设为Tile的长边,输入特征图的复用效果最好;如果把空间维设为长边,kernel的复用更好。我通常优先保通道维对齐,因为它直接影响向量化,空间维稍微差一点可以通过重叠加载来补。

4.3 尾部对齐与边界处理:补齐还是掩码

Tile尺寸很少能正好整除矩阵维度。比如M=4097,你选了Br=64,那么最后一块只剩1行。处理这种尾部有两种做法:

  • 补齐(Padding):把A和B的尾部补零到对齐宽度。缺点是多算了几个无效FLOPs;优点是控制流干净,没有分支,向量指令可以一条不差地执行。
  • 掩码(Predication):通过掩码寄存器在加载/计算阶段跳过无效部分。节省了无效计算,但引入了额外的谓词指令,而且很多架构上掩码访存的效率比正常访存低得多。

我的经验是:在AI芯片上当边界块占比不超过1/16时,一律用补齐。因为计算单元最怕分支发散,一条掩码指令带来的流水线惩罚往往比几个周期的无效计算还大。只有当矩阵维度恰好经常出现半块、残块占比很高时,才专门写一条掩码路径。

5. 一个GEMM算子的优化实录:从1.2 TFLOPS到贴近算力上限

5.1 基线版本:朴素三层循环的性能惨状

模拟一个典型案例:某推理引擎里的GEMM算子,M=N=K=4096,FP16。平台参数按前文的典型值:主存带宽100GB/s,L2带宽300GB/s,L1带宽1TB/s,L2容量32MB,L1容量256KB,FP16峰值算力20TFLOPS。

基线版本就是教科书式的三层循环:

for i in range(0, M): for j in range(0, N): for k in range(0, K): acc = 0 for kk in range(0, K): acc += A[i][kk] * B[kk][j] C[i][j] = acc

这个版本没有任何分块,而且这里的伪代码本质上是每计算一个输出点,就把A的一整行和B的一整列从主存重新读一遍。算下来实际访存量惊人,实测只有1.2 TFLOPS,计算单元利用率不到6%。这个阶段我去看profile数据,能明显看到主存读带宽被拉到极限,L1和L2的命中率都在个位数。

5.2 第一轮优化:L2级分块把数据滞留在片上

第一轮,做L2级分块。把M、N、K分别切成M0=128、N0=128、K0=128的大块,循环顺序调整为:

for m0 in range(0, M, 128): for n0 in range(0, N, 128): for k0 in range(0, K, 128): for i in range(m0, m0 + 128): for j in range(n0, n0 + 128): for k in range(k0, k0 + 128): C[i][j] += A[i][k] * B[k][j]

这一步的逻辑是:单个L2块内,A块128×128×2=32KB,B块32KB,输出块32KB,合计96KB,远小于32MB的L2容量。于是A和B在大块内可以被反复访问数百次,而不需要回到主存。

实测性能提升到了8.9 TFLOPS。瓶颈从主存带宽转移到了L2带宽。profile显示L2命中率超过95%,主存带宽占用大幅下降。

5.3 第二轮优化:L1级分块与向量化展开

第二轮,在L2大块内部再切L1级分块。取Br=32、Bc=64、Bk=32,同时把K方向按32步长展开,配合每周期能处理16×16 FP16块的向量单元。

for i0 in range(0, 128, 32): for j0 in range(0, 128, 64): for k0 in range(0, 128, 32): # 加载A[i0:i0+32][k0:k0+32]到L1 # 加载B[k0:k0+32][j0:j0+64]到L1 # 用向量指令连续累加,展开k0内部 for kk in range(0, 32): # 16x16向量乘加

这轮优化后,每个L1子块的数据量只有约10-15KB,完全驻留在L1里。K维展开后每个向量寄存器里的数据能被连续复用16次,L1带宽得到充分使用。

实测性能到了15.6 TFLOPS,计算单元利用率接近78%。再往后看瓶颈,指令等待周期大幅下降,L2带宽不再是主要矛盾,真正的限制变成了向量指令的执行吞吐和循环控制开销。

5.4 收尾优化:双缓冲和Bk扩展的取舍

最后一步微调有两个选项:一个是把Bk从32扩到64,增加K维复用长度,让每个L1块内的计算量翻倍;另一个是开双缓冲,让下一块数据预取和当前块计算重叠。

我两个都做了,但要做容量核算。Bk=64后,A子块32×64×2=4KB,B子块64×64×2=8KB,输出块32×64×2=4KB,合计16KB。按双缓冲×2算,32KB,在256KB L1里完全可行。于是L1满载预算大约60-70KB,仍然留有余裕。

最终性能稳定在17.3 TFLOPS,峰值算力利用率约86%。这个阶段再往上提升,就需要处理bank冲突和更细的地址对齐问题了。

6. 缓存冲突、对齐问题和调试复盘

6.1 性能“塌方”现场:缓存bank冲突

Tile调完后,我习惯再做一个全参数扫描,验证各个分块组合下性能是否平稳。结果发现一个诡异现象:Br=64、Bc=64、Bk=64时性能正常,改成Br=64、Bc=96、Bk=64,性能忽然掉到9 TFLOPS左右。

用profile工具看了事件计数,原因是bank冲突。某些缓存体系里,L1被划分成若干bank,同一拍里多个地址若落在同一个bank,就需要串行处理。Bc=96之后,A块的起始地址与B块的起始地址恰好在bank上模相等,导致每个加载周期都发生冲突。

解决方式很经典:给每个Tile的起始地址加上一个字节偏移(Padding)。比如在分配缓冲区时,让每一行的stride从128字节改成132字节。加Padding后,bank冲突消失,性能恢复到16.8 TFLOPS。

6.2 地址对齐:一个容易漏掉却影响巨大的细节

很多人只关注Tile尺寸,忘了检查缓冲区基地址。假设缓存行是64字节,如果Tile的起始地址不是64字节对齐,那么一次本应加载64字节的操作会跨两个缓存行,凭空多消耗一个缓存行带宽。

这个问题在动态分配场景下特别隐蔽。标准分配函数返回的地址通常满足16字节对齐,但不一定满足64字节对齐。平台自带的buffer pool一般没问题,一旦自己写malloc再传地址,就容易踩雷。排查方法很简单:打印缓冲区地址,检查是否%64==0;不对齐就换用平台提供的对齐分配接口,而不是自己手动加偏移。

6.3 profiling复盘流程:先命中率,再bank冲突,最后指令等待

我总结了一个三轮定位法,适用于大部分访存类算子的性能问题定位。

第一轮先看L1/L2命中率。如果L1命中率低于80%,优先调整L1级分块尺寸;如果L2命中率低于90%,重点检查L2级分块是否过大,导致大块数据被迫提前换出。

第二轮看bank conflict计数和非对齐访问计数。这类事件通常只在特定Tile组合下突然升高,扫描组合就能定位。如果非对齐访问和bank冲突都很低但性能还是不够,进入第三轮。

第三轮看指令等待周期和计算单元利用率。如果计算单元利用率高于85%,瓶颈主要在指令发射;如果大量周期在等待L1数据返回,说明L1内分块还需要进一步减少数据依赖,这时候可以试着调整双缓冲策略或提前预取。

6.4 精度与性能的平衡:累加顺序的坑

最后提醒一个容易被忽略的点:改变循环分块和累加顺序,等于改变浮点累加的结合顺序。矩阵乘的K方向归约本来是从0累加到4095,分块后可能变成每32个一组累加再合并,浮点舍入误差会有些微不同。

我在一次模型精度回归测试里就吃过亏:性能优化后,某个层的输出差异放大了几个百分点,定位半天,发现是K方向的分块顺序改变了累加误差的分布。解决思路不是完全放弃分块,而是固定K方向的内层累加顺序,例如始终从K=0向K=4095按顺序扫,只在前端做分块索引,不在后端改变向量内部的累加顺序。

如果精度约束很严格,还可以在最后的规约阶段用更高精度的累加寄存器,比如把FP32累加器换成FP32双精度累加,代价是部分计算单元吞吐下降,但一般只在必要时候用。

做算子访存优化,我的体感是:先理顺数据流,再调计算细节,顺序不能反。如果一开始就扎进bank冲突和向量化对齐里,很容易把复杂问题变复杂;先去把所有Tile参数和缓存容量存成一张对照表,用数据驱动的方式决定下一步,才是最快路径。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询