清华与腾讯混元联合团队使用Knight辅助完成NPU上的MoE推理优化

MLSys 2026设置MoE推理优化挑战赛，要求在指定NPU上为指定模型编写自定义算子，在输出与基线比特级对齐的前提下压低端到端解码延迟，Stanford、MIT等高校进入第二阶段。清华存储实验室与腾讯混元AI Infra联合团队夺冠。据介绍，单层MoE约占单步生成延迟82%，瓶颈在权重加载与数据搬运：团队实施E-Shard专家切分、PSUM批量读出、解码场景GEMV并发、标量引擎DMA提前启动首个专家、抑制编译器自动预取；注意力侧让QKV投影直接生成目标布局、V常驻片上、按d_head切分KV缓存更新，并把Attention、残差与MoE融合为单一算子；为满足比特级对齐，花近一周复现基线逐专家FP32转BF16累加顺序。团队还研发Agent算子优化器Knight：方案提出、实现、复盘三个Agent循环迭代，SQLite持久化经验，夜间自动优化、白天人工调优实现“人停卡不停”，并以路径检查防止修改评测脚本。据介绍最终端到端推理由14.91秒降至3.56秒，单步解码由12.63毫秒降至5.45毫秒，单层MoE由约204微秒降至约53微秒，权重加载DMA利用率提升到约80%。