GC-MoE: Overlapping Expert-Parallel Token Dispatch with FP8/FP4 Grouped GEMM via CUDA Green Contexts

GC-MoE: Overlapping Expert-Parallel Token Dispatch with FP8/FP4 Grouped GEMM via CUDA Green Contexts Yong Li, Kan Liu, and Tao Lan {brucelee.ly, liukan.lk, tao.lant}@alibaba-inc.com https://github.com/Bruce-Lee-LY/gc-moe

arXiv:2611.04222v1 [cs.DC] 3 Nov 2026

November 2, 2026 Abstract Mixture-of-experts (MoE) models are standard for large language model (LLM) serving, and expert parallelism (EP) makes all-to-all token dispatch, grouped GEMMs, and the combine step recurrent costs in every layer. In decode, each rank typically holds only 1–256 tokens, so dispatch latency and weight traffic dominate while tensor-core utilization is inherently low; because the grouped GEMM is bandwidth-bound, a decode layer does not need all SMs for arithmetic. In state-of-the-art fused operators such as DeepGEMM’s Mega MoE, dispatch and computation overlap only by sharing every SM: a few dispatch warps per SM compete with the TMA/MMA and epilogue warps for registers, shared memory, and issue slots, every phase is bound to one CTA’s combined resource budget, the synchronization domain spans the entire SM array, and every replay ends with a workspace-clean pass and a cross-rank tag-2 barrier. This paper presents GC-MoE, an open-source FP8/FP4 MoE inference operator for NVIDIA Hopper and Blackwell GPUs that uses CUDA Green Contexts to split the device at the driver level into a small communication partition and a large compute partition. The split gives each role its own kernel, resource budget, and synchronization domain, so dispatch and GEMM can be tuned independently. A three-node CUDA Graph (four nodes on Hopper with fused shared experts) launches the dispatch and grouped-GEMM kernels as concurrent roots on the two partitions; a token-ring handshake lets computation start on the first filled ring blocks, and double-buffered receive counters remove the trailing cross-rank barrier between replays. The compute partition runs operand-swapped tcgen05 block-scaled MMA on Blackwell and register-A WGMMA with in-register FP4 decoding on Hopper, with fused SwiGLU, requantization, and optional fused shared experts. On eight H20 GPUs GC-MoE accelerates DeepSeek-V4.1-Flash decode by 1.188× on average on both the routed and shared-expert suites (peaks 1.296× and 1.353×, respectively); on eight B300 GPUs it accelerates GLM-5.2 and DeepSeek-V4.1-Flash by 1.106× on average without shared experts and 1.123× with fused shared experts (peak 1.209×), outperforming DeepGEMM’s Mega MoE at all 54 measured points at matched numerical accuracy (the nine H20 shared-expert points compare against an unfused shared-expert baseline). GC-MoE reserves only 20 of 148 SMs on B300 for dispatch.

1

Introduction

Mixture-of-experts (MoE) architectures have become the dominant way to scale large language models without proportionally scaling inference cost (Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2022; Jiang et al., 2024; DeepSeek-AI, 2024a,b). In production serving, MoE layers are sharded with expert parallelism (EP): each GPU owns a subset of experts, and every layer routes each token to its top-k experts through an all-to-all token dispatch, computes the routed experts locally, and combines the partial outputs back to the source ranks. Decode dominates the serving cost of interactive applications, and in decode each rank typically processes only a handful to a few hundred tokens per layer. Under such small batch sizes, the MoE layer is not tensor-core bound: its cost is dominated by cross-rank token movement and by streaming quantized expert weights from HBM. Because of this cost structure, kernel-level MoE optimization has converged on fused, persistent operators. DeepGEMM’s Mega MoE is the state of the art on NVIDIA Blackwell (and, via a reference SM90 port, 1

on Hopper): a single fused kernel performs the expert-parallel dispatch, the FP8×FP4 grouped GEMMs with SwiGLU fused into the first epilogue, and the combine, with warp specialization inside one kernel and block-scaled tensor-core instructions (DeepSeek-AI, 2026b, 2024b). It already overlaps dispatch with computation, but only by sharing every SM (streaming multiprocessor): four dispatch warps per SM pull tokens into a per-block ring while two TMA-load warps (TMA: tensor-memory-accelerator copy engine), one MMA-issue warp (MMA: matrix multiply–accumulate), a scheduler warp, and eight epilogue warps consume ring blocks, run the GEMM, and perform the combine. The compute warps therefore never own the whole SM—the dispatch warps and their staging buffers are carved out of each SM’s register file and shared memory—and every replay ends with a workspace-clean pass followed by a cross-rank tag-2 barrier. The clean overlaps the combine, but the barrier is still a cross-rank round trip paid on every replay. The kernel’s synchronization domain is the whole device: its prologue and tail grid syncs include every SM of the device (148 on B300), so each phase waits for the slowest CTA (cooperative thread array), and the same CTA budget must host all phases, so a local change to one phase is a change to the whole kernel. These effects inflate decode latency at 1–256 tokens per rank. Two further trends sharpen the problem. First, low-precision formats have moved from FP8 activations to FP4 weights: weight traffic, not arithmetic, sets the decode lower bound, and the two GPU generations expose different instruction sets for mixed FP8×FP4 computation (tcgen05 block-scaled MMA on Blackwell, WGMMA with software FP4 decode on Hopper). Second, existing MoE systems optimize communication– computation overlap across micro-batches, layers, or experts (Li et al., 2023; He et al., 2022), or rely on stream-level concurrency, but they do not isolate SM resources inside the operator: a conventional kernel launched on the full device occupies all SMs, and concurrent kernels on separate streams compete for the same SM array with no hardware guarantee of co-residency or performance isolation. CUDA Green Contexts, introduced in CUDA 12.4, provide driver-level SM partitioning with independent workqueues, i.e., exactly the primitive needed to co-schedule a communication kernel and a compute kernel deterministically. Prior systems have used green contexts for isolation between inference processes, co-located models, or prefill/decode requests (Martín et al., 2026; Quach et al., 2025; Zhang et al., 2026); to our knowledge, none uses them to overlap the communication and computation phases within a single MoE operator. Deeper fusion alone does not address the decode bottleneck: because the grouped GEMM is bandwidth-bound, extra SMs buy little arithmetic, while a fused kernel couples the phases’ resource and synchronization budgets. This paper presents GC-MoE, an open-source MoE inference operator for Hopper (SM90) and Blackwell (SM103) GPUs that reorganizes the MoE layer as a three-node CUDA Graph (four nodes on Hopper with fused shared experts) across two CUDA Green Contexts. A small communication partition (20 of 148 SMs on B300) runs the dispatch kernel; a large compute partition runs the grouped-GEMM kernel concurrently. The two kernels are independent graph roots, so the driver schedules them simultaneously on disjoint SM sets. A token-ring protocol with generation-tagged counters lets the GEMM kernel begin consuming a block of dispatched tokens as soon as that block is filled, without waiting for dispatch to finish. A finalize kernel then aggregates the top-k and shared-expert partials into the output. Double-buffered receive counters and device-persistent synchronization state remove the trailing cross-rank barrier, so consecutive CUDA Graph replays need no cross-rank rendezvous between them. Because the grouped GEMM is bandwidth-bound at decode, ceding 20 SMs to dispatch costs little compute throughput; because each role is an ordinary kernel with its own thread, register, shared-memory, and synchronization budget, the dispatch and GEMM paths can also be tuned independently. The compute kernel implements FP8-activation×FP4-weight grouped GEMM for both architectures: an operand-swapped tcgen05 block-scaled MMA path on Blackwell that services FP4 routed and FP8 shared weights with one instruction sequence, and a register-A WGMMA path on Hopper with in-register FP4→E4M3 decoding and a window-based scale repacking. SwiGLU, routing-weight application, and FP8 requantization are fused into the first GEMM epilogue; shared experts can be fused into the same task stream. Our contributions are: • Decode-aware green-context decomposition. To our knowledge, GC-MoE is the first MoE operator that uses CUDA Green Contexts to give token dispatch and expert computation disjoint, driver-aligned SM partitions, so dispatch runs concurrently with the grouped GEMM on dedicated SMs instead of sharing every SM; balanced workqueues avoid implicit serialization between the two contexts. At decode the GEMM is bandwidth-bound, so the 20-SM dispatch partition costs little compute throughput, and each role is an ordinary kernel with its own resource budget and synchronization domain, which allows 2

the dispatch and GEMM paths to be optimized independently (Sections 3 and 4.2). • Fine-grained cross-partition synchronization. A token ring with generation-tagged full/empty counters and a 64-bit W1-completion bitmap carries the baseline’s block-granular handshake across the partition boundary, so the compute kernel starts on the first filled ring block; double-buffered receive counters eliminate the trailing cross-rank barrier between graph replays; all synchronization state survives replay without host intervention (Sections 4.3 and 4.6). • Dual-architecture FP8×FP4 grouped GEMM. Operand-swapped tcgen05 block-scaled MMA with K128 tile-prepacked FP4 weights on Blackwell, and register-A WGMMA with in-register FP4 decoding and exact window-based scale repacking on Hopper, both with fused SwiGLU, requantization, and optional fused shared experts (Sections 4.4 and 4.5). • Systematic evaluation. On 8×H20 and 8×B300, GC-MoE outperforms DeepGEMM’s Mega MoE at all 54 measured decode points: 1.188× mean on H20 (DeepSeek-V4.1-Flash, peaks 1.296× routed and 1.353× with shared experts), 1.106× mean on B300 without shared experts and 1.123× with fused shared experts (GLM-5.2 and DeepSeek-V4.1-Flash, peak 1.209×). Accuracy is matched: cosine ≥ 0.97 against BF16 everywhere and minimum cosine ≥ 0.9999997 against the baseline on both H20 suites; the nine H20 shared-expert points compare against an unfused shared-expert baseline (Section 6.3). The operator is open source.

