KOSMOS: Self-scheduling Persistent Kernels for Communication-Compute Overlap

KOSMOS: S ELF - SCHEDULING P ERSISTENT K ERNELS FOR C OMMUNICATION -C OMPUTE OVERLAP

arXiv:2611.07477v1 [cs.DC] 5 Nov 2026

Alexander Magro 1 Aristotle Martin 1 Ye Wang 1 Wen Chen 1

A BSTRACT 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.28× in BF16 and 1.10–1.25× in MXFP8. KOSMOS accelerates end-to-end training of DeepSeek-V3 and Qwen3-235B proxies by 1.29–1.83× over DeepEP. 1

I NTRODUCTION

Training frontier large language models (LLMs) requires parallelism strategies that trade memory for communication. Memory reduction strategies such as tensor and sequence parallelism (TP+SP) and expert parallelism (EP) in Mixtureof-Experts (MoE) put collectives on the critical path. TP+SP wraps every transformer block in all-gather (AG) and reducescatter (RS) collectives, whereas EP in MoE models adds all-to-alls (A2A) around each expert layer. To keep matrix cores and links saturated, overlapping communication and computation is critical to training throughput. Fusing communication into general matrix multiplication (GEMM) avoids hardware-scheduler contention and moves data at tile granularity, overlapping with the GEMM’s tilebased computation. Prior fused kernels dedicate a fixed number of compute units (CUs) to communication (Zheng et al., 2025a; Awad et al., 2025; Sul et al., 2026a) or assign each workgroup (WG) its tasks with a scheduler (Aimuyo et al., 2025; Cheng et al., 2026). However, communication demand is not fixed: In dense collectives it shifts as the kernel progresses, and in MoE it shifts with each batch’s routing. Predicting a split of communication and computation WGs does not solve this issue. A planned split from a cost model before each call is slower than a fixed split in our measurements (Table 2). Even adapting splits at runtime, as mKernel does with a kernel controller, leaves their kernel 1 Advanced Micro Devices Inc., Austin, TX, USA. Correspondence to: Alexander Magro <alex.magro@amd.com>.

Preprint.

1.18× slower than the best fixed split (Mao et al., 2026). We present KOSMOS (Kernel-Overlapped Self-scheduling Megakernels On Symmetric-memory), a library of persistent kernels in which WGs dynamically change roles as the kernel runs. A WG is either a mover, which mainly communicates, or a builder, which mainly computes GEMM tiles. A WG selects its role from data readiness, with no scheduler or per-call cost model. Scheduling is implicit: Data is produced in the order consumers claim it. Builders start on local data while movers gather; reductions are produced, ordered by groups of tiled rows (band-major), so every rank’s movers start after one step, and dispatch delivers experts in the order each GEMM consumes them. The switching policy is set by each collective-GEMM pair (Figure 1). In AG+GEMM and MoE dispatch, movers become builders once their transfers are done; in GEMM+RS (Figure 1(A)), builders join the reduction when the GEMM work queue is exhausted; in MoE combine (Figure 1(B)) WGs switch both ways, pushing whenever data is ready and building otherwise. The policy is designed around the shape of the demand: Two-way switching is slower for TP+SP and A2A dispatch collectives whose demand is monotonic; one-way switching leaves combine gaps larger than needed. KOSMOS provides broad coverage, including, to our knowledge, the first readiness-driven TP+SP fused kernels and MoE megakernels for training. Our contributions are: 1. Self-scheduling persistent kernels, with producer and consumer queues co-designed per collective around data readiness. WGs may then switch roles from local

Figure 1. CU timelines and work queues on eight GPUs for (A) fused AG and RS and (B) the MoE forward and backward. Each CU runs one persistent WG; queues hold GEMM output tiles, and hatched segments are transfers (color = rank or expert). (A) The AG queue holds local tiles first, then peer tiles interleaved row by row; the RS queue is band-major (pattern = band). AG movers (CUs 0–1 of m) become builders once their transfers finish; in RS, builders become movers when the GEMM queue empties. (B) The forward queue is expert-major, GEMM1 before GEMM2; the backward queue holds GEMM2 then GEMM1 gradient tiles, interleaving dG and dW . All CUs first route and plan (purple); dispatchers push, then build; combine movers (CUs d, d + 1) build until a combine unit is ready, then interleave pushes with GEMM2 tiles. Backward: dY is dispatched first and the dX combine overlaps the GEMM1 gradients.

KOSMOS

KOSMOS

tile-level checks. Switching lets a shape-only heuristic replace per-shape tuning: within 1% of, and up to 15% faster than, the best single fixed split. A fixed split in turn beats a split predicted per call. 2. A characterization of where idle time falls in fused communication kernels and a policy rule that reduces it: one-way switching for monotonic demand, two-way switching for interleaved demand. 3. Complete kernel coverage for training: persistent fused kernels for every TP+SP collective, forward and backward, in BF16 and MXFP8, and the first roleswitching MoE layer with a single launch per direction, from router to fold, including gradients, evaluated endto-end in Megatron-LM on two GPU models. 4. Communications optimized for bytes/CU/s: collectives for fully connected meshes that saturate the bandwidth while maintaining fine grained access, and measurements showing that on power-capped GPUs the CUs lent to communication cost the GEMM little. We demonstrate KOSMOS on MI350-series GPUs and expect many of its findings to generalize to other systems: On AMD Instinct™ MI350X and AMD Instinct™ MI355X GPUs, KOSMOS speeds up end-to-end Llama 3 training by up to 1.28× over Transformer Engine (TE), and its MoE layers speed up training by up to 1.83× over DeepEP.

2

BACKGROUND AND M OTIVATION

2.1

Parallel Strategies

2.1.1

Tensor and Sequence Parallelism

Under TP+SP (Shoeybi et al., 2020; Korthikanti et al., 2023), each transformer block contains a column-parallel and a row-parallel GEMM pair (Figure 5). Activations outside the pair are sequence-sharded. An AG assembles the sequence before a column-parallel GEMM, and an RS after the row-parallel GEMM reduces sequence shards. The backward pass adds three collectives in BF16: AGs of the output gradients and the saved input, and an RS of the input gradient. MXFP8 gathers output gradients twice, quantized along both dimensions. Each GEMM pair runs five or six collectives per training step, each on the critical path unless overlapped. 2.1.2

Fully Sharded Data Parallelism

Fully sharded data parallelism (FSDP) shards only the model states, allowing collectives to gather model states for the next layer to run off of the critical path. Intra-node FSDP (Zhao et al., 2023; Rajbhandari et al., 2020) is often been preferred to TP+SP on AMD Instinct nodes (AMD,

n.d.e; Anthony et al., 2026). However, for large-scale distributed training, FSDP is often paired with TP+SP, which then shards activation memory across sequence, head, and channel dimensions, which dominates at scale. 2.1.3

Expert Parallelism

In an MoE layer with EP, the experts are distributed throughout the world size. A router GEMM matches k experts to each token, then an A2A dispatch sends activations to the selected experts’ GEMMs. An A2A combine then returns the expert outputs to the token’s original rank, with results folded by token. Because routing is decided per pass, transfer volume and destinations change every step, with skew towards GPUs hosting hot experts. The more skewed the data, the more completion times diverge. 2.2

MXFP8 Training

MXFP8 (Rouhani et al., 2023) stores each element in FP8 and shares one power-of-two (E8M0) scale across each block of 32 consecutive elements, so a tensor moves roughly half the bytes of BF16. A block is quantized in the contracting direction. In training, as each matrix and its transpose are used, we quantize both row-wise and column-wise. Requantizing one copy from the other compounds rounding error, so both must be quantized from the higher-precision source. MXFP8 has been shown to be as accurate in training as BF16 (Mishra et al., 2025), with twice the theoretical GEMM throughput on MI350-series GPUs. 2.3

Motivation and Design Requirements

For TP+SP, the common overlap strategy pipelines chunks of data through transport and compute on two streams. This is generally inefficient, because the two kernels compete for the same CUs without coordinating. In an eight-GPU node in the MI350-series, each GPU has a theoretical bidirectional peak of 1008 GB/s of bandwidth (AMD, 2025a), split between seven 144 GB/s XGMI links. A single peer is reachable at 1/7th of the aggregate. Full bandwidth usage requires sending to and receiving from all peers simultaneously. Host initiated direct memory access engines cannot provide this fan-out at fine granularity: They lag for small transfers because of their launch and synchronization cost, and they cannot perform reductions needed by TP+SP or EP. The tile-granular overlap therefore requires in-kernel communication. This leads to our first requirement: R1: One kernel, one stream: Communication and computation must share the GPU within one kernel. In LLM training with power-capped GPUs, GEMM throughput decreases only slightly when CUs are capped. Both MI350X and MI355X (AMD, 2025a;b) run at their power limit of 1000W and 1400W, respectively, with all 256 CUs active, so when a GEMM runs on fewer CUs, the remaining

KOSMOS

ones clock higher (Figure 3). On MI350X, our GEMMs keep an average 93% of their full-GPU throughput on 224 CUs (87.5% of the GPU), and 88% on 192 CUs (75% of the GPU) (Figure 2). Reserving CUs for communication, therefore, costs the GEMM less than the CU count suggests. R2: Use the power headroom: Collectives are power cheap, so CUs should be allocated to communication up to the point where the transport saturates. Overlapping a collective with its GEMM inside one kernel means dividing the GPU’s CUs between communication and computation. Most fused kernels make that division once, leading to windows where WGs sit idle on CUs. For AG, this occurs when the gather has finished before compute is complete. On RS, we have two windows: (1) where communicators wait for the first data to send, and (2) at the end when compute WGs have finished their tiles but the reductions are still occurring. In A2A, dispatches act much like AG with front-loaded communication, but combine traffic is interleaved with each batched GEMM, leaving windows scattered throughout the kernel. Figure 4 shows these windows in traces of fused kernels that keep a fixed split of communication WGs. Retuning the split may reduce idle time, but does not solve this structural problem, as communication demand changes as the fused kernel progresses. R3: Communicators budgeted adaptively: A fixed mover count leaves either matrix cores or interconnects idle, depending on the phase of the kernel. The number of active movers should follow the work that is available. MoE layers are a chain of steps whose sizes change from iteration to iteration. Executed as separate kernels, every step requires host overhead, routing decisions must leave the device, and no step can overlap another. Additionally, WGs do not necessarily start in their ID order (Appendix Figure A.1). With cross-WG dependencies (R1, R3), a resident consumer can wait on a producer that has not started; if waiting WGs fill every CU, the producer never spawns and the kernel deadlocks. R4: Persistent and resident: Kernels must be persistent with all WGs resident. When CUs are needed for a second stream of work, kernels must adapt to a WG cap so the whole grid stays resident. One key corollary follows. R5: Self-scheduling: While kernel design may guide WGs towards certain work, kernels cannot rely on a host round trip (breaks persistence, R4), a separate scheduling kernel (a second kernel, R1), or a scheduler WG (a CU lost to communication or compute, R2). Each WG must decide its own work only from the state it observes (adapting as demand changes, R3).

3

OVERVIEW OF KOSMOS

3.1

KOSMOS Design

KOSMOS implements each fused operation as a single persistent kernel that launches one WG per CU, with each WG resident for the kernel’s lifetime (R1, R4). Each WG starts in one of two roles: A mover transports data to and from remote peers, and a builder computes GEMM tiles. Depending on the collective, either role can be the producer or the consumer of the other. WGs coordinate only through in-kernel queues and readiness flags, without the host or a scheduler (R5). Rather than launching more WGs as work appears, KOSMOS keeps this fixed resident set and lets WGs change roles as the work changes (R3). Role switching is one-way or two-way. In one-way switching, a WG takes on another role for the rest of the kernel once its own work has finished. In AG and MoE dispatch, the movers become builders after their transfers are complete. In GEMM+RS, builders join the reduce once the GEMM tiles run out (Figure 1(A)). In two-way switching, a WG with no work available in its current role works in the other role temporarily, and returns when its own work appears: In the MoE combine, a set of WGs push any combine unit that is ready and builds GEMM tiles otherwise (Figure 1(B)). 3.1.1

Readiness and Queue Order

Without a scheduler, two things decide how a WG finds work: the unit at which data is declared ready, and the order of the queues it claims from. A WG claims a tile by atomically taking the next entry from a shared queue; it then computes that tile. KOSMOS operates in bands, chunks, tiles, and steps. A tile row is a set of tiles spanning 256 rows. A rank’s chunk is the tokens/TP-size output rows it owns after RS. A band is a group of tile rows within a chunk, and each step produces every chunk’s next band. Work is tracked at the finest grain the GEMM allows. In gathering collectives, the unit is a 256-row tile row of gathered input. In reducing collectives, it is a band of rows that is sized to the output. Every output tile has exactly one completion point, so a consumer can start work on it as soon as it is finished and enqueued. Two rules fix the order of work: Upstream before downstream: A WG claims every tile of an earlier step before any tile of a later one. A WG waits only on a tile claimed earlier in queue order, to avoid deadlocks or stalls. For example, in MoE forward, GEMM2 tiles are claimed only once every GEMM1 tile has been claimed (but not necessarily finished), so a waiting tile only ever depends on the work already claimed by a resident WG (Figure 1(B)). Additionally, a two-way WG never blocks on a push (as it builds instead), so waits are always satisfied. Produce in the order consumers need: If consumers are left waiting, we lose

KOSMOS

Figure 2. GEMM-only throughput (TFLOPS) of KOSMOS GEMM kernels and hipBLASLt under a CU cap, communication off. Top: small shapes (Llama 3 70B, MBS 4); bottom: large shapes (Llama 3 405B, MBS 8); TN, NN, and NT layouts, titled M×N×K. Measured on MI350X.

Figure 3. Power and shader clock (SCLK) values over a Llama 3 70B-FC1-MBS8 GEMM stress test of about 1.1–1.3 s. Both MI350X and MI355X remain at their 1000W and 1400W power caps until dropping to 128 CUs. Both GPUs reach their SCLK peak around 96 CUs.

KOSMOS

3.2.1

Fused AG

In the AG kernels (forward AG+TN, backward AG+NN), the movers pull the gathered input directly. Each mover is pinned by its WG index to one slice of one source GPU shard, and copies it one 256-row tile row at a time from the peer into a local buffer. The mover then performs a release-add on a local arrival counter for that peer and tile row. Builders compute the local shard first, then claim peer tiles in gather order. A builder spins on the counters in gather order and performs a single acquire once the count is complete. Once their slice has been transported, the movers perform a one-way switch to the builders.

Figure 4. Advanced Thread Trace (ATT) timeline of three communication and three computation WGs for an AG+GEMM from Llama 70B FC1, with a microbatch size (MBS) of 4; a GEMM+RS from Llama 8B with MBS 8; and MoE forward for DeepSeek-V3 with balanced routing. Dark Blue: working. Light Blue: spinning on a flag or sleeping. White: exited.

overlap. In AG+GEMM, builders work on the local shard before claiming other work while movers gather the first shards. In GEMM+RS, builders must produce band-major; finishing the first band of every rank’s chunk first allows every rank’s reducers to start after one step. In MoE forward, dispatch pushes experts in the order GEMM1 consumes them, and GEMM2 works on a row block so the combine can push to its destination as soon as GEMM2 finishes it. 3.1.2

Coexistence under a CU Cap

When other work must run concurrently, such as internode FSDP, KOSMOS kernels can be CU capped by calling frameworks to allocate CUs for a second stream of work. Then, AG+GEMM and MoE kernels shrink the builder count and RS kernels shrink both roles proportionally. Every WG of the capped grid then remains resident (R4).

In MXFP8, the row-parallel backward fuses one more GEMM: Training accuracy requires the weight gradient to use dY quantized along the token dimension, so a second column-wise quantized copy of dY must also be gathered. KOSMOS gathers both quantized copies of dY , and computes the data gradient (NN) and the weight gradient (NT) in one kernel to minimize idle time. 3.2.2

The RS kernel (forward TN+RS) uses a “one-shot” pull, derived from the two-shot all-reduce (Thakur et al., 2005). Builders write their partial tiles into a local staging buffer in band-major order. Movers own a band of the output; they pull it from every rank, sum in FP32 in place in a fixed order, and round once to BF16 before storing directly to the output. Every partial crosses the fabric once and is rounded once, which has less accumulation error than RCCL’s ring, and is bitwise reproducible (Appendix Table A.6). Arrival is detected in the data itself: The staging buffer is filled with a NaN sentinel, so the landing of the data is the signal in and of itself. A BF16 signal NaN is used to avoid hangs that could arise from a matrix fused multiply-add-produced arithmetic NaN. The first builders (up to saturation) that find no GEMM tile left become movers to assist with the reduction (one-way switching). 3.2.3

3.2

TP+SP Kernels

As shown in Figure 5, TP+SP requires six GEMMs and five or six collectives. KOSMOS names GEMMs by their layout in BLAS col-major convention: “T” for “transpose,” “N” for “no transpose.” For BF16, row-parallel wgrad is able to reuse the gathered dY from row-parallel dgrad. In MXFP8 training, dY T is also gathered. To overlap the increased communication, both AGs are fused along with the pair of backward GEMMs for the row-parallel layer into a single kernel. Thus, KOSMOS supports five persistent kernels for both data types (Figure 5).

Fused RS

Bulk Kernels

In the backward pass, two collectives have no data dependency on the GEMM they run beside. The AG of X is needed only by the weight gradient, so it can run alongside the data gradient GEMM, and the RS of dX is not needed by the weight gradient. This becomes a resource optimization problem, so KOSMOS defines a heuristic that uses GB/s/CU and TFLOPS/CU curves to balance WGs. KOSMOS then employs one-way role-switching for movers and builders in bulk fusions to minimize any remaining idle time.

KOSMOS

Figure 5. Overview of TP+SP kernel fusion in KOSMOS. Top: forward. Bottom: backward. Each dashed box is one KOSMOS kernel in BF16 or MXFP8; the red box is used only for MXFP8.

Figure 6. The KOSMOS single-launch MoE forward kernel. Red and blue: local experts; grey hatching: remote experts. Green: movers (light: dispatchers, dark: combiners); purple: builders. The flow runs from Start to End; stages are the forward-pass steps of Section 3.3.1.

KOSMOS

3.3

owners. In MXFP8, dY is pushed in BF16 and quantized on the receiver in both the row-wise (dgrad2) and column-wise (wgrad2) layouts. wgrad1 also needs the column-wise X, which the movers push after dY in link time the backward leaves idle; for narrow experts where the backward is communication-bound, it is sent with the forward dispatch instead (Appendix Table B.3). After their pushes, the movers compute the routing-weight gradients dwj = ⟨dY, Yj ⟩ from the forward’s combine buffer and back-propagate it with the aux-loss gradient through the score function to the logit gradient.

MoE Megakernels

An MoE layer in EP routes each token to k experts. Then, tokens are dispatched to the GPUs that own those experts. The per-expert token counts change every step, so they are unknown until the layer runs. To avoid host overhead and enable overlap across steps, KOSMOS launches all of the steps below in a single kernel launch (Figure 6). 3.3.1

Forward Pass

1. Route. All WGs compute the router logits XWr⊤ in FP32 tiles. The last tile to finish a row block of tokens routes that block: score, (group-limited) top-k, dispatch weights, and aux-loss inputs. The last block to be routed reduces the aux loss. The routing never leaves the device and does not need a separate phase.

2. One queue. dgrad2, with the SwiGLU backward fused into its epilogue, is interleaved with wgrad2 of the previous expert. The router’s dXr and dWr tiles are queued after dgrad2/wgrad2 to hide the tail associated with the last experts to land.

2. Plan. Builders sort the routing into per-expert send lists with a stable counting sort, keeping dW bitwise reproducible, since row order sets the wgrad’s accumulation order and MXFP8 scale groups. Only per-expert counts cross ranks. Each rank waits for all peers’ counts, then builds the layout and tile queues on the device. This wait is the forward’s only cross-rank barrier. 3. Dispatch ∥ GEMM1. A crew of 16-96 movers (Appendix Table B.3) pushes token rows directly into their final, expert-grouped position in each owner’s receive arena, which at EP 8 is slightly faster than pulling (Appendix Table A.7). Push lists are destination-minor, so consecutive movers target different GPUs to keep every link busy. Per-(source, expert) flags at the destination track readiness. In MXFP8, mover quantizes their rows and pad each group to 32 rows before pushing. Each GEMM1 tile, with SwiGLU fused into its epilogue, waits only for the source whose rows it reads, and stores the pre-activations for the backward. In MXFP8, an epilogue quantizes the SwiGLU output row-wise for GEMM2 and column-wise for wgrad2. 4. GEMM2 ∥ combine. A GEMM2 tile starts as soon as GEMM1 has finished its row block. Up to 128 dualrole movers (Appendix Table B.3) push each (expert, destination) unit back to its token owner once the unit’s row blocks are done, and claim GEMM tiles between pushes (two-way switching). The combine always pushes in BF16, even in MXFP8 training. 5. Fold. Once every row owed to this rank has arrived, each token’s k rows are weighted by its FP32 router weights and summed in FP32. 3.3.2

Backward Pass

1. Dispatch. A crew of 32-64 movers (Appendix Table B.3) pushes the router-weighted dY to the expert

3. Combine ∥ wgrad1. Up to 128 dual-role movers (Appendix Table B.3) push dX row blocks back to their token owners as dgrad1 completes them, while builders run wgrad1. In BF16, wgrad1 is interleaved into dgrad1; in MXFP8, wgrad1 uses the column-wise X, which arrives late when sent in the backward. All of dgrad1 runs first to cover its arrival; wgrad1 then overlaps the combine. 4. Fold. Each token’s k gradient rows and the router’s dXr are summed in FP32 into dX. 3.4

Implementation Details

KOSMOS is written in C++20 and HIP using eight waves and 512 threads per WG; local data share and register usage are kept within single occupancy. GEMM tiles use HipKittens (Hu et al., 2026), a C++ domain-specific language for AMD GPUs, with 256x256 output tiles for all expert and TP+SP GEMMs; the router GEMM uses a 64x128 tile as it only has as many columns as experts. Shape-dependent parameters (mover crew size, transport parameters, GEMM parameters) are fixed with closed-form rules during the first training iteration (Appendix Table B.3).

4

E VALUATION

4.1

Experimental Setup

4.1.1

System Specifications

Two single-node systems were used, with eight GPUs (MI350X or MI355X (AMD, 2025a;b)) each, fully connected over XGMI. Both nodes contained two AMD Zen 5 CPUs with 64 cores each. While MI350X runs at a peak clock of 2.2 GHz at up to 1000W, MI355X runs at a peak of 2.4 GHz at up to 1400W. MI350X offers 2.31 PFLOPS

KOSMOS

of peak dense BF16 matrix multiplication, while MI355X is slightly higher at 2.52 PFLOPS. The peak performance of MXFP8 is 4.62 PFLOPS and 5.03 PFLOPS, respectively.

MoE TFLOPS count actual routed FLOPs.

The benchmarks were run on ROCm 7.14 (AMD, 2026) with RCCL 2.30.4 (AMD, n.d.c), PyTorch (Paszke et al., 2019) 2.12, ROCm Transformer Engine (AMD, n.d.f) 2.18, and ROCm Megatron-LM (AMD, n.d.d) commit b0695a0.

4.2.1

4.1.2

Integration

The microbenchmarks were run with a standalone harness in KOSMOS. For TP+SP, a RCCL collective and a GEMM are run serially as a baseline (hipBLASLt (AMD, n.d.a) in BF16, a communication-free HipKittens GEMM in MXFP8). Each configuration is run with 24 shapes arising from the Llama3 (Grattafiori et al., 2024) class of models. The 8B, 70B, and 405B model sizes are used with a microbatch size between 1 and 8 and sequence length 8192. MoE runs one layer, both forward and backward, in a standalone harness. Layers of nine MoE models are used: DeepSeek V3 (DeepSeek-AI, 2025), DeepSeek V4 Pro (DeepSeekAI, 2026), Qwen3-30B and Qwen3-235B (Yang et al., 2025), Qwen3.5-397B (Qwen Team, 2026), Kimi K2.7 Code (Moonshot AI, 2026), GLM 5.2 (Z.ai, 2026), and Llama 4 Scout and Maverick (Meta AI, 2025). The MoE models use sequence length 4096 and MBS 2. KOSMOS is tested against baselines with balanced and skewed routing to simulate different training scenarios. End-to-end runs are executed using ROCm Megatron-LM. KOSMOS TP+SP kernels are transplanted into TE and run against Megatron-LM using the default ROCm TE path. The 405B model is restricted to 24 layers to fit on a single node. For MoE, the KOSMOS adapter replaces Megatron’s router, token dispatcher, grouped multilayer perceptron (MLP), and shared expert if present. Due to resource constraints, a proxy Qwen-235B model with 24 layers and a DeepSeekV3 proxy with no dense layers and nine MoE layers are considered. TP+SP models use a global batch size (GBS) of 64, and MoE models use a GBS of 16. 4.1.3

Numerics & Validation

KOSMOS kernels are bitwise run-to-run deterministic: tile claiming changes which WG computes a tile, never the reduction or fold order. Outputs are validated against RCCL+hipBLASLt with ROCm TE’s tolerances (AMD, n.d.f), re-poisoned every iteration, and checked twice to catch races. Microbenchmarks run and discard 10 warm-up iterations before 25 timed iterations. Each is run three times with the variant order shuffled, and we report the median of the slowest rank; the median run-to-run spread is 0.4% (max 2.0%). End-to-end results report Megatron-LM’s printed iteration time and throughput. Each configuration is run three times for 20 iterations, with the median result presented.

4.2

Ablations Data Readiness Mechanisms

Table 1 isolates the design decisions behind KOSMOS, changing one mechanism at a time with everything else held fixed. The largest effects come from the order in which data is produced and consumed: local-first AG queues consumed in gather order, band-major RS production, and destinationminor A2A push lists. KOSMOS is optimized for bandwidth per CU: its 8-wave movers saturate XGMI links with half the CUs of a 4-wave kernel (Appendix Table A.4). Fusing the router and shared expert into the MoE launch is a small but consistent gain. Readiness-driven switching only pays off if readiness arrives where WGs look next: Thus, the order of production is as important as the dynamic role switching it enables. 4.2.2

Role-Switching Policy

To measure the impact of KOSMOS’ derived role switching policies, we compare them to several alternatives (Table 2). For TP+SP, as implied by the one-way relationship of AG/RS and their GEMMs, we found that one-way role switching overall outperformed two-way role switching. One-way role switching is never more than 1% slower than the best single fixed split, and up to 1.15× faster. For AG, communication-bound shapes generally do not benefit from role switching, as movers are active until the very last step: While MLP layers improve by a geometric mean of 1.05×, attention layers average 1.00×. Per-shape results for both GPUs are in Appendix Figures A.2 and A.3. For MoE, we show the benefits of our two-way role switch in KOSMOS over alternatives. KOSMOS uses an interleaved two-way role switching policy, where WGs decide their role after every tile. We compare this to a policy where dispatch is executed with one-way role switching and combine uses a dedicated WG crew, following UniEP (Zheng et al., 2026), and to a static two-way planned steal, following Weave (Huang et al., 2026), where movers finish dispatch, compute a predetermined number of tiles from a runtime cost model, then return to combine. We find the two-way role switch not only beats a fixed split in all 36 cells on both MI350X and MI355X, but is the best policy on average, with a geometric mean speedup of 1.04×/1.04× on the forward pass and 1.06×/1.08× on the backward pass (MI350X/MI355X). This gain is up to 1.14×/1.16×. Appendix Figures A.4 and A.5 show the speedups per shape on the MI350X and MI355X, respectively.

KOSMOS Table 1. Design mechanisms ablated in isolation: slowdown of the alternative relative to KOSMOS. Full results in Appendix A. M ECHANISM

A LTERNATIVE

A LTERNATIVE VS KOSMOS

AG QUEUE ORDER NO LOCAL - FIRST, CHUNK - MAJOR RS PRODUCTION ORDER ROW- MAJOR A2A PUSH ORDER DESTINATION - MAJOR 8- WAVE MOVERS 4 WAVES PER WG (7–56 WG S ) ROUTER FUSION SEPARATE ROUTER KERNELS S HARED - EXPERT FUSION SEPARATE KERNELS

1.17–1.40× SLOWER 1.41–1.45× SLOWER 1.12–2.03× SLOWER 1.23–1.96× SLOWER 1.01–1.02× SLOWER 1.02–1.04× SLOWER

Table 2. Role-policy ablation: geometric-mean speedup over a fixed communication crew. TP+SP: 24 shapes per kernel, fixed crew of 56 WGs. MoE: nine models, fixed crew is the measured peak per shape, also used by the dedicated combine crew (following UniEP (Zheng et al., 2026)). Per-step planned split following Weave (Huang et al., 2026). All policies are implemented in KOSMOS. ∗ Shipped KOSMOS policy; bold: fastest per row and GPU. MI350X

TP+SP AG+TN TP+SP TN+RS TP+SP AG+NN

One-way switch

Two-way switch

Dedicated combine

Planned split

One-way switch

Two-way switch

Dedicated combine

Planned split

1.032∗ 1.031∗ 1.018∗

0.965 1.024 0.951

– – –

– – –

1.033∗ 1.018∗ 1.017∗

0.932 1.012 0.918

– – –

– – –

– – – –

1.029∗ 1.047∗ 1.052∗ 1.064∗

1.018 1.029 1.012 1.021

0.985 0.989 0.973 0.963

– – – –

1.027∗ 1.051∗ 1.067∗ 1.093∗

1.008 1.019 1.009 1.009

