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 ```python 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: | Key | Default | When to override | |---|---|---| | `test_mode` | (not present) | **Always** set to `True` | | `num_workers` | 1 | Set from `mirage.get_configurations_from_gpu(0)` | | `num_local_schedulers` | 4 | Set from `mirage.get_configurations_from_gpu(0)` | | `max_num_batched_tokens` | 1 | Set to your test's batch size if the task kernel uses this compile-time constant | | `max_num_batched_requests` | 1 | Same as above | | `max_num_pages` / `page_size` / `max_seq_length` | 1 | Bump these so `prepare_next_batch` can fit your prefill (`max_num_pages * page_size >= prompt_length`) | | `world_size` / `mpi_rank` | 1 / 0 | For multi-GPU tests; set from `mpi4py.MPI.COMM_WORLD` | | `use_cutlass_kernel` | False | Set `True` if your layer uses CUTLASS-based kernels | | `meta_tensors` | `{}` | Auto-defaulted; **override only the entries that drive your test scenario** (typically `prompt_lengths` and/or `tokens`) — see "Meta-Tensor Defaults" below | ### `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: | Key | Default shape | Default dtype | Default content | |---|---|---|---| | `tokens` | `(1, max_seq_length)` | `int64` | zeros | | `step` | `(total_num_requests,)` | `int32` | zeros | | `prompt_lengths` | `(total_num_requests,)` | `int32` | filled with `max_num_batched_tokens` | | `input_tokens` | `(max_num_batched_tokens,)` | `int64` | zeros (filled by `prepare_next_batch`) | | `output_tokens` | `(max_num_batched_tokens,)` | `int64` | zeros | | `num_new_tokens` | `(1,)` | `int32` | zeros | | `qo_indptr_buffer` | `(max_num_batched_requests + 1,)` | `int32` | zeros (filled by `prepare_next_batch`) | | `paged_kv_indptr_buffer` | `(max_num_batched_requests + 1,)` | `int32` | zeros (filled by `prepare_next_batch`) | | `paged_kv_indices_buffer` | `(max_num_pages,)` | `int32` | zeros (filled by `prepare_next_batch`) | | `paged_kv_last_page_len_buffer` | `(max_num_batched_requests,)` | `int32` | zeros (filled by `prepare_next_batch`) | `total_num_requests` is derived from `tokens.shape[0]` (defaults to 1). **Override only what your test scenario requires.** Typical patterns: ```python # 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. ## Multi-Layer Pipeline Example Multiple layers can be chained with intermediate tensors. From the Qwen3 dense MLP pattern: ```python # Gate+Up linear → SiLU-Mul → Down+Residual # Attach weights separately, then shuffle for interleaved gate/up layout w_gate_dt = pk.attach_input(w_gate, name="w_gate") w_up_dt = pk.attach_input(w_up, name="w_up") w_gatedup_dt = pk.shuffle_tensors( inputs=[w_gate_dt, w_up_dt], shuffled_dim=0, num_groups=num_tasks // 2, name="w_gatedup", ) # Layer 1: Gate+Up fused linear pk.linear_layer(input=input_dt, weight=w_gatedup_dt, output=mlp_mid_dt, grid_dim=(num_tasks, 1, 1), block_dim=block_dim) # Layer 2: SiLU activation * element-wise multiply pk.silu_mul_layer(input=mlp_mid_dt, output=silu_out_dt, grid_dim=(num_tasks // 2, 1, 1), block_dim=block_dim) # Layer 3: Down projection + residual add pk.linear_with_residual_layer(input=silu_out_dt, weight=w_down_dt, residual=residual_dt, output=mlp_out_dt, grid_dim=(hidden_size // 64, 1, 1), block_dim=block_dim) ``` **Key pattern:** intermediate tensors (`mlp_mid`, `silu_out`) are pre-allocated and attached via `attach_input` so they can be inspected after execution if needed. For a runnable multi-task test see `tests/runtime_python/test_mode/test_diamond_fork_join_testmode.py`. ## Multi-GPU Tests Test mode supports `world_size > 1`. Each rank is independent — auto-defaults are deterministic functions of kernel params, so they produce identical values on every rank. ```python from mpi4py import MPI comm = MPI.COMM_WORLD world_size = comm.Get_size() rank = comm.Get_rank() torch.cuda.set_device(rank) params = PersistentKernel.get_default_init_parameters() params["test_mode"] = True params["world_size"] = world_size params["mpi_rank"] = rank # ... rest of setup is the same as single-GPU pk = PersistentKernel(**params) # ... attach + register layers (incl. NVSHMEM ops like pk.allreduce_layer) pk.compile(output_dir=...) pk() torch.cuda.synchronize() ```