Mirage Persistent Kernel (MPK)

1. Overview — What MPK Is

In low-batch, latency-sensitive LLM decoding, the model is read once per token and almost nothing is reused across tokens.

MPK (Mirage Persistent Kernel) is a compiler and runtime, from the Mirage group at CMU, that attacks exactly these fixed costs.

1.1 Where this page sits in “Worth the Bytes?”

Our measurement study makes one claim: transient inter-kernel L2 reuse of activations is worth less than 0.65% of step time, and MPK is the perfect adversary for that claim.

2. Why MPK Is the Right Adversary

A measurement claim is only as strong as the hardest case it survives, and MPK fuses the entire decode step into one kernel.

2.1 The paper's own attribution

MPK attributes its gain to launch elimination and intra-kernel pipelining, and bounds latency by the WEIGHT roofline — not by activation reuse.

2.2 The logical chain we test

3. Architecture — From Task Graph to One Kernel

MPK splits into a compiler that lowers a decode step into one megakernel's source, and a runtime that launches and drives it.

3.1 The compile-and-run pipeline

3.2 Baseline vs megakernel: the kernel timeline

The baseline issues a long sequence of short kernels; MPK collapses them into one persistent kernel body.

Baseline (one decode step) — many launches, each with an HBM round-trip:

L = CPU-driven kernel launch; each blue block reads activations from HBM and writes them back. Repeated per layer.

MPK (one decode step) — ONE launch, one resident kernel:

A single launch; a dispatch-only scheduler_kernel runs concurrently. No intra-step kernel boundaries.

4. Repo Analysis

We work from the mpk branch, cloned into the survey tree. The build pulls three submodules, compiles two Rust crates, runs CMake/nvcc, and cythonizes the bindings.

4.1 Directory map

4.2 Key files and their role

5. Core Mechanism (from the source)

Everything here is a verifiable static fact read directly from the cloned tree.

5.1 split_worker_scheduler is hardcoded true FACT

In init_persistent_kernel, the launch-selecting flag is set unconditionally to true at persistent_kernel.cuh:1617:

// include/mirage/persistent_kernel/persistent_kernel.cuh:1617
global_runtime_config.split_worker_scheduler = true;

That flag forks the launch. The else branch is reserved for the NVSHMEM / multi-GPU path; single-GPU always uses the split path:

// persistent_kernel.cuh:1799–1859  (launch fork)
if (global_runtime_config.split_worker_scheduler) {     // <-- single-GPU path
    // The split kernel does not support NVSHMEM ...
    worker_kernel<<<dim3(global_runtime_config.num_workers,1,1),
                    dim3(WORKER_NUM_THREADS,1,1),
                    MAX_DYNAMIC_SHARED_MEMORY_SIZE,
                    global_runtime_config.worker_stream>>>(global_runtime_config);

    scheduler_kernel<<<dim3(global_runtime_config.num_local_schedulers,1,1),
                       dim3(32,1,1), 0,
                       global_runtime_config.scheduler_stream>>>(global_runtime_config);
    // ... worker_kernel = the compute megakernel; scheduler_kernel = dispatch only
} else {
    // multi-GPU / NVSHMEM path — NOT used on single GPU
    persistent_kernel<<<dim3(num_sms_to_use,1,1), ...>>>(global_runtime_config);
    cudaDeviceSynchronize();
}

worker_kernel runs the compute; scheduler_kernel only dispatches. The megakernel of interest is worker_kernel.

// persistent_kernel.cuh:1385–1394
__global__ __launch_bounds__(WORKER_NUM_THREADS,1)
void worker_kernel(RuntimeConfig config) {
    worker_checker(config);
    execute_worker(config);          // <-- all layers' compute, stays resident
}
__global__ void scheduler_kernel(RuntimeConfig config) {
    scheduler_checker(config);
    execute_scheduler(config, 0);     // <-- dispatch only
}

5.2 num_workers = SM count, minus a few FACT

The worker count is derived from the SM count: 108 SMs (A100) → 96 workers, matching the README.

