mpk-development-norms

# MPK development norms — right place, minimal surface, clean PR These are the *change-shape* norms the maintainers actually enforce, reverse-engineered from merged PRs on `mirage-project/mpk` (see `references/exemplar-prs.md` for the cited commits and per-category file-touch tables). The sibling skills tell you HOW to add a task/model/kernel; this one tells you WHERE the change belongs and what a landable PR looks like. When in doubt, find the closest recent merged PR of the same category and mirror its footprint. The rule underneath all of them: **put a change in the file that OWNS that concern, at the smallest generic surface — not in the file that is convenient to reach from where you already are.** A diff that sprawls into shared runtime/python files to serve one model is the smell this skill exists to prevent. ## 1. Where does my change belong? (ownership map) | Change | Lives in | Must NOT touch | |---|---|---| | New GPU op / kernel | `include/mirage/persistent_kernel/tasks/<arch>/<op>.cuh` + its `tests/runtime_python/.../sm100_<op>/` unit test (+ `runtime_kernel_wrapper`) | `multigpu.py`; a model's builder | | Wiring a new task type into the runtime | C++ registration only: `runtime_header.h` (enum) → `src/kernel/{task_register,graph,runtime}.cc` → `tma.cuh` if TMA — coherently, all-or-nothing | — | | A **generic** Python op-API for that task | one `<operation>_layer` method in `python/mirage/mpk/persistent_kernel.py`, **named by the operation/algorithm, never by the model** (`moe_w13_linear_layer`, `splitk_linear_layer` — not `qwen3_*`/`deepseek_*`) | — | | Model bring-up | `python/mirage/mpk/models/<model>/builder.py` (topology, TP/EP shard rules, layer composition) + `demo/<model>/` (demo.py, HF reference, shard loader) | `persistent_kernel.py` beyond generic ops; `persistent_kernel.cuh`; `multigpu.py` | | Runtime / scheduler change | `persistent_kernel.cuh` / `runtime_header.h` / `src/kernel/runtime.cc` — as its **own** PR | a model dir; unrelated kernels | | Multi-GPU / collectives | `python/mirage/mpk/multigpu.py` — allreduce-runtime-owned (historically one owner PR) | model builders | Model composition is data, not shared code: the *order and choice* of layers, the shard-rule regexes, and the weight-name mapping are all model-specific and belong in `models/<model>/builder.py` + `demo/<model>/`. Only a genuinely reusable **operation** earns a method in `persistent_kernel.py`, and it is named for the operation. ## 2. The norms (what a reviewer checks) 1. **Right place, not convenient place.** Use the ownership map above. Ownership ≠ exclusivity: a *new task type* legitimately spans `runtime_header.h` + `src/kernel` + wrapper + `tma.cuh` (that IS its home); a runtime fix may touch a task `.cuh` when the invariant crosses the worker/task boundary. What's off-norm is reaching into a shared file to serve one model. 2. **Minimal shared-surface diff.** A model-support PR does **not** touch `persistent_kernel.cuh` or `multigpu.py`, and touches `persistent_kernel.py` only to add/fix a **generic** operation-level primitive. Shared APIs are named by operation/algorithm, never `<model>_*`. (Counter-smell this catches: `deepseek_mla_rope_q_layer`, `mla_kv_gather_unified_layer`, `dsv3_router_gate_gemv_layer` added to the shared file — those belong behind a generic API called from the model builder.) 3. **No experiment env-vars in landed code.** In `persistent_kernel.py` the only `os.environ` uses are the **5 build-path/infra vars** (`MIRAGE_HOME`, `NVSHMEM_INC_PATH`, `NVSHMEM_LIB_PATH`, `MPI_INC_PATH`, `MPI_LIB_PATH`). No `MPK_*_DBG` / `*_PROBE` / `*_GUARD` / `FASTFWD` / campaign perf toggles survive into a merged PR — a perf lever is either **hard-wired to its chosen production value** (as a named constant) or absent, and a debug/diagnostic knob is deleted with its code path. (This norm is about landed code; an *in-flight exploration* branch keeps levers env-gated default-OFF — see `mpk-lever-cleanup` for the collapse step.) Audit isn't limited to `persistent_kernel.py`: check the builder, `runtime.cc`, and C++ `getenv` debug hooks for the same residue. 4. **Runtime changes are separate, coherent PRs** — not bundled inside a model or kernel PR. `#411` (Split persistent kernel) touched exactly 2 files. If your model work needs a runtime fix, split it into its own PR so the maintainer can take/defer it independently. 5. **One PR = one coherent topic.** PRs are squash-merged (one commit each). A focused bugfix is often a single file (`#719`). Don't fold a de-cruft, a perf lever, and a new kernel into one diff. Commit granularity mirrors this even pre-squash: each commit is one reviewable idea. 6. **Comments are sparse and functional.** ~5% comment lines in kernels, ~10% in a builder — they say *what a non-obvious line does*, never a perf-campaign diary, a "we tried X" history, or narration of the obvious. No commented-out code. 7. **Tests are part of the change shape.** A new kernel/task ships its `tests/runtime_python/.../test_*_testmode.py` (usually + a `pytorch_reference.py` and, if needed, a wrapper/`setup.py`). A PR that adds a kernel with no test is off-norm. 8. **Registration/ABI is coherent.** A new task ID updates `runtime_header.h`, the `src/kernel` registration, the wrapper, and TMA/runtime glue **together** — never a dangling enum with no register/graph handler, never a handler for a deleted enum. (Fail-loud: rebuild after any enum edit — a stale enum silently mis-dispatches.) 9. **Format + no artifacts.** Run `bash scripts/format.sh` (clang-format-15, CI-enforced) before pushing. Never stage generated/local material: `scratch/`, `outputs/`, `_results/`, weight caches, generated `test.cu`/`.so`, perf logs, `PR_DESCRIPTION`/campaign notes, `.claude/` (except the sanctioned `.claude/skills/**` + `.claude/agents/**` on a skills PR). 10. **No gratuitous assertions / error-throwing.** Before adding ANY `assert` / `raise` / `throw` / `abort` / fail-loud check, ask: (a) did upstream have it? (b) is it necessary? (c) does *omitting* it have a **correctness** consequence — a silently-wrong result, not merely a later natural error? If (b)/(c) are "no", don't add it — **default to not adding.** Seemingly-correct defensive throws have caused real breakage: they fire on *valid* states and mislead debugging (a real case: `assert(params.size()==0||3)` that rejected the valid 1-param call the reader itself was written to handle). Keep a check only when it guards a real, demonstrated failure **or** a silently-wrong path (wrong-kernel selection, a BF16/FP8 fork), and even then prefer the existing/upstream idiom over a new fail-loud `abort`. A check that only pretty-prints an error the very next line would raise anyway (a `KeyError`, a dtype error) is pure surface — drop it. Config guards that merely **restate a predicate the caller already checked** are the archetype to delete. Same test for host launch/return-code checks: if upstream launched without the check and omitting it just defers to the next CUDA error, it's surface. 11. **No gratuitous renames / type-descriptors on working code.** Don't rename existing symbols (functions, params, enum symbols) or **renumber a task-type enum** or bolt on type annotations / descriptor fields / "API-parity" wrapper params to code that already runs — unless that change is itself the point. Two distinct breakages: a symbol rename breaks source/API references (external callers, imports); **renumbering** a task-type enum (changing its integer value) breaks already-**serialized** task graphs, because `task_type` is serialized numerically — a surviving TP8-only reducer keeps its upstream id, it is not re-slotted into a deleted variant's number. A `mpk: "PersistentKernel"` annotation or a `TYPE_CHECKING` import is inert at runtime and adds a dependency edge for nothing. If it ran upstream without the rename/annotation/wrapper, don't add it. Accepting-then-discarding params (`del eps, epsilon # API parity`; a `group_size` arg that only exists to be rejected when != 128) is the same smell — unused surface that only exists to be validated away. Revert to the upstream name/shape. ## 3. PR-shape checklist (run before you open/push) - [ ] Every changed file is the **owner** of its concern (ownership map). No shared-file reach for a single-model need. - [ ] `git diff --stat <merge-base>..HEAD` — is the shared-surface footprint (`persistent_kernel.cuh`, `persistent_kernel.py`, `multigpu.py`, `runtime.cc`) as **small and generic** as the closest exemplar PR? Any `<model>_*` method in `persistent_kernel.py`? - [ ] Env-var count in `persistent_kernel.py` back to the 5 build-path vars (no campaign toggles anywhere in the diff)? - [ ] Runtime/scheduler changes split into their own commit/PR? - [ ] Each commit one coherent topic; message states mechanism + (for perf) measured Δ; ends with the required `Co-Authored-By` line? - [ ] New kernels/tasks carry their test-mode test + reference? - [ ] Registration coherent (enum ⇄ register ⇄ graph ⇄ wrapper), rebuilt clean? - [ ] No **added** assert/raise/throw/abort that fails the norm-10 test (guards nothing demonstrated or silently-wrong; restates a caller predicate; pretty-prints an immediate natural error)? - [ ] No **rename** of an existing symbol / enum name-or-value / task-type ID, and no inert type-annotation / descriptor / API-parity param added to working code (norm 11)? - [ ] `scripts/format.sh` clean; no generated/local artifacts staged; sensitive-grep before push? ## References - `references/exemplar-prs.md` — the cited merged PRs per category, with their file-touch tables (the empirical basis for every claim above). Mirror the closest one. - `references/codex-checklist.md` — a self-contained, tool-agnostic review checklist (no Claude/skill framing) you can paste into `codex exec` (or hand a human reviewer) to score a diff against these norms. Feed it the diff + "review against this checklist".