MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs — L2 #

1. TL;DR #

MPK auto-compiles a multi-GPU PyTorch inference program into a single persistent mega-kernel via an SM-level task/event graph (tGraph) plus an in-kernel worker/scheduler runtime, enabling cross-task pipelining and fine-grained compute-comm overlap. Up to 1.7× lower latency vs SGLang/vLLM.

2. Q1 / Q2 / Q3 #

Q1 — 痛点 (the problem) #

The conventional kernel-per-operator execution model launches one GPU kernel per tensor operator, with an implicit kernel barrier between consecutive launches. This barrier forces all thread blocks of the prior kernel to complete before any block of the next begins. Three concrete costs follow:

1. No cross-operator software pipelining — the barrier confines pipelining (interleaving TMA loads / tensor-core compute / CUDA-core epilogue across iterations) to within a single kernel. Bubbles appear at every kernel boundary.
2. No fine-grained compute-communication overlap — a MatMul → AllReduce pair is two kernels; AllReduce waits for the entire MatMul even though each AllReduce tile depends only on one MatMul output tile.
3. Launch overhead + CUDA-Graph rigidity — hundreds-to-thousands of launches per iteration; CUDA Graphs amortize this but are static and re-instantiate poorly under dynamic shapes/control flow.

The precise regime this targets: LLM decode serving — tiny batch (1–16), 1 token per step, memory-bandwidth bound, with data-dependent attention durations (variable sequence lengths) making static scheduling ineffective.

Q2 — 方法 (the method) #

MPK represents the whole program at SM granularity instead of GPU granularity. A tGraph is a bipartite graph of tasks (a unit of compute/comm on one SM) and events (cross-task synchronization points). The compiler decomposes each operator into per-tile tasks, infers task-pair dependencies (event iff producer output region overlaps consumer input region), then fuses/normalizes/linearizes the graph into a compact device-memory structure. Per-task CUDA code is generated by the Mirage superoptimizer at thread-block level. An in-kernel runtime partitions SMs into workers (task queues, FIFO execution) and scheduler warps (event queues, dispatch), driven by an event-driven asynchronous loop, with a hybrid JIT/AOT task-launch policy.

核心技术壁垒 (the single hardest-to-replicate insight): turning an arbitrary fused mega-kernel into an efficient, indirection-free, self-scheduling structure. The chain of event fusion → tGraph normalization (every task ≤1 dependent-event and ≤1 triggering-event) → BFS linearization (tasks of one event are contiguous, so fan-out encodes as [first, last] indices) is what makes decentralized in-kernel scheduling cheap enough (only atomicAdd on circular buffers, only local scheduler state) to beat hand-tuned kernel-per-operator stacks. Anyone can fuse; making the fused graph executable at scale without a variable-length metadata explosion is the wall.

Q3 — 结果 (the results) #

• Single-batch end-to-end: 1.0–1.7× faster than SGLang/vLLM across 5 models × {A100, H100, B200}; largest wins on smaller models + newer GPUs.
• >10× over native PyTorch (hand kernels + CUDA Graphs + torch.compile).
• Multi-GPU (8× H100, tensor parallel): 1.1–1.4× over SGLang/vLLM.
• Ablations: cross-task pipelining 1.2–1.3× (final linear layer, beats cuBLAS); compute-comm overlap 1.1× per iteration.
• Near hardware limit: Qwen3-8B/A100 decode latency 14.5 ms → 12.5 ms, vs ~10 ms bandwidth roofline.

3. 架构 / 方法图 #

Target operation, shapes & precision contract #

MPK is not a single kernel but a program-level mega-kernel; still, each task is a precisely-shaped tile op:

• MatMul task: output tile of a GEMM, tiled along row and column dims of the output tensor. Precision contract: bf16 input/output, tensor-core accumulate (FP32 accumulator per the Mirage-generated implementation), RNE.
• Attention task: paged-attention decode, [B, 1, H, D] query against a paged KV cache; data-dependent duration (variable seq len) — this is the JIT-classified op.
• AllReduce task: element-wise reduce over a tile; decomposed into an NVSHMEM nvshmem_signal_wait_until data-transfer task + a local reduction task.