0.982 0.914 0.836 0.752

MoE fwd, balanced MoE fwd, skewed MoE bwd, balanced MoE bwd, skewed

4.3 4.3.1

MI355X

TP+SP Speedups TP+SP Microbenchmarks

Figure 7 compares each KOSMOS kernel to the same RCCL collective and hipBLASLt/HipKittens GEMM run without overlap. KOSMOS is faster on all 480 combinations of shape, kernel, precision, and GPU. In BF16, we see a geometric mean speedup of 1.38× on MI350X and 1.41× on MI355X. In MXFP8, the geometric mean is 1.36× on MI350X and 1.39× on MI355X. The MXFP8 GEMMs have twice the peak throughput of BF16, while an MXFP8 all-gather moves 33/64 of the BF16 bytes and the reducescatters move BF16 in both precisions, which leaves less compute to hide. Compared against unfused KOSMOS communication and GEMM kernels run without overlap gives similar geometric mean speedups: For BF16, the speedup is 1.39×/1.42×, and for MXFP8 the speedup is 1.37×/1.40×. This shows that most of the benefit over the baselines is due to the overlap mechanisms, rather than the components. 4.3.2

TP+SP End-to-End Training

Table 3 reports Llama 3 training in Megatron with TP 8, where KOSMOS replaces Transformer Engine’s TP+SP collectives and GEMMs. KOSMOS is faster for every model size, precision, and GPU model.

Figure 7. TP+SP microbenchmark: KOSMOS speedup over the baseline, geometric mean over 24 Llama 3 shapes per config.

KOSMOS Table 3. Llama 3 TP+SP end-to-end training, TP 8, sequence 8192, GBS 64: KOSMOS speedup over ROCm Transformer Engine (range over MBS 1–8, same MBS in both arms) and TFLOPS/GPU at each arm’s best MBS (TE → KOSMOS). r Full recompute in both arms; 70B at MBS 4 and 8 runs all 80 layers, and 405B runs 24 of its 126 layers. Per-MBS results: Appendix Table C.1. MI350X Model Prec.

Speedup

TFLOPS

MI355X Speedup

TFLOPS

8B

BF16 1.20–1.27 474→596 1.24–1.28 478→612 MXFP8 1.21–1.25 580→718 1.22–1.24 597→730

70B

BF16 1.15–1.21 664→769 1.22–1.28 654→834 MXFP8 1.19–1.22 893→1067 1.21–1.23 907→1119

405B BF16 1.10–1.16 570r →645r 1.17–1.22 624r →757r MXFP8 1.10–1.14 862r →985r 1.16–1.18 947r →1101r

4.4

Mixture-of-Experts Speedups

4.4.1

MoE Microbenchmarks

Figure 8 compares one KOSMOS MoE layer against RCCL A2A with TE’s grouped GEMMs. The figure covers nine models, both forward and backward, with balanced and skewed routing. KOSMOS is faster in all 144 cases, 1.28– 3.39×. Gains shrink as GEMM compute increases and top-k decreases: Qwen3-30B (narrow experts, top-8) gains up to 3.39×, while Llama 4 (wide experts, top-1) gains 1.28–1.92×. The forward pass with balanced routing gains the most, and backward pass with skewed routing is sped up the least. We also run KOSMOS components as separate launches; they are 1.25–1.51× slower than the fused layer on geometric mean which we attribute to kernel fusion. 4.4.2

MoE Proxy Training

Table 4 reports end-to-end training of two MoE proxy models. For both, balanced refers to random routing with a uniform distribution, so each expert has an equal expected number of tokens. Skewed results refer to a pure top-k routing defined by each model. For DeepSeek V3 with Megatron-generated mock data, this was seen to be up to 1.3× the data of balanced for the largest peer, whereas Qwen saw skew up to 4.2× in some iterations. KOSMOS relies on Megatron’s capacity factor flag to size shared buffers accordingly. A greater skew falls back to the default Megatron MoE layer. KOSMOS is faster than Megatron-LM with ROCm’s DeepEP (Zhao et al., 2025a; AMD, 2026) in every configuration, by 1.29–1.83× across GPU models.

5

R ELATED W ORK

Stream-level overlap. Frameworks (Li et al., 2020; NVIDIA, n.d.b) run NCCL/RCCL (NVIDIA, n.d.a; AMD, n.d.c) collectives on separate streams, where they hold CUs

Figure 8. MoE layer speedups over baselines for KOSMOS. Bars are per-model speedup for each direction, precision, and routing.

for the whole transfer and compete with the GEMM. CoCoNet (Jangda et al., 2022) adds device-side locks, chunking (Wang et al., 2022; Liang et al., 2025) shrinks the GEMMs, and copy engines (Zheng et al., 2025b) or new hardware (Pati et al., 2024) cannot reduce or signal per tile. Fused communication kernels. FLUX (Chang et al., 2024) fuses RS into a GEMM epilogue and AG into its prologue; ParallelKittens (Sul et al., 2026a), TileLink (Zheng et al., 2025a), and Iris (Awad et al., 2025) fuse TP+SP collectives per tile with a tuned number of communication streaming multiprocessors (SMs) or WGs. For MoE, COMET (Zhang et al., 2025), Mixture-of-Kittens (Sul et al., 2026b), and the MegaMoE kernels of DeepSeek-V4 (Zhao et al., 2025b; DeepSeek-AI, 2026) and Primus-Turbo (AMD, n.d.b; Zhang et al., 2026) fix communication resources at build time or per shape, before routing is known; Mixture-of-Kittens (a

KOSMOS Table 4. MoE proxy training, EP 8: KOSMOS speedup over DeepEP, and TFLOPS/GPU (DeepEP → KOSMOS). Balanced: uniform random routing; skewed: the model’s top-k router. Full results: Appendix Table C.2. MI350X Routing Prec.

MI355X

Speedup TFLOPS Speedup TFLOPS

Qwen3-235B balanced BF16 MXFP8 skewed BF16 MXFP8

1.70 1.82 1.31 1.57

359→610 372→677 249→327 245→384

1.74 1.83 1.29 1.54

387→673 400→731 275→355 268→411

DeepSeek-V3 balanced BF16 MXFP8 skewed BF16 MXFP8

1.31 1.39 1.33 1.40

423→556 446→619 414→551 437→613

1.33 1.39 1.36 1.41

480→637 500→696 465→631 488→689

second backward megakernel) and Primus-Turbo (fused operators reused in backward) compute the router outside the kernel. FlashMoE (Aimuyo et al., 2025) and MPK (Cheng et al., 2026) are persistent inference kernels driven by scheduler warps. FLUX, ParallelKittens, COMET, and Mixtureof-Kittens are NVIDIA-only, so Section 4.2 instead reproduces the UniEP and Weave policies inside the KOSMOS kernels. Adaptive communication budgets. UniEP (Zheng et al., 2026) orders one MoE task queue communication-first, with a split set by a cost model and tuner; Weave (Huang et al., 2026) sizes its communication SMs with an in-kernel cost model and lets them steal a precomputed number of GEMM tiles; mKernel (Mao et al., 2026) adapts the split with an inkernel controller but trails its best fixed split. In KOSMOS, roles follow data readiness alone, with no scheduler, task graph, per-call cost model, or steal budget, and one launch per direction covers every TP+SP collective with its GEMM, and the whole MoE layer, router and gradients included.

6

C ONCLUSION

Current state-of-the-art overlaps communication in TP+SP and EP training by running it on a separate stream, or by giving it a fixed or predicted share of the GPU inside a fused kernel. We show that on CDNA4 GPUs, both approaches are sub-optimal. Concurrent streams may contend without coordination, and a fixed split strands CUs whenever the demand for communication changes. KOSMOS addresses both by running communication and computation in one self-scheduling persistent kernel per operation or layer. Its WGs switch between moving and building as the work available changes, with in-kernel queues rather than a separate scheduler. Because collectives are limited by CUs up to saturation, giving CUs to communication costs little GEMM

throughput, if the kernel adapts. On MI350X and MI355X nodes, KOSMOS speeds up dense model training by up to 1.28× (BF16) and 1.25× (MXFP8), and MoE training by up to 1.74× (BF16) and 1.83× (MXFP8) over DeepEP. KOSMOS kernels are persistent only for intranode transports: XGMI with hipIPC memory is currently required. Remote direct memory access support and likely a hierarchical set of movers would carry the design to multinode. Thus, end-to-end results are limited to a single-node proxy model. We have evaluated with quasi-FSDP and pipeline parallelism stress tests, but complete multinode verification is future work. Finally, our structure requires one completion point per output tile. This disqualifies GEMM algorithms, such as stream-K (Osama et al., 2023), that compute a single output tile using multiple WGs. Other workloads built on fine-grained collectives, such as all-reduce as used in TP-only training and inference; SUMMA (Van De Geijn & Watts, 1997) and sparse distributed GEMMs; 3D fast Fourier transforms and spectral methods; and stencil solvers could use the same role-switching kernel design. Overlapping collectives with computation is no longer a question of how many CUs to allocate to communication, but of when. KOSMOS answers this with self-scheduling kernels that produce in the order consumers need, and lets data readiness drive role switches.

R EFERENCES Aimuyo, O., Oh, B., and Singh, R. FlashMoE: Fast distributed MoE in a single kernel. In Belgrave, D., Zhang, C., Lin, H., Pascanu, R., Koniusz, P., Ghassemi, M., and Chen, N. (eds.), Advances in Neural Information Processing Systems, volume 38, Main Conference, pp. 100676–100699. Curran Associates, Inc., 2025. doi: 10.52202/085713-3368. URL https: //proceedings.neurips.cc/paper_files /paper/2025/file/918d938bd209e5b5607 2777366f8a211-Paper-Conference.pdf. AMD. AMD Instinct MI350X GPU datasheet. https: //www.amd.com/content/dam/amd/en/doc uments/instinct-tech-docs/product-bri efs/amd-instinct-mi350x-gpu-brochure. pdf, 2025a. Product page: https://www.amd.co m/en/products/accelerators/instinct/ mi350/mi350x.html. Accessed 2026-10. AMD. AMD Instinct MI355X GPU datasheet. https: //www.amd.com/content/dam/amd/en/doc uments/instinct-tech-docs/product-bri efs/amd-instinct-mi355x-gpu-brochure. pdf, 2025b. Product page: https://www.amd.co m/en/products/accelerators/instinct/ mi350/mi355x.html. Accessed 2026-10.

KOSMOS

AMD. AMD ROCm Software, including RCCL and hipBLASLt. https://rocm.docs.amd.com/, 2026. RCCL: https://github.com/ROCm/rocm-s ystems/tree/develop/projects/rccl; hipBLASLt: https://github.com/ROCm/rocm-l ibraries/tree/develop/projects/hipbl aslt. Accessed 2026-10. AMD. hipBLASLt. https://github.com/ROCm/ rocm-libraries/tree/develop/projects /hipblaslt, n.d.a. Docs: https://rocm.docs. amd.com/projects/hipBLASLt/en/latest/. Accessed 2026-10. AMD. Mega MoE (Primus-Turbo documentation). https: //github.com/AMD-AGI/Primus-Turbo/b lob/main/docs/README_Mega_MoE.md, n.d.b. Accessed 2026-10. AMD. RCCL: ROCm Communication Collectives Library. https://github.com/ROCm/rocm-systems /tree/develop/projects/rccl, n.d.c. Docs: https://rocm.docs.amd.com/projects/r ccl/en/latest/. Accessed 2026-10. AMD. Megatron-LM for ROCm. https://github.c om/ROCm/Megatron-LM, n.d.d. Accessed 2026-10. AMD. Training a model with Megatron-LM for ROCm. https://rocm.docs.amd.com/en/docs-6.4 .1/how-to/rocm-for-ai/training/bench mark-docker/megatron-lm.html, n.d.e. ROCm 6.4.1 documentation. Accessed 2026-10. AMD. Transformer Engine on ROCm and AMDGPU. ht tps://github.com/ROCm/TransformerEng ine, n.d.f. Accessed 2026-10. Anthony, Q. et al. Training foundation models on a fullstack AMD platform: Compute, networking, and system design. In ISC High Performance 2026 Research Paper Proceedings (41st International Conference), pp. 1–18, 2026. doi: 10.23919/ISC.2026.11520488. Awad, M., Osama, M., and Potter, B. Iris: First-class MultiGPU programming experience in Triton, 2025. URL https://arxiv.org/abs/2511.12500. Chang, L.-W. et al. FLUX: Fast software-based communication overlap on GPUs through kernel fusion, 2024. URL https://arxiv.org/abs/2406.06858. Cheng, X. et al. MPK: A compiler and runtime for MegaKernelizing tensor programs. In 20th USENIX Symposium on Operating Systems Design and Implementation (OSDI 26), pp. 1909–1926, Seattle, WA, July 2026. USENIX Association. ISBN 978-1-939133-55-7. URL https://www.usenix.org/conference/os di26/presentation/cheng.

DeepSeek-AI. DeepSeek-V3 technical report, 2025. URL https://arxiv.org/abs/2412.19437. DeepSeek-AI. DeepSeek-V4: Towards highly efficient million-token context intelligence, 2026. URL https: //arxiv.org/abs/2606.19348. Grattafiori, A. et al. The Llama 3 herd of models, 2024. URL https://arxiv.org/abs/2407.21783. Hu, W. et al. HipKittens: Fast and furious AMD kernels. In Chowdhery, A. and Jia, Z. (eds.), Proceedings of Machine Learning and Systems, volume 8, pp. 1431–1443. MLSys, 2026. URL https://proceedings.mlsys.or g/paper_files/paper/2026/file/bc75fa 9843a7905bbed9d83895a88f7f-Paper-Con ference.pdf. Huang, Z. et al. Weave: Fine-grained dynamic SM scheduling in an MoE megakernel for compute-communication overlap, 2026. URL https://arxiv.org/abs/26 09.21483. Jangda, A. et al. Breaking the computation and communication abstraction barrier in distributed machine learning workloads. In Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, ASPLOS ’22, pp. 402–416, New York, NY, USA, 2022. Association for Computing Machinery. ISBN 9781450392051. doi: 10.1145/3503222.3507778. URL https: //doi.org/10.1145/3503222.3507778. Korthikanti, V. A. et al. Reducing activation recomputation in large Transformer models. In Song, D., Carbin, M., and Chen, T. (eds.), Proceedings of Machine Learning and Systems, volume 5, pp. 341–353. Curan, 2023. URL https://proceedings.mlsys.org/paper_ files/paper/2023/file/80083951326cf5 b35e5100260d64ed81-Paper-mlsys2023.p df. Li, S. et al. PyTorch distributed: Experiences on accelerating data parallel training, 2020. URL https://arxiv. org/abs/2006.15704. Liang, W. et al. TorchTitan: One-stop PyTorch native solution for production ready LLM pretraining. In Yue, Y., Garg, A., Peng, N., Sha, F., and Yu, R. (eds.), International Conference on Learning Representations, volume 2025, pp. 92161–92181, 2025. URL https: //proceedings.iclr.cc/paper_files/pa per/2025/file/e6231c5f46598cfd09ff19 70524e0436-Paper-Conference.pdf. Mao, Z. et al. mKernel: Fast Multi-GPU, multi-node fused kernels, 2026. URL https://arxiv.org/abs/26 09.13585.