2

Background

Notation. P denotes the number of ranks (GPUs), E the total number of routed experts, e = E/P the local experts per rank, H the hidden size, I the intermediate size, T the number of tokens per rank, and k the top-k routing factor. BM is the token block size of the grouped GEMM; BN = 128 and BK = 128 are the fixed MMA tile dimensions on the Blackwell path. FP8 refers to E4M3, FP4 to E2M1 packed two values per byte, and UE8M0 to the unsigned power-of-two exponent scale used by the microscaling recipes.

2.1

MoE layers under expert parallelism

A routed-expert MoE layer computes, for token t,   X (j) (j) yt = gt,j W2 SwiGLU W1 xt + W2s SwiGLU(W1s xt ), | {z } j∈TopK(t)

(1)

optional shared expert

(j)

(j)

where W1 ∈ R2I×H stacks the gate and up projections, W2 ∈ RH×I is the down projection, and gt,j is the routing weight. Under expert parallelism the E experts are distributed evenly, e = E/P per rank. Executing one layer therefore decomposes into four stages: (i) dispatch, an all-to-all exchange that moves each token, its activation scales, and routing metadata to the ranks that own its selected experts; (ii) Grouped GEMM 1 over the per-expert token groups, with SwiGLU in the epilogue; (iii) Grouped GEMM 2 producing the down projections; and (iv) combine, the weighted reduction of the k partial outputs back to the source ranks. On the Blackwell path the routed activations are FP8 with one UE8M0 scale per 32 values along the contraction dimension, the routed weights are FP4 at the same 32-value scale granularity, and the shared-expert weights are FP8; the Hopper path instead consumes E4M3FN activations with contiguous FP32 K128 scales (Section 2.3 and Table 2). In decode at small batch, each rank owns 1–256 tokens, so each local expert receives a small number of tokens and the per-expert GEMM has a tiny M dimension (often 16–32 rows after grouping). For GLM-5.2’s 256-expert, 8-rank, top-8 configuration, balanced routing gives T · k · P/E = T /4 tokens per expert; at one token per rank an expert receives 0.25 tokens on average, and the padded M tile is almost entirely padding. The grouped GEMM is then bandwidth-bound: its cost is dominated by streaming the local expert weights (FP4) and scale metadata from HBM rather than by arithmetic. At the same time, the dispatch stage is latency-bound: its critical path is the cross-rank exchange plus peer-atomic accounting, and it needs very few SMs to saturate NVLink. Because a bandwidth-bound GEMM gains little from extra SMs, the layer

3

does not need the whole device for arithmetic; yet a single fused kernel spans all SMs, and a device-wide synchronization domain is expensive when the layer itself is short. These observations—dispatch needs SMs but no tensor cores, decode GEMM needs tensor cores but is bandwidth-bound, and all-SM synchronization is costly at decode—are the premise of this work.

2.2

GPU execution model: Hopper, Blackwell, Green Contexts, and graphs

Hopper (SM90). On H20, kernels target sm_90a. The architecture provides warpgroup MMA (WGMMA), asynchronous TMA copies, and warp specialization; a warpgroup of four warps issues wgmma.mma_async instructions whose operands are read from shared memory or registers. There is no native FP4 MMA: FP4 weights must be decoded to E4M3 in software before entering the tensor core. Blackwell (SM103). On B300 (compute capability 10.3, an SM100-family part with 148 SMs), kernels target sm_103f. The fifth-generation tensor core is programmed through tcgen05 instructions: MMA accumulates into tensor memory (TMEM), a two-CTA cluster pair can issue a single MMA cooperatively (cta_group::2), and the block-scaled variant kind::mxf8f6f4.block_scale multiplies mixed FP8/FP4 operands per 32-value group using hardware scale factors. MMA completion is signaled through an mbarrier; TMEM is then read back with tcgen05.ld, so the epilogue can be software-pipelined against the next MMA without a device-wide barrier. CUDA Green Contexts. A green context is a lightweight execution context that owns a subset of the device’s streaming multiprocessors. The driver API exposes the device’s SM resource (cuDevice GetDevResource), splits it by count (cuDevSmResourceSplitByCount), and creates a green context over the resulting resource (cuGreenCtxCreate). Kernels launched into a green context’s stream are scheduled only on its SMs; two green contexts therefore run concurrently on disjoint SM sets with no MPS daemon and no process-level configuration. The split granularity and alignment are driver properties (minimum partition size and co-scheduling alignment), so the application requests a partition size and the driver determines the usable partition. Green contexts can additionally carry a workqueue configuration; the balanced sharing scope with a concurrency limit of one gives each context its own hardware workqueue, avoiding implicit serialization when two contexts’ kernels contend for one queue. CUDA Graphs. A CUDA Graph captures a kernel-node DAG once and replays it with negligible host overhead. Independent nodes without dependency edges are eligible to execute concurrently; node parameters are fixed at capture time and can be passed by value, which allows a graph to carry large descriptor tables (tensor maps, peer address tables) without per-launch pointer chasing. We use these properties to express the operator as a three-node child graph that is attached to an outer capture (Section 4.6).

2.3

Low-precision formats

On the Blackwell path the operator consumes activations quantized by the caller to FP8 E4M3 with one UE8M0 scale per 32 contiguous values along K; four exponent bytes are packed into one int32. Routed weights are FP4 (E2M1) with the same K32 microscaling granularity, packed two values per byte; the shared-expert weights are FP8. This is the (1, 1, 32) recipe: one scale per 1 × 32 value block for both operands. For shared experts we additionally support a (1, 128, 128) recipe whose weight scales cover 128 × 128 blocks; in that case the caller expands block-row scales to one packed row per output row before weight transformation. FP4 formats and their scaling conventions are shared with NVFP4/MXFP4-style training and inference work (NVIDIA, 2025; Tseng et al., 2025); on Hopper, activations use contiguous FP32 K128 scales and routed weights follow the MXFP4 convention (E2M1 with UE8M0 K32 scales). Weight transformation only rearranges already-quantized data and scales; it never quantizes floating-point weights.

4

3

Motivation: the costs of monolithic fusion at decode

Figure 1(a) shows how DeepGEMM’s Mega MoE executes a decode layer on B300. One persistent fused kernel occupies all 148 SMs. Warp specialization assigns each CTA four dispatch warps (128 threads), four non-epilogue warps (two TMA-load warps, one MMA-issue warp, and one scheduler warp; the last two run only in the leader CTA of each two-CTA cluster), and eight epilogue warps (256 threads) that also perform the combine; the kernel opens with routing accounting and a tag-1 cross-rank barrier. It then streams: the dispatch warps pull tokens, scales, and routing metadata into a per-block token ring, and the GEMM consumes a ring block as soon as that block is full. Dispatch and computation therefore do overlap—but on the same SMs. The overlap is spatial sharing: the dispatch warps and their staging buffers are carved out of the register file and shared memory of every SM, the pull is issued by only four warps per SM, and the kernel closes with workspace cleaning and a tag-2 cross-rank barrier that every rank must pass before the next replay can reuse the workspace. The clean is explicitly overlapped with the combine reduction, but the tag-2 barrier follows the clean and is a cross-rank round trip. The critical path is
shared shared Tmega = tpro + max Tpull , Tgemm + max(Tcombine , tclean + ttag2 ) , (2) where tpro is the prologue, the two middle terms execute concurrently on one shared SM array, and the tail is the longer of the final combine and the workspace clean plus the tag-2 barrier. Three costs follow. First, the compute warps never own the whole SM, so per-SM GEMM throughput is below what a dedicated SM delivers. Second, the synchronization domain is the entire device: the prologue and tail grid syncs include all 148 SMs, so each phase waits for the slowest CTA, and the tag-2 barrier adds a cross-rank round trip on every replay that the combine hides only when the combine itself outlasts the clean plus the barrier. Third, fusion couples the resource budgets: all roles share one CTA’s 64 512 registers and 227 KB of shared memory, and even an idle role keeps its share reserved, so a deeper weight pipeline, a larger epilogue, or a cheaper dispatch path must fit the combined budget and preserve every other role’s barriers and buffers; a local optimization is therefore paid for by the whole kernel. GC-MoE removes all three. Role1 runs the dispatch on a 20-SM communication context that owns the whole SM budget for the pull, role2 runs the grouped GEMM on a 128-SM compute context that carries no dispatch warps, and role3’s finalize grid is smaller still, so the grid syncs around the cross-rank tag-1 barrier span only the 20 dispatch SMs and role3’s grid sync spans only its finalize grid, rather than all 148 SMs. Parity-double-buffered receive counters let role3 flip generations instead of waiting on a tag-2 barrier. Because the roles are now separate kernels, each has its own thread, register, shared-memory, and synchronization budget: the dispatch path can be tuned for latency and the GEMM path for streaming and MMA throughput without rebalancing a shared CTA, and each role can adopt phase-specific optimizations independently. GC-MoE additionally improves the GEMM itself by prepacking FP4 weights into K128 tiles (Section 4.4). In the measured H20 configuration, GC-MoE cuts end-to-end latency by 20.9–22.8% (1.26–1.30× speedup) at 4–32 tokens per rank, the region where the grouped GEMM is still issue-bound at small M and the shared-SM penalty and the tag-2 barrier weigh most; at 256 tokens per rank weight streaming dominates the layer, both operators approach their streaming limit, and the relative gain shrinks to the roughly fixed shared-SM and tail savings (Section 6.2). Three straw-man alternatives do not solve the problem. (i) Multiple CUDA streams without isolation: a conventional kernel launched over the whole device fills every SM; a concurrently launched communication kernel must wait or steal SMs from the compute kernel, with no co-residency guarantee and no isolation from occupancy interference. (ii) One fused kernel with warp-specialized roles (the baseline): dispatch and compute can overlap, but only by sharing every SM, because a kernel cannot repartition SMs between roles without draining; the dispatch warps and their staging buffers are carved out of each SM’s budget, the pull is issued by a handful of warps, and a trailing cross-rank barrier is still needed between invocations to protect workspace reuse. Every local change is also constrained by the whole kernel: the CTA’s register and shared-memory budgets are already committed, and all roles’ barriers and buffers must remain valid, so optimizing one phase in isolation is difficult. (iii) Software-stage overlap across micro-batches or layers: this hides communication with other work but does not reduce a single layer’s critical path, and it changes numerical scheduling of the model. The design principles of GC-MoE follow directly. First, partition the SM array statically but natively, using the driver’s green-context SM split, so that dispatch and compute kernels are co-resident on disjoint SMs 5

(a) DeepGEMM Mega MoE (baseline): one fused kernel, warp roles share every SM all 148 SMs (shared)

counts + tag-1

dispatch warps (4 per SM): pull TMA/MMA + epilogue warps: GEMM, scatter

clean

tag-2

epilogue: combine

first ring block

dispatch warps and their staging buffers are carved out of every SM’s register file and shared memory t shared , T shared ) + max(T Tmega = tpro + max(Tpull combine , tclean + ttag2 ) gemm

the clean overlaps the combine; the tag-2 barrier follows the clean

(b) GC-MoE: disjoint SM sets, dispatch off the compute SMs, no tag-2 barrier 20 SMs comm

counts + tag-1

role1: pull tokens, scales, and metadata first ring block

128 SMs compute

replay N +1

role3: combine

role2: grouped GEMM, no dispatch warps consumes ring blocks as they land

saved (illustrative) t

ded , T Tgc = tpro + max(Tpull gemm ) + Tcombine

no tag-2 barrier: parity counters and role3 bookkeeping replace the cross-rank round trip

Figure 1: Execution timelines. (a) On B300 the baseline is DeepGEMM’s Mega MoE, one persistent fused kernel on all 148 SMs: after routing counts and a tag-1 cross-rank barrier, the dispatch warps stream tokens, scales, and routing metadata into a per-block token ring while the GEMM consumes each ring block as soon as it is full. The phases overlap, but on the same SMs, with the dispatch warps and their staging buffers carved out of every SM’s register file and shared memory; the tail is max(Tcombine , tclean + ttag2 )—the workspace clean overlaps the combine, but the tag-2 barrier is still a cross-rank round trip on every replay. (b) GC-MoE runs role1 (dispatch) on the 20-SM communication context and role2 (grouped GEMM) plus role3 (combine) on the 128-SM compute context. Role2 starts when the first ring block is filled (dashed trigger), and parity-double-buffered receive counters remove the tag-2 barrier, so the next replay can start as soon as role3 completes (green band). Each role is an ordinary kernel with its own register, shared-memory, and synchronization budget, so the dispatch and GEMM paths can be optimized independently. Bars are illustrative; measured end-to-end speedups span 1.03–1.30× across the routed-expert sweeps and 1.01–1.35× including shared experts (Section 6.2). Partitions are the B300 values. with hardware-level isolation. Second, realize every role as an ordinary kernel with its own thread, register, shared-memory, and synchronization budget, so that dispatch tuning and GEMM tuning are independent; the fused baseline instead charges every local change against one CTA’s combined budget. Third, carry the baseline’s block-granular token-ring handshake across the partition boundary, so the compute kernel starts on the first filled ring block rather than at the end of dispatch. Fourth, keep all synchronization state on device and persistent across graph replays, and double-buffer cross-rank counters so that no trailing barrier is needed between replays. Fifth, specialize the compute kernel per architecture, because the instruction sets for mixed FP8×FP4 differ fundamentally between Hopper and Blackwell.

4

GC-MoE design

4.1

Overview

GC-MoE executes one MoE layer as one replay of a child CUDA Graph with three kernel nodes, called roles (Table 1 and Figure 2). Role1 (dispatch) runs on the communication green context and performs routing accounting and the all-to-all pull into a local token ring. Role2 (compute) runs on the compute green context and executes the routed and shared grouped GEMMs with fused SwiGLU and requantization, scattering BF16 partials directly to the source ranks. Role3 (finalize) runs on the compute green context and combines the k (plus one shared) partials of Eq. (1) into the output. Role1 and role2 are concurrent roots of the graph—there

6

Table 1: Roles of the SM103 CUDA Graph (three nodes; on SM90, fused shared experts add a fourth node). dispatch_sms/compute_sms are the driver-aligned partition sizes. finalize_sms is min(compute_sms, 64) at up to 128 tokens per rank and compute_sms above that; with fused shared experts it is additionally tightened by the combine geometry to max(min(compute_sms, e + 1), ⌈T · combine chunks/8⌉). Role

Function

Green context

Grid

Threads

Cluster

role1 role2 role3

token dispatch (all-to-all pull) FP8×FP4 grouped GEMM + SwiGLU signal reset + combine

communication compute compute

dispatch_sms compute_sms finalize_sms

256 or 384 384 384

– (2,1,1) –

Algorithm 1 GC-MoE layer execution (one CUDA Graph replay). Require: FP8 activations x with packed UE8M0 scales, routing (idx, g), transformed FP4 routed weights W1 , W2 , symmetric buffer Ensure: BF16 output y Role1 — dispatch kernel, communication partition (concurrent graph root) 1: count local tokens per expert; publish per-expert counts to peers with packed atomics 2: NVLink barrier (tag 1); pull remote tokens, scales, and routing metadata into the token ring via TMA staging 3: release l1_full_count for each filled ring block (generation-tagged) Role2 — compute kernel, compute partition (concurrent graph root) 4: scheduler warp broadcasts W1/W2/shared tasks to both CTAs of each cluster 5: for each ring block b in generation order do 6: wait l1_full_count[b] (acquire) ▷ starts as soon as block b is filled 7: u ← FP8×FP4 GEMM W1 ; fuse SwiGLU, routing weight, and FP8 requantization 8: set l2_full_mask bits of finished W1 tiles; v ← GEMM W2 ; scatter BF16 partials to source ranks 9: end for Role3 — finalize kernel, compute partition (after role1 and role2) 10: publish combine readiness to all peers; consume double-buffered receive counts; reset ring signals 11: wait peer readiness; gather, reduce, and write y via TMA

is no dependency edge between them (Figure 2 draws them side by side)—so the driver schedules them simultaneously on disjoint SM sets. Role3 depends on both. Each role is an ordinary kernel with its own launch bounds, register and shared-memory footprint, and synchronization domain; no CTA budget is shared across roles. This is what lets role1 be tuned for communication latency and role2 for weight streaming and MMA throughput, and it is why a change to one path does not perturb the others. Algorithm 1 summarizes one replay. On the SM90 path, the same structure holds with sm_90a kernels and a default dispatch request of 4 SMs; when shared experts are fused, they execute as a fourth kernel node (role4) that role2 may wait on (Section 4.5). All cross-rank data—inputs, routing tables, the token ring, the combine buffer, and all signals—live in one symmetric-memory allocation created with PyTorch’s distributed symmetric-memory API and rendezvoused across ranks. Every rank obtains the peer virtual addresses of every other rank’s allocation and passes them to the kernels; a 592-byte by-value descriptor carries the base address and 72 relative offsets, so a kernel maps any peer’s base with a signed delta. Kernels publish and consume per-expert counters, ring counters, and combine readiness through these peer addresses over NVLink.

4.2

Green-context SM partitioning

The runtime creates a green-context pair once per CUDA context and caches it (Figure 2). Creation proceeds as follows. (1) Query the device’s SM resource, which reports the SM count, the minimum partition size, and the co-scheduling alignment. (2) Compute the compute-target SM count as ⌊(total − requested)/alignment⌋ · alignment, where requested is the application’s dispatch request (GC_MOE_DISPATCH_SMS, default 16), and

7

split the SM resource with cuDevSmResourceSplitByCount; the driver returns the compute resource and the complementary communication resource. The partition is thus driver-aligned, not application-chosen: on B300 (148 SMs) the default request yields a 20/128 communication/compute split. (3) If balanced workqueues are enabled (default), obtain a workqueue-config resource from the primary context with sharingScope=CU_WORKQUEUE_SCOPE_GREEN_CTX_BALANCED and wqConcurrencyLimit=1, and generate the context descriptor from both resources; after creation, read the resource back and verify that the balanced limit took effect. Balanced workqueues matter in practice: without them, two contexts can be placed on one hardware workqueue, which silently serializes the very kernels the design intends to overlap. (4) Create a non-blocking green stream per context and a retirement event; teardown records the event, synchronizes, and only then destroys the context. The pair is cached thread-locally per CUDA context. The partition is validated at runtime: at least 32 compute SMs, an even compute count, at least two dispatch SMs, and communication plus compute equal to the device SM count. Two consequences deserve emphasis. First, the partition is static for the lifetime of a graph, which makes replay deterministic: no kernel is ever migrated, and the dispatch partition’s SM count is a compile-time property of the launched kernels. Second, the cost of the partition is small in the regime that matters: the split assigns 20 of 148 SMs (13.5%) to dispatch and 128 to compute, but at decode batch sizes the grouped GEMM is bound by weight streaming and small-M issue latency rather than tensor-core throughput, so 128 SMs suffice; the measured effective throughput confirms that the smaller compute partition still beats the baseline’s shared-SM GEMM (Section 6.2). The partition also shrinks the synchronization domain: role1’s grid syncs, including those around the tag-1 barrier, span only the 20 dispatch SMs; role3’s grid sync spans its finalize grid (64 SMs at up to 128 tokens per rank); and the compute partition’s task scheduler is local to its own SMs, whereas the baseline’s prologue and tail grid syncs span all 148.

4.3

Overlapped dispatch and expert computation

Token ring. Dispatch does not allocate per-expert buffers. All pulled tokens land in one globally reused token ring of ring blocks of BM tokens; each ring block carries four persistent signals: l1_full_count (published by dispatch when the block is filled), l1_empty_count (released by compute when the block’s activations are consumed), l2_full_mask (a 64-bit bitmap of finished W1 output tiles), and l2_empty_count (released when the block’s L2 activations are consumed). Ring capacity is derived from a maximum-in-flight model: with per-rank capacity Tmax , the maximum routed assignments are Tmax · P · min(k, e); the pool of logical blocks is ⌈assignments/BM ⌉ + e, and the number of simultaneously live blocks is recomputed from the actual compute-partition SM count and L1/L2 cluster waves. The capacity check uses the real partition size, so a graph never starts with an undersized ring. Because the ring is reused across generations of blocks, each logical pool block maps to a ring block plus a generation number; dispatch must observe that the ring block’s consumers have released the previous generation (l1_empty_count) before overwriting it. Generation tags make block reuse safe across generations within a replay; between replays, role3 resets the ring signals in-kernel, overlapped with the combine half (Section 4.6), so no separate reset pass is needed. Producer–consumer handshake. The compute kernel’s TMA-producer warp waits with acquire semantics on l1_full_count = BM · (generation + 1) for the target ring block before publishing the block’s W1 tasks. The handshake granularity is thus one ring block: the GEMM starts on the first filled block, not at the end of dispatch. Within a block, the L2 stage does not need a separate barrier either: each W2 K128 tile consumes two W1 N128 output tiles, and the W1 epilogue sets the corresponding bits in the 64-bit l2_full_mask; the L2 producer waits for exactly its two bits. The bitmap keeps one bit per W1 N128 tile, and each W2 K128 tile waits on exactly its two bits; this is also why 2I/128 ≤ 64 is a hard design constraint (Appendix A). Scheduler. Each cluster’s leader CTA runs a scheduler warp. Tasks are phases over blocks: Linear1, Linear2, SharedLinear1, SharedLinear2. The scheduler pre-publishes a small number of L1 waves (computed from the same ring model) before the first L2 wave, then maintains a one-in-one-out discipline: every successful L2 claim publishes one more L1 task, and an L2 task is published only after all L1 tasks of its pool block have been published. Shared-expert tasks are injected before routed L2 work when small (interleaved with routed tasks) and queued after routed work when large. Task records are 32-byte descriptors broadcast to the two CTAs of a cluster through distributed shared memory with transaction-count mbarriers and double buffering. 8

child CUDA Graph (captured once, replayed once per layer) role1 · dispatch

role2 · grouped GEMM ∥

counts · all-to-all pull fills the token ring

FP8×FP4 MMA · SwiGLU requant · scatter partials

concurrent roots

role3 · finalize / combine reduces k+1 partials after peers resets ring signals · flips parity

green context A

green context B

communication partition

compute partition

20 of 148 SMs own balanced workqueue

128 of 148 SMs own balanced workqueue

symmetric-memory workspace (one allocation, peer-mapped) activations scales

routing metadata

token ring role1→role2

combine buf. role2→role3

signals counters · parity

NVLink

peer ranks P − 1 via NVLink

device-persistent, replay-aware state · no host reset between replays

graph dependency

launched into a green context

symmetric-memory access

Figure 2: Architecture of one GC-MoE layer replay. The child CUDA Graph holds three nodes: role1 (dispatch) and role2 (grouped GEMM) have no dependency edge and are concurrent roots, whereas role3 (finalize/combine) depends on both (solid arrows). Dashed arrows map each role to the green context that launches it—role1 to the 20-SM communication context, role2 and role3 to the 128-SM compute context—each with its own balanced workqueue so that the driver co-schedules them on disjoint SM sets. All cross-rank state lives in one symmetric-memory workspace whose boxes name the buffers exchanged over NVLink (double arrows): activations and scales, routing metadata, the role1→role2 token ring, the role2→role3 combine buffer, and the counters and parity signals. Every synchronization object is device-persistent and replay-aware, so a replay needs no host-side reset. Dispatch data path. Role1 proceeds in six steps. (1) File-level token counting: each warp processes 32/k tokens per iteration and accumulates per-expert counts in shared memory. (2) A packed atomic add of (1 ≪ 32) | count to the local per-expert send counter; the high word counts the number of (rank, dispatch-SM) contributors for arrival validation, the low word the token count. (3) A second scan writes each token’s routing entry to the target rank’s source-index region (peer write). (4) After a grid-wide sync, block 0 publishes the per-expert counts and prefix sums to every peer’s receive counters. (5) A two-phase NVLink barrier (tag 1) ensures all peers’ counts are visible before pulls begin; the barrier is a persistent, phase-encoded counter that survives graph replays, and it carries a configurable timeout (default 3600 s) after which the device traps instead of hanging. (6) Each warp walks the expert-concatenated, unpadded token sequence with a global cursor: lane-held per-rank count registers plus a warp reduction locate each token’s (src_rank, src_slot), and warp-private double-buffered shared-memory staging performs chunked bulk copies of the FP8 activation from the peer and the corresponding packed-scale rows, permuted to the MN-major row order the MMA path expects. Routing weight and (src_rank, src_token, slot) metadata are written alongside; the block’s full count is released when the last token lands. No trailing barrier. Early designs ended each invocation with a second NVLink barrier—DeepGEMM’s Mega MoE still ends its kernel this way, after workspace cleaning—so that the next replay could safely reuse 9

the receive counters and workspace. GC-MoE instead makes the per-expert receive counters double-buffered by parity: a replay consumes generation N mod 2, and role3 flips the parity after consuming the counts and completing its grid sync, so the next publish targets the other generation. Ordering is guaranteed by the tag-1 barrier and a release/acquire chain on combine readiness, not by a cross-rank barrier. This removes a full NVLink round trip from the inter-replay critical path. Direct partial scatter. The W2 epilogue does not write intermediate results to a local buffer for a later combine kernel. Each lane issues 16-byte stores of BF16 fragments directly to the source rank’s combine buffer at the slot (topk_slot, src_token); buffer recycling is governed by the l1_empty/l2_empty handshake. The combine kernel therefore only has to reduce, not gather from a local staging area.

4.4

FP8×FP4 grouped GEMM

Blackwell path (SM103). On B300 the operator uses the 2-CTA block-scaled MMA tcgen05.mma. cta_group::2.kind::mxf8f6f4.block_scale: two CTAs of a cluster jointly issue one MMA whose FP32 accumulator lives in TMEM and is read back with tcgen05.ld after the MMA-completion mbarrier. As in the baseline, the operands are swapped: the instruction’s A operand is the weight tile and B the activation tile, so the MMA’s N dimension is the token dimension. The instruction N is set to the tile’s valid token count rounded up to 16 rather than the padded block size, so partially filled blocks cost proportionally less. Because A carries the varying dtype, the same instruction sequence serves FP4 routed weights (E2M1, format 5) and FP8 shared weights (E4M3, format 0) with a runtime-selected format id; the CuTe DSL’s public block-scaled MMA requires matching A/B dtypes, so the instruction is emitted as an inline-PTX user op. The scale factors are moved with the dedicated UTCCP path: per K128 stage, the activation scales (BM -aligned rows, four packed K32 groups per 16 bytes) are copied to TMEM with tcgen05.cp.cta_group::2.32x128b.warpx4, 512 bytes per stage; weight scales take the same path. Each K128 stage is then issued as four K32 MMAs whose scale-factor byte selector picks the corresponding K32 group; for the (1, 128, 128) shared recipe, weight scales are staged once per four K128 blocks and the selector is driven by the K128 block index modulo 4. The Blackwell-specific contribution is therefore not the instruction pattern—operand swap, valid-M N sizing, and the UTCCP scale path follow the baseline—but the scheduling around it: K128 tile-prepacked weights, the shared-expert task stream, and the partition-level overlap. Tile geometry is fixed at BN = BK = 128, with each CTA of the pair responsible for half of the BM rows. The TMEM allocation is 2BM accumulator columns plus ⌈BM /128⌉-derived scale columns, rounded up to a power of two and capped at 512 columns. The 64-bit W1-completion bitmap caps the W1 output-tile count at 64, i.e., 2I/128 ≤ 64; the even-H/128 requirement follows from pairing N128 tiles per cluster. Weight prepacking. The weight transform rewrites each expert’s quantized FP4 matrix from rowmajor [E, rows, K/2] into K128-tile-contiguous [E, kb , rows, 64]: within a tile, all rows’ 64-byte segments are contiguous. A TMA descriptor encodes the 4-D box (128, 128, 1, 1) over (K, rows, kb , experts), so one TMA operation bursts a single 8 KB contiguous global range (16 KB for FP8 shared weights) instead of 128 segments spaced a full weight row apart; the FP4 payload is unpacked from 16U4-aligned storage into 8-bit shared memory (SMEM) by the tensor map. Gate and up rows are interleaved in 8-row groups so the SwiGLU epilogue reads pairs from one tile, and scale rows are permuted to the UTCCP consumption order; both follow the baseline’s weight transform. The K128 tile packing is the addition: it turns each weight tile load into one contiguous burst instead of 128 segments spaced a full weight row apart. The transform runs asynchronously on the caller’s stream once per weight set and is reused across token counts, capacities, and graph captures. Routed FP4 weight loads keep the baseline’s EVICT_FIRST cache hint: decode streams the weights once per layer, and polluting L2 with them hurts the activation/scales working set. W1 epilogue. Accumulators are read from TMEM in FP32, rounded to BF16, clamped, activated with the sigmoid term of SwiGLU (optionally using a polynomial exp and reciprocal for fast math), multiplied by the up projection and by the token’s routing weight (applied in the W1 stage so the L2 GEMM needs no per-token scaling), reduced to a warp-group amax, converted to a UE8M0 scale and an E4M3 payload, and stored through SMEM with a transposing stmatrix to the L2 activation ring via TMA; the requantized scale rows are written directly in the permuted order that the next stage’s UTCCP expects. At BM = 240 the 10

Table 2: SM103 (Blackwell) and SM90 (Hopper) implementation paths. Blackwell (SM103, B300)

Hopper (SM90, H20)

Compile target MMA Weight format

sm_103f tcgen05 mxf8f6f4.block_scale, 2-CTA packed FP4 with UE8M0 K32 scales

Activation format FP4 handling Scale application

E4M3 with packed UE8M0 K32 scales native E2M1 operand (format 5) UTCCP to TMEM; hardware block scale K128 tile-contiguous prepack GC_MOE_DISPATCH_SMS=16 (20 aligned) 3

sm_90a register-A WGMMA m64nNk32, E4M3 MXFP4, window-repacked with per-expert factor E4M3FN with FP32 K128 scales in-register decode FP4→E4M3 per-K32 FMA outside MMA; ×64 headroom windowed tile repack, lane-permuted GC_MOE_SM90_DISPATCH_SMS=4 3 (4 with fused shared experts)

Weight layout Dispatch request Graph nodes

epilogue double-buffers registers so the TMEM read of the next store overlaps the arithmetic of the current one. Block-size heuristic. The host√picks BM from the expected tokens per expert m̂ = T · P · k/E and its standard deviation, covering m̂ + m̂ with the number of blocks at the largest candidate; the candidate set is {16, 32, 64, 96, 128, 192, 240}, small batches pin BM = 16–32, and at least two blocks are used when the per-rank expert count is small. The chosen value balances MMA granularity against the number of live ring blocks. Hopper path (SM90). The SM90 operator is an independent implementation. Activations are E4M3FN with contiguous FP32 K128 scales; routed weights are MXFP4 (E2M1 with UE8M0 K32 scales); shared weights are FP8. Two mechanisms replace the missing hardware blockscale path. First, the host transform repacks MXFP4 weights into WGMMA-friendly tiles using a window-based rescaling: for each expert it computes a window base from the scale extremes (at most 11 exponent steps) and shifts FP4 magnitudes that fall outside the window down by exact powers of two using a lookup table over the 4-bit code space, recording the per-expert compensating factor; the tile bytes are simultaneously permuted to the lane semantics the register decoder expects. Second, the kernel decodes FP4 to E4M3 in registers with two LDS.128 loads plus PRMT/LOP3 bit manipulation, then issues register-A WGMMA m64nNk32 with E4M3 operands; block scales are applied as per-K32 fused multiply-adds outside the MMA, and a global ×64 convention on stored scales moves E4M3 values into the WGMMA e4m3 dynamic range. Tile sizes and warp specialization are chosen by a token-bound heuristic (BM ∈ {8, 16, 32, 64}, 384-thread pipelined or 256-thread non-pipelined variant), and dispatch uses epoch/counter headers with a peer table; the default dispatch request on Hopper is 4 SMs. Table 2 summarizes the two paths.

4.5

Fused shared experts

Shared experts are optional and fused into the same task stream rather than run as separate layers. The shared-expert weights are FP8 in K128 packed tiles. Three details make the fusion cheap. First, the shared W1 GEMM reads the original activation buffer by aliasing it (zero copy): shared experts see every local token, so no dispatch is needed. Second, shared activation scales must be presented in the same MN-major permuted layout that the routed path’s UTCCP consumes, although they are always K32 regardless of the shared weight recipe. The caller stages them once per token count into shared_l1_acts_sf: for token t with u = t mod BM , the destination row is row(t) = ⌊t/BM ⌋ · align(BM , 128) + ⌊u/128⌋ · 128 + (u mod 32) · 4 + ⌊(u mod 128)/32⌋ .

(3)

The row order of Eq. (3) is exactly the one the kernel itself writes for routed scales, so shared and routed experts share one TMA scale map and one UTCCP path; the buffer holds ⌈Tmax /8⌉ · 128 rows, an upper bound valid for every candidate BM . The staging can be fused into an upstream producer (e.g., the quantizer). 11

Third, shared tasks enter the unified scheduler as SharedLinear1/SharedLinear2 phases and are placed before or after routed work depending on size, and the shared output occupies lane k of the combine buffer, so the finalize kernel reduces it with no special case. On SM90, shared experts run as a fourth graph node (role4) that role2 optionally waits on.

4.6

Three-node CUDA Graph combine

Role3 is a single kernel with two halves that run concurrently. The first four warps perform the bookkeeping half: publish combine readiness to every peer with a release system atomic (required even for ranks with no routed tokens), consume this generation’s per-expert receive counts with an arrival check against dispatch_sms · P , flip recv_parity, and reset the ring signals, counters, and the W1 bitmap for the next replay, accumulating optional per-expert receive statistics along the way. Because role3 depends on both graph roots, this half starts after role2 completes, but it does not add a separate serial phase: it runs alongside the combine half, and its resets proceed while the combine warps wait for peer readiness and reduce. The remaining warps perform the combine half: wait for each peer’s readiness with a bounded timeout, then, for each token chunk assigned to the partition, prefetch and accumulate the valid top-k (and shared) slots from the possibly remote combine buffer with double-buffered bulk copies, accumulate BF16 values into FP32 registers, pack pairs back to BF16, and write y through SMEM and a bulk TMA store. The chunk geometry (1, 2, or 4 chunks per token) is chosen from the reusable SMEM budget with 16-byte per-lane alignment. The three-node formulation is deliberate. Role1 and role2 have no dependency edge, which is what makes the driver run them concurrently. Role3 depends on both and thereby absorbs the global completion condition—all W2 partials scattered—through release/acquire readiness signals rather than a second crossrank barrier. And because every synchronization object is device-persistent and replay-aware (barrier phases, ring generations, receive parity, combine readiness), a graph replay needs no host-side reset. The child graph is attached to the caller’s capture with capture-dependency semantics, or launched directly as a graph outside capture; it is cached by cubin hash, parameter signature, and resource state, so repeated layers and repeated token counts replay the same graph.

5

Implementation

GC-MoE is written in Python against NVIDIA’s CuTe DSL (CUTLASS DSL 4.8.1) and JIT-compiled per configuration to PTX, which the system assembler (ptxas 13.2) turns into cubin. The operator is a Python package; installing it builds a wheel but compiles no kernels—compilation happens on first use of a configuration. Compiler identity and caching. A compiler-identity record binds the ptxas path, size, version, and hash; the DSL frontend hash; a content hash of the operator’s own Python source tree; native-library identity; and the effective options. Options are locked to --gpu-arch <target> --opt-level 3 with a whitelist for ptxas flags (we use -O3, register-usage level 10, and fused multiply-add); relocatable device code and user overrides are rejected. Two caches sit above the assembler: a binary cache keyed by compiler identity plus kernel attributes, resources, and tensor metadata (never CUDA addresses), and a capture-key cache that additionally keys tensor-map encodings and loaded kernels by context, green partition, and all pointer/shape/timeout arguments. Cubins are loaded through the driver’s library API and reused by hash. By-value kernel parameters. CUDA Graph nodes hold kernel parameters by value, which is essential here: the compute kernel needs 18 tensor maps (TMA descriptors for activations, scales, and weights) and the peer address table. CuTe emits both as pointer formals. A post-generation PTX rewrite converts the peer descriptor into a 592-byte by-value .b8 formal (base, 72 int64 offsets, rank index) and each tensor map into a 128-byte aligned by-value formal, rewriting every param-space address in prefetch.tensormap and cp.async.bulk.tensor instructions through cvta.param.u64; the rewrite fails closed if the generated PTX shape changes, and the parameter order is normalized for the graph ABI. The SM90 path shares the JIT utilities and validates its ABI (14 pointers plus 3 scalars, and the presence of wgmma.mma_async and cp.async.bulk).

12

Distributed interface. Initialization requires one process per GPU, peer-accessible GPUs, and a torch distributed group. The buffer API allocates and rendezvouses the symmetric memory; the weight-transform API returns transformed weight pairs once per weight set and may be called outside graph capture, asynchronously on the caller’s stream. The operator call writes the output in place. Listing 1 shows the routed-expert invocation; fused shared experts extend it with shared weight pairs and the staged scale buffer. Listing 1: Minimal GC-MoE invocation (Blackwell path, routed experts). buffer = gc_moe.get_symm_buffer_for_gc_moe( group, num_experts, capacity, topk, hidden, intermediate, mma_type=’fp8xfp4’) w1, w2 = gc_moe.transform_weights_for_gc_moe(l1_weights, l2_weights) buffer.x[:tokens].copy_(x_fp8) buffer.x_sf[:tokens].copy_(x_sf) buffer.topk_idx[:tokens].copy_(topk_idx) buffer.topk_weights[:tokens].copy_(topk_weights) y = torch.empty((tokens, hidden), dtype=torch.bfloat16, device=x_fp8.device) gc_moe.fp8_fp4_gc_moe(y, w1, w2, buffer, recipe=(1, 1, 32), activation=’swiglu’, activation_alpha=1.0, activation_clamp=10.0)

Size and constraints. The device code is 1935 lines for the Blackwell (SM103) roles and 1305 lines for the Hopper (SM90) roles, plus shared scheduler, barrier, TMA, and PTX helpers. The operator’s shape constraints and their technical reasons are collected in Appendix A; the most important are I ≤ 4096 with I a multiple of 128 (64-bit W1 bitmap and K128 tiles), H a multiple of 256 (2-CTA N-tile pairing), k + 1[shared] ≤ 32 (one warp lane per combine slot), and at most 2048 experts globally (static signal sizing).

6

Evaluation

6.1

Setup

Environment. Table 3 lists hardware and software. We measure two generations: eight NVIDIA H20 GPUs (Hopper, sm_90a) and eight NVIDIA B300 GPUs (Blackwell, compute capability 10.3, sm_103f). The baseline is DeepGEMM’s Mega MoE (DeepSeek-AI, 2026b), measured with the deep-gemm package at 2.8.2+fcad603. On B300 we use the public fp8_fp4_mega_moe; when shared experts are enabled, the baseline is invoked in its fused shared-expert mode (num_shared_experts=1), so that comparison is feature-matched. The public release ships only the SM100 kernel; for H20 we built an SM90 reference port of the same design, sm90_fp8_fp4_mega_moe (one persistent fused kernel with warp-specialized dispatch, non-epilogue, and epilogue threads around the token ring, with NVLink barriers). It has no fused shared-expert path, so the H20 shared-expert baseline runs the routed Mega kernel plus a separate shared-expert kernel. Workloads. Table 4 lists the two production-scale MoE configurations. GLM-5.2 is evaluated on B300 and DeepSeek-V4.1-Flash on both hardware generations. Both use SwiGLU with α = 1.0 and clamp 10.0, the (1, 1, 32) routed recipe, fast-math sigmoid, and a per-rank token capacity of 256. Protocol. Both operators are invoked under CUDA Graph replay with identical inputs, capacities, weight transformations, and measurement settings. Each measurement window replays the graph 100 times; the per-invocation latency is the CUDA-event elapsed time divided by 100. We run 10 warmup windows and 30 measured windows, flushing L2 and synchronizing between windows. Each rank reports the windowmedian latency; the latency of a point is the maximum across the eight ranks (the slowest rank). A decode speedup is the per-point ratio latencybaseline /latencyGC ; the reported mean is the arithmetic mean of the per-point ratios (not the ratio of means). All runs are gated on idle GPUs (0% utilization, no compute processes, at least 51 200 MiB free, observed over at least three samples spanning ten seconds). We scan T ∈ {1, 2, 4, 8, 16, 32, 64, 128, 256} tokens per rank, i.e., 8–2048 tokens in total across the eight ranks.

13

Table 3: Evaluation environment. GPUs CUDA / assembler PyTorch Kernel DSL Baseline Proposed Execution

8× NVIDIA H20 (Hopper, sm_90a); 8× NVIDIA B300 (Blackwell, CC 10.3) CUDA 13.2, ptxas 13.2.78 2.11.0 (cu130) NVIDIA CUTLASS DSL (CuTe) 4.8.1 DeepGEMM 2.8.2+fcad603 (fp8_fp4_mega_moe on B300, fused shared experts when enabled; our SM90 reference port sm90_fp8_fp4_mega_moe on H20) GC-MoE 0.3.1 (commit 9bdb81f) one process per GPU; PyTorch symmetric memory over NVLink; CUDA Graph replay

Table 4: Model configurations. Routed recipe (1, 1, 32) for both models; shared-weight K group and shared recipe as shown. Model GLM-5.2 DeepSeek-V4.1-Flash

H

I

E

top-k

shared experts

shared K group

6144 5120

2048 2304

256 384

8 6

1 1

128 32

Table 5: Decode latency and speedup, DeepSeek-V4.1-Flash on 8×H20 (routed experts). Slowest-rank median per-invocation latency under CUDA Graph replay.

6.2

Tokens/rank

Total tokens

Mega (µs)

GC-MoE (µs)

Speedup

1 2 4 8 16 32 64 128 256

8 16 32 64 128 256 512 1024 2048

120.061 169.340 256.018 368.254 442.722 472.090 507.207 569.604 695.010

115.161 144.259 202.589 284.310 341.689 373.393 432.383 493.194 675.370

1.0425 1.1739 1.2637 1.2953 1.2957 1.2643 1.1731 1.1549 1.0291

Mean speedup (9 points)

–

1.1881

Decode latency and speedup

Hopper (8×H20). Table 5 and Figure 3 report DeepSeek-V4.1-Flash decode. GC-MoE is faster at all nine points, with a mean speedup of 1.188× and a peak of 1.296× at 16 tokens per rank. The shape of the curve follows the motivation: the speedup rises from 1.04× at 1 token per rank to the 1.26–1.30× plateau at 4–32 tokens per rank, where there is enough dispatch traffic for the dedicated partition to matter yet the grouped GEMM is still issue-bound at small M , so the baseline’s shared-SM carve-out and the max(Tcombine , tclean + ttag2 ) tail of Eq. (2) weigh on a short critical path; it then narrows to 1.03× at 256 tokens per rank, where weight streaming dominates and both operators approach their streaming limit, so the saving shrinks to the roughly fixed shared-SM and tail costs. Blackwell (8×B300). Table 6 and Figures 4–5 report decode for both models. GC-MoE wins all 18 points. The compute partition has 128 SMs versus the baseline’s 148, so the win is not from more arithmetic capacity: at decode the grouped GEMM is bandwidth-bound, and dedicated SMs plus the removed tail matter more than the 20 SMs ceded to dispatch. The mean speedup is 1.102× for GLM-5.2 and 1.109× for DeepSeek-V4.1-Flash (1.106× overall), and the per-point range is 1.085–1.189×. The B300 curves are flatter than the H20 curve (1.085–1.189× versus a 1.26–1.30× plateau). For DeepSeek-V4.1-Flash the absolute saving stays within 9–18 µs per layer across the sweep, so the relative gain is largest at 1 token per rank (1.189×: 11.2 µs on a 70 µs layer) and flattens as the GEMM grows (1.089× at 256). For GLM-5.2 the saving instead grows with batch (5.6 µs at 1 token/rank to 18.7 µs at 256), consistent with the dedicated-SM GEMM advantage scaling with the amount of grouped-GEMM work rather than with the fixed tail (1.095× 14

Figure 3: DeepSeek-V4.1-Flash routed-expert decode on 8×H20: latency and speedup of GC-MoE versus Mega MoE. to 1.128×). Per-rank medians are tightly clustered (at 1 token per rank for DeepSeek-V4.1-Flash, GC-MoE’s eight rank medians span 59.21–59.39 µs and the baseline’s 70.22–70.63 µs), so the speedup is not an artifact of one straggler rank.

6.3

Fused shared experts

Tables 7 and 8 and Figures 6, 7, and 8 report the shared-expert setting. On H20, fusing one shared expert is 1.031× faster on average than GC-MoE without fusion (6 of 9 points), and 1.188× faster than the unfused Mega baseline (9 of 9 points), peaking at 1.353× at 8 tokens per rank. Because the H20 reference has no fused shared-expert path, its baseline is the routed Mega kernel plus a separate shared-expert kernel; the fused-versus-unfused comparison is the within-operator ablation. At small batch the shared expert is almost free because its GEMM interleaves with routed work on the compute partition and its output joins the existing combine reduction; at 32–128 tokens per rank it adds serial compute and the ratio dips slightly below one (0.973–0.985×), which is the regime where the scheduler places shared tasks after routed work. On B300 the baseline runs its own fused shared-expert mode, so the comparison is between two feature-matched fused operators rather than a fusion ablation: GC-MoE is faster at all 18 points, with 1.131× mean for GLM-5.2 and 1.115× for DeepSeek-V4.1-Flash (1.123× overall) and a peak of 1.209× at 1 token per rank. Across the routed and shared-expert sweeps on both platforms, the evaluation covers 54 points (9 H20 routed, 9 H20 shared, 18 B300 routed, and 18 B300 shared), and GC-MoE is faster at every one.

6.4

Throughput and bandwidth analysis

Tables 9 and 10 report effective throughput. Two conventions must be stated precisely. Effective TFLOPS counts only matmul work (6×assignments×HI for routed experts), divided by the slowest-rank latency and

15

Table 6: Decode latency and speedup, 8×B300 (routed experts), both models. Slowest-rank median perinvocation latency. DSV4.1-Flash abbreviates DeepSeek-V4.1-Flash. Model

Tokens/rank

Total tokens

Mega (µs)

GC-MoE (µs)

Speedup

GLM-5.2

1 2 4 8 16 32 64 128 256

8 16 32 64 128 256 512 1024 2048

64.184 87.457 108.085 125.241 131.652 133.215 137.227 144.344 164.383

58.601 79.228 98.225 115.041 121.385 121.493 123.694 129.517 145.683

1.0953 1.1039 1.1004 1.0887 1.0846 1.0965 1.1094 1.1145 1.1284

DSV4.1-Flash

1 2 4 8 16 32 64 128 256

8 16 32 64 128 256 512 1024 2048

70.630 84.022 105.265 139.955 161.549 170.546 173.747 178.495 186.200

59.387 74.758 95.373 127.732 148.495 156.371 158.264 161.271 171.051

1.1893 1.1239 1.1037 1.0957 1.0879 1.0907 1.0978 1.1068 1.0886

Mean speedup: GLM-5.2 (9 points) Mean speedup: DeepSeek-V4.1-Flash (9 points) Mean speedup: overall (18 points)

1.1024 1.1094 1.1059

Table 7: Decode latency, fused shared experts, DeepSeek-V4.1-Flash on 8×H20. gc-se: GC-MoE with fusion; gc-unfused: GC-MoE with the shared expert run separately; mega-unfused: baseline. Ratios are against gc-se. Tokens/rank

gc-se (µs)

gc-unfused (µs)

unfused/se

mega-unfused (µs)

mega/se

1 2 4 8 16 32 64 128 256

145.171 179.286 236.353 302.667 363.640 428.722 472.348 572.091 779.019

156.866 188.566 248.009 328.262 384.421 422.456 465.376 556.513 786.387

1.0806 1.0518 1.0493 1.0846 1.0571 0.9854 0.9852 0.9728 1.0095

161.804 214.191 299.391 409.599 482.938 511.952 547.238 613.336 786.631

1.1146 1.1947 1.2667 1.3533 1.3281 1.1941 1.1585 1.0721 1.0098

Mean (9 points)

1.0307

–

1.1880

aggregated over eight GPUs. HBM working-set bandwidth and NVLink payload bandwidth are analytic models, not hardware counters: the former counts each owner’s unique active expert weights and packed scales once, the latter a one-way per-assignment payload of (H + H/32 + 2H) bytes, counting only assignments whose expert resides on a remote rank (an expected fraction (P − 1)/P of assignments). The overlap raises effective arithmetic throughput at every point. On H20, GC-MoE reaches 95.6 versus 73.8 TFLOPS at 8 tokens per rank (+29.5%, near the 1.296× peak speedup at 16) and 1287.8 versus 1251.4 at 256. On B300, GLM-5.2 reaches 8490.7 versus 7524.8 TFLOPS at 256 tokens per rank (+12.8%), and at 1 token per rank even the small-batch case improves from 75.3 to 82.5 TFLOPS. The working-set bandwidth model shows the same picture: H20 peaks at 18 931.9 versus 14 973.9 GB/s at 32 tokens per rank (+26.4%). The NVLink payload is identical by construction, and the modeled payload bandwidth rises with the speedup (239.7 to 246.7 GB/s at the largest H20 batch, +2.9%): the same bytes move in less wall time because dispatch no longer competes with the GEMM for SM resources and each replay avoids the trailing cross-rank

16

Figure 4: GLM-5.2 routed-expert decode on 8×B300: latency and speedup. barrier, so the partition does not cost cross-rank bandwidth.

6.5

Numerical accuracy

GC-MoE uses the caller’s quantized inputs and the baseline’s quantized weights and performs the same low-precision arithmetic; the difference is scheduling, layout, and instruction selection. Table 11 reports the H20 routed-expert accuracy suite: 47 cases (9 default, 23 boundary, 15 asymmetric, including empty ranks) over all eight ranks. All 376 rank-cases pass the cosine thresholds (0.97 against BF16 for GC-MoE and the baseline, 0.97 for GC-versus-Mega), and the processed-difference and graph/eager repeat checks are within 0.002 (repeat difference 0). The processed difference is 1 − 2⟨y, ŷ⟩/(∥y∥2 + ∥ŷ∥2 ) against a reference forward that dequantizes the packed weights and evaluates the layer in FP32. GC-MoE’s minimum cosine against the baseline is 0.9999998, and the two operators have identical maximum absolute error (0.109375) and identical maximum relative L2 error (0.00600836286321): at BF16 output precision they share the same maximum-error envelope, while their outputs still differ slightly. The H20 shared-expert suite likewise passes all 47 cases and 376 rank checks (cosine ≥ 0.99), with a minimum GC-versus-Mega cosine of 0.9999997 (maximum absolute error 0.0625, maximum relative L2 5.81 × 10−5 ). On B300, all 94 cases and 752 rank checks pass with cosine ≥ 0.97 (routed) and ≥ 0.99 (shared experts), with all finite checks passing. The shared-expert path stages scales in the same permuted order as the routed path, and the fused reduction accumulates BF16 values into FP32 in top-k slot order, which is why fusion does not move accuracy.

17

Figure 5: DeepSeek-V4.1-Flash routed-expert decode on 8×B300: latency and speedup. Table 8: Decode latency and speedup with fused shared experts, 8×B300. The baseline is Mega MoE in its fused shared-expert mode. Slowest-rank median per-invocation latency. DSV4.1-Flash abbreviates DeepSeek-V4.1-Flash. Model

Tokens/rank

Total tokens

Mega-se (µs)

GC-MoE (µs)

Speedup

GLM-5.2

1 2 4 8 16 32 64 128 256

8 16 32 64 128 256 512 1024 2048

68.739 93.033 113.672 130.904 136.624 138.641 143.096 150.360 170.763

59.752 79.730 99.059 115.870 121.372 123.645 128.534 136.231 152.883

1.1504 1.1669 1.1475 1.1297 1.1257 1.1213 1.1133 1.1037 1.1170

DSV4.1-Flash

1 2 4 8 16 32 64 128 256

8 16 32 64 128 256 512 1024 2048

73.541 87.140 108.908 143.151 164.977 174.223 179.868 184.979 192.473

60.816 75.478 96.132 128.697 150.251 161.154 166.320 168.524 179.688

1.2092 1.1545 1.1329 1.1123 1.0980 1.0811 1.0815 1.0976 1.0712

Mean speedup: GLM-5.2 (9 points) Mean speedup: DeepSeek-V4.1-Flash (9 points) Mean speedup: overall (18 points)

1.1306 1.1154 1.1230

18

Figure 6: DeepSeek-V4.1-Flash with shared experts on 8×H20: gc-se versus gc-unfused and mega-unfused.

Table 9: Effective throughput and analytic bandwidth models, DeepSeek-V4.1-Flash on 8×H20. TFLOPS counts matmul work only; HBM is the expert weight+scale working-set model; NVLink is the one-way assignment payload model. All values are aggregate over eight GPUs; GC denotes GC-MoE. TFLOPS

HBM (GB/s)

NVLink (GB/s)

Tokens/rank

Mega

GC

Mega

GC

Mega

GC

1 2 4 8 16 32 64 128 256

28.297 40.125 53.080 73.805 122.782 230.287 428.686 763.452 1 251.393

29.501 47.101 67.079 95.597 159.087 291.158 502.870 881.733 1 287.785

7 516.4 9 770.0 11 235.5 12 814.4 14 311.1 14 973.9 14 233.7 12 674.5 10 387.5

7 836.2 11 468.7 14 198.7 16 598.0 18 542.7 18 931.9 16 696.9 14 638.2 10 689.6

5.171 7.424 9.942 14.076 23.522 44.743 81.944 147.052 239.719

5.391 8.714 12.564 18.233 30.478 56.570 96.124 169.835 246.690

19

Figure 7: GLM-5.2 with shared experts on 8×B300: latency and speedup.

Table 10: Effective TFLOPS (matmul work only, aggregate over eight GPUs), 8×B300, routed experts. GC denotes GC-MoE. GLM-5.2

DeepSeek-V4.1-Flash

Tokens/rank

Mega

GC

Mega

GC

1 2 4 8 16 32 64 128 256

75.280 110.496 178.817 308.642 587.226 1 160.668 2 253.475 4 284.727 7 524.828

82.453 121.974 196.765 336.007 636.895 1 272.653 2 500.026 4 775.234 8 490.679

48.101 80.869 129.099 194.199 336.480 637.460 1 251.436 2 436.293 4 670.942

57.208 90.891 142.488 212.781 366.062 695.245 1 373.857 2 696.491 5 084.628

Table 11: Numerical accuracy, DeepSeek-V4.1-Flash on 8×H20, routed experts. Thresholds: cosine ≥ 0.97; processed diff and eager/graph repeat diff ≤ 0.002. Cosine columns are per-category minima. Category

Cases

GC/BF16

Mega/BF16

GC/Mega

Max proc. diff

default boundary asymmetric

9/9 23/23 15/15

0.977039 0.977672 0.977837

0.977039 0.977672 0.977837

0.99999976 0.99999976 0.99999976

1.80 × 10−5 9.95 × 10−6 1.80 × 10−5

GC: max absolute error 0.109375; max relative L2 0.00600836286321; repeat diff 0. Mega: max absolute error 0.109375; max relative L2 0.00600836286321.

20

Figure 8: DeepSeek-V4.1-Flash with shared experts on 8×B300: latency and speedup.

21

7

Related Work

MoE models and systems. Sparsely gated MoE introduced conditional computation with top-k routing (Shazeer et al., 2017); GShard and Switch Transformers established expert parallelism and all-to-all dispatch at scale (Lepikhin et al., 2020; Fedus et al., 2022); recent models use fine-grained experts and low precision (Jiang et al., 2024; DeepSeek-AI, 2024a,b). Systems work optimizes MoE training and inference at the framework level: DeepEP provides low-latency all-to-all kernels for expert parallelism (DeepSeek-AI, 2026a); MegaBlocks casts MoE as block-sparse GEMM with no token dropping (Gale et al., 2023); Tutel adapts parallelism and pipelining (Hwang et al., 2023); FasterMoE models and overlaps communication with computation across experts (He et al., 2022); Lina’s non-blocking all-to-all overlaps communication with expert computation across batches (Li et al., 2023); SonicMoE optimizes MoE kernels for high sparsity and IO (Guo et al., 2025). These systems overlap communication with compute across units of work (micro-batches, layers, experts). GC-MoE instead places dispatch and compute on disjoint SM sets inside one layer, so the overlap costs the compute SMs no dispatch warps or staging buffers; this is orthogonal and composable with the above. Beyond overlap, the decomposition gives each phase its own resource and synchronization budget, so dispatch and GEMM can be tuned independently; this matters most when the layer is bandwidth-bound, as in decode. Efficient kernels and serving. FlashAttention established IO-aware tiling (Dao et al., 2022), its successors exploit parallelism and Hopper asynchrony with FP8 (Dao, 2023; Shah et al., 2024); Triton and CUTLASS provide the tiled-kernel substrate (Tillet et al., 2019; NVIDIA, 2026b). DeepGEMM provides the FP8/FP4 grouped-GEMM family, including the Mega MoE baseline, which already overlaps dispatch with GEMM by sharing every SM; GC-MoE restructures that shared-SM, warp-specialized organization into disjoint SM partitions (DeepSeek-AI, 2026b); FlashInfer contributes inference kernels (FlashInfer Team, 2026). Serving engines schedule requests and memory (Kwon et al., 2023; Agrawal et al., 2023); GC-MoE is an operator that such engines can call under CUDA Graph capture. Low-precision formats. FP8 formats and their scaling behavior were characterized early (Kuzmin et al., 2022), and FP8 training/inference stacks established per-group scaling (Peng et al., 2023; DeepSeek-AI, 2024b); NVFP4 and MXFP4 define the FP4 microscaling conventions used here (NVIDIA, 2025; Tseng et al., 2025). GC-MoE’s contribution is not a new format but the dual-architecture realization of FP8×FP4 grouped GEMM: hardware block scaling with K128 weight prepacking and partition-level scheduling on Blackwell, and in-register FP4 decoding with window-based scale repacking on Hopper (Section 4.4). GPU partitioning and execution control. CUDA Green Contexts are a driver-level SM-partitioning primitive (NVIDIA, 2026a); MPS and MIG provide coarser isolation. Prior work uses green contexts and spatial partitioning for isolation: characterizing green-context, MPS, and MIG isolation (Martín et al., 2026), co-locating models on spatially partitioned GPUs (Quach et al., 2025), and isolating prefill/decode slots in an agentic serving system (Zhang et al., 2026). To our knowledge, GC-MoE is the first operator to use green contexts for intra-operator communication–computation overlap, with co-resident kernels synchronized at ring-block granularity.

8

Conclusion

GC-MoE organizes the expert-parallel MoE layer as a three-node CUDA Graph (four on Hopper with fused shared experts) across two CUDA Green Contexts: a small communication partition streams tokens into a ring while a large compute partition consumes ring blocks with FP8×FP4 tensor-core work, and a finalize node reduces the outputs without a trailing cross-rank barrier. At decode the grouped GEMM is bandwidth-bound, so the 128-SM compute partition loses little by ceding 20 SMs to dispatch, while each role gains its own resource and synchronization budget and can be optimized independently. The design required three ingredients beyond the partition itself—generation-tagged ring synchronization, double-buffered cross-rank counters, and architecture-specific low-precision MMA paths for Blackwell and Hopper—and it yields 1.188× mean decode speedup on 8×H20 and 1.106–1.123× on 8×B300 across two production MoE

22

configurations, with accuracy that matches the baseline’s error envelope at BF16 precision (minimum cosine 0.9999997 against the baseline on both measured H20 suites). Limitations. The operator targets decode; prefill-sized token counts are supported but the compute partition, not overlap, becomes the binding resource, and at the largest measured batch (256 tokens per rank) the gain narrows to 1.01–1.13×. The shape constraints of Appendix A reflect hard design decisions (64-bit completion masks, static signal sizing, a 72-rank by-value peer table). A controlled ablation of partition size and an overlap-on/off comparison are not included in this version; the fused-versus-unfused shared-expert comparison is the only within-operator ablation reported. The workspace is single-instance, so calls that share it must be serialized. Future work. We plan controlled experiments over partition sizes and dispatch SM counts, extension to prefill-heavy regimes and additional model families, adaptive partitioning when the driver permits finer splits, and integration into a graph-capturing serving stack.

References Amey Agrawal, Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav S. Gulavani, and Ramachandran Ramjee. SARATHI: Efficient LLM inference by piggybacking decodes with chunked prefills. arXiv preprint arXiv:2308.16369, 2023. Tri Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691, 2023. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memoryefficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems (NeurIPS), 2022. DeepSeek-AI. DeepSeek-V2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434, 2024a. DeepSeek-AI. DeepSeek-V3 technical report. arXiv preprint arXiv:2412.19437, 2024b. DeepSeek-AI. DeepEP: An efficient expert-parallel communication library. deepseek-ai/DeepEP, 2026a. Accessed 2026-11-02.

https://github.com/

DeepSeek-AI. DeepGEMM: A high-performance FP8 GEMM library for MoE and dense model inference. https://github.com/deepseek-ai/DeepGEMM, 2026b. Version 2.8.2, accessed 2026-11-02. William Fedus, Barret Zoph, and Noam Shazeer. Switch Transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1–39, 2022. FlashInfer Team. FlashInfer: Kernel library for LLM serving. https://github.com/flashinfer-ai/ flashinfer, 2026. Accessed 2026-11-02. Trevor Gale, Matei Zaharia, Cliff Young, and Erich Elsen. MegaBlocks: Efficient sparse training with mixture-of-experts. In Machine Learning and Systems (MLSys), 2023. Wentao Guo et al. SonicMoE: Accelerating MoE with IO and tile-aware optimizations. arXiv preprint arXiv:2512.14080, 2025. Jiaao He, Jidong Zhai, et al. FasterMoE: Modeling and optimizing training of large-scale dynamic pre-trained models. In ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming (PPoPP), 2022. Changho Hwang, Wei Cui, Yifan Xiong, et al. Tutel: Adaptive mixture-of-experts at scale. In Machine Learning and Systems (MLSys), 2023.

23

Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, et al. Mixtral of experts. arXiv preprint arXiv:2401.04088, 2024. Andrey Kuzmin, Mart Van Baalen, Yuwei Ren, and Markus Nagel. FP8 quantization: The power of the exponent. arXiv preprint arXiv:2208.09225, 2022. Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with PagedAttention. In ACM Symposium on Operating Systems Principles (SOSP), 2023. Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, et al. GShard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.16668, 2020. Jiamin Li, Yimin Jiang, Yibo Zhu, Cong Wang, and Hong Xu. Accelerating distributed MoE training and inference with Lina. In USENIX Annual Technical Conference (ATC), 2023. Juan José Martín et al. Performance isolation for inference processes in edge GPU systems. arXiv preprint arXiv:2601.07600, 2026. NVIDIA. Pretraining large language models with NVFP4. arXiv preprint arXiv:2509.25149, 2025. NVIDIA. CUDA C++ programming guide (Green Contexts). cuda-c-programming-guide/, 2026a. Accessed 2026-11-02.

https://docs.nvidia.com/cuda/

NVIDIA. CUTLASS: CUDA templates for linear algebra subroutines. https://github.com/NVIDIA/ cutlass, 2026b. Version 4.8.1, accessed 2026-11-02. Houwen Peng, Kan Wu, Yixuan Wei, Guoshuai Zhao, et al. FP8-LM: Training FP8 large language models. arXiv preprint arXiv:2310.18313, 2023. Ryan Quach et al. ECLIP: Energy-efficient and practical co-location of ML inference on spatially partitioned GPUs. arXiv preprint arXiv:2506.12598, 2025. Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, and Pradeep Ramani. FlashAttention-3: Fast and accurate attention with asynchrony and low-precision. arXiv preprint arXiv:2407.08608, 2024. Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538, 2017. Philippe Tillet, H. T. Kung, and David Cox. Triton: An intermediate language and compiler for tiled neural network computations. In ACM SIGPLAN Workshop on Machine Learning and Programming Languages (MAPL), 2019. Albert Tseng et al. Training LLMs with MXFP4. arXiv preprint arXiv:2502.20586, 2025. Yuning Zhang et al. AgentServe: Algorithm-system co-design for efficient agentic AI serving on a consumergrade GPU. arXiv preprint arXiv:2603.10342, 2026.

24

A

Operator constraints

Table 12 lists the supported envelope and the technical reason for each constraint. Table 12: Supported envelope and rationale (SM103 path unless noted). Constraint

Technical reason

I ≤ 4096, I mod 128 = 0

W1 output tiles 2I/128 must fit the 64-bit l2_full_mask; K128 tiles complete. 2-CTA clusters pair N128 tiles, requiring H/128 even. One warp lane carries one combine slot; dispatch reads 32/k tokens per warp. Static per-expert signal region; lane-held per-expert counters.

H mod 256 = 0 k + 1[shared] ≤ 32, k ≥ 1 E ≤ 2048 global, ≤ 512 per rank, E mod P = 0 P ≤ 72 ranks Routed recipe (1, 1, 32); SwiGLU BM ∈ {16, 32, 64, 96, 128, 192, 240} Shared recipe (1, 1, 32) or (1, 128, 128) One process per GPU, peer-accessible Serialized calls per workspace

592-byte by-value peer descriptor with 72 offsets. MMA scale-factor path hard-codes K32 UE8M0; SwiGLU fused in the W1 epilogue. 2-CTA MMA alignment; TMEM/SMEM budget. Scale staging and selector logic implemented for these K granularities. Symmetric memory, green contexts, and captured graph nodes. Persistent single-instance ring and signal state.

25

Related documents

• CascadeEP: Asynchronous Expert Execution for MoE Prefill under Attention Imbalance#1117964
• CascadeEP: Asynchronous Expert Execution for MoE Prefill under Attention Imbalance#1108732
• MonoMoE: An Efficient Fused Mega-kernel for Quantized MoE Decoding#927545
• A Multi-Engine Dataflow for MoE Decoding on AWS Tranium#1242549
• A Multi-Engine Dataflow for MoE Decoding on Scratchpad-Based Tensor Accelerators#1050941