Mpk Internals

MPK Internals: Compilation-to-Runtime Pipeline

This document traces the full lifecycle of an MPK megakernel from Python graph construction through CUDA compilation to persistent kernel execution.

Pipeline Overview

Phase 1: Python Graph Building
  PersistentKernel.compile()
  → layer methods build KNGraph/TBGraph
  → kn_graph.generate_task_graph()
        |
        v
Phase 2: C++ Code Generation  (runtime.cc)
  Graph::generate_task_graph()
  → register_mugraph()     — builds task/event lists
  → print_task_graph()     — emits CUDA code + JSON
        |
        v
  Two artifacts:
    test.cu              — _init_persistent_kernel(), _execute_task(), Python C ext
    task_graph.json      — task descriptors, events, dependencies
        |
        v
Phase 3: CUDA Compilation
  nvcc test.cu → test.so  (Python extension module: __mirage_launcher)
        |
        v
Phase 4: Runtime Initialization
  init_persistent_kernel()
  → loads JSON, allocates GPU queues, builds RuntimeConfig
        |
        v
Phase 5: Runtime Execution
  launch_persistent_kernel()
  → prepare_kernel (reset queues)
  → worker_kernel + scheduler_kernel (persistent loop)
  → workers fetch tasks, wait on events, call _execute_task()
  → schedulers process events, enqueue tasks to workers

Phase 1: Python Graph Building

Key file: python/mirage/mpk/persistent_kernel.py

Entry point: PersistentKernel.compile()

The compilation method does the following in order:

1. Generate task graph — calls self.kn_graph.generate_task_graph(num_gpus, my_gpu_id) which bridges through Cython (python/mirage/_cython/core.pyx, generate_task_graph()) into C++. Returns {"cuda_code": str, "json_file": str}.

2. Write files — writes test.cu (CUDA code + HARD_CODE Python extension wrapper) and task_graph.json to a temp directory.

3. Compile — builds the nvcc command via get_compile_command() and calls subprocess.check_call().

4. Load module — uses importlib.util.spec_from_file_location() to dynamically load the compiled .so as Python module __mirage_launcher. Extracts init_func, launch_func, init_request_func, finalize_func.

5. Initialize runtime — calls init_func(...) with meta-tensor pointers, worker/scheduler counts, and serving config.

How layers build the graph

Each layer method (e.g., rmsnorm_layer, linear_layer, moe_w13_fp8_layer) does:

1. Create a TBGraph with CyTBGraph(grid_dim, block_dim, forloop_range, reduction_dimx)
2. Call tb_graph.new_input(dtensor, partition, forloop_dim, store_in_dmem) for each input and output
3. Call self.kn_graph.customized([tensors...], tb_graph) to register the operator
4. Call self.kn_graph.register_task(tb_graph, "task_name") which dispatches to C++ Graph::register_task()

HARD_CODE: the Python C extension wrapper

The HARD_CODE constant (top of persistent_kernel.py) is a C string appended to the generated CUDA code. It defines a Python extension module with four functions:

• init_func — parses Python args, calls C++ init_persistent_kernel()
• launch_func — takes a CUDA stream pointer, calls launch_persistent_kernel(stream)
• init_request_func — calls init_request_resources() (for online serving)
• finalize_func — calls finalize_persistent_kernel()

Layer API: TBGraph Partition Scheme

Each layer method (e.g., rmsnorm_layer, linear_layer, moe_w13_linear_layer) builds a TBGraph that describes how the global tensors are sliced into per-task tiles. This section explains every parameter.

CyTBGraph constructor

tb_graph = TBGraph(CyTBGraph(grid_dim, block_dim, forloop_range, reduction_dimx))

tb_graph.new_input() — registering a tensor

tb_graph.new_input(dtensor, input_map, forloop_dim, store_in_dmem)

Called for every tensor the task touches — both inputs and outputs. The first num_inputs calls register inputs; the remaining register outputs. This ordering must match num_inputs/num_outputs in graph.cc's task_config tuple.

input_map: the partition tuple

A 3-element tuple (mx, my, mz) that maps grid dimensions → tensor dimensions:

input_map.y and input_map.z work identically for grid_dim.y and grid_dim.z.

