大语言模型做生成推理时最消耗算力的不是“模型有多大”而是每一层、每一个 Token 都要执行的大量矩阵乘法。很多情况下这些矩阵乘法里存在明显的输入相关冗余某一部分输出通道对当前 Token 并没有贡献却仍然被完整算了一遍。Reduced Matrix Multiplication 这个名字对应的正是这一类优化思路在输入信息驱动下把一个完整的矩阵乘积削减成更小规模的乘积从而降低 LLM 推理阶段的无效计算。相比操作符融合、CUDA Kernel 优化这类“改底层实现”的方式它更偏向从算法层面决定“这一次乘法到底要算哪些部分”。这篇文章会围绕 LLM 推理中的矩阵乘法削减做一次技术拆解先讲清楚这类方法想解决什么再分析输入自适应机制可能的实现路径然后用 PyTorch 写一个概念原型帮助理解最后给出资源观察、实验验证与工程落地建议。1. 核心能力速览能力项说明技术方向LLM 推理加速、矩阵乘法计算量削减、输入自适应路由与稀疏计算目标算子Attention 中的 QKV 投影与加权求和、FFN/Gated MLP 中的上下投影矩阵关键机制根据当前输入 Token 或上下文判断哪些乘法输出通道相对重要只保留这些通道参与计算优化收益降低矩阵乘法的 FLOPs可能带来更低延迟、更高吞吐也可能因为索引开销导致收益缩小适用模型以 Transformer 架构为代表的解码器模型理论上可推广到编码器模型部署方式属于算法/系统协同优化思路需要落到 vLLM、TensorRT-LLM 等推理框架中才能形成服务API 能力不直接提供 HTTP API取决于具体接入的推理框架批量任务需要额外设计连续 Token 的索引合并策略否则动态裁剪会破坏批量矩阵乘法的连续性开源状态当前公开材料未给出明确实现仓库需以原始论文或项目发布页为准这里需要先说明从已公开的材料看这个主题更像是一类方法描述而不是某个可以直接拉下来的完整推理引擎。本文后续会按“优化机制分析 概念验证 工程落地观察”的顺序展开。2. 适用场景与使用边界Reduced Matrix Multiplication 适合解决的问题是 LLM 推理中“大量矩阵乘法存在结构性和输入性冗余”的场景。比较典型的是 Feed-Forward 层。Transformer 的 FFN 会把输入从 hidden_dim 先映射到一个很大的中间维度再用激活函数做非线性变换。ReLU、GELU 这类激活会让一部分中间神经元输出趋近于零或饱和。从单个输入看最终起作用的往往不是全部中间通道。如果能在计算前就知道哪些通道不需要参与理论上就可以把两次大矩阵乘法变成两次窄矩阵乘法。Attention 里同样存在类似空间。例如某些注意力头可能主要对句法信息敏感某些头主要对局部位置信息敏感。输入不同真正起作用头的集合也不同。如果按输入动态挑选头或剪掉部分 Value 通道Attention 的矩阵乘积规模也可以降低。但这类方法不是没有代价。第一矩阵乘法削减会引入额外计算。需要用一个打分器去预测每个输入对应的重要性通道。打分器本身也是神经网络也是一次前向计算。如果打分器的规模没有控制好省下来的算力反而会被打分器吃掉。第二动态索引会破坏稠密矩阵乘法的高效性。GPU 上的矩阵乘法之所以快是因为数据排布连续可以充分利用 Tensor Core。如果每个 Token 选择的通道索引不同权重矩阵的切片可能不连续需要 gather、索引、重排这些操作在 GPU 上经常比直接算一次大矩阵乘法还要慢。第三评估不能只看单次延迟。对 LLM 推理来说更重要的是端到端吞吐和显存带宽。减少 FLOPs 不代表减少显存访问。如果削减后的权重仍然以稠密格式存储实际推理速度可能没有明显变化。所以适用场景主要集中在输入间差异明显、固定剪枝损失较大的模型FFN 中间激活稀疏性较强的模型能够配合自定义 Kernel 实现连续 Gather 和分段矩阵乘法的推理系统对延迟不敏感、但对吞吐极度敏感的服务端批量推理场景。不适合的场景包括没有足够校准数据、无法接受效果波动、推理框架不支持动态 shape 变化的场景。涉及安全合规时还需要注意模型压缩和推理优化的实验材料应使用合法授权的模型权重与数据不得用未经授权的内容做校准集也不要在生产环境直接使用未经效果验证的裁剪策略。3. LLM 推理矩阵乘法瓶颈到底在哪要理解 Reduced Matrix Multiplication 的价值先看 LLM 推理中的矩阵乘法分布。3.1 Prefill 与 Decode 的矩阵乘法形态LLM 推理可以粗略分成两个阶段。Prefill 阶段处理整段 Prompt输入是一个比较长的序列。这个阶段矩阵乘法通常是 [batch, seq_len, hidden_dim] 乘以 [hidden_dim, output_dim] 的形态计算密度比较高GPU 利用率相对容易做上去。Decode 阶段每次只生成一个 Token输入序列长度为 1。矩阵乘法的形态变成 [batch, 1, hidden_dim] 乘以 [hidden_dim, output_dim]。虽然单次计算量不大但整个回复过程要循环执行很多次而且每次都要读取大量权重。Decode 阶段对显存带宽的依赖往往比对 FLOPs 的依赖更强。在这样的背景下矩阵乘法削减有两种作用路径在 Prefill 阶段通过减少中间通道数量来降低整体 FLOPs在 Decode 阶段不仅减少 FLOPs还减少需要从显存读取的权重列数从而缓解带宽压力。第二种路径往往更有吸引力。因为 Decode 阶段如果能把参与计算的权重矩阵列数减半显存读取量也约等于减半实际延迟可能获得接近线性的改善。3.2 矩阵乘法中的冗余来自哪里冗余可能来自三个层面。第一层是激活函数带来的通道级冗余。很多模型使用 GELU 或 SiLU它们会产生负区间饱和。对于某些输入很大一部分中间通道输出几乎为零。一个输入自适应系统可以提前预测哪些通道会被激活然后只计算这些通道对应的矩阵乘积。第二层是样本级多样性带来的统计冗余。不同领域的输入激活的高响应通道往往不同。固定剪枝会剪掉全局重要性低的通道但这种方法可能误伤某些领域输入的关键通道。输入自适应机制的好处是保留了每一类输入各自需要的通道。第三层是任务级冗余。在问答、摘要、代码生成等不同任务下模型内部不同 Attention Head 和 FFN 神经元的参与程度也不同。如果输入侧能提供任务信息那么矩阵乘法的削减粒度可以更细。4. Input-Adaptive 矩阵乘积削减的机制拆解从名字看Reduced Matrix Multiplication 包含三个关键词Reduced削减矩阵乘法的规模Input-Adaptive削减策略由当前输入决定Matrix-Product Reduction削减对象是矩阵乘积而不是简单减少 Token 数或量化位数。4.1 核心链路打分、路由、裁剪、恢复一个通用的输入自适应矩阵乘积削减流程可以分成四个步骤对当前输入计算通道重要性分数根据分数选出需要保留的通道索引只在保留通道上执行矩阵乘法把结果映射回原始输出维度供后续层继续使用。这里最关键的是第一步。重要性打分器不能太复杂否则额外前向计算会抵消收益。常见打分方式包括使用一个低秩线性层预测通道分数复用模型内部已有的 Gate 输出作为近似分数按照最近若干 Token 的激活统计量估计通道重要性通过聚类把输入映射到预计算的通道子集。4.2 矩阵乘积削减的数学抽象假设某一层需要计算Y X W其中 X 是输入 [B, T, D] 或 [B*T, D]W 是权重 [D, N]。完整计算需要做 B*T 行与 N 列的矩阵乘法。输入自适应削减希望找到一组索引集合 I(x)只计算Y[:, I] X W[:, I]其中 W[:, I] 表示 W 中由输入 x 决定的列子集。如果 |I| 远小于 N则计算量显著下降。更进一步Attention 中的加权求和可以写成O softmax(Q K^T / sqrt(d)) V这里的输入自适应削减可以在两个层面执行在 Q K^T 之前对部分 Head 的输出通道做降维在 softmax 之后对 V 的参与行做选择或加权。对于 FFN 类结构比如H activation(X W_up) Y H W_down可以对 H 的通道做输入自适应选择只保留响应较高的通道再与 W_down 中对应的行做乘积。这样 W_down 参与计算的行数就减少了。4.3 与传统剪枝和稀疏方法的区别传统结构化剪枝通常是离线完成的在验证集上统计每个通道的重要性剪掉固定比例后导出模型。所有输入共享同一套通道子集。输入自适应方法不同它允许每个 Token 拥有不同的通道子集。这种灵活性理论上能保留更多有效信息但也带来两个工程难点索引集合不固定无法像普通稀疏矩阵那样预先压缩存储需要对每个 Token 单独做 Gather批量推理时需要特殊 Kernel 支持。因此Reduced Matrix Multiplication 既不是简单的“稀疏矩阵乘法”也不是固定结构剪枝更接近“输入相关的动态稠密子矩阵乘法”。5. 概念原型用 PyTorch 看矩阵乘法削减效果为了保证这部分可执行、可观察下面写一个概念验证代码。它不代表原始论文实现而是用于演示输入自适应矩阵乘法的数据流。读者可以在普通开发机上运行观察矩阵规模变化带来的计算量差异。5.1 模拟一层 FFN 的通道重要性先给一个简单的离线通道重要性分析工具import torch def channel_importance_by_activation(activations, ratio0.25): 统计一组激活中每个通道的平均响应幅度。 activations: [num_samples, hidden_dim] ratio: 保留通道比例 importance activations.abs().mean(dim0) importance importance / (importance.max() 1e-6) keep_num max(1, int(activations.shape[1] * ratio)) topk_index torch.topk(importance, keep_num).indices.sort().values return importance, topk_index # 模拟 512 条样本、4096 维中间激活 fake_activations torch.randn(512, 4096) importance, topk_index channel_importance_by_activation(fake_activations) print(保留通道数:, topk_index.numel()) print(前 16 个保留通道索引:, topk_index[:16].tolist())这个工具是为了说明给定一批输入激活可以计算哪些通道值得保留。真实场景中输入自适应方案不会等激活出来后才决定而是在矩阵乘法之前用一个打分器预测。5.2 输入自适应的简化 Forward 原型下面这段代码演示的是“打分 - 裁剪 - 子空间乘法”的数据流。它使用了循环写法便于理解逻辑不代表高效实现。import torch import torch.nn as nn class AdaptiveFFNBlock(nn.Module): 一个最小化的输入自适应 FFN 原型 1. 用打分器预测中间通道的重要性 2. 对每个 Token 保留 top_k 通道 3. 在降维后的子空间执行输出投影。 def __init__(self, in_dim, hidden_dim, top_k128): super().__init__() self.top_k top_k # 常规 FFN 权重 self.up nn.Linear(in_dim, hidden_dim, biasFalse) self.down nn.Linear(hidden_dim, in_dim, biasFalse) # 输入自适应打分器尽量轻量 self.scorer nn.Sequential( nn.Linear(in_dim, 64), nn.ReLU(), nn.Linear(64, hidden_dim), ) def forward_single(self, x): # x: [in_dim] scores self.scorer(x) topk_index torch.topk(scores, self.top_k).indices topk_index topk_index.sort().values # 只计算 top_k 个通道的 up 投影 up_weight self.up.weight.index_select(0, topk_index) # [top_k, in_dim] h_topk torch.relu(up_weight x) # [top_k] # down 投影也只使用 top_k 行 down_weight self.down.weight.index_select(1, topk_index) # [in_dim, top_k] y down_weight h_topk return y def forward(self, x): # x: [batch, seq, in_dim] B, T, _ x.shape outs [] for b in range(B): for t in range(T): outs.append(self.forward_single(x[b, t])) return torch.stack(outs, dim0).view(B, T, -1) # 随便构造一组输入检查能否跑通 model AdaptiveFFNBlock(in_dim128, hidden_dim1024, top_k128) x torch.randn(1, 4, 128) y model(x) print(输出形状:, y.shape)上面这段代码里top_k 与 hidden_dim 恰好相等是为了验证结构能跑通。实际使用中 top_k 应小于 hidden_dim。5.3 观察矩阵乘法规模对耗时的影响完整 LLM 推理测试需要较大显存可以先在普通矩阵上观察规模带来的耗时差异import time import torch def bench_matmul(input_shape, weight_shape, repeats20): device cuda if torch.cuda.is_available() else cpu x torch.randn(*input_shape, devicedevice) w torch.randn(*weight_shape, devicedevice) for _ in range(5): _ x w if torch.cuda.is_available(): torch.cuda.synchronize() start time.perf_counter() for _ in range(repeats): _ x w if torch.cuda.is_available(): torch.cuda.synchronize() return (time.perf_counter() - start) / repeats full_time bench_matmul((32, 128, 4096), (4096, 4096)) reduced_time bench_matmul((32, 128, 4096), (4096, 1024)) print(f完整矩阵乘法耗时: {full_time * 1000:.2f} ms) print(f削减 75% 列后的耗时: {reduced_time * 1000:.2f} ms)这个基准只观察矩阵规模差异不能代表端到端加速。实际推理中裁剪带来的 Gather 开销和稀疏索引计算很可能消掉一部分收益。6. 接入推理框架与批量任务设计的工程观察如果把这类算法真正接入 LLM 推理服务需要考虑的就不是单个 Python 类而是与推理框架的深度配合。6.1 自定义 Kernel 的必要性动态按 Token 选择通道后常规矩阵乘法 Kernel 无法直接复用。PyTorch 中使用index_select或者 Gather 操作并不能高效利用 Tensor Core。工程化时通常要做两件事先把同一批次中具有相同或相似通道索引的 Token 分组再让每组 Token 执行独立的稠密矩阵乘法。这个逻辑可以用一个简化例子来说明def group_by_channel_index(token_scores, num_groups8): 把通道索引相近的 Token 聚到同一组便于执行批量矩阵乘法。 这里只演示“分组”的数据结构不代表推理框架内部实现。 # 假设每个 Token 打分后得到 top_k 索引 # 这里用一个哈希片段代表索引模式 token_group (token_scores.argmax(dim-1) % num_groups) groups {} for token_id, group_id in enumerate(token_group.tolist()): groups.setdefault(group_id, []).append(token_id) return groups # 模拟 16 个 Token 的通道模式和分组结果 token_scores torch.randn(16, 4096) groups group_by_channel_index(token_scores, num_groups8) print(分组结果:, groups)实际框架中这类分组策略需要基于通道索引的相似度而不是单一 argmax 值并且要平衡分组的数量与每组内 Token 的个数。6.2 批量任务设计从批量任务角度需要关注三个问题。第一批大小越大通道索引的差异可能越大。如果批内 Token 来自完全不同的领域每个 Token 保留的通道子集交集很小分组后每组 Token 数少矩阵乘法的利用率反而下降。第二请求调度层需要感知这种动态通道选择。服务端通常使用 Continuous Batching 动态拼接请求。如果每个请求的通道索引不固定KV Cache 和权重索引的管理会变得更复杂。第三校准数据的选取非常关键。如果输入自适应打分器只在校准集上观测过某种风格的文本线上出现分布外输入时通道选择可能出现明显偏差最终表现为输出质量下降。7. 显存占用与性能观察方法由于公开材料没有给出明确的测试环境本文不引用固定的显存数字。这里提供一套可用于自行评估的观察方法。7.1 观察 GPU 显存使用运行推理实验时可以开启 PyTorch 的显存统计import torch def print_gpu_memory(): if not torch.cuda.is_available(): print(未检测到 CUDA 设备) return print(已分配显存: {:.2f} GB.format(torch.cuda.memory_allocated() / 1024**3)) print(缓存显存: {:.2f} GB.format(torch.cuda.memory_reserved() / 1024**3)) print_gpu_memory()一个常见的误判是“显存占用没变说明削减没用”。实际上推理服务的显存占用主要由模型权重、KV Cache 和中间激活决定。如果权重没有稀疏存储即使只计算一部分通道显存占用也可能不变真正变化的是计算时的带宽读取量。7.2 性能指标观察清单评估矩阵乘法削减方案至少记录以下指标指标观察内容说明单 Token 延迟Prefill 后首个 Token 的生成延迟Decode 阶段延迟受权重读取量影响显著吞吐量单位时间完成的请求数关注批量大小上升后的变化趋势FLOPs实际参与计算的乘加次数使用 profiler 或理论估算显存带宽利用率从 HBM 到 SM 的数据读取效率动态裁剪可能导致带宽利用率下降输出质量在验证集上的困惑度或任务指标需要与压缩前对比打分器开销打分器单独的前向耗时防止“省了乘法多了额外网络”7.3 降低现象开销的方向如果发现端到端延迟没有明显下降优先检查索引计算和数据重排开销。可选择的方向包括降低打分器的计算频率不为每个 Token 都计算一次而是每隔若干 Token 或在一个块内共享索引减少动态选择频率对同一请求内的连续 Token 使用同一组通道设置通道索引的保底下界防止某些 Token 只保留极少通道导致输出质量波动把索引结果缓存下来同一个前缀被重复请求时可以直接复用历史通道选择结果。8. 常见问题与排查思路以本地实验和框架接入中的常见情况来做排查表问题现象可能原因排查方式解决方案输出质量明显下降保留通道过少或打分器不准确对比不同 top_k 下的困惑度提高保底通道数增加打分器容量延迟没有下降索引/Gather 开销占比过高单独基准打分器与 Gather 耗时降低打分频率使用分组执行显存占用没有下降权重仍是稠密存储观察权重加载量和实际读取量使用稀疏分块存储或索引连续化批量推理吞吐不升反降批内 Token 索引差异大分组后每组 Token 少查看批内通道索引分布按相似输入聚类请求或减少分桶数打分器增加额外耗时打分器结构偏重profiler 观察打分器耗时减小打分器维度或共享多个层同一打分器在 CPU 上运行效果不明显CPU 对索引重排开销不敏感但收益受限对比 CPU 与 GPU 的浮点峰值利用率CPU 场景优先做算子融合而不是动态裁剪CUDA Kernel 报 shape 错误动态索引导致维度不固定打印各阶段张量形状固定小组内 Token 的索引强制对齐 shape量化叠加后效果崩溃动态裁剪与量化误差相互放大分别测试裁剪与量化效果先固定裁剪通道再做量化或联合校准9. 工程落地最佳实践从概念验证到推理服务落地有几点经验值得提前沉淀。第一先做冗余度分析不要直接做动态裁剪。保留一份校准集对每一层统计中间通道激活幅度分布。如果绝大多数通道对大多数输入都有较大响应说明该层本身没有太多可削减空间盲目套用输入自适应方法意义不大。第二控制打分器的额外开销。打分器只承担“决定性不强但要相对准”的任务应尽量采用低秩结构。一个更省力的方案是复用已有的 Gate 输出或 RMSNorm 前的统计量避免为每个模块单独加一套全连接打分器。第三给通道选择设置纪律性约束。例如限制每个 Token 最多只能从某个预定义权重分组中选择固定数量的通道。这样虽然损失了一部分灵活性但更容易生成连续 kernel批量矩阵乘法也能保持较高利用率。第四分层单独评估。不要把所有层一次性换成动态裁剪而是选择 FFN 占比高、冗余明显的层逐层输入相关观察每一层替换后的累积质量损失。第五离线加上安全校验。涉及人物肖像、语音、版权文本等内容时需要先确定授权边界。推理优化实验也应使用可合法使用的数据避免用未经授权的内容做通道重要性的校准集合。第六如果准备对外提供接口服务最好在裁剪策略上保留一个开关。在生产环境出现异常请求时可以先回退到完整矩阵乘法避免因为动态通道选择引入输出质量风险。10. 总结与下一步Reduced Matrix Multiplication 的核心价值在于把矩阵乘法的削减从“全局静态”推进到“输入自适应”。它对 FFN 激活冗余较高、批次内请求差异较大的 LLM 推理场景有实际意义但工程实现难度明显高于固定剪枝。如果要从零开始验证这个方向第一步不是去改推理框架而是运行一套激活冗余度分析脚本观察目标模型各层有多少中间通道真实有效。第二步再实现一个打分网络验证是否能用较低开销预测通道重要性。第三步再结合自定义 Kernel、分组策略和批量调度做端到端评测。最容易踩的坑是“只减 FLOPs不看实际延迟”。动态通道选择和分组带来的访存压力可能抵消削减收益。后续值得扩展的方向包括把打分器与投机解码合并、利用前缀缓存复用通道索引、把动态裁剪扩展到 KV Cache 的稀疏读取。这类方法适合对推理系统有深入优化需求的团队。先把冗余度分析做扎实再把动态选择的判定逻辑做简单才有可能在真实服务中得到稳定收益。