“两个矩阵相减之前,‘有多少子矩阵明明可以不乘’——这个问题,往往到模型后端优化到后期才被认真问起。”
如果把目光只放在 CUDA Kernel 和算子融合上,能提的优化会越来越薄。矩阵乘法的每个浮点都被 tensor core 吃掉的时候,再想省出 20% 的端到端延迟,就得换一种思路:不要总想着把“上一次乘法”算得更快,而要想着如何让这一次矩阵乘法本身更小。这正是标题中 “Matrix-Product Reduction” 想表达的意图:在输入(input)的引导下,把矩阵乘积的有效计算量降下来,而不是只在 GEMM 的循环展开和分块上做文章。
按照标题来看,这属于 LLM 推理侧一个很值得捋清楚的新优化方向。很多读者可能第一反应会问:这跟稀疏矩阵乘法有什么区别?跟 MoE 的路由、跟剪枝又有什么边界?看完这篇文章,你会得到一套相对清晰的判断框架,也能在自己的模型服务里走一遍“输入自适应 → 选择子矩阵 → 做缩减后的 GEMM”的基本流程。文章不以“某个框架多快”作为结论,而把核心放在原理、适用性、代价和工程落地上。
1. LLM 推理为什么绕不开“矩阵乘法变少”这个问题
Decoder-only 结构的大模型,在生成阶段每一次只产生一个 token,但模型内部要执行的矩阵乘法和训练阶段没有本质区别。把一层 Transformer 运算拆开看,普通 Dense LLM 至少包含这几组矩阵乘法:
- QKV 投影:把输入向量从隐藏层维度线性映射到 Query、Key、Value 空间;
- Attention Score:Query 与 Key 做矩阵乘积,得到注意力权重;
- Attention 输出投影:把加权后的 Value 结果重新映射回隐藏层维度;
- 第一个 FFN 线性层:将隐藏层升维到中间维度;
- 第二个 FFN 线性层:把中间维度降维回到隐藏层维度。
这还只是单层。层数一多,LLM 推理的计算量就高度集中在“高维矩阵乘法”上。当模型支持 128K、256K 长上下文时,Attention 部分的矩阵张量还会按照序列长度的平方或者与 KV Cache 规模成正比地膨胀。服务侧经常看到的“GPU 算力没有打满但延迟很高”,往往不是单个 Kernel 算得不够快,而是模型内部存在大量低信息密度的矩阵运算,它们把显存带宽和算子发射周期占住了。
如果只看 Kernel 层面,CUDA 生态里已经有非常成熟的 GEMM 优化框架,比如 CUTLASS、CUBLAS 以及各种融合 Attention 实现。这类优化的核心目标非常朴素:对于一个形状确定的矩阵乘法,尽可能利用 tensor core、共享内存、向量化访存把计算压到极限。问题是,它默认“这次乘法所有参与计算的行、列、块都有必要算”。
RMM 这类思路的切入点是:先别默认这个假设成立。同一个权重矩阵在训练完成后是固定的,但推理阶段的输入是变化的。某些输入 token 激活的神经元分布完全不同,某些中间向量经过激活函数后已经大量趋近于零,某些历史 KV 位置与当前 query 的相关性极低。如果能在矩阵乘之前有效判断“哪些块可以被安全跳过”,那么需要的就不再是“更快的 GEMM”,而是“更少的 GEMM”。把这两个层次分开思考,是理解 LLM 推理优化的一个关键进步。
2. “矩阵乘积归约”到底归约了什么
对“Reduced Matrix Multiplication”产生误解,通常是从英文单词 Reduction 开始的。在并行计算里,Reduction 常常指用加法和比较把一组数据聚合成更小的结果,比如数组求和、求最大值。但这篇文章标题里的 Matrix-Product Reduction,并不是 CUDA 里那种 block-reduce 的归约。
如果把它放在大模型推理语境里,更准确的理解是:减少矩阵乘积本身的规模。也就是原本要算一个比较大的矩阵乘法,经过输入条件判断后,只取其中一部分子矩阵进行计算,最后仍然得到等价或近似等价的输出。
举个例子。设模型某一层线性映射为:
Y = X @ W^T其中 X 的形状为[1, H]或[T, H],是当前输入激活;W 的形状为[M, H],是层权重。传统做法直接调用 GEMM 计算整个 Y。RMM 的思路则是先找到一个与输入相关的判断条件,例如基于前一阶段激活的阈值、top-k 排序或低秩打分,认为权重 W 中只有部分行对应的中间神经元在当前 token 上是有意义的。于是只取出这些行构成子权重W_sub,执行:
Y ≈ X @ W_sub^T如果选出的行数远小于 M,矩阵乘法的主循环次数就变少了。这样的做法可以发生在 FFN 的升维、降维层,也可以发生在对 Attention Score 分块的行、列压缩上。核心特征是“减”和“择”发生在矩阵乘之前或中间,而且选择结果依赖当前输入,不是训练后固定不变的。
因此,“Input-Adaptive”这个词非常关键。预训练模型的权重虽然固定,但不同输入产生的激活差异很大。若把“该减哪些块”变成一个编译期静态决定,那就是常规结构化剪枝;若把“该减哪些块”交给当前 token 的激活分布实时决定,才是标题里的输入自适应。也正是这个特性,让 RMM 和静态剪枝在工程上的挑战完全不同:它无法像剪枝一样提前把权重重新排布,只能在运行时快速生成索引,再交给后续 Kernel 去按紧凑内存布局执行。
3. 先看清冗余在哪里,再谈归约
不是所有矩阵乘法都值得做输入自适应归约。要判断哪些层适合,先要弄清楚一层计算中的冗余从何而来。
3.1 Token 或序列方向上的冗余
在生成阶段,自回归模型需要反复对历史 token 的 KV 向量做 attention。序列越长,KV Cache 越大,Attention Score 对应的 Key 侧矩阵也越大。可是对很多已生成的 token,它们的语义与当前生成位置相关性并不高;如果每一轮都把全部历史 token 同等对待,矩阵乘积里就会包含大量低贡献计算。
如果把 Key 向量按 token 分块,先通过一个低成本打分判断出哪些历史 token 块可能存在高注意力,哪些块基本不会获得权重,就可以只在少部分 Key 块上执行完整的 QK^T 矩阵乘法。上下文越长,这种序列方向上的可削减空间越明显。
3.2 隐藏层神经元方向上的冗余
FFN 中间层常常是“输入自适应稀疏激活”的重灾区。很多模型使用 SiLU、ReLU 这类非线性激活。实际推理时,输入一个 token,经过 gate 分支后,会有相当比例的中间神经元激活值非常小,甚至等于零。
如果某个中间神经元在当前输入下激活为零,那么它对应的权重行在后续 down-projection 矩阵乘法中参与计算的贡献也接近零。提前把这一列/行索引出来剔除,后续矩阵乘法就可以从全量中间维度退化为有效活跃维度。严格选择还是近似选择,取决于算法设计,但总思想都是让矩阵乘积的活动区间跟着输入走。
3.3 Attention Head 和矩阵块方向上的冗余
多头注意力中,不同 head 在不同输入上关注模式差异很大。某些 head 对当前输入的信息增益很低,但这部分是否值得做输入自适应归约需要谨慎。因为跳过一个 head 不只是减少一次矩阵乘法,还可能改变多头拼接结果的结构和语义。更工程化的做法是在每个 head 内部再做分块或沿序列方向做裁剪,而不是直接丢弃整个 head。
整体上,RMM 的归约单元可以是 block、row group、column group,也可以是 token group。粒度越细,理论上能省掉的浮点运算越多,但运行时索引和重排也越复杂。粒度太粗则难以达到收益。实际项目里,既要让跳过的块足够“整”,又要尽量接近真实信息密度边界。
4. 它和剪枝、稀疏、MoE 到底有什么区别
很多读者会把“输入自适应矩阵乘积归约”和其他热门优化放在一起比较,这里值得认真做一次区分,避免后续选型时张冠李戴。
| 优化方法 | 决策是否依赖输入 | 权重是否长期固定 | 一次推理中的开销形态 | 典型代表 |
|---|---|---|---|---|
| 静态权重剪枝 | 否 | 是 | 部署前已经把无效参数去掉 | 结构化稀疏、N:M 稀疏 |
| MoE 专家路由 | 是 | 是 | 每个 token 选择少数专家 | Mixtral、各类 MoE |
| ReLU/SiLU 稀疏激活 | 是 | 是 | 激活后很多值为 0,但稀疏度不规整 | ReLU 系列 FFN |
| KV Cache 淘汰/压缩 | 是 | 是 | 通过策略删除部分历史 token | H2O、StreamingLLM |
| RMM 输入自适应矩阵乘积归约 | 是 | 是 | 根据输入生成子矩阵选择并在 GEMM 前缩减计算 | 本文讨论的主题 |
静态剪枝是在模型层面把权重矩阵的某些位置永久去掉,它与输入无关;MoE 是在“多个独立专家”之间做路由,权重本身分属不同参数集;而 RMM 更接近“同一份权重矩阵内部,按当前激活裁出一个子矩阵来做计算”。它和稀疏激活关注的问题有交叉,但要求更明确——稀疏激活只是说让矩阵中出现零值,RMM 则是希望把有零值的位置直接变成更小的稠密 GEMM,让底层算子计算时不触碰零值数据。
所以,RMM 的优势不是创造出新的数学结构,而是为“激活稀疏”和“上下文动态相关”找到一套更适合 executor 的形态。它把“值接近零”从数据面上的特征,翻译成“子矩阵选择”这种算子层面的动作,最终受益的还是 GEMM 循环层。
5. 一次输入自适应矩阵乘积归约的核心流程
把前面的思想落地成代码前,先把流程拆成几个阶段,这样看代码时会更容易理解每步在做什么。
5.1 计算轻量 gate,生成候选活跃度
RMM 不会对每一处权重都无中生有地去猜哪些 block 有用,而是依赖一个计算量很小的 gate 或打分器。常用的做法是复用模型自身前一层算子产生的激活。比如对 FFN,可以先算 gate 分支的中间结果,然后得到一个维度为中间层大小的向量。把这个中间向量按 block 聚合,得到每个 block 的平均绝对值,作为该 block 的“信息重要度”。
5.2 把重要度转化为块掩码
连续的 block 选择比零散元素选择更适合底层 GEMM。假设中间层有 16384 个神经元,按每 64 个神经元分成 256 个 block。对 256 个 block 做 top-k 或阈值筛选,得到保留下来的 block 编号。这时我们可以构造一个长度为 256 的 bool mask,再把 mask 扩展成与原始神经元一一对应的索引向量。
5.3 使用 gather 得到紧凑子矩阵
选中 block 后,原始中间激活和共享权重要做两个 gather 操作。激活矩阵按 block 抽取,权重的中间维度也按同一组 block 抽取。此时两个矩阵可以在连续内存里组成更小的一次 GEMM。带来的性能收益主要来自这里——不再是 M 维度全量计算,而是实际的active_M维度计算。
5.4 结果放回与准确性校验
缩减后的矩阵乘法输出的是某些位置上的值,对于被跳过的位置,往往需要补零或做 mask。另一种更简单的形式是只把最终输出 y 计算出来,因为 y 的矩阵形状并没有变。如果中间步骤需要和其他分支相加,则要保持被跳过位置不影响数值逻辑。工程上建议在实现每一层 RMM 后,对比全量计算与缩减计算的输出余弦相似度或均方误差,合格后再接入上层。
6. 一个最小可运行的代码示例
在这个示例中,我会模拟一个带 gate 的 FFN 层,并用 block 级别的输入自适应选择缩减 down projection 的矩阵乘法。
6.1 环境准备
代码依赖 PyTorch,不需要特殊硬件也能先跑通逻辑。若想看到时间收益,建议在 NVIDIA GPU 上运行并反复 warmup。Python 版本、PyTorch 版本以实际可用为准,本文重点是演示思路。
# 文件路径:demonstrate_rmm_ffn.py import torch import torch.nn.functional as F def ffn_full(x, W_gate, W_up, W_down): """ 传统 Dense LLM 中常见的 FFN 计算。 W_gate, W_up: [M, H] W_down: [H, M] x: [1, H] """ gate = F.linear(x, W_gate) # [1, M] up = F.linear(x, W_up) # [1, M] h = F.silu(gate) * up # [1, M] y = F.linear(h, W_down) # [1, H] return y, h def ffn_block_reduce(x, W_gate, W_up, W_down, block=64, keep_ratio=0.5): """ 输入自适应的 down projection 缩减版本。 先根据 gate/up 激活结果选择需要保留的 block,再做一次紧凑 GEMM。 """ M = W_gate.size(0) assert M % block == 0, "示例代码要求 M 能被 block 整除" gate = F.linear(x, W_gate) # [1, M] up = F.linear(x, W_up) # [1, M] h = F.silu(gate) * up # [1, M] # 把中间维度切成 block,计算每个 block 的平均绝对激活 num_blocks = M // block h_block = h.view(1, num_blocks, block) # [1, num_blocks, block] block_score = h_block.abs().mean(dim=2) # [1, num_blocks] # 输入自适应:保留 block_score 最大的前 keep_ratio 个 block keep_num = max(1, int(num_blocks * keep_ratio)) _, top_block_idx = torch.topk(block_score, k=keep_num, dim=1) top_block_idx, _ = torch.sort(top_block_idx, dim=1) # 保持顺序稳定 # 从原始激活中 gather 出保留 block h_keep = h.view(1, num_blocks, block)[:, top_block_idx[0], :] # [1, keep_num, block] h_keep = h_keep.reshape(1, -1) # [1, keep_num * block] # W_down 形状是 [H, M],同样按 block 做 gather W_down_block = W_down.view(W_down.size(0), num_blocks, block) # [H, num_blocks, block] W_keep = W_down_block[:, top_block_idx[0], :] # [H, keep_num, block] W_keep = W_keep.reshape(W_down.size(0), -1) # [H, keep_num * block] # 紧凑矩阵乘:此时实际参与运算的中间维度变小 y_reduce = h_keep @ W_keep.T # [1, H] return y_reduce, top_block_idx[0]代码里有两点需要解释。第一,top_block_idx是输入相关的,同一个权重在不同 token 下会选择出完全不同的 block,这正是“输入自适应”的朴素体现。第二,为了能真正在 GPU 上受益,选择粒度必须是 block,而不是单个神经元或单个元素,因为零散的元素排列会让底层矩阵乘法失去稠密 tensor core 的优势。
再看一段验证函数,用来对比全量结果、缩减后的时间以及输出差异。验证逻辑不追求零误差,因为它本质上是一个启发式近似;重点是把误差暴露出来,方便阅读者判断阈值或 keep_ratio 是否合理。
# 文件路径:demonstrate_rmm_ffn.py(续) def run_demo(): torch.manual_seed(42) H = 1024 M = 4096 block = 64 keep_ratio = 0.5 x = torch.randn(1, H) W_gate = torch.randn(M, H) * 0.02 W_up = torch.randn(M, H) * 0.02 W_down = torch.randn(H, M) * 0.02 y_full, h_full = ffn_full(x, W_gate, W_up, W_down) y_reduce, keep_idx = ffn_block_reduce( x, W_gate, W_up, W_down, block=block, keep_ratio=keep_ratio, ) print("full output shape:", y_full.shape) print("reduce output shape:", y_reduce.shape) print("kept block ratio:", keep_idx.numel() / (M // block)) # 不追求完全一致,但观察两者的大致差距 diff = (y_full - y_reduce).abs().mean().item() print("mean abs diff:", diff) if __name__ == "__main__": run_demo()运行方式很简单:
python demonstrate_rmm_ffn.py预期输出中,两个输出 shape 都是[1, H],这说明缩减后的矩阵乘并没有改变字段层的输入输出协议。mean abs diff 的大小由随机初始化的权重和 keep_ratio 决定。如果 keep_ratio 调到 1.0,理论上 reduce 后的结果应该与 full 完全一致。
6.2 算力对比角度
可以把“满计算”和“缩减计算”的浮点运算量做一个粗略估算:down projection 这一步的主要浮点量约等于2 * H * active_M。若 active_M 是 M 的一半,理论矩阵乘法浮点量就减少一半。但实际端到端收益不会等于这个比例,原因是 gate 与 up 的线性映射没有缩减,仍要全量计算。
矩阵乘法压缩是一层一层累积出来的:attention 的 QK 部分可以按 context block 压缩,FFN 的 down projection 可以按中间神经元 block 压缩,后面的输出投影又能继续依赖前面的算子结果做输入自适应。只有当多层都有机会压缩时,端到端效果才足够明显。单层单独压缩,很可能被额外索引和重排成本抵消。
从工程角度,建议先给模型层铺上 profiling,记录每层、每个矩阵乘的输入活跃度分布,再决定实现哪些层的 RMM。不要一开始就把所有层都改造成动态索引,“全都要”的后果往往是 profiling 复杂度和数值风险同时上升。
7. LLM 推理中哪些层更容易受益
不同矩阵乘法的可缩减空间差异很大。下表可以作为初步选型参考。
| 矩阵乘法位置 | 输入自适应依据 | 可缩减维度 | 工程难度 | 收益潜力 |
|---|---|---|---|---|
| FFN down projection | 激活后的中间神经元活跃度 | 中间维度 M | 中等 | 高 |
| FFN gate 与 up projection | 前一层输出对若干神经元的驱动 | H 维度 | 较低,但收益有限 | 中 |
| Attention QK^T | 历史 token/KV block 与当前 token 的相关性 | 序列长度 T | 较高 | 长上下文时极高 |
| Attention 输出投影 | 是否真正需要融合大量 Value 向量 | Value/序列方向 | 高 | 中高 |
| Embedding/LM Head | 极少直接做动态选择 | 词表维 | 较低 | 视词表稀疏性而定 |
对于 FFN,中间层维度通常远大于隐藏层维度,所以只要能把 block 级别稀疏性用起来,理论降幅会非常明显。对于 Attention,收益主要体现在长上下文场景。短上下文时,序列本身不长,固定开销、索引计算和 kernel launch 反而可能淹没收益。
这里要特别提醒:KV Cache 淘汰与 Attention 里的矩阵削减并不相同。KV Cache 淘汰是删除一些 Key 和 Value,之后根本不参与计算;而基于矩阵乘积归约的 Attention 优化,更像是在做 QK^T 之前先不计算低相关部分的矩阵乘。如果采用硬删除策略,需要保证被删除的位置不会因后续层交互重新产生重要影响。对于有安全性要求的任务,建议先用软性掩码做离线实验,再决定是否采用硬删除。
8. 做得不好会有什么代价
任何优化都有代价。RMM 的代价主要在四个方面。
第一是索引和重排开销。每次输入不同,子矩阵选择就不同,需要将活跃 block 的数据复制到连续内存。这个复制过程本身占用显存带宽。如果一块数据在原始矩阵里本来就内存连续,但被 block gather 后变成非连续访问,复制开销会抵消掉一部分 GEMM 收益。
第二是数值误差风险。基于 top-k 或阈值截断的 RMM 天然是近似计算。被跳过的 block 可能在某些样本上信息量不高,但在极端样本上突然重要。生产环境只有先做充分离线测评,才能回答“近似误差是否可接受”。对输出 logits 极敏感的任务,可能还要把阈值调得非常宽松,从而导致收益变小。
第三是 Kernel 复杂度上升。真正把 RMM 用到高性能 GEMM 上,通常需要针对 block 数量设计多个 kernel 路径:当活跃 block 较多时走全量路径;当活跃 block 数量中等时走紧凑 GEMM;当特别稀疏时还要考虑转成稀疏 kernel。每一种路径都要验证,复杂度比单纯调一个 CUBLAS 调用要高得多。
第四是动态 shape 与批处理冲突。如果一次推理服务的 batch size 很大,每个输入 token 选择的 block 不同,为了保持矩阵乘形状规整,往往要按 block 做求并集或按公共 mask 计算。这会让“每个 token 节省 40%”变成“整个 batch 只能节省 15%”。这并不代表思路无效,而是提醒你在批处理场景下需要更精细的调度。
9. 常见问题与排查思路
| 问题现象 | 可能原因 | 排查方式 | 解决方案 |
|---|---|---|---|
| 缩减后速度反而变慢 | 活跃 block 太少时索引、复制开销高于 GEMM 节省 | 用 profiling 统计 kernel 耗时和 memcpy 耗时 | 增大 block 粒度,减少动态索引次数;先固定 block mask 跑一版 |
| 输出与全量结果差异过大 | top-k 或阈值没有贴合模型激活分布 | 统计每层激活强度分布与方差 | 调整 keep_ratio,改用分位数阈值,或引入校准集 |
| block 不规整导致 Kernel 报错 | M 无法被 block 整除 | 检查权重形状和维度 | 对 padding 后权重做掩码,或选择能整除 M 的 block 值 |
| 批处理时收益大幅下降 | 不同样本保留的 block 不一致,出现求并集膨胀 | 统计 batch 内 block mask 的交并比 | 对相似样本分桶,或对每个 block 使用计数路由 |
| 长上下文场景不稳定 | 只按位置或简单分数丢弃历史 token,造成重要信息丢失 | 用局部敏感度测试逐层观测 attention score | 引入第二道回退检查,先做低风险淘汰 |
10. 工程落地中几条比较务实的建议
第一,先做 block 级 profiling,再做算子改造。在设想任何节省比例之前,先用一个离线脚本统计每层激活的 block 活跃率、KV block 的相关度分布。如果发现某层即使全量计算也只有极少数 token 会走到边缘 block,那就先考虑把它