In short: the value tells you which tensor dimension that grid axis splits. -1 means "don't split by this grid axis."

forloop_dim (vestigial in MPK)

In the Mirage superoptimizer, forloop_dim and forloop_range together control tiled reduction loops within a TBGraph. However, in MPK forloop_range is always 1, which makes forloop_dim a no-op — the dimension division (dim / 1) and stride multiplier (* 1) have no effect regardless of what value you pass.

MPK task kernels handle their own internal tiling and reduction directly in CUDA (e.g., looping over the K dimension in a matmul). The TBGraph forloop mechanism is not used. You'll see various forloop_dim values in existing layer methods (e.g., 1, 2, -1), but they're all equivalent when forloop_range=1. By convention, existing code sets forloop_dim to the "reduction dimension" of the operation, but this is cosmetic.

store_in_dmem

• True — the per-task tensor slice lives in device (global) memory. Should be set to True for all MPK tensors.

Annotated example: moe_w13_linear_layer

def moe_w13_linear_layer(self, input, weight, moe_routing_indices,
                         moe_mask, output, grid_dim, block_dim):

    # input:              (batch_size, hidden_size)                        2D bf16
    # weight:             (num_experts, 2*intermediate_size, hidden_size)  3D bf16
    # moe_routing_indices:(num_experts, batch_size)                        2D int32
    # moe_mask:           (num_experts + 1,)                               1D int32
    # output:             (batch_size, num_experts_per_tok, 2*inter_size)  3D bf16

    tb_graph = TBGraph(CyTBGraph(grid_dim, block_dim, 1, 64))

    #                       tensor,              input_map,    forloop_dim*, store_in_dmem
    #  (* forloop_dim is vestigial in MPK — has no effect when forloop_range=1)

    tb_graph.new_input(input,               (-1, -1, -1),  1,           True)
    #  → No partition on any grid axis. Every task sees full (batch, hidden).

    tb_graph.new_input(weight,              (-1,  1, -1),  2,           True)
    #  → grid_dim.y partitions dim 1 (the 2*intermediate_size axis).
    #    Each task handles 2*inter_size / grid_dim.y rows of the weight matrix.

    tb_graph.new_input(moe_routing_indices, (-1, -1, -1), -1,           True)
    #  → No partition. Every task sees the full routing table.

    tb_graph.new_input(moe_mask,            (-1, -1, -1), -1,           True)
    #  → No partition. Every task sees the full mask.

    tb_graph.new_input(output,              (-1,  2, -1), -1,           True)
    #  → grid_dim.z partitions dim 2 (the 2*intermediate_size axis of the output).
    #    Each task writes to its slice of output columns.

    self.kn_graph.customized([input, weight, moe_routing_indices, moe_mask, output], tb_graph)
    self.kn_graph.register_task(tb_graph, "moe_w13_linear_sm100")

How partitioning connects to task pointers

At runtime, the partition tuple is resolved during task graph generation (src/threadblock/graph.cc). For each task instance (one grid coordinate), the code generator computes a byte offset from the tensor's base pointer:

per_task_ptr = base_ptr
             + blockIdx.x * stride_for(input_map.x)
             + blockIdx.y * stride_for(input_map.y)
             + blockIdx.z * stride_for(input_map.z)

These offsets are baked into the TaskDesc at init time (via JSON → FullTaskDesc → TaskDesc). The task kernel receives pre-offset pointers in task_desc->input_ptrs[i] and task_desc->output_ptrs[i] — this is why tasks are blockIdx-agnostic.

Phase 2: C++ Code Generation

Key file: src/kernel/runtime.cc

Entry point: Graph::generate_task_graph()

This function orchestrates all code generation:

1. register_mugraph() — walks the KNGraph operators and converts each into FullTaskDesc entries. For each KN_CUSTOMIZED_OP, it queries task_config[op] (a tuple of num_inputs, num_outputs, TaskType, variant_id set by Graph::register_task()) to determine the task type and variant. It also creates EventDesc entries for inter-task dependencies and populates first_tasks (the initial ready tasks).

