KOSMOS：用于通信与计算重叠的自调度持久化内核

arXiv:2611.07477v1 公告类型：交叉发布 摘要：张量与序列并行（TP+SP）以及专家并行（EP）使集合通信处于 Transformer 层的关键路径上。最先进的融合内核通过固定角色或调度方案来缓解这一问题，但这种刚性设计仍会留下大量空闲时段。预测通信与计算的资源划分也无法消除这些空闲时段，因为预先确定的划分无法随数据就绪状态动态调整。我们提出 KOSMOS，这是一个自调度持久化内核库，其工作组使用按消费者需求排序的队列，在数据就绪时切换通信与计算角色，无需调度器、逐次调用的成本模型或经调优的资源划分。KOSMOS 将 Transformer 层中的每个 TP+SP 集合通信操作与其 GEMM 融合，并在每个传播方向上仅通过一次内核启动，就能执行从路由器到折叠聚合的整个专家层。在 AMD Instinct MI350X 和 AMD Instinct MI355X GPU 上的评估表明，KOSMOS 将 Megatron-LM 中 TP+SP 端到端训练的速度在 BF16 下提升至 1.10–1.28 倍，在 MXFP8 下提升至 1.10–1.25 倍。与 DeepEP 相比，KOSMOS 将 DeepSeek-V3 和 Qwen3-235B 代理模型的端到端训练速度提升至 1.29–1.83 倍。