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 ```python tb_graph = TBGraph(CyTBGraph(grid_dim, block_dim, forloop_range, reduction_dimx)) ``` | Parameter | Meaning | |---|---| | `grid_dim` | `(x, y, z)` — number of task instances in each dimension. Total tasks = `x * y * z`. | | `block_dim` | `(threads, 1, 1)` — threads per task. Must be `(128,1,1)` Ampere, `(256,1,1)` Hopper/Blackwell. | | `forloop_range` | Number of forloop iterations (always **1** in MPK — see note below). | | `reduction_dimx` | Tile size for the reduction dimension (always **64** in MPK). | ### `tb_graph.new_input()` — registering a tensor ```python 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.x` value | Meaning | |---|---| | `-1` | `grid_dim.x` does **not** partition this tensor. Every task sees the full extent of every dimension. | | `0` | `grid_dim.x` partitions **tensor dimension 0**. Task at grid position `gx` sees the slice `[gx * dim[0]/grid_x : (gx+1) * dim[0]/grid_x]` along dim 0. | | `1` | `grid_dim.x` partitions **tensor dimension 1**. Same slicing logic on dim 1. | | `2` | `grid_dim.x` partitions **tensor dimension 2**. | `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` ```python 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**: