“两个矩阵相减之前‘有多少子矩阵明明可以不乘’——这个问题往往到模型后端优化到后期才被认真问起。”如果把目光只放在 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 ScoreQuery 与 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、各类 MoEReLU/SiLU 稀疏激活是是激活后很多值为 0但稀疏度不规整ReLU 系列 FFNKV Cache 淘汰/压缩是是通过策略删除部分历史 tokenH2O、StreamingLLMRMM 输入自适应矩阵乘积归约是是根据输入生成子矩阵选择并在 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, block64, keep_ratio0.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(dim2) # [1, num_blocks] # 输入自适应保留 block_score 最大的前 keep_ratio 个 block keep_num max(1, int(num_blocks * keep_ratio)) _, top_block_idx torch.topk(block_score, kkeep_num, dim1) top_block_idx, _ torch.sort(top_block_idx, dim1) # 保持顺序稳定 # 从原始激活中 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, blockblock, keep_ratiokeep_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那就先考虑把它