KOSMOS: Self-scheduling Persistent Kernels for Communication-Compute Overlap
Abstract
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.
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
KOSMOS introduces self-scheduling persistent kernels that dynamically switch workgroup roles between communication and computation based on data readiness, achieving significant speedups in TP+SP and MoE training on AMD GPUs. The paper presents a rigorous systems-level contribution that addresses a fundamental bottleneck in distributed training, offering a novel alternative to static resource partitioning with strong empirical validation on modern hardware.
Methodology
The paper introduces KOSMOS, a library of persistent kernels that fuse communication and computation for Tensor/Sequence Parallelism (TP+SP) and Expert Parallelism (EP). The core innovation is a "self-scheduling" mechanism where workgroups (WGs) dynamically switch between "mover" (communication) and "builder" (computation) roles based on data readiness queues, rather than relying on fixed resource splits or host-side schedulers. This approach effectively addresses the idle time inherent in static fused kernels by allowing the hardware to adapt to the dynamic nature of collective operations (e.g., All-Gather, Reduce-Scatter, All-to-All). The design is deeply integrated with the AMD CDNA4 architecture (MI350X/MI355X), utilizing specific memory access patterns (band-major, local-first) to maximize bandwidth utilization and minimize latency. The methodology is rigorous, providing a clear theoretical justification for why dynamic role switching outperforms static or cost-model-based approaches, particularly in the context of MoE layers where routing decisions change per batch.
Results & evaluation
The evaluation is comprehensive, covering both microbenchmarks (individual kernel performance) and end-to-end training (Megatron-LM integration). The paper reports significant speedups: 1.10-1.28x for TP+SP training and 1.29-1.83x for MoE training compared to state-of-the-art baselines like DeepEP and Transformer Engine. The experiments are conducted on two different AMD GPU models (MI350X and MI355X) and include both BF16 and MXFP8 precisions. Ablation studies are provided to isolate the impact of specific design choices, such as queue ordering and role-switching policies, demonstrating that the self-scheduling mechanism is the primary driver of performance gains. The use of proxy models for large-scale MoE training (DeepSeek-V3, Qwen3-235B) is a reasonable approach given hardware constraints, though full-scale multi-node validation is lacking.
Reproducibility
The paper provides detailed implementation specifics, including the use of HIP, C++20, and specific tile sizes. It references open-source frameworks like Megatron-LM and Transformer Engine, which aids in understanding the integration. However, the code for KOSMOS itself is not explicitly linked in the provided text (though it is likely available given the venue and industry context). The heuristics for determining mover/builder counts are documented, which is crucial for reproducibility. The bitwise determinism of the kernels is a strong point for reproducibility and debugging.
Limitations
The primary limitation is the evaluation scope: results are limited to single-node (8-GPU) setups. The paper acknowledges that extending to multi-node requires RDMA support and hierarchical scheduling, which is left as future work. Additionally, the kernels are optimized for AMD CDNA4 architecture; while the concepts may generalize, the specific implementation details (e.g., wavefront sizes, memory hierarchy) are tightly coupled to AMD hardware. The paper also notes that the design disqualifies certain GEMM algorithms like Stream-K, which could limit applicability in some edge cases.
Broader impact
This work has significant implications for the efficiency of large-scale LLM training, particularly for models utilizing MoE architectures. By reducing the critical path of communication, KOSMOS can lower the cost of training frontier models. The self-scheduling paradigm could inspire similar approaches in other domains of parallel computing where dynamic workloads require efficient resource allocation. The focus on AMD GPUs also contributes to the diversification of high-performance computing ecosystems, providing a competitive alternative to NVIDIA-centric solutions. KOSMOS introduces self-scheduling persistent kernels that dynamically switch workgroup roles between communication and computation based on data readiness, achieving significant speedups in TP+SP and MoE training on AMD GPUs. The paper presents a rigorous systems-level contribution that addresses a fundamental bottleneck in distributed training, offering a novel alternative to static resource partitioning with strong empirical validation on modern hardware.
Author of this paper? Send a correction.