最近为了把一个本地模拟的 decoder-only 大模型推理速度提上来，我把某国产 NPU 异构计算架构的算子库从上层接口一直翻到了底层 kernel。翻完最大的感受是：大模型上层那些 Attention 变体、KV Cache 优化都很重要，但真正兜底性能的，还是那个最不起眼的 MatMul 算子——大语言模型矩阵运算的核心组件，几乎每一层都在反复调用它。这篇文章不算什么教科书讲解，而是结合我实际排障过程，把这个算子在算子库里的实现骨架、性能边界、优化取舍拆开讲清楚。想快速定位大模型推理性能瓶颈的人、做算子开发或者框架性能优化的人，应该都能从里面找到可以立刻拿去用的判断方法。
先说明一下，文中的算子库并不特指某一个商业产品，而是我在本地环境中研究过的、常见的 NPU 异构计算架构里的算子集合，接口逻辑和主流算子库大同小异。为了避免歧义，下文统一用“NN 算子库”代称。MatMul 在所有深度学习框架里都是高频算子，但由于 NPU 的硬件调度方式和通用 GPU 不同，它的内部设计侧重点很不一样。这篇文章会尽量把“为什么这样设计”讲透。
1. 先理解为什么 LLM 每一层都在碰 MatMul
1.1 从 Attention 和 FFN 看矩阵乘的落点
很多人觉得矩阵乘不过是线性代数课本里的一个运算，但放到生成式大模型里，它承担了几乎全部的可学习权重。Decoder-only 模型的每个 Transformer Block，常规结构可以拆成两部分：Attention 和 FFN，而这两部分的主要计算全是 MatMul。
比如 Attention 的第一步是 QKV 投影。输入矩阵 X 的形状是 [seq_len, hidden_size]，权重矩阵 W_Q、W_K、W_V 的形状都是 [hidden_size, hidden_size]，于是要做三个矩阵乘得到 Q、K、V。接下来计算注意力分数，要把 Q 和转置后的 K 相乘；拿到分数之后，又要用分数和 V 相乘，得到注意力输出；最后还有一个输出投影矩阵，把 Attention 结果再映射回 hidden_size。光一个 Attention，已经出现了六次 MatMul。
FFN 也不遑多让。常见的做法是先把 hidden 状态投影到一个更大的中间维度，过激活函数，再投影回 hidden_size，这就是两次大矩阵乘。再加上 Embedding 后的映射、最后一层输出 logits 的投影，一个 12 层的模型，单次生成一个 token 要执行的 MatMul 次数轻松超过几十次。如果 QKV 投影没有被融合，这个数字还会更高。
所以我在看 NN 算子库时，第一个反应就是：只要把 MatMul 这个算子的性能吃透，整个大模型推理项目的性能画像就完成了一大半。这不是夸张，因为大模型绝大部分算力消耗都落在这个算子上。
1.2 GEMM、BMM、量化 MatMul：同一算子的几种面孔
算子库里的 MatMul 并不是只有一种实现。从接口上看，它通常接收两个输入 tensor，a 和 b，再通过 transpose_a、transpose_b 这类参数决定是否转置。根据输入维度和数据格式，至少能分出三种常见形态：
- 标准 GEMM：即通用矩阵乘，对应线性层、投影层。输入是二维矩阵，公式为 C = alpha * A * B + beta * C，alpha 和 beta 默认一般是 1 和 0。
- 批量矩阵乘 BMM：输入带有 batch 维，比如 Attention 里对多个 head 的 Q 和 K 做批量乘。形状可以是 [B, M, K] 乘 [B, K, N]，输出 [B, M, N]。这种形态对并行调度的要求更高，因为 batch 维天然适合分发给不同计算单元。
- 量化 MatMul：输入可能是 int8 或 int4 的权重，激活可能仍是 fp16。计算时需要在硬件矩阵乘单元里做反量化，或者走混合精度通道，累加器保持 fp32，最终再写回目标 dtype。
为什么算子库要区分这些形态？因为底层硬件对不同数据布局、不同维度的支持差别很大。比如 BMM 的 batch 维可以当作并行度；量化 MatMul 则受限于硬件是否原生支持低比特乘加指令。如果算子库只用一套朴素实现，性能很可能在某种形态下崩掉。
1.3 推理场景对 MatMul 的特殊要求
大模型推理和训练对 MatMul 的压力完全不同。训练时通常是大的 batch size，M 很大，形状相对规整，追求的是吞吐；推理时有两个截然不同的阶段：prefill 阶段要处理整段输入，M 等于输入序列长度，可能很大；decode 阶段每次只生成一个 token，M 等于 1。
这种“M 一会儿大一会儿小”的动态 shape，会让同一个 MatMul 算子呈现出完全不同的性能特征。M=1 时，计算量极小，权重搬运占主导；M=1024 时，计算量骤增，矩阵乘法单元是否饱和成为关键。NN 算子库在设计 MatMul 时，必须能根据 M、N、K 和硬件资源动态选择不同 kernel，而不是静态绑定一个最优配置。这也是我后面调试时踩坑最多的地方。
2. 矩阵乘的性能边界：先算清楚 Roofline 再谈优化
2.1 一次 MatMul 的计算量与访存量
在动手调优之前，我习惯先做一次“纸上计算”。MatMul 的计算量非常固定：C 的每个元素要做 K 次乘加，C 本身有 M 乘 N 个元素，所以总乘加次数是 MNK，如果按浮点运算数 FLOPs 算，还要乘以 2，因为一次乘加算两个运算：
FLOPs = 2 * M * N * K访存量则取决于输入数据要从主存搬运多少字节。最粗略的估算是不考虑缓存，假设 A、B、C 都从主存走一遍：
访存字节数 ≈ (M*K + K*N) * sizeof(type) + M*N * sizeof(acc)真正关键的不是这两个绝对值，而是它们的比值，也就是算术强度：
arithmetic_intensity = FLOPs / 访存字节数用 Roofline 模型的话来说，算术强度低，就是访存受限；算术强度高，才是计算受限。我举个具体例子。假设 hidden_size=4096，decode 阶段 M=1，权重是 fp16：
- FLOPs = 2 * 1 * 4096 * 4096 ≈ 33.6 MFLOPs
- 访存字节数主要是读取权重 B，4096*4096 个 fp16 约等于 33.6MB
- 算术强度 ≈ 1 FLOP/Byte
这个强度低得吓人。换句话说，每从主存读一个字节，只做了一次浮点运算。再对比 prefill 阶段 M=1024：
- FLOPs = 2 * 1024 * 4096 * 4096 ≈ 34.4 GFLOPs
- A 的大小变成 1024*4096 个 fp16，大约 8.4MB，B 仍然是 33.6MB，C 大约 8.4MB
- 算术强度 ≈ 34.4e9 / 50e6 ≈ 700 FLOP/Byte
同样一个权重矩阵，M 变大之后，算术强度提升了几百倍，瓶颈属性完全反过来。这就是为什么大模型推理里“decode 慢”那么普遍：它一直在访存受限区里跑。
2.2 同一个算子，在 Roofline 两侧反复横跳
MatMul 比较特殊的地方在于，它既会被归为访存密集，又会被归为计算密集，关键看 shape。NN 算子库里的 MatMul 为了兼顾两种场景，往往会准备两套优化路线。
访存受限时，优化目标不是把矩阵乘法单元打满，而是减少主存访问。手段包括：把权重尽量压在缓存里、多请求合并 batch、使用低比特权重、避免中间结果来回写。计算受限时，优化目标变成提高矩阵乘单元利用率，比如选择合适的分块、把数据搬运和计算重叠起来、用低精度乘加指令。
这两种思路如果搞混，会出现很典型的性能问题：一个人拼命调分块大小，想把计算流水线加满，结果 M=1 时根本没有足够多的计算任务，所有的努力都在给一个“等待权重搬运”的过程做装饰；另一个人拼命做 kernel 融合，减少中间结果访存，但大 M 场景下其实计算早就饱和，减少的那点访存根本改变不了瓶颈。所以每次优化前，我都要求自己先算出算术强度，再决定动作方向。
2.3 分块到底改变了什么
分块是矩阵乘优化里最基础也最核心的手段。朴素实现里，计算 C 逐个元素时，每个 C 元素都要遍历 A 的一行和 B 的一列，导致 A 和 B 的数据在内存里被反复读取。分块之后，C 被切成 M_tile × N_tile 的子块，A 和 B 也按对应的 K 方向切成小条，让每个子块都能装进片上存储，被多次复用。
一个典型的分块伪代码大概长这样：
for (mo = 0; mo < M; mo += MT) { for (no = 0; no < N; no += NT) { for (ko = 0; ko < K; ko += KT) { load A_tile = A[mo:mo+MT, ko:ko+KT]; load B_tile = B[ko:ko+KT, no:no+NT]; C_tile += A_tile * B_tile; } } }MT、NT、KT 怎么选，直接影响性能。分块太大装不进片上存储，分块太小复用率不够。NPU 上还得对齐矩阵乘单元的形状，比如某些硬件矩阵乘单元一次处理 16×16、32×16 这样的块。算子库里的 MatMul 之所以看起来复杂，很大一部分就是在处理这些“对齐”和“切分”的细节。
3. 算子库中 MatMul 的实现骨架：从接口到 Kernel
3.1 算子描述、参数校验和 Shape 推导
NN 算子库里任何一个算子对外都有一个标准描述：算子名、输入 tensor、属性、输出 tensor。MatMul 的属性通常包括 transpose_a、transpose_b，以及可能存在的 alpha、beta。在真正进入计算前，算子框架会先做 shape 推导和参数校验。
这一步看起来简单，但直接影响图编译阶段能不能静态分配内存。比如动态 shape 下，如果 shape 推导失败，整个图的执行计划可能退化成运行时推断，带来额外开销。我在调试模拟项目时就遇过一个问题：输入的 seq_len 在 prefill 和 decode 之间切换，某些绑定静态 shape 的 kernel 被强行兜底到通用的慢版本，profile 里一眼就能看出来。后来把 shape 推导修好，让调度器正确识别 M=1 的小形状，MatMul 耗时立刻降了一截。
参数校验还包括 dtype 检查。fp16、bf16、int8、fp32 混用是常事，算子的实现必须明确累加器精度和输出 dtype，否则精度问题会一层层传上去。
3.2 Tiling 策略是性能好坏的分水岭
分块听上去简单，做起来全是细节。同一个 MatMul，M=1 和 M=1024 的最优分块可能完全不同。NN 算子库一般会先按 M、N、K 和 dtype 查一张经验配置表，如果表里没有，再跑一个代价模型去估算不同 tile 方案的耗时。
举例来说，如果目标硬件的片上 SRAM 约 256KB，权重和激活都是 fp16，那么一个 192×64×64 的 tile 可能需要 64KB，再加上双缓冲的副本，可能已经占满。而如果权重是 int8，同样容量能放更大的 tile。tile 不是越大越好，还需要考虑多核并行的粒度。常见做法是外层按 M 和 N 切给多个核，内层在 K 方向做流水，这样核间通信最少。
在实际调试中，我自己不会一上来就调 tile 参数，而是先用 profile 看硬件的矩阵乘单元利用率。如果利用率高但整体耗时高，问题往往在数据搬运；如果利用率低，再看是不是 tile 大小导致计算任务喂不饱。
3.3 多核并行与流水线重叠
NPU 上跑 MatMul，通常不是单核硬算，而是把输出 C 切成多个子块，分给多个计算核心。切分原则是沿 M 和 N 方向切，尽量让每个核独立算完自己的 C 块，避免跨核共享 A 或 B 带来的同步开销。
除了多核，还有流水线。MatMul 的完整数据流是“主存→片上存储→矩阵乘单元→写回主存”。如果不做流水，计算单元会一直等数据搬运；做了双缓冲之后，当前 tile 在计算时，下一个 tile 已经在搬运，单位时间能处理的数据量大幅提升。这个过程很像餐厅后厨：炒菜和备菜同时进行，而不是炒完一道再去备下一道。
算子库里的调度代码通常会显式管理这些 buffer：主存地址、片上 buffer index、计算任务队列。不同 NPU 的 buffer 大小和搬移指令差别很大，所以这部分代码很少能跨平台复用，也是 MatMul 算子实现里最难维护的部分。
3.4 不通用的底层 Kernel 和统一的分发逻辑
同一个 MatMul，算子库往往需要维护多套底层 kernel。CPU 上可以用 SIMD 向量化，NPU 上要用专用的矩阵乘单元指令，如果还支持通用 GPU，还得有另一套基于通用 GPU 指令的实现。所以围绕 MatMul 会有一个策略分发层：根据硬件类型、dtype、shape、是否转置、是否 batch，选择具体 kernel。
分发层的存在很有价值。它可以做到“上层算子只写一次，底层按平台扩展”。我在本地模拟环境里主要看的是 NPU 路径，但很多优化思路对 CPU 同样适用，只是分块大小和指令选择不同。
4. 排查 MatMul 性能瓶颈的一段真实经历
4.1 模拟场景：为什么 decode 阶段一个 token 要 90 毫秒
我在本地搭了一个模拟项目 X，结构是 12 层 decoder-only transformer，hidden_size=4096，32 个 head，head_dim=128，权重全部 fp16。prefill 阶段一次性输入 1024 个 token，表现还过得去；真正让人头疼的是 decode 阶段，每生成一个 token，端到端耗时要 90 毫秒左右。
第一件事是打开算子级 profiling。结果里 MatMul 相关算子（QKV 投影、注意力分数、注意力输出、输出投影、FFN 上下投影、logits 投影）加起来占了总耗时的 70% 以上。这个占比完全符合预期：12 层模型，每层一堆矩阵乘，每次生成 token 时所有权重基本都要从头读一遍。
当时我大概估算了一笔账：每层 QKV 三个 4096×4096 的 fp16 权重，加上输出投影和 FFN 的几组大权重，全部加起来有数 GB。假设我的模拟环境有效带宽约 200GB/s，光是读取一遍权重就要二十多毫秒，90 毫秒端到端说明除了 MatMul 本身，还有很多中间搬运和调度浪费可以压缩。
4.2 用算术强度判断该往哪个方向调优
Profile 数据里有每个 MatMul 的 M、N、K，我用前面那个公式手算了一遍。decode 阶段几乎每个 MatMul 的 M=1，算术强度在 1 到几个 FLOP/Byte 之间，清一色访存受限。
结论很明确：再去纠结矩阵乘单元的利用率没有意义，M=1 的任务根本喂不饱计算单元。真正的目标是让权重访问变得“便宜”：要么一次少读一点，要么让读过来的数据多算一点，要么干脆别反复读同样的权重。
顺着这个思路，我做了三件事。第一，检查调度器是否给 M=1 选用了小形状专用 kernel，避免大 tile 带来的无谓循环开销；第二，尝试把部分权重从 fp16 压到 int8，同时用 per-channel scale 做量化补偿；第三，把 QKV 投影、FFN 这类多个 MatMul 和中间激活合并成融合算子，减少中间变量写回主存的次数。
4.3 动刀之后，90 毫秒降到了 42 毫秒
这三件事不是一次全上的，而是一步步验证。
先开 M=1 专用 kernel，端到端从 90 毫秒降到 71 毫秒，效果明显，因为调度开销和无效的并行等待被砍掉不少。然后把注意力相关的几个权重换成 int8，权重读取字节数下降，decode 耗时又降了十几毫秒，但这里要特别小心精度，我是先用一小段测试集做对比，确认生成质量没有明显回退才继续的。最后做融合算子，把 QKV 三个 MatMul 加 bias 和残差合并，把 FFN 两个 MatMul 加激活合并，端到端来到 42 毫秒左右。
整个排查过程复盘一下，真正让性能改善的不是某个“黑科技 kernel”，而是：先意识到瓶颈是访存，然后顺着访存路径把所有浪费捡干净。这里我建议任何做推理优化的人都养成一个习惯：动手调代码前，先花五分钟算一下当前算子的算术强度，再决定优化方向，否则很容易南辕北辙。
5. 从 MatMul 往外看：融合、量化与大模型落地的关键
5.1 算子融合把多个 MatMul 变成一次循环
在做 MatMul 优化时，绕不开“融合”这两个字。Attention 里的 QK^T 和后续的 softmax、乘 V，如果按原始算子拆开做，QK^T 这个中间矩阵要先写回主存，做完 softmax 再读出来，再和 V 乘。这个中间矩阵哪怕只写一次，访存成本也相当可观。
目前主流的分块注意力算法就是把 Q、K、V 按小块加载到片上，Q 和 K 的子块相乘得到分数子块，softmax 也在片上做，然后直接拿结果子块去乘 V 的子块，最后累加。整个过程中间矩阵不落主存，访存压力大幅下降。NN 算子库里的融合版 MatMul，本质上就是让这种“嵌套计算”在同一个 kernel 里完成，而不是靠一次次算子 launch 拼起来。
这种思路不仅限 Attention。FFN 里的 MatMul、激活函数和下一个 MatMul 也可以融合。融合后少了好几次全局内存 round-trip，decode 场景下收益尤其大。
5.2 低比特 MatMul 为什么在推理里那么香
量化是另一个大杀器，原因也藏在 Roofline 里。访存受限时，推理延迟大约正比于需要读取的权重字节数。fp16 换成 int8，字节数直接减半；int4 如果硬件支持得好，再减半。对着一个几乎完全访存受限的 decode 阶段，这等于直接把“卡脖子”的通道拓宽。
但低比特不是白来的。int8 和 int4 的矩阵乘通常需要 scale 和 zero point，计算时累加器还得用 fp32 来保证精度。如果一个模型没有做过量化感知训练，直接截断权重，生成质量很可能崩。我的经验是：先在离线数据上小批量验证，观察困惑度或输出分布，再决定要不要上量化。算子库里的量化 MatMul 实现通常会封装反量化逻辑，但模型侧校准仍要自己做。
5.3 可迁移到日常开发的三条判断准则
把 MatMul 的优化经验收拢一下，其实可以抽象出三条通用准则，放到任何性能敏感的开发里都成立：
- 动手前先算算术强度。判断瓶颈是访存还是计算，不同结论对应的优化动作完全不同。
- 用 profile 数据验证假设。不要凭猜，每一步改动都要能解释为什么变量在变。
- 尊重 shape 的多样性。一套静态最优解很难通吃所有场景，调度器最好能根据输入形状动态切换策略。
这三条听上去很朴素，但我在实际项目里反复吃过亏，才把它们变成肌肉记忆。尤其是第一条，很多优化方案看起来很高级，一放进 Roofline 框架里就原形毕露：要么是在给不存在的计算瓶颈加速，要么是绕过了一个根本不是瓶颈的访存。
最后分享一个小技巧。如果你也想研究这类算子库，别一上来就扎进汇编级的 kernel 代码。先把算子的接口、shape 推导、tiling 配置、多核调度这些“外围逻辑”看完，再结合自己的模型 profile 数据去理解 kernel。这样读代码的效率会高很多，而且更容易定位到实际问题。我自己就是在把 MatMul 算子拆完之后，才发现推理瓶颈背后的真正逻辑。
