Test Mode

MPK Test Mode

Test mode compiles and runs an MPK task graph exactly once and exits. It exercises the full pipeline — Python layer API, task registration, C++ code generation, nvcc compilation, runtime dispatch, and the persistent runtime's metadata setup (init_kernel + prepare_next_batch) — making it the primary tool for validating that a new layer or task works end-to-end.

Test mode is selected by setting params["test_mode"] = True at construction time. Internally this defines -DMPK_TEST_MODE for the launcher build, which:

1. Auto-allocates any meta tensors the test author didn't pass (so paged-attention / embedding / sampling layers see valid qo_indptr_buffer, paged_kv_*, input_tokens, etc.).
2. Lets init_request_resources() and prepare_next_batch run normally — the same code paths production uses.
3. Forces prepare_next_batch's always-finalize shortcut on iter 1, which returns false and terminates the scheduler after exactly one task-graph pass.

Required: PyTorch Reference Comparison

Every test mode file must include a PyTorch reference implementation that computes the same operation, and must compare the MPK output against it numerically. A test that only runs the kernel without checking correctness is not a valid test — it only proves the kernel doesn't crash.

The reference should:

• Use plain torch ops (or torch.nn.functional) to implement the same math the layer performs.
• Run on the same input tensors as the MPK kernel (cast to a higher precision like float32 if needed for a trustworthy reference).
• Be compared with a tolerance appropriate to the dtype: bf16 typically atol=1e-2, rtol=1e-2; fp16 similar; fp32 much tighter.

Use torch.testing.assert_close(out, ref, atol=..., rtol=...) and/or print (out - ref).abs().max() so failures surface immediately rather than silently producing wrong numbers.

Where the reference lives: pytorch_reference.py

Per-layer test_mode files must import their PyTorch reference from pytorch_reference.py in the same folder, not redefine it inline. The folder layout is tests/runtime_python/<arch>/sm100_<layer>/, with one pytorch_reference.py per folder containing one function per in-scope layer. Both the new test_mode test (test_<layer>_testmode.py) and the existing kernel-wrapper test (test_<layer>.py) import from the same file, so they stay aligned on a single canonical reference.

If pytorch_reference.py does not yet exist for the layer, create it. If a kernel-wrapper test already exists with an inline reference, extract that reference into pytorch_reference.py and refactor the kernel-wrapper test to import from it.

Quick Start

import torch
import mirage
from mirage.mpk.persistent_kernel import PersistentKernel

# 1. Configure
num_workers, num_schedulers = mirage.get_configurations_from_gpu(0)
params = PersistentKernel.get_default_init_parameters()
params["test_mode"] = True
params["num_workers"] = num_workers
params["num_local_schedulers"] = num_schedulers
pk = PersistentKernel(**params)

# 2. Create tensors and attach (both inputs AND outputs)
x   = torch.randn(16, 4096, dtype=torch.bfloat16, device="cuda")
w   = torch.randn(4096, dtype=torch.bfloat16, device="cuda")
out = torch.zeros(16, 4096, dtype=torch.bfloat16, device="cuda")

x_dt   = pk.attach_input(x, name="x")
w_dt   = pk.attach_input(w, name="w")
out_dt = pk.attach_input(out, name="out")

# 3. Build layer(s)
block_dim = (256, 1, 1) if pk.target_cc >= 90 else (128, 1, 1)
pk.rmsnorm_layer(input=x_dt, weight=w_dt, output=out_dt,
                 grid_dim=(16, 1, 1), block_dim=block_dim)

# 4. Compile and run — same call as production; MPK_TEST_MODE is baked into the .so
pk.compile(output_dir="./test_output")   # saves .cu and .json for debugging
pk()
torch.cuda.synchronize()

# 5. Compare against a PyTorch reference
def torch_rmsnorm(x, w, eps=1e-6):
    x_f32 = x.to(torch.float32)
    rms = x_f32.pow(2).mean(dim=-1, keepdim=True).add(eps).rsqrt()
    return (x_f32 * rms * w.to(torch.float32)).to(x.dtype)

