先给结论:这个主题如果落地成工程实验,核心不是再发明一种新的矩阵乘法算子,而是回答一个问题——同一个 LLM 推理过程里,矩阵乘法到底是按固定形状被反复执行,还是可以跟着输入变成更小的形状、更短的序列、更窄的通道范围去执行。这个问题解决得好,收益不一定出现在峰值算力上,而是出现在推理吞吐、显存占用和长上下文场景里更明显。
我看到的标题关键词是 Reduced Matrix Multiplication、Input-Adaptive Matrix-Product Reduction、LLM Inference。从这三个词里能拼出两条主线:一条是对矩阵乘法本身做“减少”,比如减少乘了多少次、少算哪些位置、把一串矩阵乘积的中间结果压缩掉;另一条是“输入自适应”,也就是减少的策略不能写死,要随输入变化。这里说的输入不只是用户发来的 prompt 文字,还包括 batch 里句子长度、KV cache 状态、激活值的分布、当前 token 位置这些动态信息。
换句话说,它不是一篇单纯讲低秩矩阵近似的文章,更像是在探讨:在做 LLM 推理时,能不能按当前输入条件临时决定哪一次矩阵乘积可以被约减,并且保证模型输出不崩。
下面我会把这个思路拆成几条可执行的技术路径讲,同时也把容易翻车的地方拉出来逐个过一遍。
1. LLM 推理的矩阵乘法开销到底集中在哪,Reduced Matrix Multiplication 能切哪里
1.1 从一次 LLM 前向计算看乘法密度
大语言模型推理过程里,绝大多数计算都落在矩阵乘法上,这不是夸张。一个普通的 decoder-only transformer,在生成每个 token 时都会执行同一套线性变换:词嵌入查找后要过 query、key、value 投影,注意力分数算完之后要做 output 投影,然后 MLP 里还有 up 和 down 两个投影。新一点的模型还会用 gated 结构,比如先算 gate、up 两组结果,再做点乘或合并,再进 down 投影。
这不是一次就结束。生成第几个 token,就要把这套流程从头到尾再走一遍。即使优化过的推理框架会做 KV cache,避免重复计算前面 token 的 key 和 value,query 投影、attention output、MLP 还是每步都要算。随着序列变长,序列维度的计算量会继续上升,显存占用也集中在 KV cache 和中间激活上。
这意味着,如果真要对矩阵乘法做“减少”,可下手的位置很多:
- 注意力里的 Q、K、V 投影矩阵;
- 注意力分数乘 V 之后的 output 投影;
- MLP 里的大矩阵;
- 以及不同矩阵连乘时产生的中间结果。
但这里有个很容易踩的误区:不要一看“约减矩阵乘法”就以为只是把矩阵变小一点。把 4096 维投影砍成 2048 维,那是模型压缩,是另一个领域。模型结构固定之后,真正能想办法减少的,是“这次输入条件下到底有没有必要把全部维度都计算完”。
1.2 从算子层面看,什么是矩阵乘积的约减
矩阵乘法如果用最朴素三层循环来看,是拿两个矩阵做乘加,最后还需要做 reduction 把中间乘积累加到一起。比如 A 是 m×k,B 是 k×n,C 里每个元素都要对 k 维做累加。这个 k 维就是最典型的 reduction dimension。
普通情况下,GPU 会把这项工作交给高度优化过的 BLAS 算子,算子内部有分块、tile、寄存器复用、向量化,几乎不需要普通应用层去干预。
但如果标题里的 Reduced Matrix Multiplication 指的是“把这个 reduction 过程从固定 k 变成输入可变的 k'”,那就进入了新问题:k' 的选择不能按全局固定,而要取决于当前输入向量在哪些分量上有主要能量。这和稀疏矩阵、低秩矩阵的思路不完全一样。稀疏矩阵是静态表明某些位置为零;低秩近似是找一个固定的小维度替代大维度。输入自适应约减更像是在运行时评估:在这一次的输入分布下,有哪些运算可以被剪掉或提前结束,以不显著影响输出分布为准。
理解了这个层次,就会明白为什么不能指望随便实现一个内核就一步到位。真正难的不是怎么把矩阵变小,而是在什么条件下变小不会伤害生成质量。
1.3 输出对输入自适应的要求
LLM 有大量非线性层,比如 GELU、SiLU、softmax、LayerNorm。这些非线性会让参数矩阵的有效利用程度随输入改变。
例如一个 batch 里有一条很长的请求和一条很短的请求。如果框架把它们强行 pad 到相同长度再统一走矩阵乘法,短句子的很多 padding 区域都在浪费计算。如果框架按长度分桶,或者把长句和短句拆到不同步数处理,同样的计算量就能更集中地用在有效 token 上。
这种“按 batch 和序列长度输入的矩阵乘法”也属于一种明显的 input-adaptive 操作。
类似的,在 gated MLP 里,gate 分支决定哪部分激活要保留。如果某一段激活经 SiLU 后接近零,它对最终输出的贡献就很小。理论上有条件对 down 投影的输入列做裁剪,把这个 token 里某些通道强制置零或跳过,从而在 down 投影里少算一部分矩阵乘积。实际做的时候要谨慎:跳跃通道会不会导致最后一层 logits 翻盘,需要实验验证。
所以这一个主题真正值得研究的一条线索是:LLM 推理的所有矩阵乘法里,哪些计算可以用低风险的规则临时减少,哪些计算即使冗余也必须保留。
2. 输入自适应不等同于低秩近似或静态稀疏,必须分清边界
2.1 低秩矩阵不用做运行时判断,输入自适应需要
低秩近似在模型推理中已经非常常见。把全连接层的权重 W 分解成 U、V 两个更小矩阵,或者用低秩 adapter 去适应下游任务,都可以减少矩阵乘法量。早年间很多模型压缩工作都围绕这个思路展开。
但低秩近似有一个明显特点:它把权重矩阵的 rank 看成相对稳定,压缩关系一旦训完就固定了,不会因为用户输入变长或变短而发生变化。
Input-Adaptive 思路不满足这种静态假设。它强调系统应当根据给定输入状态动态地改变矩阵乘积的形式。例如:
- 对当前 token 来说,某些通道的输入激活值普遍很低,就可以选择跳过这些通道参与的乘积计算;
- 对整段序列来说,如果上下文长度很短,就可以用更小的中间矩阵尺度去执行某一部分计算;
- 对重复出现多次的矩阵乘法链,可以在中间结果已经足够收敛时提前终止累加。
执行这类逻辑时,应用层要拿到算子的调度权,所以也不是简单地换一个 torch.matmul 能解决。你需要在整个推理管线里增加采样、判断和分支入口。
2.2 静态稀疏算不好,动态稀疏才难
静态稀疏基本可以留到模型训练或转模型时处理。例如把不重要的权重剪枝成稀疏矩阵,推理框架按稀疏格式存储和计算。只要权重不动,稀疏模式不动,算子很容易优化。
输入自适应的约减更像动态稀疏或条件计算。这里有两个难点:
- 每个输入 batch 的稀疏模式不一样,意味每层都要重新判断哪些计算要减掉;
- 判断本身也要付出代价。如果判断一次要读入全部激活统计,再做矩阵裁剪,这部分代价很可能超过真正省下的矩阵乘法开销。
因此,真正的 RMM 系统不能只用“减少的 FLOPs”衡量收益,还要把输入分析、决策、内存拷贝、kernel 切换这些 overhead 全部加起来看真实延迟。我的建议是,每一步都测端到端 token 生成时间,不要只看单个 matmul 峰值对比。
2.3 需要区分“减少矩阵乘法”和“减少矩阵乘法的执行次数”
标题写成 Matrix-Product Reduction,有一种解释是:LLM 推理存在很多连续矩阵乘积,比如先算 XW,再算后面一层矩阵。理想情况是把几个线性变换合并成单次 更高维度的线性变换?但从数值上看,两个矩阵相乘后接非线性层,再把结果传给下一层,中间夹了激活函数,不是简单能提取公因式。
比较实际的样子反而可能是:XW1 出来是一个大矩阵,接着要乘以 W2。矩阵乘积的顺序是先做第一个乘积,存下中间结果,再做第二个乘积。若我们能把两个矩阵乘法的外层或内层维度裁剪,让第二次只需要用到中间结果的一部分,整体上多个矩阵乘积都会被“reduce”。
这就要说到执行顺序的重要性。比如线性层加线性层,中间无激活时可以提前合并两个权重矩阵,这也是权重融合的一种做法;但 LLM 中大多数层之间有激活或归一化,不容易这样合并。此时能做的是找到矩阵乘积链中的瓶颈维度,比如第一个矩阵输出的列数过大,又只有少量列会被第二个矩阵大权重行使用。这种结构允许做输入自适应的列裁剪。
不过这种判断不一定每次成立。层的输入是激活值,激活值来自上一层非线性,本身已经和输入产生复杂关系。真正是否可裁剪,必须放到一个固定模型上实测,不能只从理论上推。
3. 跑一个可验证的最小实验,先看到“输入自适应”有没有收益
3.1 实验环境怎么搭
我建议先找一个小模型或中等模型。做算子实验时不建议一上来就跑了几百亿参数的大模型,那样很难定位现象。可以用开源社区容易拿到的中等规模模型,也可以直接用自己已经部署过的模型。关键是能方便地替换 forward 函数,能够注入额外的输入分析和分支逻辑。
如果只是验证概念,默认环境大致需要:
| 条件 | 建议值 | 说明 |
|---|---|---|
| GPU 显存 | 至少能容纳模型和 batch 为 1 的中间激活 | 不需要做大 batch,但最好能跑长序列实验 |
| 推理框架 | PyTorch 或兼容的自定义运行环境 | 方便在 forward 里插入 hook |
| 数据 | 代表性 prompt 集合 | 不要只用一条,需要长短句混合 |
| 记录工具 | 日志 + 简单数值统计脚本 | 记录每层输出差、采样耗时、输出 token |
我倾向于先拿几十条不同长度的 prompt,固定 seed,先把模型原始输出完整记下来,生成一条“参照答案”。后面做任何约减实验,都要拿结果和这个“参照答案”对比。
3.2 选择模型的哪些层做切入口
前面提过,可以从三类矩阵乘法入手:
- attention 中的 qkv 投影;
- attention output 投影;
- MLP 的 up/gate/down 矩阵。
每一步矩阵乘法都可以看作一个 hook 点。在 PyTorch 这类框架里,可以用 forward hook 拿到输入张量,也能在计算前自定义替代实现。
一个具体做法是,把某一次矩阵乘法拆成主路径和次要路径来观察:
- 先算出当前输入张量每个通道的绝对均值;
- 把通道按贡献排序,只取前 80% 能量的通道;
- 删掉其余通道在第二个矩阵里的对应行,重新计算输出;
- 比较真实输出和裁剪后输出的差异。
这里示例代码如下,只是做概念演示,不是某个现成框架里的库函数:
def cut_matmul_by_channel_energy(input_tensor, weight): # 假设输入形状是 [seq_len, hidden_dim] # 按绝对均值估计激活能量,权重维度第一个维度对应 hidden_dim energy = input_tensor.abs().mean(dim=0) keep = energy > energy.mean() * 0.2 # 这是个示例阈值 reduced_input = input_tensor[:, keep] reduced_weight = weight[keep, :] return reduced_input @ reduced_weight这类代码有个明显问题:选出 keep 通道后要执行 gather 和切片,操作本身会带来内存拷贝。如果 hidden_dim 不大,预取释放出的算力不足以抵消管理开销。所以把这类思路放到需要优化的场景之前,你先做一次“足够大的矩阵乘法才值得切”的门禁。小维度切了只会更慢。
3.3 对比任务设计
最少要有三组实验:
- 基线组:原始注意力 + 原始 MLP,跑 20 条 prompt,记录输出 token、logits、耗时。
- 通道约减组:选定某些层做动态列裁剪,保证裁剪比例先从 5% 或 10% 开始。
- 输入长度分组组:按 sequence length 分桶,对长序列走更激进优化,短序列保持原样。
每组最好重复跑 3 次。模型推理本身有随机性,即使不开采样,GPU 上矩阵计算也可能由于并行结果略有差异。固定 seed 后普通情况差异很小,但仍建议多跑几次看方差。
判断收益时,重点不单是 PPL 或 loss,还要看:
- 同一条 prompt 生成的 top-1 token 序列是否一致;
- 裁剪后 logits 分布和原 logits 分布的具体差距;
- 推理的总时延和显存峰值变化;
- 越是长输入,收益是否越明显。
我自己做这类实验会先只看前 5 层或单层。因为整个模型 30 层以上,普通做法是把每层都套上逻辑。如果单层结果都不可控,整个模型更不可控。先从单层入手,能省下大量调错时间。
4. 关键指标与判定标准:怎么才算“减少成功”
4.1 不要只看速度,要看输出保真度
任何对 LLM 推理过程的约减都会威胁输出。你需要三个层面的指标:
| 指标 | 含义 | 通过标准建议 |
|---|---|---|
| 最大 logits 绝对差 | 裁剪后 logits 与原 logits 每个位置的最大差别 | 随模型规模与阈值而定,先用最小偏差 |
| top-k 翻转率 | 裁剪后 top-k token 是否改变或顺序变化 | 对分类/greedy 场景尽量为 0 |
| 平均生成 token 一致率 | 用同一 prompt 跑完整解码,对比逐 token 候选序列 | 刚开始接近 100%,后面再放宽 |
| 相对延迟变化 | 端到端耗时相对 baseline 是否变慢 | 至少不能变慢,否则无落地意义 |
模型数值天生带有噪声,所以不是每个位置都必须严格相等。但只要解码阶段采用 greedy 采样,任何一步造成 top-1 token 翻转,就可能让后面整段输出完全不同。这是我在做这类实验时最先接受的约束。
因此我建议设计两套判定标准:
- 轻量判定只在单次 forward 中验证误差;
- 完整判定需要跑采样或 greedy 解码,比较输出 token 序列。
轻量判定通过不代表整体能过;完整判定一旦没过,就可能要降低裁剪比例或把裁剪限制在特定层。
4.2 资源指标怎么写清楚
不要只说“耗时降低 xx%”。这类实验要写成可复现记录:
- 被裁剪矩阵在那个输入 batch 的实际 shape;
- 每次选择裁剪维度所耗时;
- 裁剪后余下矩阵乘法耗时;
- 端到端 batch 从输入到首 token 的总时间;
- 单 token 生成平均时间;
- 显存占用峰值。
只有把这些完整记录在案,才能判断收益到底来自哪一步,是不是换了 kernel 带来的偶然变化。
还有一个很容易踩的点:GPU 有没有统一内存或 CUDA graph。如果框架本身已经在用 CUDA graph 做静态图加速,动态增加分支会使 kernel 无法完全预录,图模式可能退化。此时看似省了 FLOPs,实际却慢了。
所以我给出的实际经验是:如果目标部署环境是高吞吐 batch 推理,别盲目引入太多动态 if 逻辑;先用静态 batch 大小、固定 prompt 长度把基准落稳,再考虑动态分支。
4.3 正常实验记录怎么排
下面是我会使用的日志字段:
layer_name input_shape weight_shape energy_threshold reduce_ratio original_output_norm reduced_output_norm max_abs_diff time_original time_reduced cuda_memory_delta运行结束后拿这些字段做汇总。不要只记录打印到 stdout 的对数,最好直接落成 CSV,后面好画曲线。若发现 reduce_ratio 高但输出误差也高,就要换掉当前策略。
5. 落到实际部署前,能扩展成哪些动态减少方案
5.1 基于 batch 和序列长度做组内矩阵乘积裁剪
实际推理时,一个 batch 内每个样本的序列长度可能差异很大。如果把所有样本 padding 到同一个最大长度后执行矩阵乘法,会有大量 padding token 参与计算。这类冗余完全不依赖模型输出质量,纯属输入长度不均衡带来的浪费。
按长度分组或分桶,是最容易落地的“input adaptive”方案。这种方案不要求你改动每层网络结构,只是改变同一时间进算子的输入形状。分桶粒度可以控制在 64、128、256 的倍数,尽量降低碎片化损失。
更好的方式是让推理框架支持“batch 内动态长度矩阵乘法”,用不规则调度去执行。不过很多框架目前还是先 bucket 再跑。对实验者来讲,先做 bucket 相对容易,收益也能量化。
5.2 基于激活稀疏性的列裁剪
这种方案要对激活值做统计。如果某一列在 batch 内激活都很小,那这一列在后续权重矩阵内整列对乘加结果贡献就更小。
一个比较稳的思路不是硬裁掉,而是先做“软裁剪”验证:
- 把激活值小的通道随机 mask 掉,跑一遍任务;
- 再把同样通道 mask 掉并观察 logits diff;
- 收集多批数据后,看哪些层/哪些位置经常能被mask。
反复多次后,有些层可能呈现出稳定的稀疏现象,就可以在部署时导入规则。这比一次性完全动态判断要安全一些。
但这个方法会有一个缺陷:激活值与 Batch 内容高度绑定。今天这批 prompt 下某通道不激活,明天换领域数据后可能又激活了。想作为生产方案,必须有足够多的样本去估计通道被裁剪的风险,不能只凭一次实验拍板。
5.3 对矩阵乘积的维度顺序做优化
LLM 中很多层可以写成几个矩阵乘积的组合。比如 MLP 里先算 XW_gate 和 XW_up,再做逐元素乘,最后乘 W_down。从算子角度看,等于 X 先分别和两个权重矩阵乘,产生了两个大中间矩阵。
某些实现会用变量让 Gate 在 down 投影前参与 pointwise 乘法。若 Gate 中部分元素数值在激活后接近零,理论上能跳过对应 down 投影的某些列或某些行。但困难仍然在判断和内存布局。
还有一种做法是换乘法结合顺序。如果第二矩阵的某个中间维度远小于输入维度,可以先乘第二个矩阵,再回头做第一次乘。这需要线性代数结构允许交换顺序,LLM 多层非线性层通常不允许。
所以实际项目能尝试的通常有两个方向:一是在可压缩矩阵乘积链中减少参与计算的维度;二是用近似算子代替完整矩阵乘积,比如更大粒度的稀疏块计算。目标不是零错误,而是让错误在可接受范围。
6. 最容易出问题和被高收益预期误导的地方
6.1 显存减少和 FLOPs 减少不一定成正比
要特别注意,很多约减方案是因为显存瓶颈才被提出来。比如减少中间激活后,显存可能下降。但 GPU 上单次矩阵的耗时不一定和 FLOPs 严格成正比,因为小矩阵无法打满算力,反而会让 GPU 利用率下降。
在低显存卡上跑大模型时,你会觉得计算很慢。但如果直接在算子层面把矩阵切小不一定变快,因为 GPU 要处理更多小 kernel,启动时间和带宽占用都会扩大。
不要只根据理论 FLOPs 就去推算收益。标准做法是先在真实 GPU 上跑 minibatch 和单条长序列样例。长上下文场景显存收益可能远大于短文本场景,矩阵乘法的形状切换成本又很干扰。要分开统计“短 prompt”和“长 prompt”两组结果。
6.2 把判断规则做进图编译层,而不是每层都用 Python 判断
如果按我前面说的,直接在模型 forward 里写 Python 判断,每次推理都要解释执行,耗时极高。生产化时必须把判断逻辑下放到自定义算子或编译图里。
一个可以接受的中间方案是:
- 使用 kernel 内按线程块做条件判断;
- 或使用静态 shape specialization,为不同输入长度生成不同 kernel;
- 不要在前向 Python 层做细粒度矩阵切片。
也就是说,Input-Adaptive Matrix-Product Reduction 一旦要落地,还必须考虑编译器和运行时对动态 shape 的支持。如果算子库只能处理静态 shape,动态列裁剪会导致每次 kernel 重新编译或多次内存复制。这个成本通常在短 batch 里是灾难。
6.3 注意输出 token 的连锁偏差
有些实验只测单层 forward 输出误差。问题在于单层误差小,不代表几十层累积后误差还小。LayerNorm 会对数值缩放敏感,后面的 MLP 又可能放大局部误差。即便是很小的最大绝对差,只要改变了一个 token 的排序,最后生成的文本就可能完全不同。
如果你想对完整模型使用这种思路,至少应设置一组“长链稳定性”测试:跑 20 条长 prompt,用 greedy 解码生成 50 到 200 个 token,对比原始和裁剪后的输出。任何一次 token 不一致,都是潜在风险,不能简单归因于模型随机性而忽略。
6.4 什么时候不要用约减
这是我一直保留的“边界清单”:
- 输入 batch 本身很小,矩阵单次执行时间低于约 10 微秒时,不约减;
- 判断和切片开销大于省下的矩阵计算时,不约减;
- 需要严格保证输出一致性的场景,比如某些审计或复现任务,不约减;
- 没有做好日志和可回退机制前,不约减;
- prompt 任务本身很短且模型不大时,约减收益很低,别增加复杂度。
一个系统真正可用,不只是“快了一点”,而是出错时可快速回退,可观测性够用。
7. 当你顺着这套思路做实验,遇到问题从这条链路排查
我会按这个顺序排查,而不是一头扎进改参数:
- 先看你到底改没改对矩阵乘积。很多应用层以为自己在裁矩阵,实际在裁维度时选错了轴,比如把 weight 里的行和列裁反了,这样输出 shape 都会变。
- 接着确认输入自适应是否真的生效。要看每个样本的输入长度、激活统计是否如预期变化。很多情况下判断写得没问题,但前面有 padding mask 干扰统计,导致裁掉的全是有效 token。
- 然后看数值误差分布。不要只看平均误差。如果某几个位置 top-1 logits 翻转,要把这个位置的上下文和裁剪通道记录摆出来。
- 再观测时间。用 profile 工具看端到端矩阵乘积总耗时,而不是只盯着某个操作打印的时间。可能节省在算子 A,却增加在算子 B。
- 最后才调阈值和层选择。从低裁剪比例、单层开始,达标后再扩多层。
- 每次只改一个变量。同一时间改掉列裁剪、batch分桶、kernel替换等三种逻辑后,出了问题你根本定位不出来是哪种逻辑在影响结果。
最后写几条实操心得
很多人看到这类题目会先问“用哪个库实现的、有没有公开代码”。但从工程取舍来看,先搞清楚“在什么输入条件下少算哪些乘法是安全”更关键。没有可复现的误差判定标准,就算拿到代码也未必能用于自己的场景。
我的建议是:把整个项目拆成三段来验证。第一段在模型内部插入 hook 记录每个矩阵乘法的输入统计和输出一致性;第二段实现简单的通道裁剪或者序列长度分组逻辑;第三段把动态判断整合进自定义算子中,并在生产级推理环境下压测。
如果第一段已经发现某些层对裁剪极其敏感,那就不要硬裁。换成只做长度分组或只做浅层裁剪就可能稳很多。如果第一段表现还稳定,再逐步扩展裁剪层数和 reduce 比例,每一步都保留原始模型作为参照,把日志落全。
最后一点,这个方向真正上线时,收益也要用两种视角衡量:单次前向性能视角会看矩阵乘法耗时少了多少,而产品视角看的是吞吐、最大并发、长序列能力、输出一致性。后面一种才决定能否长期使用。
Reduced Matrix Multiplication 能不能成为 LLM 推理里的常规优化手段,最终不取决于“约减”这一步多聪明,而取决于判断逻辑能不能像普通算子和显存管理一样被可靠地纳入整个系统。只要这一点还没被验证,就先把收益预期控制住,拿真实 benchmark 说了算。