KOSMOS

Meta AI. The Llama 4 herd: The beginning of a new era of natively multimodal AI innovation. https://ai.m eta.com/blog/llama-4-multimodal-intel ligence/, 2025. Accessed 2026-10. Mishra, A., Stosic, D., Layton, S., and Micikevicius, P. Recipes for pre-training LLMs with MXFP8, 2025. URL https://arxiv.org/abs/2506.08027. Moonshot AI. Kimi K2.7 Code. https://huggingf ace.co/moonshotai/Kimi-K2.7-Code, 2026. Model card. Accessed 2026-10. NVIDIA. NVIDIA Collective Communications Library (NCCL). https://github.com/NVIDIA/nccl, n.d.a. Accessed 2026-10. NVIDIA. Transformer Engine. https://github.com /NVIDIA/TransformerEngine, n.d.b. Accessed 2026-10. Osama, M., Merrill, D., Cecka, C., Garland, M., and Owens, J. D. Stream-K: Work-centric parallel decomposition for dense matrix-matrix multiplication on the GPU. In Proceedings of the 28th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming, PPoPP ’23, pp. 429–431, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400700156. doi: 10.1145/3572848.3577479. URL https://doi.org/10.1145/3572848.3577 479. Paszke, A. et al. PyTorch: An imperative style, highperformance deep learning library. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurip s.cc/paper_files/paper/2019/file/bdb ca288fee7f92f2bfa9f7012727740-Paper.p df. Pati, S., Aga, S., Islam, M., Jayasena, N., and Sinclair, M. D. T3: Transparent tracking & triggering for fine-grained overlap of compute & collectives. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, ASPLOS ’24, pp. 1146–1164, New York, NY, USA, 2024. Association for Computing Machinery. ISBN 9798400703850. doi: 10.1145/3620665.3640410. URL https://doi.org/10.1145/3620665.3640 410. Qwen Team. Qwen3.5: Towards native multimodal agents. https://qwen.ai/blog?id=qwen3.5, 2026. Accessed 2026-10.

Rajbhandari, S., Rasley, J., Ruwase, O., and He, Y. ZeRO: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1–16, 2020. doi: 10.1109/SC41405.2020.00024. Rouhani, B. D. et al. Microscaling data formats for deep learning, 2023. URL https://arxiv.org/abs/ 2310.10537. Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-LM: Training multibillion parameter language models using model parallelism, 2020. URL https://arxiv.org/abs/19 09.08053. Sul, S. H., Arora, S., Spector, B. F., and Ré, C. ParallelKittens: Systematic and practical simplification of Multi-GPU AI kernels. In Chowdhery, A. and Jia, Z. (eds.), Proceedings of Machine Learning and Systems, volume 8, pp. 1076–1089. MLSys, 2026a. URL https: //proceedings.mlsys.org/paper_files/ paper/2026/file/ff997469ac66cf893c41 83efeb22212a-Paper-Conference.pdf. Sul, S. H., Brown, N., Wildermuth, H., Lin, W., Cassano, F., and Ré, C. Mixture-of-Kittens: MoE megakernel for NVL72s, 2026b. URL https://arxiv.org/abs/ 2609.36070. Thakur, R., Rabenseifner, R., and Gropp, W. Optimization of collective communication operations in MPICH. The International Journal of High Performance Computing Applications, 19(1):49–66, 2005. doi: 10.1177/10943420 05051521. URL https://doi.org/10.1177/10 94342005051521. Van De Geijn, R. A. and Watts, J. SUMMA: Scalable universal matrix multiplication algorithm. Concurrency: Practice and Experience, 9(4):255–274, 1997. Wang, S. et al. Overlap communication with dependent computation via decomposition in large deep learning models. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1, ASPLOS 2023, pp. 93–106, New York, NY, USA, 2022. Association for Computing Machinery. ISBN 9781450399159. doi: 10.1145/3567955.3567959. URL https: //doi.org/10.1145/3567955.3567959. Yang, A. et al. Qwen3 technical report, 2025. URL https: //arxiv.org/abs/2505.09388. Z.ai. GLM-5.2: Built for long-horizon tasks. https: //z.ai/blog/glm-5.2, 2026. Accessed 2026-10.

KOSMOS

Zhang, L. et al. Enabling DeepSeek-V4-Flash training on AMD Instinct MI355X GPUs with Primus. https: //rocm.blogs.amd.com/software-tools-o ptimization/primus-deepseek-v4/README .html, 2026. ROCm Blogs. Accessed 2026-10. Zhang, S. et al. COMET: Fine-grained computationcommunication overlapping for Mixture-of-Experts. In Zaharia, M., Joshi, G., and Lin, Y. (eds.), Proceedings of Machine Learning and Systems, volume 7. MLSys, 2025. URL https://proceedings.mlsys.org/pa per_files/paper/2025/file/e27ea0cd50 b798ff8942caf9203f0992-Paper-Confere nce.pdf. Zhao, C. et al. DeepEP: an efficient expert-parallel communication library. https://github.com/deepsee k-ai/DeepEP, 2025a. Zhao, C. et al. DeepGEMM: clean and efficient BLAS kernel library on GPU. https://github.com/dee pseek-ai/DeepGEMM, 2025b. Zhao, Y. et al. PyTorch FSDP: Experiences on scaling fully sharded data parallel, 2023. URL https://arxiv. org/abs/2304.11277. Zheng, S., Zheng, X., Chang, L.-W., and Zhai, J. UniEP: Unified expert-parallel MegaKernel MoE for LLM training. In Proceedings of the 35th International Symposium on High-Performance Parallel and Distributed Computing, HPDC ’26, pp. 387–401, New York, NY, USA, 2026. Association for Computing Machinery. ISBN 9798400726408. doi: 10.1145/3806645.3807818. URL https://doi.org/10.1145/3806645.3807 818. Zheng, S. et al. TileLink: Generating efficient computecommunication overlapping kernels using tile-centric primitives. In Zaharia, M., Joshi, G., and Lin, Y. (eds.), Proceedings of Machine Learning and Systems, volume 7. MLSys, 2025a. URL https://proceedings.ml sys.org/paper_files/paper/2025/file/ c6ee784cbe46d854843e4c883a3321ef-Pap er-Conference.pdf. Zheng, S. et al. Triton-distributed: Programming overlapping kernels on distributed AI systems with the Triton compiler, 2025b. URL https://arxiv.org/abs/ 2504.19442.

KOSMOS

A

A BLATIONS

A.1

WGs out of Order

While workgroups of a persistent grid are all resident, they do not start in ID order. The dispatcher has tendencies to fill sorted by XCD rather than ID, but this is not guaranteed either. With oversubscription, the order depends on per-WG work. Variable work per WG can lead to highly unpredictable start order.

Figure A.1. Workgroup launch order and time sorted plotted against WG ID number. For persistent kernels, all WGs are resident but order is not guaranteed, and instead leans towards an accelerator complex die (XCD) ID-sorted launch order. With oversubscription, 1024 WGs exhibit different behavior based on the workload assigned to each WG.

A.2

Role Switching Policies

These figures give the per-shape results behind Table 2. For TP+SP, one-way switching is never significantly slower than the best fixed split, but up to 15% faster. For MoE, two-way switching beats a fixed split in all 72 cells, and is faster than a dedicated split in 63 of 72 cells and within 2% in the rest. Planned splits appear to lose where routing is skewed within a step, as a fixed crew size is determined before data arrives.

KOSMOS

Figure A.2. MI350X: speedup of one-way and two-way role switching over a best single fixed split of 56 communication WGs (8 per peer) for the three fused formats in training: AG+TN GEMM, TN GEMM+RS, and AG+NN GEMM. Switching arms use the KOSMOS heuristic crew (Table B.3).

KOSMOS

Figure A.3. MI355X: speedup of one-way and two-way role switching over a best single fixed split of 56 communication WGs (8 per peer) for the three fused formats in training: AG+TN GEMM, TN GEMM+RS, and AG+NN GEMM. Switching arms use the KOSMOS heuristic crew (Table B.3).

KOSMOS

Figure A.4. MI350X: speedup over a fixed split of different role-switching policies in MoE. KOSMOS two-way role switching with data readiness role switching, a UniEP-style one-way role switching for dispatch with a dedicated combine crew, and a Weave-style planned split. Crew size for non-KOSMOS policies is the smallest that reaches 98% of the peak dispatch and combine bandwidth (24, 32, or 48 WGs)

Figure A.5. MI355X: speedup over a fixed split of different role-switching policies in MoE. KOSMOS two-way role switching with data readiness role switching, a UniEP-style one-way role switching for dispatch with a dedicated combine crew, and a Weave-style planned split. Crew size for non-KOSMOS policies is the smallest that reaches 98% of the peak dispatch and combine bandwidth (16, 24, 32, or 48 WGs)

KOSMOS

A.3

Fused vs Separate Router and Shared Expert

Running the router or the shared expert as separate kernels is 1.01–1.02× and 1.02–1.04× slower, respectively, than the fused launch (Table A.1). Table A.1. Fusing the router and the shared expert into the MoE launch. Fused (KOSMOS): one launch. Separate: The router (GEMM, top-k, aux loss) or the shared expert runs as its own kernels (hipBLASLt for the shared-expert GEMMs) on the same stream. Without: the same MoE layer with routing precomputed (router rows) or routed experts only (shared-expert rows). Component

Model

Pass

Fused (ms)

Separate (ms)

Without (ms)

Speedup

Router Router Router Router

dsv3 dsv3 qwen235 qwen235

forward backward forward backward

7.557 13.773 3.622 5.808

7.612 14.005 3.679 5.918

7.168 13.711 3.398 5.752

1.007 1.017 1.016 1.019

Shared expert Shared expert

dsv3 dsv3

forward backward

7.655 15.533

7.990 15.797

7.218 13.865

1.044 1.017

A.4

Local First AG Queue vs Naive Ordering

A naive rank-order queue without local-first is 1.17–1.40× slower than the shipped queue (Table A.2). Table A.2. AG consumption order, configs 1 (AG+TN) and 3 (AG+NN), TP=8, atomic tile claim (queue order is the schedule). Shipped: local shard first, then peer tiles in the order movers gather them. Chunk-major: local first, peers chunk by chunk (mismatches the gather order). Naive: chunk-major in rank order, no local-first. Time: slowest rank, median of 3 runs × 25 calls; slowdown vs. shipped in parentheses. Config

Shape

Shipped (ms)

Chunk-major

Naive

TN TN

70B-FC1-MBS8 405B-QKV-MBS8

7.579 5.552

8.040 (+6.1%) 7.314 (+31.7%)

8.873 (+17.1%) 7.742 (+39.4%)

NN NN

70B-Down-MBS8 405B-OProj-MBS8

4.038 5.184

4.821 (+19.4%) 7.022 (+35.5%)

5.109 (+26.5%) 7.240 (+39.7%)

A.5

Band-Major vs. Row-Major RS Queue

Row-major production is 1.41–1.45× slower than band-major production (Table A.3). Table A.3. RS production order, config 2 (GEMM → RS), TP=8. Band-major (shipped) produces band b of every rank’s chunk before band b+1, so reducers start on the first band; row-major produces each chunk completely before the next. Shape 8B-OProj-MBS8 70B-Down-MBS8

A.6

Band-major (ms)

Row-major (ms)

Row / band

1.366 4.156

1.930 6.038

1.41× 1.45×

8 vs. 4 Waves

With four waves per WG, transport is 1.23–1.96× slower measured by bytes/CU/s at 7–56 WGs (Table A.4).

KOSMOS Table A.4. Mover waves vs. request depth, 1 GiB transfer (70B-MBS8), TP=8. The shipped 8-wave mover keeps one 16-byte remote load in flight per wave (two waves per single instruction, multiple data unit). 4-wave variants: same loop (depth 1), two loads in flight per wave (depth 2), or load/store interleave (pipelined, AG only). (Bold: fastest) Op

WGs

8 waves (shipped)

4 waves

4 waves, depth 2

4 waves, pipelined

AG AG AG AG AG

7 14 28 56 112

9.772 5.056 2.734 2.217 2.278

19.131 9.791 5.044 2.727 2.230

10.525 5.451 2.899 2.220 2.273

19.460 9.920 5.063 2.729 2.237

RS RS RS RS RS

7 14 28 56 112

15.900 7.988 4.068 2.288 2.325

30.378 15.205 7.627 3.878 2.236

18.490 9.263 4.683 2.513 2.259

– – – – –

A.7

Dst-Minor vs. Dst-Major Push

Destination-major push lists are 1.12–2.03× slower than destination-minor lists (Table A.5). Table A.5. Push order for EP transport, dsv3, EP=8, routed experts only (router and shared expert excluded). Destination-minor (KOSMOS): Consecutive WGs push to different destinations, one destination per WG. Destination-major: Consecutive WGs push to the same destination. Same kernels and data. Routing Phase Destination-minor (ms) Destination-major (ms) Slowdown Balanced Balanced Skewed Skewed

A.8

8.148 13.789 18.686 30.077

Forward Backward Forward Backward

16.532 20.567 24.841 33.653

2.03× 1.49× 1.33× 1.12×

RS Error vs. RCCL Ring, against a FP64 Reference

While KOSMOS accumulates all eight partials in FP32 and rounds to BF16 or MXFP8 once, ring algorithms often keep BF16 partials and round at every hop. The result is 2.21× lower reduction error for KOSMOS, and 1.71-1.73× lower error against a FP64 GEMM. KOSMOS is also bitwise run-to-run deterministic. Table A.6. RS error, TP=8. Every arm reduces the same BF16 partials. Reduction: vs. the FP64 sum of the partials. end-to-end: vs. an FP64 GEMM (32,768-element sample). Reduction Shape

KOSMOS

8B-OProj-MBS8 70B-Down-MBS8

A.9

−3

1.67 × 10 1.67 × 10−3

end-to-end

RCCL

Ratio −3

3.68 × 10 3.68 × 10−3

2.21 2.21

KOSMOS −3

2.34 × 10 2.35 × 10−3

RCCL

Ratio −3

4.04 × 10 4.02 × 10−3

1.73 1.71

Push vs Pull

While pushing rows into the destination’s arena is slightly faster than pulling them this requires more memory than a pull algorithm. As such, for larger EP sizes it may be required to use a pull algorithm instead.

KOSMOS

Table A.7. Push vs. pull EP transport in the full MoE layer (EP=8). Push (KOSMOS): The producer writes rows into the destination’s buffer. Pull: The consumer reads them from the source in the forward dispatch; combine and the backward push in both arms, so the backward rows are a control. Speedup: pull over push time. Model