ref = torch_rmsnorm(x, w)
print("Max diff:", (out - ref).abs().max().item())
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)

# 6. Cleanup
pk.finalize()

What Test Mode Actually Does

The launcher's first scheduler event is always EVENT_END_OF_TASK_GRAPH, so prepare_next_batch fires before iter 0. Concretely:

init_kernel:           zero step / request_ids / qo_indptr / paged_kv_indptr; seed page_queue
1st END_OF_TASK_GRAPH (iter_num=0): prepare_next_batch fills meta tensors for iter 0 → returns true
iter 0:                the test — layer-under-test runs with valid meta tensors
2nd END_OF_TASK_GRAPH (iter_num=1): prepare_next_batch finalizes (MPK_TEST_MODE always-finalize)
                                    → no new prefills (next_request_id == total_num_requests)
                                    → returns false → terminate

So tests of meta-tensor-dependent layers (paged attention, MoE routing, embedding, sampling) work — they read the values that prepare_next_batch wrote.

API Reference

PersistentKernel.get_default_init_parameters() (classmethod)

Returns a dict with safe defaults for test mode. You must set params["test_mode"] = True — it is not in the defaults.

Commonly overridden keys:

mirage.get_configurations_from_gpu(rank)

Returns (num_workers, num_schedulers) tuned for the GPU at the given rank. Always use this rather than hardcoding — the values depend on SM count and architecture.

pk.attach_input(tensor, name)

Registers a PyTorch CUDA tensor with the computation graph. Returns a DTensor for use in layer calls.

• Call for every tensor — inputs, weights, AND outputs.
• Output tensors are modified in-place when the kernel runs.
• Tensor must be contiguous (row-major / C-order).
• Name must be unique across all attached tensors.

pk.compile(output_dir=None)

Generates CUDA code, compiles with nvcc, loads the resulting .so module.

• Set output_dir to save test_rank0.cu and task_graph_rank0.json — essential for debugging compilation errors or incorrect results.
• Compilation can be slow (1–10+ minutes) depending on which task kernels are instantiated.

pk() — Launch the kernel

Same call as production. In test mode the launcher was compiled with -DMPK_TEST_MODE so it terminates after one task-graph pass. The previous pk.run_test_mode() method has been removed; use pk() directly.

• Must be called after compile().
• Does not synchronize — call torch.cuda.synchronize() before reading output tensors.
• Optional default_stream=stream kwarg if you don't want the current stream.
• Profiler export: pass params["profiler_tensor"] and optional params["trace_name"] before compile(). After pk() returns, both <trace_name>.perfetto-trace and <trace_name>.csv are written. See "Profiling" below.

pk.finalize()

Frees GPU resources (queues, events, task/event storage). Call when done.

Meta-Tensor Defaults

Test mode auto-allocates any of the 10 meta tensors that you don't pass:

total_num_requests is derived from tokens.shape[0] (defaults to 1).

Override only what your test scenario requires. Typical patterns:

# Single prefill of length N (controls qo_indptr_buffer / paged_kv_* via prepare_next_batch)
params["meta_tensors"] = {
    "prompt_lengths": torch.tensor([N], dtype=torch.int32, device="cuda"),
}

# Specific prompt content (e.g. for embedding-layer tests that read input_tokens)
params["meta_tensors"] = {
    "prompt_lengths": torch.tensor([N], dtype=torch.int32, device="cuda"),
    "tokens": torch.tensor([[101, 7592, 2088, ...]], dtype=torch.int64, device="cuda"),
}

# Multi-request batch — total_num_requests inferred from tokens.shape[0]
params["meta_tensors"] = {
    "tokens": torch.zeros((4, max_seq_length), dtype=torch.int64, device="cuda"),
    "prompt_lengths": torch.tensor([16, 8, 32, 4], dtype=torch.int32, device="cuda"),
}

The shape/dtype assertions that production runs through (e.g. tokens.shape[1] == max_seq_length, prompt_lengths.dtype == int32) all run in test mode too — defaults satisfy them by construction; user overrides will fail loudly if they don't match.