# python/mirage/utils.py:33–53
def get_configurations_from_gpu(rank):
    props = torch.cuda.get_device_properties(rank)
    sm_cnt = props.multi_processor_count
    if   sm_cnt >= 160: worker = 144
    elif sm_cnt >= 132: worker = 128   # H100 = 132 SMs -> 128 workers
    elif sm_cnt >= 108: worker = 96    # A100 = 108 SMs -> 96 workers (README config)
    elif sm_cnt >= 68:  worker = 64
    ...
    return worker, get_scheduler(sm_cnt, worker)

5.3 The headline number lives in the demo FACT

Both modes time with CUDA events. Baseline loops model.forward (many kernels); MPK calls mpk() once. Speedup = baseline_ms / mirage_ms.

# demo/qwen3/demo.py — baseline (:786–824, many kernel launches per step)
for cur_pos in range(prompt_len, decode_limit):
    logits = model.forward(input_ids=..., step=step, stream=stream)  # N kernels/step
    next_token = logits.argmax(dim=-1)
    ...
run_time = starter.elapsed_time(ender)
per_tok_ms = run_time / max(prompt_len + tokens_generated, 1)

# demo/qwen3/demo.py — MPK (:843–865, ONE megakernel call)
starter.record()
mpk()                                     # <-- the whole decode, one persistent kernel
ender.record(); torch.cuda.synchronize()
run_time = starter.elapsed_time(ender)
per_tok_ms = run_time / max(prompt_lengths[0].item() + tokens_generated, 1)

6. Claims vs Implementation

The table checks the paper's structural claims against the cloned source; the speedup is measured at 2.733× (Section 9).

7. Our Reproduction — Hypothesis & Method

7.1 Hypothesis — how we expect MPK's gain to decompose

Our a-priori prediction (written before the run) follows directly from the architecture. At batch 1 the decode step streams ~16 GB of weights once, which sets a weight-bandwidth floor neither engine can beat. The baseline pays that floor PLUS the launch/dispatch of N kernels per step PLUS the HBM round-trips of intermediate activations. MPK removes the launches and the activation round-trips and pipelines the resident work, but still pays the weight floor. So the speedup is the ratio of (floor + removable overhead) to (floor + residual). Within the removable part, the launch elimination should dominate; the activation/L2 reuse slice should be a sliver. A single hidden vector passed between two kernels at batch 1 is KB-scale (~hidden_size × 2 bytes ≈ a few KB); aggregated over all kernel boundaries in a step, the inter-kernel activation traffic is only O(MB) per token — either way orders of magnitude below the GB-scale weight stream. The H100 run bore this out: MPK lands at 1.21× the 4.84 ms/token weight roofline while the baseline sits at 3.3× it (Section 9.3).

Predicted attribution of the MPK speedup PRED

Illustrative a-priori hypothesis. The red L2 slice is bounded below 0.65% by the Section 9.3 roofline byte-ratio argument — not by ncu counters (not collected; profiler crash). The realized attribution is the Section 9.3 roofline table.

7.2 Method — what we build and run on PACE

Hardware (as run): one PACE H100 (sm_90, 132 SMs → 128 workers), built with CUTLASS kernels (sm_90a, TMA). Model: Qwen3-8B fp16, batch 1, greedy, max_num_batched_tokens=8. Reported metric is steady-state per-token decode latency.

The two headline runs

# cd $REPO/demo/qwen3  (demo imports models.* from cwd)
# BASELINE — native Triton/FlashInfer, many kernel launches per step
# both runs decode to natural EOS (baseline stops at 64 tokens, MPK at 326) —
# do NOT cap with --max-new-tokens, or you cannot reproduce the 326-token MPK run
python demo.py --prompt "<fixed prompt>" \
       --save-tokens $RESULT_DIR/outputs/baseline.json

# TREATMENT — single fused megakernel; mpk() does the whole decode
python demo.py --use-mirage --prompt "<fixed prompt>" \
       --save-tokens $RESULT_DIR/outputs/mirage.json

# speedup = baseline.latency_ms_per_token / mirage.latency_ms_per_token  (measured 2.733x on H100)

