KOSMOS — Self-scheduling persistent kernels — What does it mean?

The news. On November 5, 2026, AMD researchers posted the preprint KOSMOS: Self-scheduling Persistent Kernels for Communication-Compute Overlap. On AMD Instinct MI350X and MI355X GPUs it reports 1.10–1.28× faster end-to-end Megatron-LM training with tensor plus sequence parallelism in BF16, and 1.29–1.83× faster DeepSeek-V3 and Qwen3-235B proxy training than DeepEP. Read the paper →

Start with the problem. When a model is split across GPUs with tensor and sequence parallelism, every transformer layer must gather pieces of its input from the other GPUs before a matrix multiplication, and scatter and add the partial results afterward. Mixture-of-experts layers add an all-to-all: tokens travel to the GPU that holds their expert, and results travel back. If the GPU waits for those transfers, its matrix cores sit idle, so training frameworks try to overlap the transfers with the multiplication. The data moves over the GPU-to-GPU links, which are separate hardware from the matrix cores, so in principle both can be busy at once.

Picture a restaurant kitchen at the dinner rush. The cooks are builders: each makes one dish, which is one tile of the matrix multiplication. The waiters are movers: they carry ingredients in from the pantry across the building, which is the other GPUs. The usual fix is a manager who assigns three waiters and five cooks at opening time and never changes it. At 6 pm the waiters are swamped while the cooks wait for ingredients; at 9 pm every ingredient has arrived and the waiters stand idle while the cooks fall behind. A fixed split is right for exactly one moment of the evening. Predicting the split with a cost model before each call does not fix this either: in the paper's measurements, a planned split ran slower than a plain fixed one (Table 2).

KOSMOS removes the manager. Each fused operation is one persistent kernel: it launches once with one workgroup resident per compute unit (the AMD counterpart of an SM) and avoids the cost of a launch per step. Workgroups claim work by atomically taking the next entry from a shared queue, and switch roles when their own work runs out. Switching is one-way in the all-gather and reduce-scatter kernels: movers become builders once their transfers finish, and builders join the reduction when the multiplication tiles run out. It is two-way in the mixture-of-experts combine step, where a workgroup pushes any ready result and otherwise builds a tile. This is operator fusion taken one step further: the communication is fused into the multiplication, not only the arithmetic.

Role switching only works if work shows up where workers look next, so the second ingredient is queue order: data is produced in the order its consumers will claim it. In the all-gather, builders compute their own local shard first while movers fetch remote shards in the order the multiplication needs them. In the reduce-scatter, output is produced band by band across every GPU's chunk, so every GPU's reducers can start after the first band is ready. In the kitchen, the ticket rail is sorted by which table is waiting, so whoever is free takes the most urgent ticket. Nobody schedules; the order of the rail is the schedule. The ablations back this up: producing reduce-scatter output row by row instead of band by band made the kernel 1.41–1.45× slower, a bigger loss than any role-switching policy gains.

How do the policies compare? The paper measured each against a fixed communication crew, as geometric-mean speedups on one GPU model (MI350X); these are one lab's preprint results.

A worked example shows where the gain comes from (illustrative numbers, built on one measured fact). The paper measures that its multiplications keep about 93% of full throughput on 224 of 256 compute units, because the idle units free power headroom and the active ones clock higher. Reserve 32 units as a fixed communication crew for a multiplication that takes 10 ms on the full GPU: it now takes 10 / 0.93 ≈ 10.75 ms. Suppose the reduction still needs 1.5 ms after the last tile, with every builder idle (illustrative): 12.25 ms total. If those builders join the reduction and the tail shrinks to 0.5 ms (illustrative), the total is 11.25 ms, so 12.25 / 11.25 ≈ 1.09× faster with the same amount of work. That sits inside the up-to-1.15× the paper reports for one-way switching over the best fixed split.

The limits matter as much as the gains. Communication-bound shapes gain little from switching, because movers stay busy until the last step: attention layers average 1.00× in the paper. Every measurement is on AMD MI350X and MI355X GPUs, and the DeepSeek-V3 and Qwen3-235B results come from proxy models, not full runs, so whether the same gains hold on NVIDIA hardware is not measured. The durable idea is separable from the benchmark: when demand for communication moves during a kernel, let readiness assign roles instead of predicting them.

Goes deeper in: GPU & CUDA → Operator Fusion & FlashAttention → Operator Fusion