Routing

Phase

Push (ms)

Pull (ms)

Speedup

dsv3 dsv3 dsv3

balanced balanced balanced

forward backward step

8.117 13.783 21.043

8.158 13.725 21.025

1.005 0.996 0.999

dsv3 dsv3 dsv3

skewed skewed skewed

forward backward step

18.628 30.055 45.144

19.004 30.052 45.506

1.020 1.000 1.008

qwen235 qwen235 qwen235

balanced balanced balanced

forward backward step

3.783 5.710 9.276

3.849 5.699 9.322

1.017 0.998 1.005

KOSMOS

B

A DDITIONAL KOSMOS D ETAILS

B.1

Fused Communication Framework Comparison

Table B.1. Fused communication–computation systems. Scope: what one launch fuses (a collective–GEMM pair or a whole MoE layer). Crew sized by: how the number of communication units is chosen. Roles/Switch trigger: whether and when a unit changes between communication and computation. Rtr/wgrad: router and weight gradients inside the launch. MoE launches F/B Crew sized by

System

Scope

Train

FLUX (Chang et al., 2024) TileLink (Zheng et al., 2025a) ParallelKittens (Sul et al., 2026a) Iris (Awad et al., 2025) COMET (Zhang et al., 2025) Primus MegaMoE (AMD, n.d.b) UniEP (Zheng et al., 2026)

pair pair pair pair pair pair pair

✓ ✗ ✗ ✗ △ ✓ ✓

– – – – ∼2 / ? 3 / 5a 3 / 6a

FlashMoE (Aimuyo et al., 2025) layer† DeepSeek MegaMoE (Zhao et al., 2025b) layer† Weave (Huang et al., 2026) layer†

✗ ✗ ✗

1/– 1+aux / – 1+aux / –

Mixture-of-Kittens (Sul et al., 2026b) MPK (Cheng et al., 2026)

layer model

✓ ✗

KOSMOS

pair + layer

✓

– launch config launch config launch config offline, per shape autotuned per shape perf-model tuner

scheduler launch config prologue cost modelb 1+aux / 1+aux launch config 1/– scheduler 1/1

shape model

Roles

Switch trigger

own tile fixed fixed fixed fixed fixed one-way sched. tasks fixed warps planned phases fixed sched. tasks

– – – – – – comm tasks exhausted, queue scheduler warp – precomputed steal countc – scheduler SMs

one-way + two-way

readiness at every claime

Switches /unit

Rtr wgrad TP+SP

0 0 0 0 0 0 1

– – ✗ – ✗ ✗ ✗

– – ✗ – ? part.a ✗a

✓ ✓ ✓ △ ✗ ✗ ✗

per task 0 ≤2

✓ ✗ ✗

– – –

✗ ✗ ✗

0 per task

✗d ✓

✓ –

✗ ✗

unbounded ✓f

✓

✓

†

Forward pass only (inference). a From public code. b Per layer and per GPU, from routing counts and calibrated bandwidth/throughput curves; fixed for the rest of the kernel. Communication SMs dispatch, steal a precomputed number of GEMM tiles, then combine; all SMs join the final combine. d Router runs before the launch; routing-score gradient fused. e Movers hand off when their share is done (AG, dispatch, bulk); combine movers check for ready work before every tile claim, in forward and backward. f Router GEMM, top-k and aux loss forward; router gradients (logits, dX, dWr ) backward.

c

B.2

Supported TP+SP Configs Table B.2. TP+SP KOSMOS kernels per collective-GEMM pair (†: MXFP8 only).

B.3

Heuristics

Pass

Collective

GEMM

Kernel

fwd fwd bwd bwd bwd bwd

AG X RS Y AG dY AG dY , AG dY ⊤ † AG X RS dX

col-parallel row-parallel row-parallel dgrad row-parallel dgrad + wgrad col-parallel dgrad col-parallel wgrad

fused AG+TN fused TN+RS fused AG+NN fused AG+NN(+NT) † bulk AG ∥ NN bulk RS ∥ NT

KOSMOS

Table B.3. Per-shape heuristics, evaluated once per shape on the host. Rules were fitted at TP 8 and EP 8; at other TP sizes mover count is scaled. Mover counts are per remote peer; MoE crews are WGs per GPU. All counts are then fitted to the CU cap so that every mover runs beside at least one builder. Parameter

Kernel

Input

Rule

Movers per peer

AG+GEMM, fprop AG+GEMM, dgrad bulk AG

col.

gathered width N

BF16: 6 if N is at least 3584, else 8. MXFP8: 8.

row

gathered width N

BF16: 6 if N is at least 4608, else 8. MXFP8: 8.

contraction K

BF16: 2, 4, 6 if K is at least 7168, 3584, 2304; else 8. MXFP8: 2, 4 if K is at least 3584, 1792; else 8.

GEMM+RS, row fprop

local K, waves w (tiles per cap)

bulk RS

output width n

K at least 4096: 4 (MXFP8 5). K at least 3072: 6 (MXFP8 8 if w is at least 16, else 7). K at most 768: 14, 10, 8 if w is at least 16, at least 8, or less. Otherwise 8 (MXFP8 10 if w is at least 16). BF16: 2, 4 if n is at least 8192, 3584; else 8. MXFP8: 8 if n is 5120 to 8191, else 4.

Dispatch crew

MoE forward

expert width I

Combine crew

MoE forward

CU cap

Dispatch crew

MoE backward

Combine crew

MoE backward

total bytes, GEMM work total bytes, GEMM work

Push burst depth

MoE dispatch, combine

hidden size H

Largest of 8, 7, 6, 4 that divides H/512.

Split-K factor

AG+GEMM row dgrad, bulk AG bulk RS

output tiles

BF16: 4 or 2 if the FP32 partials fit the workspace (at most 64 MiB), else 1. MXFP8: 1. BF16: 4 if at most 64 tiles and the partials fit the workspace, else 1. MXFP8: 1.

AG+GEMM (all)

tiles, tile rows

Reducers per peer

XCD bucketing

output tiles

total total

BF16: 32, 24, 16 if I is below 1024, below 3072, or larger. MXFP8: 24576/I rounded to a multiple of 8 between 16 and 48, doubled if the forward also sends the column-wise input. Minimum over ranks. The smaller of 128 and the cap, rounded down to a multiple of 8. 256 WGs times push time over push plus GEMM time, rounded up to a multiple of 8, between 32 and 64. 256 WGs times push time over GEMM time, rounded up to a multiple of 8, between 24 and 128 (MXFP8 32 and 128).

On if at most 1024 tiles and at least 5 tile rows.

KOSMOS

C

F ULL P ER -S HAPE AND P ER -GPU TABLES BEHIND THE S UMMARY F IGURES

Table C.1. Llama 3 TP+SP end-to-end training, TP 8, sequence 8192, GBS 64: median over three runs of Megatron-LM’s printed numbers (each the mean over iterations 3–19). Speedup is TE baseline ms/iter over KOSMOS ms/iter. OOM: out of memory. r Full recompute in both arms. Layer-reduced proxies: The 405B model runs 24 of its 126 layers with full recompute. 70B at MBS 4 and 8 runs with full recompute. Three runs per point, 20 iterations each. BF16

MXFP8

Model MBS Layers TE baseline KOSMOS KOSMOS Speedup TE baseline KOSMOS KOSMOS Speedup ms/iter ms/iter TFLOPS ms/iter ms/iter TFLOPS MI350X 8B 8B 8B 8B 70B 70B 70B 70B 405B 405B 405B 405B

1 2 4 8 1 2 4 8 1 2 4 8

32 32 32 32 80 80 80r 80r 24r 24r 24r 24r

10056.2 7727.6 7119.4 7117.2 47371.0 44371.2 62232.3 62377.7 56475.7 57459.5 57912.5 58049.1

7913.1 6435.7 5811.7 5658.0 41279.4 38308.3 51524.3 51415.7 51236.1 49923.4 49899.2 49998.7

426.6 524.4 580.7 596.4 713.3 768.6 571.4 572.6 628.1 644.6 644.9 643.6

1.271 1.201 1.225 1.258 1.148 1.158 1.208 1.213 1.102 1.151 1.161 1.161

8708.3 6729.2 5939.6 5813.0 37514.8 32956.2 43348.7 43643.7 37706.9 37438.2 37347.7 37655.0

6967.7 5454.8 4902.3 4702.4 30841.8 27603.6 36146.8 36258.4 34361.3 32840.9 32671.8 OOM

484.9 619.1 688.5 717.6 954.7 1066.6 814.5 812.0 936.5 979.9 985.0 –

1.250 1.234 1.212 1.236 1.216 1.194 1.199 1.204 1.097 1.140 1.143 –

MI355X 8B 8B 8B 8B 70B 70B 70B 70B 405B 405B 405B 405B

1 2 4 8 1 2 4 8 1 2 4 8

32 32 32 32 80 80 80r 80r 24r 24r 24r 24r

9996.1 7876.5 7220.7 7059.0 48695.2 45029.7 59942.3 58645.8 54176.8 52786.9 52032.3 51565.3

7915.3 6351.8 5709.9 5515.6 39918.9 35307.9 47738.9 46949.1 46203.8 43727.3 42743.1 42513.7

427.0 531.3 591.0 611.8 737.6 833.9 616.7 627.1 696.5 735.9 752.9 756.9

1.263 1.240 1.265 1.280 1.220 1.275 1.256 1.249 1.173 1.207 1.217 1.213

8647.1 6664.8 5893.8 5649.1 37191.5 32447.1 42303.4 41509.8 37137.2 35251.2 34474.2 33989.8

6971.3 5368.7 4833.4 4619.6 30146.4 26317.1 34680.9 34187.4 31882.3 29995.2 29407.8 29226.6

485.5 628.9 698.5 730.5 976.7 1118.8 848.9 861.2 1009.4 1072.9 1094.3 1101.1

1.240 1.241 1.219 1.223 1.234 1.233 1.220 1.214 1.165 1.175 1.172 1.163

KOSMOS

Table C.2. MoE proxy training, EP 8: median over three runs of Megatron-LM’s printed numbers, 20 iterations each. KOSMOS speedup is the DeepEP ms/iter over KOSMOS’s. DeepEP: Megatron-LM with ROCm’s DeepEP. Balanced: uniform random routing; skewed: the model’s top-k router. BF16 Model

Routing

Backend

MXFP8

ms/iter TFLOPS KOSMOS speedup ms/iter TFLOPS KOSMOS speedup

MI350X Qwen3-235B

balanced KOSMOS 545.1 DeepEP 925.6 Qwen3-235B skewed KOSMOS 1022.6 DeepEP 1337.7 DeepSeek-V3 balanced KOSMOS 612.7 DeepEP 804.8 DeepSeek-V3 skewed KOSMOS 617.1 DeepEP 822.7

MI355X Qwen3-235B

balanced KOSMOS 494.3 DeepEP 859.7 Qwen3-235B skewed KOSMOS 941.6 DeepEP 1211.8 DeepSeek-V3 balanced KOSMOS 534.3 DeepEP 709.7 DeepSeek-V3 skewed KOSMOS 539.7 DeepEP 732.1

610 359 327 249 556 423 551 414

– 491.4 1.698 894.7 – 868.1 1.308 1359.2 – 550.1 1.314 762.8 – 555.4 1.333 778.5

677 372 384 245 619 446 613 437

– 1.821 – 1.566 – 1.387 – 1.402

673 387 355 275 637 480 631 465

– 454.9 1.739 831.1 – 810.5 1.287 1244.2 – 488.7 1.328 680.0 – 493.8 1.356 697.9

731 400 411 268 696 500 689 488

– 1.827 – 1.535 – 1.391 – 1.413

Table C.3. TP+SP microbenchmark, AG+GEMM, column fprop, TP 8, MI350X: per iteration the slowest rank, median. Base: the RCCL collective and the GEMM run back to back on one stream with no overlap (BF16: torch.mm (hipBLASLt); MXFP8: a communication-free HipKittens MXFP8 GEMM). Speedup is base over KOSMOS; vs unf. is KOSMOS’s communication and GEMM WGs run back to back over the fused kernel. BF16 Shape

M

N

MXFP8

K Base ms KOSMOS ms Speedup vs unf. Base ms KOSMOS ms Speedup vs unf.

405B-FC1-MBS1 8192 13312 16384 405B-QKV-MBS1 8192 2304 16384 70B-FC1-MBS1 8192 7168 8192 70B-QKV-MBS1 8192 1280 8192 8B-FC1-MBS1 8192 3584 4096 8B-QKV-MBS1 8192 768 4096 405B-FC1-MBS2 16384 13312 16384 405B-QKV-MBS2 16384 2304 16384 70B-FC1-MBS2 16384 7168 8192 70B-QKV-MBS2 16384 1280 8192 8B-FC1-MBS2 16384 3584 4096 8B-QKV-MBS2 16384 768 4096 405B-FC1-MBS4 32768 13312 16384 405B-QKV-MBS4 32768 2304 16384 70B-FC1-MBS4 32768 7168 8192 70B-QKV-MBS4 32768 1280 8192 8B-FC1-MBS4 32768 3584 4096 8B-QKV-MBS4 32768 768 4096 405B-FC1-MBS8 65536 13312 16384 405B-QKV-MBS8 65536 2304 16384 70B-FC1-MBS8 65536 7168 8192 70B-QKV-MBS8 65536 1280 8192 8B-FC1-MBS8 65536 3584 4096 8B-QKV-MBS8 65536 768 4096 geometric mean

3.994 1.128 1.091 0.516 0.502 0.377 8.108 2.170 1.963 0.993 0.766 0.468 16.169 4.268 4.664 1.832 1.348 0.873 31.274 9.082 10.035 3.528 2.751 1.668

3.320 0.926 0.995 0.469 0.310 0.234 6.908 1.569 1.734 0.756 0.544 0.398 13.525 2.832 3.530 1.353 0.938 0.693 26.651 5.578 7.678 2.532 1.829 1.289

1.203 1.218 1.097 1.101 1.619 1.611 1.174 1.383 1.132 1.313 1.409 1.174 1.196 1.507 1.321 1.354 1.436 1.260 1.173 1.628 1.307 1.394 1.505 1.294 1.316

1.188 1.366 1.072 1.028 1.241 1.041 1.137 1.434 1.165 1.262 1.334 1.036 1.123 1.515 1.237 1.288 1.460 1.194 1.131 1.612 1.257 1.344 1.536 1.219 1.249

1.903 0.740 0.613 0.308 0.259 0.192 3.943 1.349 1.092 0.583 0.444 0.269 8.418 2.493 2.174 1.066 0.822 0.489 16.921 4.993 5.006 2.074 1.621 0.919

1.651 0.521 0.539 0.281 0.209 0.165 3.476 0.883 0.966 0.442 0.346 0.248 7.455 1.589 1.894 0.786 0.619 0.411 14.983 3.116 4.238 1.479 1.225 0.766