Primary architecture figure #



Paper's Figure 1. The pipeline: a PyTorch tensor program plus an inference configuration and request stream feed the MPK Compiler (§4), which emits an SM-level graph (tGraph) of MM/AT/E/AR nodes, executed by the in-kernel parallel runtime (§5) inside one persistent kernel. Reader should notice the entire loop below the compiler lives inside a single launch — no per-op kernel boundaries.



Paper's Figure 4. (a) is the coarse operator DAG (MatMul→Attention→MatMul→AllReduce); (b) is the decomposed tGraph with blue compute tasks, orange comm tasks, and green event circles alternating; (c) is a functionally-equivalent but suboptimal tGraph where events only capture operator-level dependencies — i.e. behaves like kernel barriers. The load-bearing contrast is (b) vs (c): the same computation admits many tGraphs and only fine-grained events unlock overlap.

The two barriers this removes are best seen in the background figures:



Paper's Figure 2 (stored as fig3.png). (a) Kernel A then Kernel B with pipeline bubbles across the TMA / Tensor-Core / CUDA-Core rows; (b) MPK interleaves Task A and Task B stages with no bubble. This is the intra- vs cross-task pipelining distinction that motivates the paged-shared-memory design.

4. 作者证明 #

无严格性能定理 — 正确性有形式化 (event fusion + normalization + linearization), 性能为实证。 The formal content is graph-rewrite correctness, not a throughput bound. Notation and the 6 minimum checks follow.

Notation table #

Symbol	Meaning
t, t1, t2	tasks (SM-level units of compute/comm)
e, e1, e2, e′	events (cross-task sync points)
InTasks(e)	set of producer tasks that trigger e
OutTasks(e)	set of consumer tasks that depend on e
e.counts	# triggers required to activate e
∪	set union

方程物理意义 (the two fusion rewrites) #