Profiling artifacts (repo's own entrypoint)

# --profiling attempts a tg4perfetto trace + task_graph_0.json + kernel_0.cu
#  in our run perfetto export crashed, so no usable trace was collected (attribution is roofline-only)
# (a short --max-new-tokens here is only an illustrative cap for a cheap profiled run,
#  NOT the headline runs above, which decode to natural EOS)
python demo.py --use-mirage --profiling --max-new-tokens 4 \
       --prompt "<fixed prompt>" --trace-name $RESULT_DIR/mpk_trace

8. L2 Attribution — Isolating the Activation-Reuse Slice (Planned ncu Method)

The plan was to split the measured MPK speedup into five buckets and isolate the fifth: (i) kernel-launch/dispatch elimination, (ii) SRAM / shared-memory staging, (iii) prefetch / latency hiding, (iv) reaching the weight-HBM-traffic roofline (weight-roofline-bound — MPK does NOT reduce weight traffic), and (v) inter-kernel L2 ACTIVATION reuse — and to show that (v) is < 0.65% of step time. Because the single-GPU megakernel-of-interest is worker_kernel (Section 5.1), and the repo's analyze.sh sudo-wraps ncu in a way PACE forbids, we would invoke Nsight Compute directly with explicit metrics that replicate its memory sections. (This counter-based pass did not complete; see the banner above.)

8.1 The Nsight metrics that separate the buckets

8.2 The two profiling runs we would have issued

# C1: MPK — profile the single compute megakernel worker_kernel (one launch)
ncu --target-processes all --kernel-name "regex:worker_kernel" --launch-count 1 \
    --metrics "$NCU_METRICS" --csv --export $RESULT_DIR/ncu/mpk_worker_kernel.ncu-rep \
    python demo.py --use-mirage --max-new-tokens 4 --prompt "<fixed prompt>"

# C2: baseline — capture one steady decode step's full Triton/FlashInfer kernel mix
ncu --target-processes all --launch-count 60 --launch-skip 40 \
    --metrics "$NCU_METRICS" --csv --export $RESULT_DIR/ncu/baseline_kernels.ncu-rep \
    python demo.py --max-new-tokens 4 --prompt "<fixed prompt>"

8.3 The attribution arithmetic the counters would feed

9. Results (updating as PACE runs land)

9.1 Headline reproduction

9.2 Attribution (the headline of this study)

9.3 The roofline arithmetic (H100) FACT

The 2.733× comes from launch-overhead elimination + intra-kernel pipelining moving the step down to the weight roofline — not from L2 activation reuse.

Baseline (16.047 ms/tok) is ~3.3× roofline → launch/dispatch-bound; MPK (5.872 ms/tok) is ~1.21× the weight roofline. Even 100%-L2-served inter-kernel activations save < act/total ≪ 1% — so the gain is launch elimination + reaching the weight roofline, not L2 reuse, corroborating our <0.65% bound.

10. How to Reproduce

10.1 Build (one-time, on a compute node with network)

# shallow clone leaves deps/ empty — submodule fetch is MANDATORY
cd $REPO    # mirage, mpk branch
git submodule update --init --recursive       # cutlass, z3, nlohmann/json
module load cuda/12.6.1                        # fallback 12.1.1; provides nvcc + ncu
conda create -y -p $ENV python=3.11 && conda activate $ENV
pip install "torch>=2.4" && pip install -r requirements.txt
export CUDACXX=$(which nvcc)
pip install -e . -v                            # Rust crates + CMake/nvcc + Cython
python -c "import mirage as mi; print('mirage OK')"

10.2 Run (headline + profiling + attribution)

cd $REPO/demo/qwen3
# headline runs decode to natural EOS (baseline 64 tokens, MPK 326) — no --max-new-tokens cap
python demo.py --prompt "$P" --save-tokens $OUT/baseline.json
python demo.py --use-mirage --prompt "$P" --save-tokens $OUT/mirage.json
# speedup = baseline.latency_ms_per_token / mirage.latency_ms_per_token
# profiled / ncu runs use a short --max-new-tokens 4 cap (cheap; illustrative, not the headline)
python demo.py --use-mirage --profiling --trace-name $OUT/mpk_trace --max-new-tokens 4 --prompt "$P"
ncu --kernel-name "regex:worker_kernel" --launch-count 1 --metrics "$NCU_METRICS" \
    --csv --export $OUT/ncu/mpk_worker_kernel.ncu-rep python demo.py --use-mirage --max-new-tokens 4 --prompt "$P"

10.3 Key configuration