1.152 1.420 1.137 1.097 1.242 1.167 1.134 1.529 1.132 1.318 1.282 1.083 1.129 1.568 1.148 1.357 1.328 1.189 1.129 1.603 1.181 1.402 1.324 1.200 1.252

1.100 1.343 1.134 1.044 1.180 1.040 1.086 1.413 1.181 1.240 1.265 1.048 1.067 1.463 1.190 1.277 1.295 1.177 1.085 1.485 1.159 1.317 1.280 1.199 1.205

KOSMOS

Table C.4. TP+SP microbenchmark, AG+GEMM, column fprop, TP 8, MI355X: per iteration the slowest rank, median. Base: the RCCL collective and the GEMM run back to back on one stream with no overlap (BF16: torch.mm (hipBLASLt); MXFP8: a communication-free HipKittens MXFP8 GEMM). Speedup is base over KOSMOS; vs unf. is KOSMOS’s communication and GEMM WGs run back to back over the fused kernel. BF16 Shape

M

N

MXFP8

K Base ms KOSMOS ms Speedup vs unf. Base ms KOSMOS ms Speedup vs unf.

405B-FC1-MBS1 8192 13312 16384 405B-QKV-MBS1 8192 2304 16384 70B-FC1-MBS1 8192 7168 8192 70B-QKV-MBS1 8192 1280 8192 8B-FC1-MBS1 8192 3584 4096 8B-QKV-MBS1 8192 768 4096 405B-FC1-MBS2 16384 13312 16384 405B-QKV-MBS2 16384 2304 16384 70B-FC1-MBS2 16384 7168 8192 70B-QKV-MBS2 16384 1280 8192 8B-FC1-MBS2 16384 3584 4096 8B-QKV-MBS2 16384 768 4096 405B-FC1-MBS4 32768 13312 16384 405B-QKV-MBS4 32768 2304 16384 70B-FC1-MBS4 32768 7168 8192 70B-QKV-MBS4 32768 1280 8192 8B-FC1-MBS4 32768 3584 4096 8B-QKV-MBS4 32768 768 4096 405B-FC1-MBS8 65536 13312 16384 405B-QKV-MBS8 65536 2304 16384 70B-FC1-MBS8 65536 7168 8192 70B-QKV-MBS8 65536 1280 8192 8B-FC1-MBS8 65536 3584 4096 8B-QKV-MBS8 65536 768 4096 geometric mean

3.705 1.200 1.183 0.540 0.499 0.337 7.427 2.335 2.289 1.008 0.795 0.483 14.025 4.648 4.490 1.926 1.534 0.891 27.347 9.179 9.132 3.703 2.968 1.720

2.929 0.959 0.958 0.486 0.308 0.243 5.848 1.639 1.682 0.800 0.536 0.421 11.265 2.924 3.206 1.433 0.943 0.734 21.930 5.396 6.475 2.686 1.808 1.375

1.265 1.252 1.235 1.113 1.621 1.386 1.270 1.424 1.361 1.260 1.482 1.149 1.245 1.590 1.401 1.344 1.626 1.215 1.247 1.701 1.410 1.378 1.641 1.251 1.360

1.239 1.431 1.261 1.031 1.328 1.042 1.208 1.544 1.345 1.318 1.481 1.039 1.229 1.581 1.371 1.347 1.624 1.186 1.234 1.668 1.347 1.359 1.624 1.205 1.323

1.958 0.754 0.656 0.308 0.258 0.190 3.735 1.340 1.226 0.580 0.461 0.271 7.511 2.525 2.393 1.050 0.847 0.497 14.561 4.892 4.776 2.036 1.661 0.926

1.625 0.524 0.532 0.291 0.210 0.169 3.176 0.910 1.000 0.455 0.347 0.257 6.332 1.664 1.923 0.816 0.620 0.430 12.447 3.205 3.738 1.551 1.250 0.802

1.205 1.438 1.233 1.059 1.229 1.125 1.176 1.472 1.227 1.276 1.328 1.053 1.186 1.518 1.245 1.288 1.366 1.155 1.170 1.526 1.278 1.313 1.329 1.155 1.258

1.202 1.419 1.251 1.051 1.192 1.030 1.159 1.491 1.247 1.270 1.319 1.045 1.154 1.538 1.271 1.316 1.360 1.172 1.157 1.545 1.266 1.319 1.349 1.182 1.255

KOSMOS

Table C.5. TP+SP microbenchmark, GEMM+RS, row fprop, TP 8, MI350X: per iteration the slowest rank, median. Base: the RCCL collective and the GEMM run back to back on one stream with no overlap (BF16: torch.mm (hipBLASLt); MXFP8: a communication-free HipKittens MXFP8 GEMM). Speedup is base over KOSMOS; vs unf. is KOSMOS’s communication and GEMM WGs run back to back over the fused kernel. BF16 Shape

M

N

MXFP8

K Base ms KOSMOS ms Speedup vs unf. Base ms KOSMOS ms Speedup vs unf.

405B-Down-MBS1 8192 16384 6656 405B-OProj-MBS1 8192 16384 2048 70B-Down-MBS1 8192 8192 3584 70B-OProj-MBS1 8192 8192 1024 8B-Down-MBS1 8192 4096 1792 8B-OProj-MBS1 8192 4096 512 405B-Down-MBS2 16384 16384 6656 405B-OProj-MBS2 16384 16384 2048 70B-Down-MBS2 16384 8192 3584 70B-OProj-MBS2 16384 8192 1024 8B-Down-MBS2 16384 4096 1792 8B-OProj-MBS2 16384 4096 512 405B-Down-MBS4 32768 16384 6656 405B-OProj-MBS4 32768 16384 2048 70B-Down-MBS4 32768 8192 3584 70B-OProj-MBS4 32768 8192 1024 8B-Down-MBS4 32768 4096 1792 8B-OProj-MBS4 32768 4096 512 405B-Down-MBS8 65536 16384 6656 405B-OProj-MBS8 65536 16384 2048 70B-Down-MBS8 65536 8192 3584 70B-OProj-MBS8 65536 8192 1024 8B-Down-MBS8 65536 4096 1792 8B-OProj-MBS8 65536 4096 512 geometric mean

2.015 1.140 0.752 0.541 0.354 0.293 4.411 2.252 1.389 0.971 0.593 0.503 9.581 4.385 2.837 1.886 1.118 0.864 18.697 9.350 6.004 3.621 2.171 1.704

1.703 0.819 0.528 0.375 0.234 0.205 3.646 1.581 1.002 0.712 0.381 0.356 7.571 3.087 1.880 1.446 0.714 0.698 15.043 6.196 4.158 2.921 1.444 1.367

1.183 1.392 1.424 1.441 1.509 1.431 1.210 1.424 1.386 1.364 1.555 1.413 1.265 1.421 1.509 1.304 1.567 1.238 1.243 1.509 1.444 1.240 1.504 1.247 1.380

1.252 1.460 1.316 1.297 1.317 1.246 1.291 1.487 1.453 1.371 1.477 1.309 1.287 1.528 1.601 1.298 1.607 1.307 1.284 1.602 1.527 1.260 1.548 1.350 1.390

1.382 0.967 0.576 0.462 0.271 0.239 2.728 1.882 1.105 0.877 0.502 0.440 5.713 3.674 2.167 1.700 0.960 0.837 11.873 7.322 4.271 3.313 1.863 1.622

1.253 0.766 0.441 0.381 0.222 0.210 2.401 1.417 0.792 0.718 0.384 0.379 4.763 2.723 1.436 1.350 0.711 0.721 9.270 5.297 2.894 2.629 1.341 1.361

1.103 1.261 1.307 1.215 1.217 1.138 1.136 1.328 1.395 1.221 1.309 1.161 1.199 1.349 1.509 1.259 1.350 1.161 1.281 1.382 1.476 1.260 1.390 1.191 1.271

1.396 1.367 1.306 1.220 1.238 1.203 1.464 1.381 1.560 1.333 1.327 1.240 1.538 1.398 1.650 1.280 1.469 1.257 1.597 1.434 1.566 1.278 1.434 1.306 1.379

KOSMOS

Table C.6. TP+SP microbenchmark, GEMM+RS, row fprop, TP 8, MI355X: per iteration the slowest rank, median. Base: the RCCL collective and the GEMM run back to back on one stream with no overlap (BF16: torch.mm (hipBLASLt); MXFP8: a communication-free HipKittens MXFP8 GEMM). Speedup is base over KOSMOS; vs unf. is KOSMOS’s communication and GEMM WGs run back to back over the fused kernel. BF16 Shape

M

N

MXFP8

K Base ms KOSMOS ms Speedup vs unf. Base ms KOSMOS ms Speedup vs unf.

405B-Down-MBS1 8192 16384 6656 405B-OProj-MBS1 8192 16384 2048 70B-Down-MBS1 8192 8192 3584 70B-OProj-MBS1 8192 8192 1024 8B-Down-MBS1 8192 4096 1792 8B-OProj-MBS1 8192 4096 512 405B-Down-MBS2 16384 16384 6656 405B-OProj-MBS2 16384 16384 2048 70B-Down-MBS2 16384 8192 3584 70B-OProj-MBS2 16384 8192 1024 8B-Down-MBS2 16384 4096 1792 8B-OProj-MBS2 16384 4096 512 405B-Down-MBS4 32768 16384 6656 405B-OProj-MBS4 32768 16384 2048 70B-Down-MBS4 32768 8192 3584 70B-OProj-MBS4 32768 8192 1024 8B-Down-MBS4 32768 4096 1792 8B-OProj-MBS4 32768 4096 512 405B-Down-MBS8 65536 16384 6656 405B-OProj-MBS8 65536 16384 2048 70B-Down-MBS8 65536 8192 3584 70B-OProj-MBS8 65536 8192 1024 8B-Down-MBS8 65536 4096 1792 8B-OProj-MBS8 65536 4096 512 geometric mean

2.164 1.194 0.807 0.539 0.353 0.281 4.309 2.404 1.520 0.995 0.620 0.483 8.764 4.688 3.047 1.972 1.154 0.892 16.904 9.329 5.955 3.815 2.280 1.768

1.795 0.840 0.559 0.393 0.246 0.219 3.320 1.530 1.035 0.735 0.401 0.383 6.467 2.949 1.911 1.456 0.740 0.735 12.565 5.796 3.731 2.906 1.495 1.440

1.206 1.421 1.445 1.371 1.432 1.283 1.298 1.571 1.468 1.354 1.545 1.261 1.355 1.590 1.594 1.354 1.559 1.214 1.345 1.610 1.596 1.313 1.525 1.228 1.408

1.271 1.433 1.378 1.307 1.308 1.220 1.344 1.545 1.506 1.329 1.493 1.267 1.351 1.549 1.630 1.308 1.579 1.295 1.380 1.612 1.613 1.278 1.510 1.335 1.404

1.514 1.016 0.609 0.476 0.275 0.244 2.973 1.986 1.181 0.914 0.517 0.452 5.808 3.893 2.315 1.777 1.002 0.870 11.491 7.697 4.547 3.472 1.959 1.695

1.184 0.770 0.442 0.383 0.229 0.220 2.198 1.454 0.792 0.721 0.388 0.384 4.365 2.776 1.458 1.393 0.722 0.744 8.499 5.408 2.861 2.719 1.379 1.433

1.279 1.320 1.377 1.242 1.200 1.107 1.352 1.366 1.491 1.268 1.334 1.176 1.331 1.402 1.588 1.276 1.388 1.169 1.352 1.423 1.589 1.277 1.420 1.183 1.324

1.525 1.338 1.383 1.212 1.218 1.173 1.603 1.374 1.572 1.286 1.324 1.225 1.596 1.400 1.588 1.274 1.417 1.234 1.615 1.433 1.595 1.285 1.425 1.294 1.384

KOSMOS

Table C.7. TP+SP microbenchmark, AG+GEMM, row dgrad, TP 8, MI350X: per iteration the slowest rank, median. Base: the RCCL collective and the GEMM run back to back on one stream with no overlap (BF16: torch.mm (hipBLASLt); MXFP8: a communication-free HipKittens MXFP8 GEMM). Speedup is base over KOSMOS; vs unf. is KOSMOS’s communication and GEMM WGs run back to back over the fused kernel. BF16 Shape

M

N

MXFP8

K Base ms KOSMOS ms Speedup vs unf. Base ms KOSMOS ms Speedup vs unf.

405B-Down-MBS1 8192 6656 16384 405B-OProj-MBS1 8192 2048 16384 70B-Down-MBS1 8192 3584 8192 70B-OProj-MBS1 8192 1024 8192 8B-Down-MBS1 8192 1792 4096 8B-OProj-MBS1 8192 512 4096 405B-Down-MBS2 16384 6656 16384 405B-OProj-MBS2 16384 2048 16384 70B-Down-MBS2 16384 3584 8192 70B-OProj-MBS2 16384 1024 8192 8B-Down-MBS2 16384 1792 4096 8B-OProj-MBS2 16384 512 4096 405B-Down-MBS4 32768 6656 16384 405B-OProj-MBS4 32768 2048 16384 70B-Down-MBS4 32768 3584 8192 70B-OProj-MBS4 32768 1024 8192 8B-Down-MBS4 32768 1792 4096 8B-OProj-MBS4 32768 512 4096 405B-Down-MBS8 65536 6656 16384 405B-OProj-MBS8 65536 2048 16384 70B-Down-MBS8 65536 3584 8192 70B-OProj-MBS8 65536 1024 8192 8B-Down-MBS8 65536 1792 4096 8B-OProj-MBS8 65536 512 4096 geometric mean

2.172 1.069 0.743 0.511 0.411 0.368 4.209 2.052 1.430 0.897 0.575 0.463 9.128 3.983 2.595 1.731 1.112 0.820 18.304 8.486 5.649 3.390 2.049 1.566

1.976 0.940 0.666 0.421 0.273 0.220 3.696 1.529 1.160 0.764 0.421 0.380 7.253 2.788 2.070 1.351 0.736 0.694 14.027 5.205 4.066 2.511 1.370 1.281

1.099 1.137 1.115 1.214 1.506 1.674 1.139 1.343 1.233 1.175 1.366 1.217 1.258 1.429 1.253 1.281 1.512 1.181 1.305 1.630 1.389 1.350 1.496 1.222 1.305

1.127 1.070 1.074 1.056 1.025 1.152 1.215 1.299 1.195 1.055 1.256 1.096 1.260 1.415 1.351 1.180 1.423 1.058 1.274 1.629 1.417 1.252 1.501 1.130 1.220

2.314 1.171 0.806 0.660 0.465 0.312 4.916 2.267 1.531 1.074 0.737 0.484 10.103 4.554 2.993 1.952 1.295 0.864 20.140 9.607 6.696 3.682 2.418 1.641

