最近为了把一个本地模拟的 decoder-only 大模型推理速度提上来我把某国产 NPU 异构计算架构的算子库从上层接口一直翻到了底层 kernel。翻完最大的感受是大模型上层那些 Attention 变体、KV Cache 优化都很重要但真正兜底性能的还是那个最不起眼的 MatMul 算子——大语言模型矩阵运算的核心组件几乎每一层都在反复调用它。这篇文章不算什么教科书讲解而是结合我实际排障过程把这个算子在算子库里的实现骨架、性能边界、优化取舍拆开讲清楚。想快速定位大模型推理性能瓶颈的人、做算子开发或者框架性能优化的人应该都能从里面找到可以立刻拿去用的判断方法。先说明一下文中的算子库并不特指某一个商业产品而是我在本地环境中研究过的、常见的 NPU 异构计算架构里的算子集合接口逻辑和主流算子库大同小异。为了避免歧义下文统一用“NN 算子库”代称。MatMul 在所有深度学习框架里都是高频算子但由于 NPU 的硬件调度方式和通用 GPU 不同它的内部设计侧重点很不一样。这篇文章会尽量把“为什么这样设计”讲透。1. 先理解为什么 LLM 每一层都在碰 MatMul1.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 并不是只有一种实现。从接口上看它通常接收两个输入 tensora 和 b再通过 transpose_a、transpose_b 这类参数决定是否转置。根据输入维度和数据格式至少能分出三种常见形态标准 GEMM即通用矩阵乘对应线性层、投影层。输入是二维矩阵公式为 C alpha * A * B beta * Calpha 和 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 sizeM 很大形状相对规整追求的是吞吐推理时有两个截然不同的阶段prefill 阶段要处理整段输入M 等于输入序列长度可能很大decode 阶段每次只生成一个 tokenM 等于 1。这种“M 一会儿大一会儿小”的动态 shape会让同一个 MatMul 算子呈现出完全不同的性能特征。M1 时计算量极小权重搬运占主导M1024 时计算量骤增矩阵乘法单元是否饱和成为关键。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_size4096decode 阶段 M1权重是 fp16FLOPs 2 * 1 * 4096 * 4096 ≈ 33.6 MFLOPs访存字节数主要是读取权重 B4096*4096 个 fp16 约等于 33.6MB算术强度 ≈ 1 FLOP/Byte这个强度低得吓人。换句话说每从主存读一个字节只做了一次浮点运算。再对比 prefill 阶段 M1024FLOPs 2 * 1024 * 4096 * 4096 ≈ 34.4 GFLOPsA 的大小变成 1024*4096 个 fp16大约 8.4MBB 仍然是 33.6MBC 大约 8.4MB算术强度 ≈ 34.4e9 / 50e6 ≈ 700 FLOP/Byte同样一个权重矩阵M 变大之后算术强度提升了几百倍瓶颈属性完全反过来。这就是为什么大模型推理里“decode 慢”那么普遍它一直在访存受限区里跑。2.2 同一个算子在 Roofline 两侧反复横跳MatMul 比较特殊的地方在于它既会被归为访存密集又会被归为计算密集关键看 shape。NN 算子库里的 MatMul 为了兼顾两种场景往往会准备两套优化路线。访存受限时优化目标不是把矩阵乘法单元打满而是减少主存访问。手段包括把权重尽量压在缓存里、多请求合并 batch、使用低比特权重、避免中间结果来回写。计算受限时优化目标变成提高矩阵乘单元利用率比如选择合适的分块、把数据搬运和计算重叠起来、用低精度乘加指令。这两种思路如果搞混会出现很典型的性能问题一个人拼命调分块大小想把计算流水线加满结果 M1 时根本没有足够多的计算任务所有的努力都在给一个“等待权重搬运”的过程做装饰另一个人拼命做 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:moMT, ko:koKT]; load B_tile B[ko:koKT, no:noNT]; C_tile A_tile * B_tile; } } }MT、NT、KT 怎么选直接影响性能。分块太大装不进片上存储分块太小复用率不够。NPU 上还得对齐矩阵乘单元的形状比如某些硬件矩阵乘单元一次处理 16×16、32×16 这样的块。算子库里的 MatMul 之所以看起来复杂很大一部分就是在处理这些“对齐”和“切分”的细节。3. 算子库中 MatMul 的实现骨架从接口到 Kernel3.1 算子描述、参数校验和 Shape 推导NN 算子库里任何一个算子对外都有一个标准描述算子名、输入 tensor、属性、输出 tensor。MatMul 的属性通常包括 transpose_a、transpose_b以及可能存在的 alpha、beta。在真正进入计算前算子框架会先做 shape 推导和参数校验。这一步看起来简单但直接影响图编译阶段能不能静态分配内存。比如动态 shape 下如果 shape 推导失败整个图的执行计划可能退化成运行时推断带来额外开销。我在调试模拟项目时就遇过一个问题输入的 seq_len 在 prefill 和 decode 之间切换某些绑定静态 shape 的 kernel 被强行兜底到通用的慢版本profile 里一眼就能看出来。后来把 shape 推导修好让调度器正确识别 M1 的小形状MatMul 耗时立刻降了一截。参数校验还包括 dtype 检查。fp16、bf16、int8、fp32 混用是常事算子的实现必须明确累加器精度和输出 dtype否则精度问题会一层层传上去。3.2 Tiling 策略是性能好坏的分水岭分块听上去简单做起来全是细节。同一个 MatMulM1 和 M1024 的最优分块可能完全不同。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 transformerhidden_size409632 个 headhead_dim128权重全部 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 的 M1算术强度在 1 到几个 FLOP/Byte 之间清一色访存受限。结论很明确再去纠结矩阵乘单元的利用率没有意义M1 的任务根本喂不饱计算单元。真正的目标是让权重访问变得“便宜”要么一次少读一点要么让读过来的数据多算一点要么干脆别反复读同样的权重。顺着这个思路我做了三件事。第一检查调度器是否给 M1 选用了小形状专用 kernel避免大 tile 带来的无谓循环开销第二尝试把部分权重从 fp16 压到 int8同时用 per-channel scale 做量化补偿第三把 QKV 投影、FFN 这类多个 MatMul 和中间激活合并成融合算子减少中间变量写回主存的次数。4.3 动刀之后90 毫秒降到了 42 毫秒这三件事不是一次全上的而是一步步验证。先开 M1 专用 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-tripdecode 场景下收益尤其大。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 算子拆完之后才发现推理瓶颈背后的真正逻辑。