小米MiMo-V2.5系列模型推理全链路优化:释放Hybrid SWA架构极致效率

2026-05-30 18:59:53
来源:AI观察
分享

据Xiaomi Mimo消息,小米AI团队近期针对MiMo-V2.5系列大模型(包含MiMo-V2.5、MiMo-V2.5-Pro等)完成了一次覆盖推理全链路的深度工程优化,成功将模型架构的理论效率优势,尤其是Hybrid Sliding Window Attention(Hybrid SWA)的低成本潜力,全面兑现到生产环境。通过系统性地重构KVCache管理、分级缓存系统、调度策略及Prefill/Decode执行链路,并结合MoE与多模态场景的专项优化,实现了在长上下文与多模态场景下显著的性能提升与成本降低。

一、Hybrid SWA架构的推理效率优势

Hybrid SWA通过混合局部窗口注意力(SWA)与全局注意力(Full Attention),在保持长程依赖建模能力的同时,大幅降低了计算与存储开销。

1.1 计算量分析

以MiMo-V2.5-Pro为例,模型共70层,其中60层为SWA(窗口大小128),10层为Full Attention。相比纯Full Attention架构,Hybrid SWA的计算量降至约1/7,在近似计算受限的Prefill阶段带来显著成本缩减。

1.2 KVCache存储分析

SWA层仅需保留滑动窗口内的KV,使得整体KVCache存储占用同样下降至接近Full Attention的1/7。在近似内存带宽受限的Decode阶段,这直接降低了长序列场景下的推理延迟。经估算,MiMo-V2.5系列在KVCache存储效率上位列国产模型前列。

二、KVCache系统重构

为充分发挥Hybrid SWA的存储优势,团队对KVCache系统进行了深度重构。

2.1 SWA KVCache管理

* KVCache双池:将KVCache拆分为Full Attention与SWA两个独立存储池。SWA池严格按窗口大小配置容量,实现O(W)存储约束,使整体KVCache容量效率提升约7倍。

* KVCache按层异步拉取:优化后,SWA层只需预取极少量KVCache,实现了与计算过程近乎完美的重叠,使缓存读取成本趋近于零。

* SWA-aware前缀缓存树:改造传统前缀树匹配规则,引入“窗口安全长度”校验,确保SWA模式下前缀复用的正确性。同时优化淘汰机制与索引管理,在降低存储占用的同时,实际提升了缓存命中率。

* KVCache命中率提升优化:针对分布式缓存各层级状态不一致的问题,实施了主动同步、高频序列保护、中短序列SWA留存等一系列优化,将缓存容量优势转化为高有效命中长度。优化后,服务端KVCache命中率平均可达93%,高强度用户场景下可超过95%。

2.2 GCache:高性能分布式缓存基础设施

GCache是小米自研的高性能通用缓存,作为L3 KVCache,具备以下特点:

* 架构设计:采用非中心化元数据管理,支持内存/磁盘多级缓存与平滑扩缩容,提供多语言SDK与异步回调。

* 网络优化:优先使用GPU网卡通信,并通过NUMA绑定等优化,单进程RDMA读吞吐可达170 GB/s,延迟仅280 us。

* 存储成本优化:采用在GPU机器上混布的方式,利用闲置内存与NVMe SSD,实现额外存储成本为零。

* 稳定性保障:通过代码加固、逻辑分组、自动化故障处理与快速超时机制,在混布单副本状态下保障了高可用性。

2.3 针对缓存命中率的讨论

SWA的低存储占用与GCache的大容量支持,显著延长了KVCache的留存时间(TTL),从而大幅提升了缓存命中率。更小的带宽需求也降低了多级存储间的数据搬运开销。

三、调度优化

3.1 KVCache与负载亲和调度

自研无状态调度器LLM-Router使用Redis作为中心化存储,避免单点故障导致的KVCache调度回退。调度策略优先选择已缓存请求前缀且负载较低的节点,上线后将L2缓存命中率提升约25%,单机输入吞吐提升约30%。

3.2 TTFT优化

在等待队列中,优先调度实际需计算token数更少的请求进行Prefill,并引入等待时间惩罚机制避免饥饿。该策略最高可将TTFT的P90指标降低30%,改善了用户首token等待时间。

四、Prefill优化

4.1 分布式配置

SWA KVCache优化后,所需显存减少,使得专家并行(EP)规模得以缩减至原先的1/2,带来端到端性能约40%的提升。

4.2 长度分桶策略

为缓解长短请求混合调度导致的负载不均衡,采用三级长度分桶策略(0–64K / 64K–256K / 256K–1M),将负载特征相近的请求聚合计算,显著提升了线上Prefill的平均吞吐。

4.3 MoE负载均衡

得益于预训练阶段的优化,模型各层平均专家负载度已达约0.85的较优水平,目前未引入额外负载均衡策略。

4.4 解决NUMA冲突

禁用系统内核中与推理框架冲突的numa_balancing参数,消除了计算间隙,使端到端性能提升约10%。

五、Decode优化

5.1 显存优化

通过完整支持Decode阶段的SWA KVCache、优化PD分离架构中的预分配策略、调优CUDA Graph参数等措施,有效扩充了单节点可承载的并发量,提升了Decode吞吐。

5.2 MTP优化

在Prefill阶段引入MTP支持并优化缓存适配,使Decode前期令牌预测加速效果大幅提升:第0–128 token加速比达2.3倍,第128–256 token加速比达1.5倍,有效降低了Agentic场景的Decode成本。

六、多模态推理优化

基于SGLang社区方案,团队对多模态Encoder进行了深度优化,在延时基本不变的情况下,将Encoder吞吐提升至2倍。

6.1 架构优化

支持多模态Embedding数据复制与推理并行、Encoder数据并行(TP=1)以及跨请求组Batch,提升了GPU利用率。

6.2 预处理优化

将大图预处理迁移至GPU、实现图片并行下载、解耦多模态数据下载与推理过程、采用视频并行解码,显著降低了端到端延迟,例如1小时视频的Encoder延时从156秒降至23秒。

6.3 缓存优化

通过一致性哈希路由提升多模态缓存命中率30%,并通过共享内存实现机内多卡间缓存数据共享。

七、后记

本次优化是Hybrid SWA、MoE与多模态复合架构首次全面的大规模工程落地实践。通过多维度协同优化,最终将架构的理论效率优势转化为生产环境中的真实性能与成本收益。相关优化成果已通过API降价回馈用户,并将部分工作以PR形式回馈SGLang开源社区。

原文:MiMo-V2.5 系列推理全链路优化:将 Hybrid SWA 效率推向极致(来源:Xiaomi Mimo)

免责声明:风险提示:本文内容仅供参考,不代表同花顺观点。同花顺各类信息服务基于人工智能算法,如有出入请以证监会指定上市公司信息披露平台为准。如有投资者据此操作,风险自担,同花顺对此不承担任何责任。
homeBack返回首页
不良信息举报与个人信息保护咨询专线:10100571违法和不良信息涉企侵权举报涉算法推荐举报专区涉青少年不良信息举报专区

浙江同花顺互联信息技术有限公司版权所有

网站备案号:浙ICP备18032105号-4
证券投资咨询服务提供:浙江同花顺云软件有限公司 (中国证监会核发证书编号:ZX0050)
AIME
举报举报
反馈反馈