1.840 0.871 0.582 0.382 0.269 0.238 4.165 1.742 1.088 0.680 0.455 0.397 8.679 3.665 2.160 1.329 0.843 0.722 17.485 7.415 5.148 2.670 1.670 1.395

1.258 1.345 1.384 1.728 1.730 1.307 1.180 1.301 1.407 1.579 1.619 1.220 1.164 1.243 1.386 1.469 1.536 1.197 1.152 1.296 1.301 1.379 1.448 1.176 1.357

1.163 1.269 1.377 1.240 1.386 1.101 1.145 1.264 1.397 1.442 1.513 1.083 1.181 1.205 1.417 1.469 1.515 1.184 1.205 1.316 1.329 1.467 1.551 1.201 1.302

KOSMOS

Table C.8. TP+SP microbenchmark, AG+GEMM, row dgrad, TP 8, MI355X: per iteration the slowest rank, median. Base: the RCCL collective and the GEMM run back to back on one stream with no overlap (BF16: torch.mm (hipBLASLt); MXFP8: a communication-free HipKittens MXFP8 GEMM). Speedup is base over KOSMOS; vs unf. is KOSMOS’s communication and GEMM WGs run back to back over the fused kernel. BF16 Shape

M

N

MXFP8

K Base ms KOSMOS ms Speedup vs unf. Base ms KOSMOS ms Speedup vs unf.

405B-Down-MBS1 8192 6656 16384 405B-OProj-MBS1 8192 2048 16384 70B-Down-MBS1 8192 3584 8192 70B-OProj-MBS1 8192 1024 8192 8B-Down-MBS1 8192 1792 4096 8B-OProj-MBS1 8192 512 4096 405B-Down-MBS2 16384 6656 16384 405B-OProj-MBS2 16384 2048 16384 70B-Down-MBS2 16384 3584 8192 70B-OProj-MBS2 16384 1024 8192 8B-Down-MBS2 16384 1792 4096 8B-OProj-MBS2 16384 512 4096 405B-Down-MBS4 32768 6656 16384 405B-OProj-MBS4 32768 2048 16384 70B-Down-MBS4 32768 3584 8192 70B-OProj-MBS4 32768 1024 8192 8B-Down-MBS4 32768 1792 4096 8B-OProj-MBS4 32768 512 4096 405B-Down-MBS8 65536 6656 16384 405B-OProj-MBS8 65536 2048 16384 70B-Down-MBS8 65536 3584 8192 70B-OProj-MBS8 65536 1024 8192 8B-Down-MBS8 65536 1792 4096 8B-OProj-MBS8 65536 512 4096 geometric mean

2.229 1.152 0.782 0.516 0.396 0.328 4.190 2.238 1.505 0.931 0.596 0.458 8.425 4.366 2.910 1.815 1.137 0.835 16.544 8.603 5.655 3.516 2.162 1.613

1.919 0.966 0.630 0.435 0.279 0.226 3.409 1.622 1.119 0.804 0.438 0.401 6.467 2.924 2.067 1.433 0.766 0.741 12.430 5.359 3.903 2.675 1.453 1.370

1.162 1.192 1.241 1.184 1.416 1.450 1.229 1.380 1.345 1.157 1.360 1.140 1.303 1.493 1.408 1.267 1.485 1.128 1.331 1.605 1.449 1.314 1.487 1.177 1.314

1.220 1.124 1.209 1.096 1.086 1.073 1.268 1.336 1.352 1.092 1.299 1.070 1.307 1.465 1.390 1.209 1.455 1.059 1.338 1.577 1.443 1.281 1.464 1.115 1.255

2.438 1.233 0.827 0.667 0.457 0.311 4.806 2.403 1.615 1.101 0.732 0.499 9.360 4.724 3.157 2.023 1.304 0.878 18.272 9.597 6.479 3.833 2.454 1.687

1.782 0.907 0.569 0.396 0.262 0.248 3.604 1.858 1.046 0.720 0.445 0.414 7.264 3.674 1.981 1.420 0.828 0.759 14.587 7.376 4.341 2.846 1.646 1.472

1.368 1.360 1.453 1.684 1.746 1.255 1.333 1.293 1.544 1.528 1.646 1.206 1.289 1.286 1.594 1.424 1.574 1.156 1.253 1.301 1.492 1.347 1.491 1.146 1.398

1.267 1.310 1.462 1.260 1.457 1.101 1.268 1.305 1.519 1.445 1.565 1.092 1.286 1.276 1.595 1.464 1.584 1.171 1.298 1.318 1.500 1.470 1.655 1.197 1.360

KOSMOS

Table C.9. TP+SP microbenchmark, bulk AG, column dgrad, TP 8, MI350X: per iteration the slowest rank, median. Base: the RCCL collective and the GEMM run back to back on one stream with no overlap (BF16: torch.mm (hipBLASLt); MXFP8: a communication-free HipKittens MXFP8 GEMM). Speedup is base over KOSMOS; vs unf. is KOSMOS’s communication and GEMM WGs run back to back over the fused kernel. BF16 Shape

M

N

MXFP8

K Base ms KOSMOS ms Speedup vs unf. Base ms KOSMOS ms Speedup vs unf.

405B-FC1-MBS1 8192 16384 13312 405B-QKV-MBS1 8192 16384 2304 70B-FC1-MBS1 8192 8192 7168 70B-QKV-MBS1 8192 8192 1280 8B-FC1-MBS1 8192 4096 3584 8B-QKV-MBS1 8192 4096 768 405B-FC1-MBS2 16384 16384 13312 405B-QKV-MBS2 16384 16384 2304 70B-FC1-MBS2 16384 8192 7168 70B-QKV-MBS2 16384 8192 1280 8B-FC1-MBS2 16384 4096 3584 8B-QKV-MBS2 16384 4096 768 405B-FC1-MBS4 32768 16384 13312 405B-QKV-MBS4 32768 16384 2304 70B-FC1-MBS4 32768 8192 7168 70B-QKV-MBS4 32768 8192 1280 8B-FC1-MBS4 32768 4096 3584 8B-QKV-MBS4 32768 4096 768 405B-FC1-MBS8 65536 16384 13312 405B-QKV-MBS8 65536 16384 2304 70B-FC1-MBS8 65536 8192 7168 70B-QKV-MBS8 65536 8192 1280 8B-FC1-MBS8 65536 4096 3584 8B-QKV-MBS8 65536 4096 768 geometric mean

3.794 1.178 1.055 0.548 0.493 0.364 8.014 2.232 1.991 1.008 0.753 0.499 16.086 4.425 4.335 1.920 1.382 0.908 31.460 9.445 9.304 3.791 2.661 1.751

3.386 0.669 0.879 0.324 0.254 0.176 6.959 1.393 1.815 0.602 0.493 0.323 13.411 2.824 3.719 1.218 0.933 0.617 26.567 5.904 7.733 2.456 1.918 1.242

1.121 1.761 1.200 1.694 1.942 2.062 1.152 1.603 1.097 1.676 1.529 1.546 1.199 1.567 1.166 1.576 1.480 1.472 1.184 1.600 1.203 1.544 1.387 1.410 1.443

1.257 1.743 1.606 1.582 1.511 1.342 1.354 1.647 1.577 1.578 1.569 1.388 1.374 1.660 1.639 1.556 1.720 1.403 1.390 1.680 1.636 1.569 1.638 1.386 1.528

1.813 0.665 0.598 0.319 0.253 0.185 4.015 1.279 1.139 0.576 0.430 0.296 8.535 2.552 2.289 1.094 0.797 0.531 17.226 5.114 4.825 2.164 1.549 1.002

1.639 0.391 0.478 0.197 0.173 0.121 3.670 0.751 0.904 0.355 0.286 0.205 7.590 1.478 1.666 0.710 0.533 0.385 14.820 3.031 3.955 1.445 1.076 0.754

1.106 1.701 1.251 1.621 1.468 1.524 1.094 1.702 1.260 1.620 1.500 1.443 1.125 1.727 1.374 1.542 1.495 1.378 1.162 1.687 1.220 1.498 1.440 1.328 1.414

1.145 1.722 1.355 1.631 1.583 1.419 1.133 1.770 1.380 1.547 1.698 1.498 1.183 1.811 1.537 1.553 1.675 1.357 1.221 1.773 1.330 1.521 1.649 1.364 1.480

KOSMOS

Table C.10. TP+SP microbenchmark, bulk AG, column dgrad, TP 8, MI355X: per iteration the slowest rank, median. Base: the RCCL collective and the GEMM run back to back on one stream with no overlap (BF16: torch.mm (hipBLASLt); MXFP8: a communication-free HipKittens MXFP8 GEMM). Speedup is base over KOSMOS; vs unf. is KOSMOS’s communication and GEMM WGs run back to back over the fused kernel. BF16 Shape

M

N

MXFP8

K Base ms KOSMOS ms Speedup vs unf. Base ms KOSMOS ms Speedup vs unf.

405B-FC1-MBS1 8192 16384 13312 405B-QKV-MBS1 8192 16384 2304 70B-FC1-MBS1 8192 8192 7168 70B-QKV-MBS1 8192 8192 1280 8B-FC1-MBS1 8192 4096 3584 8B-QKV-MBS1 8192 4096 768 405B-FC1-MBS2 16384 16384 13312 405B-QKV-MBS2 16384 16384 2304 70B-FC1-MBS2 16384 8192 7168 70B-QKV-MBS2 16384 8192 1280 8B-FC1-MBS2 16384 4096 3584 8B-QKV-MBS2 16384 4096 768 405B-FC1-MBS4 32768 16384 13312 405B-QKV-MBS4 32768 16384 2304 70B-FC1-MBS4 32768 8192 7168 70B-QKV-MBS4 32768 8192 1280 8B-FC1-MBS4 32768 4096 3584 8B-QKV-MBS4 32768 4096 768 405B-FC1-MBS8 65536 16384 13312 405B-QKV-MBS8 65536 16384 2304 70B-FC1-MBS8 65536 8192 7168 70B-QKV-MBS8 65536 8192 1280 8B-FC1-MBS8 65536 4096 3584 8B-QKV-MBS8 65536 4096 768 geometric mean

3.519 1.246 1.168 0.558 0.493 0.347 7.095 2.406 2.264 1.042 0.792 0.502 13.862 4.701 4.385 2.000 1.504 0.935 26.956 9.374 8.699 3.880 2.950 1.787

3.047 0.673 0.897 0.339 0.269 0.187 6.006 1.394 1.709 0.646 0.489 0.344 11.584 2.777 3.456 1.292 0.928 0.662 22.603 5.497 6.880 2.596 1.875 1.318

1.155 1.850 1.302 1.645 1.838 1.850 1.181 1.727 1.325 1.613 1.620 1.459 1.197 1.693 1.269 1.548 1.621 1.411 1.193 1.705 1.264 1.495 1.573 1.356 1.479

1.345 1.776 1.619 1.511 1.498 1.336 1.405 1.725 1.689 1.518 1.625 1.344 1.421 1.732 1.651 1.521 1.725 1.344 1.460 1.742 1.667 1.568 1.685 1.355 1.546

1.937 0.689 0.647 0.320 0.252 0.180 3.800 1.341 1.238 0.584 0.446 0.294 7.633 2.651 2.411 1.126 0.840 0.532 14.849 5.235 4.793 2.228 1.641 1.025

1.629 0.401 0.506 0.203 0.169 0.127 3.120 0.792 0.911 0.374 0.286 0.212 6.282 1.561 1.758 0.712 0.521 0.403 12.191 3.045 3.504 1.465 1.070 0.764

1.189 1.718 1.279 1.577 1.495 1.423 1.218 1.693 1.358 1.560 1.556 1.388 1.215 1.699 1.371 1.581 1.613 1.320 1.218 1.719 1.368 1.521 1.533 1.342 1.447

1.232 1.748 1.376 1.581 1.570 1.405 1.242 1.745 1.467 1.517 1.792 1.476 1.253 1.762 1.463 1.582 1.812 1.345 1.282 1.775 1.461 1.545 1.721 1.375 1.511

KOSMOS

Table C.11. TP+SP microbenchmark, bulk RS, column wgrad, TP 8, MI350X: per iteration the slowest rank, median. Base: the RCCL collective and the GEMM run back to back on one stream with no overlap (BF16: torch.mm (hipBLASLt); MXFP8: a communication-free HipKittens MXFP8 GEMM). Speedup is base over KOSMOS; vs unf. is KOSMOS’s communication and GEMM WGs run back to back over the fused kernel. BF16 Shape

M

N

MXFP8

K Base ms KOSMOS ms Speedup vs unf. Base ms KOSMOS ms Speedup vs unf.

405B-FC1-MBS1 13312 16384 8192 405B-QKV-MBS1 2304 16384 8192 70B-FC1-MBS1 7168 8192 8192 70B-QKV-MBS1 1280 8192 8192 8B-FC1-MBS1 3584 4096 8192 8B-QKV-MBS1 768 4096 8192 405B-FC1-MBS2 13312 16384 16384 405B-QKV-MBS2 2304 16384 16384 70B-FC1-MBS2 7168 8192 16384 70B-QKV-MBS2 1280 8192 16384 8B-FC1-MBS2 3584 4096 16384 8B-QKV-MBS2 768 4096 16384 405B-FC1-MBS4 13312 16384 32768 405B-QKV-MBS4 2304 16384 32768 70B-FC1-MBS4 7168 8192 32768 70B-QKV-MBS4 1280 8192 32768 8B-FC1-MBS4 3584 4096 32768 8B-QKV-MBS4 768 4096 32768 405B-FC1-MBS8 13312 16384 65536 405B-QKV-MBS8 2304 16384 65536 70B-FC1-MBS8 7168 8192 65536 70B-QKV-MBS8 1280 8192 65536 8B-FC1-MBS8 3584 4096 65536 8B-QKV-MBS8 768 4096 65536 geometric mean

3.971 1.145 1.242 0.559 0.444 0.330 8.329 2.151 2.386 1.014 0.763 0.525 16.695 4.193 4.731 1.922 1.396 0.932 32.763 9.099 10.576 3.724 2.695 1.736

3.256 0.660 0.771 0.336 0.238 0.190 7.155 1.279 1.513 0.619 0.494 0.364 14.465 2.573 3.673 1.195 0.967 0.651 29.801 5.821 8.094 2.384 1.956 1.226

1.220 1.735 1.611 1.663 1.864 1.730 1.164 1.682 1.577 1.638 1.545 1.441 1.154 1.630 1.288 1.608 1.444 1.433 1.099 1.563 1.307 1.562 1.378 1.416 1.476

1.589 1.793 1.610 1.558 1.870 1.509 1.443 1.821 1.681 1.585 1.722 1.249 1.453 1.885 1.488 1.602 1.759 1.276 1.419 1.774 1.417 1.613 1.798 1.302 1.582

2.262 1.011 0.825 0.479 0.321 0.299 4.827 1.949 1.567 0.886 0.590 0.552 10.298 3.780 3.151 1.713 1.115 1.037 20.770 7.988 6.813 3.315 2.188 2.017

