KOSMOS：自调度持久内核实现通信与计算重叠
KOSMOS: Self-scheduling Persistent Kernels for Communication-Compute Overlap
一篇很硬核的内核优化论文：不靠预设切分，让 GPU 自己在通信和计算间切换，DeepSeek-V3 训练比 DeepEP 快 1.29-1.83 倍。
KOSMOS 是一个自调度持久内核库，workgroup 按消费者需求排序的队列在通信与计算角色间切换，无需调度器、按调用代价模型或手工调参。它将 transformer 层的所有 TP+SP 集合通信与 GEMM 融合，并把整个专家层从 router 到 fold 单次 launch 完成。在 AMD Instinct MI350X 和 MI355X 上测试，Megatron-LM 的 TP+SP 端到端训练 BF16 加速 1.10-1.28x、MXFP8 加速 1.10-1.25x。相比 DeepEP，KOSMOS 将 DeepSeek-V3 和 Qwen3-235B 代理的端到端训练加速 1.29-1.83x。
KOSMOS: Self-scheduling Persistent Kernels for Communication-Compute Overlap
Tensor and sequence parallelism (TP+SP) and expert parallelism (EP) put collectives on the critical path of transformer layers. State-of-the-art fused kernels mitigate this with fixed roles or schedules, but this rigidity leaves large idle windows. Predicting a communication-computation split does not remove them either; a predetermined split cannot follow data readiness. We present KOSMOS, a library for self-scheduling persistent kernels whose workgroups use queues ordered by consumer demand to switch between communication and computation roles as data become ready, with no scheduler, per-call cost model, or tuned split. KOSMOS fuses every TP+SP collective of a transformer layer with its GEMM, and runs an entire expert layer, from router to fold, in one launch per direction. Evaluated on AMD Instinct MI350X and AMD Instinct MI355X GPUs, KOSMOS speeds up TP+SP end-to-end training in Megatron-LM by 1.10-1.28x in BF16 and 1.10-1.25x in MXFP8. KOSMOS accelerates end-to-end training of DeepSeek-V3 and Qwen3-235B proxies by 1.29-1.83x over DeepEP.
