百度在昆仑芯上适配MiMo Flash V2推理

小米 MiMo Flash V2 为 MoE 架构（3090 亿总参数、150 亿活跃参数），采用 SWA+Sink（滑动窗口注意力+锚定令牌）与 Full Attention 混合注意力机制，其 SWA+Sink 机制中 Key 头 192 个、Value 头 128 个，非对称结构导致模型无法直接在硬件上部署。 且该模型的部署能力当时仅存在于 vLLM 开发分支，vLLM 最新 v0.13.0 正式版本未包含相关适配。百度百舸和昆仑芯的技术团队参考社区 MiMo Flash V2 PR，在 vLLM-Kunlun Plugin 中完成两步：一是继承 vLLM v0.11.0 版本 Linear 模块中的 QKVParallelLinear 类并修改其核心逻辑，为 Key 权重与 Value 权重分别创建对应维度的张量，从根源解决维度适配问题。 二是在 vLLM v0.11.0 不支持该模型的情况下自定义实现 MiMoV2FlashForCausalLM_KUNLUN 模型类，实现上层模型与底层 P800 算子的对接。性能侧，团队在 2 天内基于已有的 SWA attention 算子增量开发 Sink 功能并整合进高性能算子库，开发者升级算子库即可在 P800 上高性能运行。 随后用插件配套的 torch_xray 精度对齐工具自动比对 GPU 与昆仑芯 P800 的逐层输出以定位并优化数值偏差，用 xpu_profiler 生成算子调用时序图识别性能瓶颈与计算气泡，最终使 MiMo Flash V2 在 P800 上的推理效果与 GPU 平台保持一致，实现无损适配与高效运行。 此外团队已提前完成 MTP（多 token 预测）的底层逻辑铺设：本次开发的 SWA+Sink 算子在 decode 阶段支持的 query 长度最大可达 32，远超该模型 MTP（N=3）所需的 4，开源社区修复相关问题后 P800 可立即启用该加速模式。正文未给出吞吐、时延等量化指标。