1.653 0.625 0.552 0.333 0.183 0.179 3.696 1.193 1.052 0.613 0.335 0.325 7.812 2.339 2.129 1.174 0.622 0.608 16.225 4.630 4.902 2.294 1.207 1.162

1.369 1.619 1.496 1.436 1.752 1.674 1.306 1.633 1.489 1.447 1.760 1.699 1.318 1.616 1.480 1.458 1.792 1.706 1.280 1.725 1.390 1.445 1.813 1.735 1.551

1.505 1.516 1.472 1.328 1.670 1.544 1.453 1.525 1.431 1.329 1.644 1.568 1.383 1.524 1.404 1.336 1.675 1.590 1.346 1.551 1.326 1.339 1.703 1.611 1.486

KOSMOS

Table C.12. TP+SP microbenchmark, bulk RS, column wgrad, TP 8, MI355X: per iteration the slowest rank, median. Base: the RCCL collective and the GEMM run back to back on one stream with no overlap (BF16: torch.mm (hipBLASLt); MXFP8: a communication-free HipKittens MXFP8 GEMM). Speedup is base over KOSMOS; vs unf. is KOSMOS’s communication and GEMM WGs run back to back over the fused kernel. BF16 Shape

M

N

MXFP8

K Base ms KOSMOS ms Speedup vs unf. Base ms KOSMOS ms Speedup vs unf.

405B-FC1-MBS1 13312 16384 8192 405B-QKV-MBS1 2304 16384 8192 70B-FC1-MBS1 7168 8192 8192 70B-QKV-MBS1 1280 8192 8192 8B-FC1-MBS1 3584 4096 8192 8B-QKV-MBS1 768 4096 8192 405B-FC1-MBS2 13312 16384 16384 405B-QKV-MBS2 2304 16384 16384 70B-FC1-MBS2 7168 8192 16384 70B-QKV-MBS2 1280 8192 16384 8B-FC1-MBS2 3584 4096 16384 8B-QKV-MBS2 768 4096 16384 405B-FC1-MBS4 13312 16384 32768 405B-QKV-MBS4 2304 16384 32768 70B-FC1-MBS4 7168 8192 32768 70B-QKV-MBS4 1280 8192 32768 8B-FC1-MBS4 3584 4096 32768 8B-QKV-MBS4 768 4096 32768 405B-FC1-MBS8 13312 16384 65536 405B-QKV-MBS8 2304 16384 65536 70B-FC1-MBS8 7168 8192 65536 70B-QKV-MBS8 1280 8192 65536 8B-FC1-MBS8 3584 4096 65536 8B-QKV-MBS8 768 4096 65536 geometric mean

3.620 1.226 1.228 0.566 0.446 0.304 7.231 2.377 2.329 1.042 0.805 0.509 14.160 4.596 4.537 2.010 1.546 0.926 27.920 9.101 9.363 3.918 3.025 1.764

3.015 0.691 0.841 0.357 0.234 0.206 6.188 1.361 1.671 0.659 0.505 0.383 12.013 2.625 3.368 1.278 0.993 0.690 23.735 5.212 6.846 2.509 1.967 1.301

1.201 1.774 1.460 1.583 1.907 1.471 1.169 1.747 1.393 1.582 1.594 1.328 1.179 1.751 1.347 1.573 1.556 1.341 1.176 1.746 1.368 1.562 1.538 1.355 1.474

1.586 1.925 1.637 1.518 1.926 1.432 1.513 1.914 1.613 1.552 1.819 1.277 1.508 1.902 1.586 1.563 1.848 1.302 1.515 1.943 1.520 1.562 1.836 1.326 1.617

2.323 1.065 0.852 0.490 0.323 0.298 4.638 2.058 1.622 0.918 0.597 0.556 9.387 4.026 3.215 1.778 1.137 1.056 18.367 8.139 6.592 3.466 2.231 2.057

1.612 0.667 0.533 0.356 0.194 0.194 3.348 1.276 1.008 0.658 0.354 0.354 6.753 2.484 1.983 1.262 0.660 0.656 13.417 4.912 4.424 2.478 1.278 1.262

1.441 1.598 1.598 1.375 1.664 1.536 1.385 1.613 1.610 1.394 1.686 1.570 1.390 1.621 1.621 1.409 1.723 1.610 1.369 1.657 1.490 1.399 1.745 1.629 1.543

1.602 1.501 1.581 1.289 1.594 1.472 1.491 1.512 1.560 1.291 1.617 1.499 1.463 1.528 1.558 1.298 1.635 1.512 1.435 1.554 1.425 1.305 1.653 1.524 1.492

KOSMOS

Table C.13. MoE microbenchmark, BF16, forward pass, EP 8: per call the slowest rank, median. Baseline: RCCL A2A with TE’s grouped GEMMs. Speedup is baseline over KOSMOS; vs KOSMOS serial is the same KOSMOS kernels as separate launches over the fused layer. Routing

Model

MI350X balanced DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean skewed DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean MI355X balanced DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean skewed DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean

Baseline ms KOSMOS ms KOSMOS TFLOPS Speedup vs KOSMOS serial 17.041 11.026 17.443 18.219 14.753 8.742 4.158 3.358 4.373

7.670 3.791 8.531 8.117 6.562 3.405 1.555 2.381 2.710

752.6 543.8 761.2 711.1 754.0 726.6 397.7 865.8 760.8

31.357 19.649 30.492 32.022 27.324 16.608 7.093 7.586 7.679

15.792 7.848 17.083 16.005 13.640 7.547 3.126 5.737 5.651

365.5 262.7 380.1 360.7 362.7 327.8 197.9 359.4 364.8

15.856 11.215 16.249 17.165 13.696 8.421 4.103 3.183 4.090

6.806 4.008 7.091 7.081 5.882 3.463 1.644 2.287 2.423

848.1 514.4 915.8 815.2 841.2 714.4 376.1 901.6 850.8

30.078 19.235 29.068 30.608 26.230 15.958 7.048 6.574 7.143

14.522 8.481 14.683 14.620 12.631 7.338 3.368 5.123 4.966

397.5 243.1 442.3 394.8 391.7 337.1 183.6 402.4 415.2

2.222 2.908 2.045 2.244 2.248 2.568 2.674 1.410 1.614 2.164 1.986 2.504 1.785 2.001 2.003 2.201 2.269 1.322 1.359 1.898

1.463 1.610 1.314 1.438 1.476 1.568 1.508 1.099 1.212 1.400 1.324 1.402 1.286 1.335 1.333 1.373 1.216 1.078 1.223 1.282

2.330 2.799 2.292 2.424 2.329 2.432 2.495 1.392 1.688 2.199 2.071 2.268 1.980 2.094 2.077 2.175 2.093 1.283 1.438 1.911

1.592 1.570 1.479 1.571 1.591 1.601 1.456 1.158 1.277 1.469 1.379 1.314 1.394 1.401 1.376 1.388 1.198 1.115 1.241 1.308

KOSMOS

Table C.14. MoE microbenchmark, BF16, backward pass, EP 8: per call the slowest rank, median. Baseline: RCCL A2A with TE’s grouped GEMMs. Speedup is baseline over KOSMOS; vs KOSMOS serial is the same KOSMOS kernels as separate launches over the fused layer. Routing

Model

MI350X balanced DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean skewed DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean MI355X balanced DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean skewed DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean

Baseline ms KOSMOS ms KOSMOS TFLOPS Speedup vs KOSMOS serial 25.939 16.245 29.207 26.821 23.205 13.479 6.068 6.675 9.309

14.337 5.918 16.807 14.994 12.294 6.040 1.793 4.383 6.198

805.2 696.7 772.8 770.0 804.9 819.1 690.0 940.8 665.2

48.415 28.235 57.860 50.908 42.471 24.938 10.382 14.418 15.242

30.535 11.970 34.486 30.922 26.627 13.154 4.152 10.122 11.350

378.1 344.4 376.6 373.4 371.6 376.1 297.9 407.3 363.3

24.104 15.765 27.974 25.323 21.445 12.770 5.834 5.651 8.226

11.493 5.162 13.883 12.049 9.930 5.176 1.777 3.888 5.327

1004.5 798.8 935.6 958.2 996.6 956.0 696.0 1060.4 774.1

46.070 27.346 55.182 49.450 40.206 23.386 10.072 11.729 13.823

25.534 10.560 28.365 25.540 21.712 11.339 3.772 8.816 9.801

452.1 390.5 457.9 452.0 455.8 436.3 328.0 467.7 420.7

1.809 2.745 1.738 1.789 1.888 2.232 3.385 1.523 1.502 1.996 1.586 2.359 1.678 1.646 1.595 1.896 2.501 1.424 1.343 1.744

1.237 1.479 1.187 1.249 1.252 1.408 1.703 1.111 1.113 1.293 1.196 1.467 1.127 1.220 1.186 1.356 1.574 1.101 1.095 1.248

2.097 3.054 2.015 2.102 2.160 2.467 3.283 1.453 1.544 2.169 1.804 2.590 1.945 1.936 1.852 2.062 2.670 1.330 1.410 1.909

1.403 1.647 1.270 1.398 1.410 1.559 1.763 1.162 1.140 1.403 1.297 1.578 1.228 1.331 1.325 1.460 1.702 1.152 1.135 1.345

KOSMOS

Table C.15. MoE microbenchmark, MXFP8, forward pass, EP 8: per call the slowest rank, median. Baseline: RCCL A2A around Transformer Engine MXFP8 grouped GEMMs. Speedup is baseline over KOSMOS; vs KOSMOS serial is the same KOSMOS kernels as separate launches over the fused layer. Routing

Model

MI350X balanced DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean skewed DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean MI355X balanced DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean skewed DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean

TE MXFP8 ms KOSMOS ms KOSMOS TFLOPS Speedup vs KOSMOS serial 15.554 10.323 14.399 15.916 13.561 8.661 4.285 2.731 3.450

5.326 3.202 5.516 5.514 4.637 2.846 1.591 1.749 1.899

1083.9 643.9 1177.4 1046.9 1067.1 869.1 388.6 1178.8 1085.9

29.825 19.397 27.322 30.105 25.839 16.442 7.770 5.586 6.052

11.790 6.396 11.497 11.811 10.215 6.002 2.839 3.905 3.794

489.6 322.4 564.9 488.7 484.4 412.2 217.9 527.9 543.3

14.836 10.013 13.527 15.154 12.950 8.353 4.174 2.713 3.163

5.325 3.364 5.140 5.363 4.664 2.870 1.687 1.698 1.838

1084.0 612.8 1263.5 1076.4 1060.9 862.0 366.7 1214.2 1121.5

28.632 18.822 26.344 28.936 24.851 15.819 7.634 5.361 5.682

11.643 6.836 10.897 11.726 10.066 6.031 2.964 3.620 3.455

495.8 301.6 595.9 492.3 491.5 410.2 208.7 569.6 596.7

2.921 3.224 2.611 2.887 2.925 3.043 2.693 1.562 1.817 2.566 2.530 3.033 2.377 2.549 2.530 2.739 2.737 1.430 1.595 2.328

1.498 1.464 1.424 1.505 1.483 1.418 1.421 1.111 1.240 1.390 1.375 1.389 1.364 1.394 1.386 1.416 1.320 1.080 1.211 1.322

2.786 2.976 2.632 2.826 2.777 2.911 2.474 1.598 1.721 2.468 2.459 2.753 2.418 2.468 2.469 2.623 2.576 1.481 1.645 2.276

1.489 1.439 1.476 1.515 1.477 1.442 1.388 1.136 1.242 1.395 1.380 1.342 1.399 1.390 1.391 1.406 1.287 1.093 1.256 1.323

KOSMOS

Table C.16. MoE microbenchmark, MXFP8, backward pass, EP 8: per call the slowest rank, median. Baseline: RCCL A2A around Transformer Engine MXFP8 grouped GEMMs. Speedup is baseline over KOSMOS; vs KOSMOS serial is the same KOSMOS kernels as separate launches over the fused layer. Routing

Model

MI350X balanced DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean skewed DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean MI355X balanced DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean skewed DeepSeek-V3 Qwen3.5-397B DeepSeek-V4-Pro Kimi-K2.7-Code GLM-5.2 Qwen3-235B Qwen3-30B Llama 4 Scout Llama 4 Maverick geometric mean

TE MXFP8 ms KOSMOS ms KOSMOS TFLOPS Speedup vs KOSMOS serial 23.684 14.637 25.560 25.133 20.706 12.346 5.557 5.197 8.289

10.519 5.149 12.512 11.021 8.911 4.681 1.699 2.704 4.477

1097.5 800.8 1038.1 1047.6 1110.5 1056.9 728.0 1524.5 920.9

44.307 26.233 45.664 45.325 38.530 23.169 10.217 10.547 12.913

21.928 9.945 23.454 21.599 18.280 10.053 3.659 6.415 7.873

526.5 414.6 553.8 534.5 541.3 492.2 338.1 642.7 523.7

21.941 13.748 23.446 23.152 19.171 11.553 5.281 4.766 7.377

9.032 5.246 10.559 9.394 7.663 4.268 1.732 2.616 3.968

1278.2 786.0 1230.1 1228.9 1291.3 1159.3 714.0 1576.4 1039.0

41.350 24.885 42.728 42.354 36.189 21.759 9.708 9.539 11.684

19.153 9.154 20.454 18.686 15.886 9.105 3.730 5.921 6.847

602.8 450.4 635.0 617.8 622.9 543.4 331.6 696.3 602.2

2.251 2.843 2.043 2.281 2.324 2.637 3.270 1.921 1.851 2.343 2.021 2.638 1.947 2.098 2.108 2.305 2.792 1.644 1.640 2.101

1.430 1.631 1.316 1.476 1.452 1.531 1.508 1.317 1.290 1.435 1.387 1.691 1.355 1.472 1.434 1.520 1.442 1.225 1.172 1.403

2.429 2.621 2.220 2.465 2.502 2.707 3.048 1.822 1.859 2.378 2.159 2.718 2.089 2.267 2.278 2.390 2.603 1.611 1.706 2.173

1.586 1.599 1.431 1.628 1.607 1.681 1.488 1.311 1.308 1.510 1.548 1.860 1.469 1.645 1.607 1.654 1.412 1.253 1.250 1.510

Related documents

• Beaver: Elastic GPU Sharing between ML and Latency-Critical vRAN Workloads#1186836
• Beaver: Elastic GPU Sharing between ML and Latency-Critical vRAN Workloads#1181594
• Joint Effects of GPU Server Topology, Parallelism, and Congestion Control on MoE Inference: A Controlled Simulation Study#1142136
• Joint Effects of GPU Server Topology, Parallelism, and Congestion Control on MoE Inference: A Controlled Simulation Study#1122251
• Poseidon: DAG-Guided Parallelism Search for LLM Pre-Training on Heterogeneous Clusters#973693