2. print_task_graph() — generates two outputs:

   Output 1: CUDA code containing three generated functions:

   • construct_task_graph() — loads task_graph.json at runtime, parses it into FullTaskDesc/EventDesc vectors, and creates TMA descriptors for Hopper/Blackwell tasks.
   • _init_persistent_kernel() — sets up tensor pointers from io_configs (torch tensors, cudaMalloc buffers, shuffled tensors, NVSHMEM buffers). Called once during initialization.
   • _execute_task() — a giant if/else dispatcher that maps (task_type, variant_id) pairs to the actual kernel function calls. Each branch contains the code string generated by the corresponding TaskRegister::register_*_task() function.

   Output 2: JSON task graph — serializes all tasks, events, and dependencies (see JSON Schema section below).

Key file: src/kernel/graph.cc

Graph::register_task() maps task name strings to registration functions:

"moe_w13_fp8_sm100" → register_moe_fp8_sm100_task() → TASK_MOE_W13_FP8_SM100

Each registration function (in src/kernel/task_register.cc) reads tensor dimensions from the TBGraph, generates a CUDA code string calling the templated kernel with those dimensions, and returns a variant_id via register_task_variant(). Same code string → same variant_id (deduplication).

Phase 3: CUDA Compilation

Key function: get_compile_command() in persistent_kernel.py

Builds the nvcc command with:

• Includes: Python headers, Mirage headers, CUTLASS, JSON library
• Architecture flags: -gencode=arch=compute_90a,code=sm_90a (Hopper), compute_100a,code=sm_100a (Blackwell)
• Feature defines: -DMPK_ENABLE_TMA (Hopper/Blackwell), -DMIRAGE_GRACE_HOPPER or -DMIRAGE_GRACE_BLACKWELL
• Runtime defines: -DMODE_OFFLINE, -DMPK_MAX_NUM_BATCHED_REQUESTS=N, -DMPK_MAX_NUM_BATCHED_TOKENS=N, -DMPK_MAX_NUM_PAGES=N, -DMPK_PAGE_SIZE=N, -DMPK_MAX_SEQ_LENGTH=N
• Scheduler config: -DMAX_WORKER_PER_SCHEDULER=N (computed from worker/scheduler ratio)
• Output: shared library (.so) as a Python extension module

For multi-GPU (NVSHMEM): adds -rdc=true, NVSHMEM/MPI includes and libraries.

Phase 4: Runtime Initialization

Key file: include/mirage/persistent_kernel/persistent_kernel.cuh

init_persistent_kernel() sets up the full runtime state:

1. Meta-tensor mapping — stores 10 meta-tensor pointers in global_runtime_config (step, tokens, input_tokens, output_tokens, num_new_tokens, prompt_lengths, qo_indptr, paged_kv_indptr, paged_kv_indices, paged_kv_last_page_len).

2. NVSHMEM init (if multi-GPU) — calls nvshmemx_init_attr(), creates NVSHMEM teams for cross-GPU communication.

3. Call generated _init_persistent_kernel() — this loads the JSON task graph via construct_task_graph(), allocates GPU memory for intermediate tensors, and populates the all_tasks, all_events, first_tasks vectors.

4. Allocate runtime queues on GPU:

   • worker_queues[2 * num_workers] — per-worker task queues (local + remote). Each is a circular buffer of TaskId with length per_worker_queue_len (1024).
   • sched_queues[num_schedulers + 1] — per-scheduler event queues + one global broadcast queue. Circular buffers of EventId.
   • worker_queue_last_ready_task_id[2 * num_workers] — atomic counters for queue tail.
   • sched_queue_last_ready_event_id[num_schedulers + 1] — atomic counters for event queue tail.
   • all_event_counters[num_events] — atomic counters tracking how many times each event has been triggered.
   • all_event_num_triggers[num_events] — how many triggers each event needs before it's considered "ready".

5. Copy task/event data to GPU — all_tasks, all_events, first_tasks are copied to device memory.

6. Set kernel attributes — sets cudaFuncAttributeMaxDynamicSharedMemorySize for worker and scheduler kernels.

7. Create streams and events — separate CUDA streams for workers and schedulers (split mode), plus synchronization events.

8. Call init_request_resources() — launches init_kernel which initializes per-request state (step counters, page queues for MODE_OFFLINE/MODE_ONLINE).