在推理优化圈子里，我最近一段时间基本都泡在 MiMo-V3 的 prefill 性能实验里。项目组决定用 HySparse2 来做稀疏化加速，实验配置单上明确写着“前 25 层 prefill”，不少同事第一反应都是：为什么是 25 层？不是应该跑整个模型吗？这问题其实挺有代表性的，因为它不只是一个数字选择，背后涉及到模型结构、显存峰值、算子特性、稀疏化验证策略等多层因素。我把整个分析和踩坑过程整理出来，希望能给正在做长上下文推理优化、稀疏化加速的朋友一些可直接参考的经验。
1. 先弄清楚 MiMo-V3 和 HySparse2 的关系
1.1 MiMo-V3 到底是什么样的模型
MiMo-V3 不是那种层数统一、结构上下一致的规整 Transformer。它是我们在内部推进的混合架构实验模型，总深度 42 层，但真正吃计算、吃显存的核心部分，是前 25 层标准 Transformer block。这 25 层负责对输入 token 做语义对齐、局部关系建模和初步特征抽取，每一层都包含完整的 QKV 投影、注意力计算、MLP 扩展和 LayerNorm 之类的基础组件，计算密度高、中间激活大，是典型的“计算密集区”。
从第 26 层开始，模型切换成线性注意力结构，配合共享 KV cache 的机制，把原来和时间序列长度强相关的 KV 存储解耦，注意力计算的复杂度从平方级降下来。也就是说，模型后面的 17 层虽然也参与 prefill，但它们计算量小、显存占用低，不是瓶颈。瓶颈就在前面这 25 层。
所以如果你拿到 MiMo-V3 的 prefill 性能报告，第一眼看到“25 层”时不要觉得这是测试偷懒，这其实是刻意聚焦到主要矛盾上。模型设计时就把最强、最重的计算都压在前面的稠密注意力阶段，而 HySparse2 要做的，就是把这个关键阶段的冗余计算消掉。
1.2 HySparse2 在 prefill 里扮演什么角色
HySparse2 是我们内部迭代到第二版的稀疏化推理加速框架，核心思路一句话概括：在保证端到端精度可控的前提下，把 prefill 阶段的计算稀疏掉，减少无效 FLOPs，压缩 KV cache 和中间激活，从而换取更低的延迟和更高的吞吐。
第一版 HySparse1 主要做的是权重剪枝和 KV cache 量化，效果有，但问题也不少：剪枝粒度太粗，导致精度波动大；量化只是降低了存储，没有真正砍掉计算量。HySparse2 在架构上做了三件关键事：
- 支持结构化稀疏模式，包括 2:4 半结构化稀疏和 block-sparse 块稀疏，目标是让稀疏后的矩阵计算真正落到 GPU Tensor Core 上，而不是变成空转的“假稀疏”。
- 针对 prefill 做了算子融合，把 QKV 投影、RoPE 位置编码、FlashAttention 入口、MLP 激活等环节的 kernel 边界重新划分，减少中间张量的写回和读取。
- KV cache 层做了多级压缩，包括低精度存储、冗余头合并和动态驱逐策略，在长序列场景下显存峰值能明显降下来。
用生活化的类比来说，HySparse1 是“把书里没用的字涂掉”，而 HySparse2 更像是“把整本书重新排版，只保留必要的内容，并且每页都做了摘要”。前者看起来省了内容，后者才是真正省了阅读时间。
我一开始接触 HySparse2 时也有个误区，以为它主要优化 decode 阶段。其实在长 prompt 场景下，prefill 才是真正的耗时大头。特别是输入 8K 甚至 32K token 的业务数据，prefill 阶段要一口气并行处理所有输入，生成完整的 KV cache 和中间状态，计算密度极高。而 decode 阶段一次只生成一个 token，虽然是串行的，但单步计算量小、内存带宽敏感，优化思路完全不同。
2. 为什么选 25 层做 prefill 实验
2.1 从架构上看，25 层是一个天然分界线
很多人在模型里见到“25 层”这种数字，本能会以为是从总层数里随机截一段。但在 MiMo-V3 里，这个数字是结构决定的。
我把 MiMo-V3 的层配置贴出来给大家看，一目了然：
- 第 1 到第 25 层：标准 Transformer block，完整自注意力 + FFN，参与全量 token 交互。
- 第 26 到第 42 层：线性注意力层，使用固定的平行扫描机制，复杂度与序列长度近似线性。
前 25 层和后 17 层的计算模式完全是两套东西。prefill 实验如果直接跑全模型，HySparse2 的算子库需要同时兼容两种注意力实现，调试和定位问题的难度会翻倍。而先只跑前 25 层，相当于把实验范围限制在“标准注意力 + 稠密 FFN”这一套最通用的组件里，这是稀疏化最容易产生收益、也最容易出问题的地方。
而且从推理开销的实际分布来看，MiMo-V3 在长输入场景下，前 25 层的计算时间占总 prefill 时间的 80% 以上。后面的线性注意力层虽然层数不少，但单层 FLOPs 很低。所以 25 层就是这个模型的“性能主战场”，实验当然选主战场先打。
2.2 从显存和计算量看，25 层是性价比拐点
除了架构原因，25 层的选择还有一个非常现实的理由：显存峰值和计算量在 25 层附近出现了明显的性价比拐点。
我做了个不同层数的 prefill 预实验，使用单张 H800 80GB，输入长度 8192，batch size 8，KV cache 采用 fp16 存储。实际测出来的数据大概是这样：
| prefill 层数 | 理论 KV cache 占用 | 中间激活峰值 | prefill 耗时 |
|---|---|---|---|
| 16 层 | 约 6.4GB | 约 18GB | 1.21s |
| 25 层 | 约 10.0GB | 约 27GB | 1.87s |
| 32 层 | 约 12.8GB | 约 34GB | 2.36s |
| 42 层全量 | 约 16.8GB | 约 44GB | 3.02s |
从 16 层到 25 层，计算量增加了大约 55%，但显存和耗时都在可控范围内；从 25 层到 32 层，计算量只增加 28%，中间激活峰值却涨得很快。如果直接跑全量 42 层，即便显存勉强装下，留给 batch 扩展和后续 decode 生成的空间就非常少了，实验的灵活性会大打折扣。
所以 25 层是我们在“覆盖主要计算瓶颈 + 保留显存余量 + 等待时间合理”三者之间找到的平衡点。这不是拍脑袋定的，是根据实际显存预算反推出来的结论。如果你的模型结构不同，我建议也做一次这样的层数扫描，而不是照抄别人的 25 层。
2.3 稀疏化验证本身需要分层推进
做稀疏化有个重要的工程原则：不能一次性把整个模型全部打上稀疏，然后期待一切正常。风险太高，定位问题太痛苦。
我在做 HySparse2 实验时，第一步是逐层做敏感性分析，也就是把每一层的权重单独稀疏化，观察对最终 loss 和下游指标的影响。前 25 层里，不同层的敏感性差异很大，靠近输入的层对精度极其敏感，稀疏率稍微一高，loss 立刻飙升；而靠近第 25 层的部分相对宽容一些。
如果一开始就全模型稀疏化，你会面对一个很不幸的局面：整个模型精度掉了，但你根本不知道是哪一层、哪一个模块拖垮的。拆分到 25 层里逐层排查，至少能把问题范围缩小到一个 block 内部。
另外，稀疏化算子的正确性也需要分层验证。HySparse2 的 kernel 替换了原本的线性层计算路径，我们需要对比稀疏前后每一层的输出，用余弦相似度和最大绝对误差来评估。这个工作放到 25 层上，每层的调试成本还能接受；如果放全模型 42 层，每次实验的编译时间、前向计算时间、日志分析时间都会呈指数级增加。
3. 用 HySparse2 跑 25 层 prefill 的完整操作
3.1 环境准备与模型裁剪
实验的硬件环境是单机 8 卡 H800 80GB，软件栈包括 CUDA 12.4、PyTorch 2.5，以及 HySparse2 配套的算子库和运行时。要做 25 层 prefill 实验，第一步不是写稀疏化代码，而是先把模型裁剪出来，确保实验只加载前 25 层。
我遇到的一个常见错误是：直接在原始模型上设置一个层数开关，却没有处理好 embedding 和后续层的依赖关系。正确的做法是构造一个只包含前 25 层 block 的子模型，输入输出保持不变，但内部不再有第 26 层之后的参数和计算路径。
裁剪之后，还要固定随机种子，确保每次实验的输入一致。我习惯的做法是使用固定的 prompt 语料，长度分布覆盖 512、2048、4096、8192、16384 这几个档位，而不是只用单一长度。有些朋友做 prefill 实验只测一种长度，结果换成长 prompt 后显存直接 OOM，这就是输入分布没覆盖到位。
3.2 稀疏化配置与参数选择
HySparse2 的配置通过一个 JSON 文件控制，核心参数大概长这样：
{ "model": { "name": "mimo_v3_head25", "layers": [0, 25], "precision": "bf16" }, "sparsity": { "weight_sparsity": { "pattern": "2:4", "target_layers": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_up_proj", "down_proj"], "sparsity_ratio": 0.5 }, "kv_cache_compression": { "enable": true, "quant_bit": 8, "head_prune_ratio": 0.25 } }, "runtime": { "kernel_fusion": true, "flash_attention": true, "prefill_cuda_graph": true } }这里最关键的是pattern选 2:4 结构化稀疏。为什么不用非结构化稀疏？因为 GPU 上非结构化稀疏的矩阵乘法很难真正利用硬件加速，数据排布不规则，访存效率极差。2:4 结构化稀疏意味着每 4 个连续权重里固定保留 2 个，形状规则，Tensor Core 有专门的硬件指令支持，实测可以有效加速。
权重稀疏率我一开始设的是 0.5，也就是 50%。但这个值不是所有层都能承受的。靠近输入的第一层，稀疏率降到 0.25 才让精度满足要求；而中后段层可以到 0.6。所以实际配置里需要对每一层单独指定稀疏率，而不是用一个全局值一刀切。
KV cache 压缩方面，fp16 转到 8bit 量化，加上 head prune 之后，实测 KV cache 体积能减少 50% 左右。但这里有个容易踩的坑：KV cache 压缩会直接影响注意力分数的计算精度，特别是长序列上，量化误差会被累积放大，所以压缩后一定要跑一遍长文本生成验证，确认没有出现注意力崩塌或者重复生成的问题。
3.3 性能测试流程与数据记录
性能测试不是简单跑一次前向就算数。我采用的流程是先做 warmup 再测多轮取中位数，因为 GPU kernel 首次执行时会有初始化、编译和显存分配的开销，直接测第一轮的耗时完全不能反映真实情况。
预热跑 10 轮，然后正式测 20 轮，取 P50 和 P99 两个分位数。P50 反映典型性能，P99 反映极端抖动。如果 P99 比 P50 高出 30% 以上，通常说明存在显存碎片、kernel 调度不稳定或者电源功耗限制的问题，这时候要先解决稳定性再谈优化。
25 层 prefill 在完成 HySparse2 优化后的实测数据，我记录如下：
| 输入长度 | 未稀疏耗时 | HySparse2 耗时 | 加速比 | 显存峰值下降 |
|---|---|---|---|---|
| 2048 | 0.53s | 0.34s | 1.56x | 28% |
| 8192 | 1.87s | 1.16s | 1.61x | 35% |
| 16384 | 3.72s | 2.31s | 1.61x | 38% |
| 32768 | 7.55s | 4.68s | 1.61x | 40% |
可以看到加速比在长序列下稳定在 1.6 倍左右，没有随着序列变长而衰减，这个结果说明稀疏化算子没有引入额外的复杂度瓶颈，整体扩展性是不错的。但我也要实话实说，1.6 倍距离理论上的 2 倍还有差距，主要原因是 2:4 稀疏虽然减少了 FLOPs，但内存读取、kernel launch、LayerNorm 等非稀疏部分的开销没有同步减少。
4. 一路踩过的坑和排查方法
4.1 精度偏差直接爆掉
第一次把稀疏化应用到 25 层 prefill 时，我 meet 到一个非常尴尬的现象：中间层的输出余弦相似度前 20 层都很好，但到第 22 层之后突然剧烈下降。一开始我怀疑是深层对稀疏更敏感，排查了半天，最后发现是norm层和稀疏化线性层之间的操作顺序出了问题。
具体来说，HySparse2 在融合 QKV 投影时，把 LayerNorm 之后的缩放值重新映射到了稀疏矩阵的 scale 上。对于浅层，这个映射误差很小；到了深层，累积的均值和方差偏移变大，映射误差被放大，精度就崩了。解决方法是改为在每层单独做校准，把 scale 从权重统计中直接计算出来，而不是继承上一层的统计值。
这个坑给我的教训是：稀疏化不是简单替换算子，每替换一个组件，都要重新验证它和相邻组件之间的数值接口。特别是涉及归一化层、残差连接这种会改变数值分布的地方，最容易出隐蔽问题。
4.2 稀疏率达标但速度没变
还有个让人抓狂的情况：权重矩阵确实变成了 50% 稀疏，但 prefill 耗时几乎没有变化。查了半天才发现，瓶颈根本不在权重计算，而在注意力部分。
当输入长度较长时，prefill 的计算热点会从 MLP 和 QKV 投影转移到注意力分数计算和 softmax 上。我当时只稀疏化了线性层，注意力部分用的是普通的 FlashAttention，却忘了 HySparse2 支持对注意力分数矩阵做动态阈值截断。也就是说，权重稀疏只解决了模型参数的计算量，没有解决激活层面的计算量。
后来我在配置里打开注意力稀疏开关，设置了一个自适应的 top-k 阈值，对小数值注意力分数做截断，prefill 耗时才真正降下来。这个问题的本质是：稀疏化必须跟着计算热点走，不是在所有层均匀撒胡椒面。
4.3 KV cache 峰值算错导致 OOM
KV cache 显存的计算公式看起来很简单：2 * num_layers * num_kv_heads * head_dim * seq_len * bytes_per_element，但实际使用时我照这个公式算出来的值和实际显存占用差了近一倍。
原因在于公式只算了 KV cache 本身，没有算上分页表的开销、缓存池对齐、以及 beam search 时的多个序列副本。HySparse2 采用类似 PagedAttention 的分页管理方式，页面大小是 16 或者 32，如果序列长度不是页面大小的整数倍，最后一页会整体分配，产生浪费。
现在我用的是更保守的估算公式：总显存 = KV cache 理论值 * 1.3 + 中间激活峰值 + 权重显存。实验设计时还要给 CUDA context 和 fragment 预留至少 8GB 的余量。否则很容易出现测试时单 batch 没问题，一加大 batch 就 OOM 的尴尬情况。
4.4 回归测试里最容易漏掉的点
做 25 层 prefill 实验时，我一度沉迷于优化算子性能，差点漏掉一个重要问题：稀疏化对不同 padding 策略的兼容性。输入序列如果做了 padding，注意力 mask 会把 padding 位置的分数置为无穷小负值。但稀疏化后，注意力分数的动态截断逻辑如果不认识 mask 的位置，可能会把 padding 位置的有效信息误删掉，或者保留无意义的 padding 分数，导致生成结果异常。
回归测试时，我强烈建议加入“变长输入 + padding”的场景，而不是只用整齐对齐的固定长度序列。另外还要测试 batch 中同时存在短序列和超长序列的混合情况，这种场景最容易暴露索引错位和访存越界的问题。
5. 这个实验实际带给我的东西
5.1 prefill 和 decode 的优化思路真的不一样
以前我做推理优化，习惯性把 prefill 和 decode 放在一起调优，总觉得加速一个阶段总归对另一个阶段也有帮助。这次用 HySparse2 做 25 层 prefill 实验，最大的认知刷新是：两个阶段的优化重心完全不一样。
prefill 阶段是计算密集的，所有输入 token 同时参与计算，GPU 的算力利用率高，内存带宽相对不是瓶颈。所以优化 prefill 的核心是减少无效计算，权重稀疏、激活稀疏、算子融合都能直接产生收益。而 decode 阶段是内存带宽密集的，GPU 算力往往跑不满，卡在权重加载和 KV cache 读取上，这时候稀疏权重反而可能因为不规则的访存拖慢速度，更适合的做法是权重压缩和更好的缓存布局。
理解了这一点，再回头看“为什么只跑 25 层 prefill”这个问题，就变得非常清晰了。我们是想在最短的迭代周期内，验证 MiMo-V3 前面这段计算密集区能不能通过 HySparse2 实现稳定加速，而不是急着把全模型都改造完。
5.2 25 层实验结论如何扩展到全部层
25 层 prefill 实验跑通后，向 42 层全模型扩展其实没有想象中那么复杂，但也没有那么简单。关键在于分清楚哪些经验可以直接平移，哪些需要重新验证。
可以直接平移的是算子层面的优化，比如 QKV 投影的 2:4 稀疏、MLP 的稀疏化、算子融合方案。这些在前 25 层验证过的问题，在后 17 层线性注意力结构中同样适用，只是 kernel 需要适配新的数据布局。
需要重新验证的是 KV cache 压缩策略。后 17 层使用共享 KV cache，本来 cache 的尺寸就小很多，压缩的收益没有前 25 层那么明显，但风险更高，因为共享意味着误差会被多个位置复用和放大。所以扩展时，我是先保持后 17 层不稀疏，等前 25 层精度完全稳定后，再逐步打开后面的稀疏开关，一步步做敏感性确认。
根据我个人做这轮实验的体会，最关键的还是那句话：稀疏化实验一定要分层推进、逐层校准、按热点优化，不要迷恋某一个整体数字。25 层这个数字本身不重要，重要的是它帮助我们把 MiMo-V3 最核心的计算瓶颈完整地暴露出来，并且在一个可控的规模内跑通了整个优化链路。下次你再看到实验报告里出现奇怪的层数，建议先问问自己：这个数字对应的计算代价和显存代价长什么样？模型结构在哪里分界？性能热点分布在哪里？把这三个问题回答清楚，你就能判断这个层数定得合理不合理了。
