MiMo v2.5推理优化：极限提升混合SWA效率

背景

MiMo-V2.5 系列模型（包括 MiMo-V2.5 和 MiMo-V2.5-Pro）在架构上融合了多项设计选择：Hybrid Sliding Window Attention（混合滑动窗口注意力，Hybrid SWA）将 KVCache 存储压缩至 Full Attention 的约 1/7；稀疏 MoE 激活在保持模型容量的同时降低了每 token 计算量；多模态编码器使模型能够理解视觉、音频和视频。这些特性赋予了 MiMo-V2.5 系列在长上下文和多模态场景中显著的性能与效率潜力。

从设计之初，目标就很明确：训练一个既强大又高效的长上下文推理模型。然而，这两个目标存在固有矛盾。强推理能力需要建模长程依赖，通常需要更大规模的注意力计算和更高的 KVCache 开销。在传统的 Full Attention 架构中，注意力计算量和 KVCache 存储量都随上下文长度快速增长，导致长上下文训练和推理成本极高。Hybrid SWA 通过在相邻层中交错使用局部 Sliding Window Attention（SWA）和全局 Full Attention 来解决这一问题：大部分层仅在局部窗口内计算注意力，少数关键层保留全局视图。理论上，这种结构将注意力复杂度降低到接近线性，同时保留了建模长程依赖的能力。

然而，理论上的架构优势并不会自动转化为生产效率。Hybrid SWA 引入了管理 KVCache 命中率、前缀匹配以及维持 Full Attention 层与 SWA 层之间双重语义一致性的新复杂性。实际工程系统还面临更多挑战——跨多级存储的数据移动、异步预取与调度的错位、分布式缓存状态同步困难——这些都阻碍了理论优势的直接实现。

核心内容

理论效率边界量化

以 MiMo-V2.5-Pro 为例，该模型共有 70 层，其中 10 个 Full Attention 层和 60 个 SWA 层，滑动窗口大小为 128。与 Full Attention 相比，Hybrid SWA 的计算成本如下图所示（原文配有图，此处省略）。SWA 层占总层数的 6/7，因此 Hybrid SWA 架构的总计算量约为 Full Attention 的 1/7。在 Chunked Prefill 场景下（预填充主要由计算瓶颈限制），这直接转化为预填充成本的等比降低。

由于 SWA 层只需保留滑动窗口内的 KV（而非整个序列），KVCache 内存使用量同样降至接近 1/7。解码阶段主要由内存带宽限制，其延迟与模型参数和 KVCache 的总读取字节数成正比。对于长序列，KVCache 体积可能远超模型参数，因此 KVCache 存储的减少在长序列场景中几乎直接转化为解码成本的降低。

不同模型架构的 KVCache 存储差异很大，访问模式也不同。如下图所示，MiMo-V2.5-Pro 和 MiMo-V2.5 在 KVCache 效率上排名第二，仅次于 DeepSeek-V4-Pro 和 DeepSeek-V4-Flash。

值得注意的是，实际成本差异并不严格与 KVCache 大小比例对应，因为存在与序列长度无关的固定计算和内存访问成本。但在长上下文场景中，总体趋势成立：短序列收益微乎其微，序列越长，推理成本优势越明显。

工程挑战：从理论到生产

MiMo-V2 和 MiMo-V2.5 系列是最早采用 Hybrid SWA 架构的模型之一，但当时主流开源推理框架和缓存系统都未完整支持 SWA。在推出 MiMo API 时，团队选择 SGLang v0.5.5 作为服务后端代码库，并立即面临严峻挑战。在该版本中，SGLang 的 HiCache 不支持 SWA，或者说，早期的 SWA 支持是通过存储完整 KVCache 来维持兼容性的。尽管有一些变通方法可以让 SWA 更可用，但团队希望构建一个具有更高性能上限和更好可用性的 KVCache 系统。

KVCache 分离设计：Full KV Pool 与 SWA KV Pool

Hybrid SWA 引入了一个根本性的存储冲突：Full Attention 层需要存储完整的序列 KV（O(N)），而 SWA 层只需维护滑动窗口内的 KV（O(W)）。在传统的单一 KV 池设计中，系统必须为所有层分配 O(N) 的 GPU 内存，无法利用 SWA 的窗口稀疏性——实际上退化成了接近全量 KVCache 的实现。

一个自然的解决方案是将 KVCache 拆分为两个独立的池，分别用于 Full Attention 和 SWA，并在系统层面进行统一抽象：

• 物理层：维护独立的 Full KV 池和 SWA KV 池。SWA 池的大小仅按窗口大小设计，并支持基于窗口的独立驱逐策略，严格将 SWA 存储限制在 O(W)。该机制同样扩展到 L2 和 L3 存储层级。
• 逻辑层：向上层（前缀树、调度器、传输协议）暴露单一的序列视图，以 Full Attention 索引为权威参考，并维护一个 Full→SWA 映射以实现透明的分层存储。
• 调度约束：系统在接受请求时同时验证 Full KV 和 SWA KV 的容量约束，避免单维度检查导致的资源错配。
• 数据移动：跨层级传输仅基于 SWA mask 进行，确保只移动有效的窗口数据，避免冗余带宽消耗。

通过这一设计，SWA KVCache 在系统层面实现了严格的 O(W) 存储约束，整体 KVCache 容量效率提升了约 7 倍，从而释放了 Hybrid SWA 的结构优势。主流推理框架也采用了类似的实现方式。

SWA 感知的前缀缓存树（Prefix Cache Tree）

在 SWA KVCache 存储优化完成后，SWA 层只需预取极小量的 KVCache。通过按层调度，可以实现 Host-to-Device KVCache 预取与计算近乎完美的重叠，使推理过程中缓存读取的成本接近于零。

传统的 RadixAttention 命中规则建立在一个简单假设之上：相同的 token 序列 → 相同的 KV。这一假设在 Full Attention 下成立——只要两个请求共享相同的 token ID，它们对应的 KV 保证仍在池中且可直接复用。

但在 SWA 下，这一假设被打破。原因是前缀树的逻辑生命周期与 SWA KV 的物理生命周期出现了错位。前缀树节点的长度不受 SWA 窗口约束——一个节点的序列长度可以短于窗口，也可以远长于窗口，节点还会通过请求的合并、拆分和删除而不断变化。因此，一个前缀树节点在逻辑上可能仍然代表完整的 token 序列，但其对应的 SWA KV 可能只保留了尾部部分，甚至可能已被完全驱逐。如果前缀树仍然提供复用……

（原文在此处中断，后续内容缺失。以上基于已有部分完整呈现。）

关键要点

• Hybrid SWA 的理论效率：通过将 6/7 的层替换为局部滑动窗口注意力，计算量降至 Full Attention 的约 1/7，KVCache 存储同样降至约 1/7，且长序列下收益更显著。