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

# MiMo-V2.5 系列全流水线推理优化：将混合 SWA 效率推向极致 来源：https://mimo.xiaomi.com/blog/mimo-v2-5-inference logo 2026 年 5 月 30 日 中文版› (https://mimo.xiaomi.com/zh/blog/mimo-v2-5-inference) V2.5 模型系列，包括 MiMo-V2.5 和 MiMo-V2.5-Pro，融合了多项架构设计选择：混合滑动窗口注意力（Hybrid SWA）将 KVCache 存储压缩至全注意力的约 1/7；稀疏 MoE 激活在保持模型容量的同时降低了每 token 的计算量；多模态编码器实现了跨视觉、音频和视频的跨模态理解。这些特性共同赋予了 MiMo-V2.5 系列在长上下文和多模态场景中显著的性能与效率潜力。 从一开始，我们的目标就很明确：训练一个既强大又高效的长上下文推理模型。这两个目标本质上存在张力。强推理需要建模长距离依赖，这通常需要更大规模的注意力计算和更高的 KVCache 开销。在传统的全注意力架构中，注意力计算和 KVCache 存储都随着上下文长度快速增长，使得长上下文训练和推理成本过高。Hybrid SWA 的工作原理是在各层间交替使用局部滑动窗口注意力（SWA）和全局全注意力：大多数层仅在局部窗口内计算注意力，而少数关键层保留全局视角。理论上，这种结构将注意力复杂度降低到接近线性，同时保留了建模长距离依赖的能力。 然而，理论上的架构优势并不会自动转化为生产效率。Hybrid SWA 在管理 KVCache 命中率、前缀匹配以及在 Full Attention 和 SWA 层之间保持双语义一致性方面引入了新的复杂性。真实的工程系统还面临更多挑战——跨多级存储的数据移动、未对齐的异步预取和调度、分布式缓存状态同步困难——这些都阻碍了理论优势的直接实现。 除了 Hybrid SWA，MoE 对分布式调度和负载均衡提出了严格要求，而多模态编码器在大图像和长视频场景中仍然是吞吐瓶颈。调度策略以及 Prefill/Decode 执行流水线也需要仔细优化。本文介绍了 MiMo-V2.5 系列推理系统的端到端工程实践，涵盖 KVCache 管理、分层缓存系统、SWA 感知前缀缓存树、调度策略、Prefill/Decode 执行流水线以及多模态优化——系统性地将架构的理论效率潜力（尤其是 Hybrid SWA）在生产环境中实现。 ## 1. Hybrid SWA：推理效率优势 在深入具体优化之前，让我们先量化 Hybrid SWA 的理论效率边界——这是设计选择的架构原理，也是后续所有优化的基准线。 ### 1.1 计算分析 以 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 场景中，prefill 很大程度上是计算受限的，这直接转化为 prefill 成本的同比降低。 ### 1.2 KVCache 存储分析 由于 SWA 层只需要保留滑动窗口内的 KV（而非完整序列），KVCache 内存使用同样下降到接近 1/7。Decode 阶段主要是内存受限的，其延迟与模型参数和 KVCache 的总读取字节数成正比。对于长序列，KVCache 体积可能远超模型参数，因此在长序列场景中，KVCache 存储的减少几乎直接转化为 decode 成本的降低。 不同模型架构的 KVCache 存储差异很大，访问模式也不同。如下所示，MiMo-V2.5-Pro 和 MiMo-V2.5 在 KVCache 效率上排名第二，仅次于 DeepSeek-V4-Pro 和 DeepSeek-V4-Flash。 值得注意的是，实际成本差异并不严格对应 KVCache 大小比例，因为存在与序列长度无关的固定计算和内存访问成本。但在长上下文场景中，总体趋势是成立的：**短序列时增益有限，而序列越长，推理成本优势越大**。 ## 2. KVCache 系统重构 MiMo-V2 和 MiMo-V2.5 系列是最早采用 Hybrid SWA 架构的模型之一，但当时主流开源推理框架和缓存系统都未提供完整的 SWA 支持。当我们推出 MiMo API 时，我们选择 SGLang v0.5.5 作为服务后端代码库——并立即遇到了严峻挑战。在那个版本中，SGLang 的 HiCache 不支持 SWA，或者更确切地说，早期的 SWA 支持是通过存储完整 KVCache 来保持兼容性实现的。虽然有一些变通方法可以使 SWA 更可用，但我们希望构建一个具有更高性能上限和更好可用性的 KVCache 系统。 ### 2.1 SWA KVCache 管理 #### KVCache 双池设计 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 的结构优势。主流推理框架也已采用类似的实现方法。 #### 逐层 KVCache 预取 在 SWA KVCache 存储优化到位后，SWA 层只需要预取极少量的 KVCache。这使得通过逐层调度，Host-to-Device 的 KVCache 预取与计算之间可以实现近乎完美的重叠，将推理过程中缓存读取的成本降至接近零。 逐层 KVCache 预取时间线：(a) 计算流因等待 KVCache 加载而停顿；(b) 逐层调度使回载与计算重叠，GPU 无需等待即可运行 #### SWA 感知前缀缓存树 传统的 RadixAttention 命中规则基于一个简单假设：**相同 token 序列 → 相同 KV**。在 Full Attention 下这个假设成立——只要两个请求共享相同的 token ID，其对应的 KV 保证仍在池中，可直接复用。 但这个假设在 SWA 下失效了。原因是前缀树的逻辑生命周期与 SWA KV 的物理生命周期发生了错位。前缀树节点的长度不受 SWA 窗口约束——一个节点的序列长度可以小于窗口，也可以远大于窗口，而且节点通过请求合并、拆分和移除不断变化。因此，一个前缀树节点可能仍然逻辑上代表完整的 token 序列，但其对应的 SWA KV **可能只剩下尾部部分，或者已经完全被逐出**。如果前缀树仍然基于“token 相等 → 命中”规则提供复用长度，调度器可能会接收到一个带有已逐出尾部 KV 的伪命中——后续注意力计算将读取无效或已被覆盖的槽位，直接降低模型正确性。 为了在 SWA 下保持前缀复用的正确性和高效性，必须从三个方面修订前缀树语义： 1. **匹配规则升级为“窗口安全长度”：** 除了 token 相等外，尾部 W 个 token 必须在 SWA 池中仍然具有有效槽位。匹配长度被截断到这一新边界——超出部分视为缺失。这确保从命中段检索到的 KV 始终有效。 2. **逐出绑定到请求生命周期：** 长 prefill 中每个 chunk 完成时、请求终止时，以及 decode 期间每生成 N 个 token 时，都触发超出窗口的 SWA 释放。这使得在长上下文/长输出任务中，SWA 池的使用量保持在 W 或 chunk 级别，而不是随序列长度增长。 3. **节点携带双索引：** 每个前缀树节点记录两组信息——Full Attention 段索引（决定逻辑顺序，参与 Full Attention 层计算）和 SWA 段映射（决定窗口安全性）。逐出分别管理：窗口外的 SWA 段可以独立逐出，同时保留 Full Attention 段（使 Full Attention 层可复用该前缀），或者整个段都可以逐出。 SWA 将 KV 体积压缩到 1/7 是**容量级别**的收益，而命中率是**复用级别**的收益。两者共同决定了实际的 prefill 计算成本曲线。在引入“窗口安全长度”匹配规则后，给定 token 容量的原始命中率略有下降——但在相同存储预算内可容纳的 token 数量增长了几倍。以固定存储预算衡量，有效命中率显著提高。 SWA 感知前缀缓存树：每个节点携带每 token 的 Full Attention 状态和 SWA 状态，窗口大小为 4；节点追踪哪些尾部 token 仍具有有效的 SWA 槽位 #### KVCache 命中率优化 在所有三层 HiCache 都重构为 SWA 感知后，设备、主机和存储后端各自维护了“哪些位置的 SWA 有效”的状态。然而，HiCache 的数据移动流水线是异步的，不同部署的缓存状态不同，跨会话的共享前缀长度也不同；各层之间的 Full Attention Cache 和有效 SWA 索引很容易不同步。根据 SWA 感知前缀缓存树的匹配规则，如果序列在 Full Attention Cache 上命中但在 SWA Cache 上缺失，会发生严重的匹配长度截断：截断越多，需要重新计算的量越大，SWA Cache 的优化效果越低。因此，我们针对不同场景优化了分布式一致性和缓存命中率： - **设备完整，主机不足：** 当 L3→L2 预取由于带宽-延迟权衡仅拉取尾部段，或者当 L1 前缀树重组未同步到 L2/L3 时，会出现这种情况。我们在前缀树节点合并和 prefill 完成等时间点主动检查设备与主机之间 SWA 占用的差异，在主机 SWA 池中分配补充槽位，并通过 D2H 传输异步写入设备 SWA KV。 - **主机完整，设备不足：** 在下一次 H2D 传输时自然对齐——无需主动修复。 - **高频序列 L3 前缀逐出：** 长序列头部由于高频访问而保留在 L1/L2 中，缓存亲和性将相同前缀的请求路由到同一节点。L3 缓存由于长时间未被直接访问，可能被存储逐出策略逐出——过早释放全局高频序列的 L3 Cache，严重降低跨机器复用。我们在访问 L1/L2 Cache 时定期查询 L3 Cache，以防止过早逐出。 - **中/短序列 SWA 保留策略：** 基于用户请求模式，我们为中等/短序列在固定长度位置保留相对密集的 SWA KV Cache。虽然增加 SWA 密度会提高 SWA 在整体 KVCache 中的占比，但它直接有利于多用户共享系统提示等场景。 通过这些优化，我们将 KVCache 容量扩展转化为更长的有效命中长度，使跨会话的长前缀复用成为可能——这对于长 agent 会话、多用户共享系统提示以及重复调用同一代码库的工具特别有益。 ### 2.2 GCache：高性能分布式缓存基础设施 GCache 是小米存储团队开发的高性能通用缓存系统，构成统一训练-推理存储架构的关键部分。早期在训练场景中，存储团队就认识到某些开源缓存项目对分布式文件系统的加速有限，无法充分发挥性能潜力，因此开始研发内部解决方案。后来随着 MiMo 大模型的发布和推理服务的推出，团队将 GCache 改造为独立的存储产品，用于模型分发以及作为推理引擎的 L3 KVCache。 GCache 支持文件与 KV 双语义，跨内存/磁盘/远程的多级缓存，共享内存持久化和全路径零拷贝，高并发非阻塞 IO 和 RDMA 通信，在满足上层服务高吞吐低延迟需求的同时，保持了良好的扩展性。 #### 架构设计 GCache 架构：用户线程（多语言 SDK）→ gcache SDK（管理器、切片分发/收集、工作线程、本地缓存）→ gcache 集群（gcache-servers，带有元数据、基于一致性哈希的内存/磁盘缓存排列），Raft+rocksdb Master 用于服务发现和心跳，后端对象存储（Ceph/HDFS） GCache 具有几个关键特性： 1. 去中心化元数据管理，实现集群无限扩展： - 基于 key 的一致性哈希确定存储位置。 - Master 采用基于 Raft 的高可用部署，但仅管理心跳和服务发现——IO 路径不经过 Master。 2. 服务端同时支持内存和磁盘缓存： - 内存中的冷数据被逐出到磁盘；磁盘上的热数据被提升到内存。这种方法对推理场景非常有利，自动保证活跃会话的性能，同时降低长时间空闲会话的成本。 - 缓存条目持久化到共享内存——服务重启不丢失缓存。 - 支持平滑扩容或缩容，不丢失缓存。 3. 多语言 SDK 配备专用线程进行请求切片和分发： - 这些线程不消耗用户线程资源；切片提高并发性，并将 IO 大小保持在 RDMA 友好的范围内。 - 线程使用异步回调，回调粒度灵活——可以是单 KV 级别、批级别。