mpk-development-norms

name

mpk-development-norms

description

The MPK team's "where does a change belong + what a clean PR looks like" norms, extracted from mirage-project/mpk merged-PR history. Read FIRST — before starting any MPK change, opening/shaping a PR, deciding which file a change goes in, reviewing a diff's shape, or cleaning up a branch that grew messy/off-norm ("改得太乱/不符合开发规范"). Complements add-mpk-model / add-mpk-task / v2-model-support (the HOW) with the WHERE + the PR-shape gate.

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".