• Successor-set fusion (Def 4.1): if $OutTasks(e_1) = OutTasks(e_2)$, replace with $e'$ where $InTasks(e') = InTasks(e_1) \cup InTasks(e_2)$, $OutTasks(e') = OutTasks(e_1)$. Physically: two events whose consumers are identical are the same barrier — consumers wait on both anyway, so union the producers.
• Predecessor-set fusion (Def 4.2): if $InTasks(e_1) = InTasks(e_2)$, replace with $e'$ where $InTasks(e') = InTasks(e_1)$, $OutTasks(e') = OutTasks(e_1) \cup OutTasks(e_2)$. Physically: two events fired by the same producers fire simultaneously — union the consumers.

6 minimum checks #

1. Dependency preservation: an event exists for pair $(t_1,t_2)$ iff $t_1$'s output region overlaps $t_2$'s input region — so every producer-consumer edge is retained; fusion only merges events with identical consumer or producer sets, which cannot drop an edge. ✔ consistent.
2. Fusion is semantics-preserving: merging equal-OutTasks (or equal-InTasks) events keeps the wait condition unchanged (consumers still wait for all original producers). ✔
3. Normalization termination/equivalence: each fan-out/fan-in rewrite adds one event $e'$ + $k$ empty tasks and reduces a task's event degree to 1; empty tasks do no compute, so semantics are preserved. Overhead only when graph is "wide". ✔
4. Linearization contiguity (Alg 1): BFS enqueues each task once, each event once; lines 5–7 append all tasks depending on a dequeued event consecutively, so an event's fan-out is a contiguous [first,last] index range. ✔ (this is the encoding claim, not a bound).
5. Roofline placement — see below; the target regime is memory-bound, so the proof of value is bandwidth utilization, not FLOP peak.
6. % peak / lower-bound derivation — see below; the paper's "close to hardware limits" claim is checked against a bandwidth roofline, not a compute roofline.

Roofline placement (kernel-specific) #

Decode serving at batch 1–16 with 1 token/step is memory-bandwidth bound: the arithmetic intensity of loading model weights once per token is roughly $\text{AI} \approx \frac{2 \cdot P}{2 \cdot P} = 1$ FLOP/byte for a bf16 weight matrix (2 FLOPs per MAC, 2 bytes per bf16 weight loaded once) — far below the ridge point of any of A100/H100/B200 (H100 ≈ 990 bf16 TFLOP/s ÷ 3.35 TB/s ≈ 295 FLOP/byte). The op therefore lands deep on the left (memory-bound) side of the roofline. MPK's wins are consistent with this: they come from removing overheads that inflate wall-time above the bandwidth floor (launches, CPU scheduling, pipeline bubbles), not from higher FLOP utilization.

% peak / lower-bound derivation #

The paper's own bound for Qwen3-8B on A100: weights ≈ 16 GB, A100 HBM ≈ 1.6 TB/s, so the pure weight-load lower bound is $16\text{ GB} / 1.6\text{ TB/s} = 10\text{ ms}$ per token. Measured: MPK 12.5 ms → ~80% of the bandwidth roofline ($10/12.5$); the tuned baselines sit at 14.5 ms → ~69%. So MPK closes about a third of the remaining gap to the memory-bandwidth ceiling but is still ~25% above it (KV-cache traffic, non-overlapped comm, residual scheduling all account for the rest).

Why this launch config (Table 1) is chosen #

Schedulers are pinned at exactly 4 SMs × 4 warps = 16 scheduler warps on all three GPUs regardless of total SM count (A100 108, H100 132, B200 148); the rest are workers (104/128/144). Shared memory is paged at 32 KB/page → 5/7/7 pages per SM. A larger scheduler budget would steal SMs from compute with no benefit (scheduling is atomicAdd-light and decentralized); a smaller one would bottleneck JIT dispatch for the data-dependent attention wave. Page size trades pipelining depth (more pages = deeper cross-task prefetch) against per-page waste; 32 KB is the point where H100/B200's ~228 KB smem yields 7 usable pages.

Design space & constraint derivation #

Axis	Value chosen	Rejected alternatives	Constraint that blocks them
Task granularity	~#tasks ∝ #SMs (per operator)	1 task = whole op (kernel-per-op)	coarse tasks re-introduce the barrier; too-fine tasks explode scheduling atomicAdd traffic
Event metadata	≤1 dep + ≤1 trig event per task (normalized)	variable-length event lists	max-fanout preallocation → large device-memory + indirect indexing cost
Event fan-out storage	[first,last] contiguous index range (linearized)	explicit dependent-task index list	per-event list storage grows with fan-out, costly to read in-kernel
Task launch mode	hybrid JIT (data-dependent ops) + AOT (rest)	pure JIT / pure AOT	pure JIT = 2 syncs/task (latency); pure AOT = static assignment can't absorb attention skew
Scheduling	decentralized (local state)	globally-coordinated	global coordination adds cross-SM comm/sync overhead
Shared memory	paged (32 KB pages)	monolithic per-kernel smem	monolithic smem forbids cross-task prefetch (both tasks need smem)

5. 实验与数据 #



Paper's Figure 9. Numbers above each MPK bar are its speedup over the best existing system (1.0–1.7×). Notice the gradient: speedup rises for smaller models and newer GPUs — the mega-kernel advantage is dominated by overhead removal, which shrinks in relative terms as compute grows. The 1.0× floor (parity) shows up on larger models / older hardware.



Paper's Figure 12. MPK-Pipe vs MPK-No-Pipe on the final linear layer: 1.2–1.3× speedup, and it even beats cuBLAS-compiled kernels. This is the load-bearing evidence that the paged-shared-memory + pre-load/compute phase split actually removes bubbles rather than just relocating them.



Paper's Figure 13. Overlap is disabled by capturing only coarse operator-level dependencies (the Fig 4c suboptimal graph). Enabling fine-grained overlap cuts per-iteration latency by 1.1× — modest, but it is the only isolable contribution of the compute-comm overlap mechanism, so it bounds how much the AllReduce-decomposition strategy buys.



Paper's Figure 10. MPK-Hybrid-MoE vs SGLang-MoE in microseconds; the hybrid balancer beats purely-static partitioning across all batch sizes. The gather step alone was up to 11% of SGLang's MoE time — fusing it into the GEMM data-loading phase removes that standalone kernel entirely.



Paper's Figure 11. Up to 10× over PyTorch and 1.1–1.4× over SGLang/vLLM at 8 H100. The multi-GPU win is where compute-comm overlap matters most, since AllReduce is now on the critical path.

Optimization techniques inventory #

Technique	Target bottleneck	HW primitive	Measured contribution
Cross-task software pipelining	pipeline bubbles at kernel boundaries	TMA async copy + tensor cores	1.2–1.3× (Fig 12)
Paged shared memory	smem lifetime tied to kernel blocks prefetch	on-chip shared memory	enabler for pipelining (no standalone number)
Fine-grained compute-comm overlap	comm latency behind compute	NVSHMEM signal-wait	1.1× per iter (Fig 13)
Hybrid JIT/AOT launch	dispatch latency vs load imbalance	device-memory semaphores / atomicAdd	qualitative (Fig 8: 2 vs 1 sync)
Fused gather-GEMM (MoE)	standalone gather kernel (≤11% MoE time)	async token-level copy in GEMM load phase	folds gather into GEMM (Fig 10)
Task-description prefetch	352-byte descriptor device-memory latency	smem prefetch	overhead reduction

Numerical considerations #

• Precision: all systems run bf16; per-task tensor-core GEMMs accumulate in FP32 (Mirage-generated). No FP8/FP4 path in this work.
• Accuracy vs reference: MPK is semantics-equivalent by construction (fusion/normalization preserve producer-consumer regions); the paper reports no numerical-divergence table — accuracy comes directly from the reference PyTorch task implementations, not re-derived.
• Edge cases (overflow/denormal) are not analyzed; the contribution is scheduling/fusion, not new numerics.

6. 论证链 #

#	Step	Support
1	Kernel-per-operator imposes an implicit barrier between launches, enforcing full-operator completion.	§2.1: thread blocks scheduled independently → no cross-block sync → barriers inserted by runtime.
2	Barriers block cross-task pipelining, compute-comm overlap, and add launch/CUDA-Graph overhead.	§2.1 Fig 2/3; §1 three limitations.
3	Representing dependencies at SM (task) granularity instead of operator granularity exposes the fine-grained overlap the barrier hides.	§3 tGraph: event per overlapping producer-consumer tile pair (Fig 4b vs 4c).
4	A naive SM-level graph has variable-length event metadata that is too costly to read in-kernel.	§4.1 para 7: max-fanout preallocation → memory + indirect indexing overhead.
5	Event fusion + normalization + linearization compress the graph so each task stores ≤1 dep/trig event and each event stores a [first,last] task range.	§4.1 Defs 4.1/4.2, Fig 6 rewrites, Alg 1 BFS.
6	A compact graph enables a cheap decentralized in-kernel runtime (atomicAdd queues, local scheduler state, hybrid JIT/AOT).	§5–§6.1: workers/schedulers, semaphores, 40K C++ / 84K CUDA / 10K Py.
7	End-to-end this yields 1.0–1.7× over tuned kernel-per-operator baselines, near the bandwidth roofline.	§6.3 Fig 9; Qwen3-8B 14.5→12.5 ms vs ~10 ms bound.

7. 实现 cross-reference #

Public repo: https://github.com/mirage-project/mirage (MPK is the persistent-kernel component). Specific line-level citations were not extractable from the L1, so per-file entry points are marked as required-but-unverified locations rather than fabricated line numbers.

• Kernel entry / host stub: the PyTorch backend path — torch.compile(backend=MPK) invokes the compiler and returns a callable issuing one mega-kernel launch. [实现未公开 at line level] — see mirage persistent-kernel backend.
• Inner loop: worker loop dequeues a task (circular buffer + atomicAdd), runs the Mirage-generated device function, notifies its triggering event; scheduler warps poll event queues and dispatch. [实现未公开 at line level]
• Epilogue / write-back: per-task epilogue + nvshmem_signal_wait_until for AllReduce data-transfer tasks feeding local reduction tasks. [实现未公开 at line level]

核心技术壁垒 (dedicated paragraph) #

The replication wall is not the mega-kernel idea (FlashDMoE and Spector et al. already hand-built ones) but the automatic, indirection-free execution substrate. The normalize-then-linearize invariant — every task has at most one dependent and one triggering event, and every event's dependent tasks occupy a contiguous index range — is what lets the in-kernel runtime schedule with nothing but atomicAdd counters and [first,last] bounds, i.e. no per-task variable-length metadata reads on the hot path. Reproducing MPK means reproducing this compiler-side graph canonicalization and the 84K lines of CUDA runtime that assumes it. Miss the invariant and the in-kernel scheduler's memory-indirection cost erases the overlap gains it was meant to unlock.

关键实现细节 (easy-to-miss tricks) #

1. Empty tasks as degree-reduction glue: normalization inserts no-compute tasks (Fig 6 T1..Tk) purely to force fan-in/fan-out to 1. They look like overhead but keep the descriptor fixed-size; overhead stays <1% because real models are "deep not wide".
2. 352-byte task descriptors prefetched into smem: because tasks are far finer than kernels, per-task descriptor loads from device memory would dominate; MPK hides them by prefetching upcoming descriptors into shared memory before dequeue.
3. AllReduce is not a collective: it is compiled into NVSHMEM signal-wait data-transfer tasks + local reduction tasks — an ordinary pair of tGraph nodes — which is what makes async compute-comm overlap expressible at all.

Portability analysis #

• GPU families: evaluated on sm_80 (A100), sm_90 (H100), sm_100 (B200). Depends on TMA async copy (Hopper+ for the fused gather-GEMM path) and NVSHMEM for multi-GPU.
• Porting to CDNA (gfx942/gfx950): the graph-level machinery (tGraph, fusion, normalization, linearization, worker/scheduler runtime, atomicAdd queues) is architecture-neutral and would survive. What would not survive unmodified: TMA-dependent gather-GEMM (needs an mfma + async buffer_load/ds_read rewrite) and NVSHMEM comm (needs ROCm SHMEM / RCCL equivalent). The per-task code is Mirage-generated, so a CK/Triton retarget is a Mirage-backend problem, not an MPK-runtime problem.

Software → Hardware reverse implication #

MPK reveals what the mega-kernel paradigm wants from hardware:

1. A cheap cross-SM synchronization primitive: MPK synthesizes cross-SM events out of device-memory semaphores + atomicAdd because the ISA offers no first-class sub-kernel cross-block barrier. A hardware "event register" with SM-to-SM notification would shed the entire scheduler-warp layer (4 SMs × 4 warps) and cut JIT dispatch to one hop.
2. A persistent, kernel-lifetime-independent shared-memory scratchpad: MPK reinvents this as a paged abstraction because smem is tied to block lifetime. Hardware that let shared memory persist across "tasks" would delete the paging bookkeeping and deepen pipelining for free.
3. Larger / bankable smem to raise page count: only 5–7 pages/SM at 32 KB caps pipeline depth. More on-chip SRAM (or a dedicated async-copy staging buffer) directly increases cross-task prefetch overlap.
4. Concrete proposal: if the next-gen ISA exposed a hardware SM-to-SM event/notify instruction and lifetime-decoupled shared memory, MPK could drop the decentralized scheduler-warp partition and the paged-smem layer — shedding a large slice of the 84K CUDA runtime lines and closing more of the residual ~25% gap to the bandwidth roofline (from 80% toward the 10 ms bound).