V2 Model Support

V2 Model Support — compute graph → working Runtime-V2 demo

This is the GENERIC pipeline for putting ANY model on MPK Runtime-V2. It was distilled from the campaign that took DeepSeek-V3 decode from a v1-only model to a full-61-layer Runtime-V2 e2e run at TP8 EP2 bs=1 (commit e31b34dd, opt-in --use-v2, default build byte-identical). DSv3 examples below are clearly labeled worked-example material — the recipe does not depend on them; the staged plan that drove that campaign is archived at references/V2_DSV3_DECODE_MASTER_PLAN.md (M0→M5 ladder). A SECOND, smaller worked example — Qwen3-8B on v2, dense, single-GPU-capable — is in §"Worked example #2" below (it is the closest starting point for a dense-model campaign). It is a context+phased-recipe skill: architectures vary, the PHASES and GATES do not.

Read mpk-development-norms FIRST. This skill is the HOW (graph→plan→demo); that one is the WHERE + the PR-shape gate that decides what lands cleanly on mpk when the campaign is done — model code in models/<model>/builder.py + demo/<model>/, only GENERIC ops (never <model>_*) in shared persistent_kernel.py, no experiment env-vars in landed code, runtime fixes as separate PRs. During exploration keep levers env-gated default-OFF (mpk-lever-cleanup); before opening the PR, conform to the norms.

Environment prerequisites (what must exist on the machine)

In-repo (travels with every clone): this skill's references/, the v2 runtime + kernels (include/mirage/persistent_kernel/, tasks/blackwell_v2/), the harness (tests/runtime_python/blackwell_v2/), both worked-example demos (demo/deepseek_v3/, demo/qwen3/), the repo agents (.claude/agents/*.md), and the sibling skills. Machine-local (keep working without them, as noted):

• Remote multi-GPU box — only for multi-rank gates + verdict configs. The box CLI (~/nebius_box.sh) and the machine inventory in references/box-orchestration.md §1-2 are SITE-SPECIFIC (placeholders in-repo; actual IPs/users/keys/paths stay in operator-local notes — never committed); the structural rules there (§3-8: rsync/build, session discipline, safety, testing tiers) transfer to any box. No box ⇒ single-GPU phases still run end-to-end (Qwen3-class models need no box at TP1).
• Local GPU(s) — needed from the first harness gate onward (graph-build/test-mode gates are 0-GPU). Torch-probe cards before use.
• Model checkpoints — site-specific paths; the input contract below is what matters.
• User-level agents (~/.claude/agents/: mpk-perf-analyzer, ablation-logic-reviewer, codex-task-dispatcher) — present only on the same user account. The repo-level roster (.claude/agents/mpk-*, v2-*, ferret-*) travels with the clone. If ablation-logic-reviewer is missing, run the review discipline with a general-purpose subagent given its first-principles brief.
• Codex MCP (mcp__codex__codex) — the cross-check second engine; .mcp.json is machine-local (git-ignored). If unconfigured, reviews degrade to subagent-only — say so.
• Personal memory (~/.claude/projects/-home-muhengl-mirage/memory/) — optional context, same-account only. The load-bearing lessons are already distilled into this suite's docs.
• experiment_history/ — git-ignored, so a FRESH CLONE STARTS EMPTY. That is expected: create INDEX.md + the journal on first use (contract in Phase d); the anti-loop evidence that must survive clones lives in v2-kernel-writing/references/m1-decode-evidence.md.

Input contract (what you need before starting)

1. Model config — hidden size, layer count/types (dense vs MoE, attention variant), head geometry, vocab, dtype/quantization (FP8 block-scale? BF16?), norm eps.
2. HF checkpoint — safetensors; know which weights need conversion/absorption (e.g. MLA absorbed q_b_proj) and which need requantization.
3. Parallelism plan — TP degree, EP degree (routed_tp_size = world/EP), which linears are Column vs Row parallel, where the AllReduces land, vocab-parallel or replicated lm_head.
4. The compute graph — a draw.io file (ops as nodes labeled name+shape+dtype, edges = tensor deps, TP-sharding annotations). Format + parsing: references/graph-to-plan.md. If none is supplied, derive the graph from the HF modeling_*.py and WRITE the plan doc as if you had one — the plan doc is the contract, the drawio is just its serialization.
5. Target machine(s) — which box runs the multi-GPU verdict configs (references/box-orchestration.md). Single-GPU/local for micro-gates only.

Why phased (the one-sentence history)

DSv3-on-v2 succeeded because it NEVER registered a live-path task without a complete v2 consumer body (a bodyless consumer silently deadlocks the whole box — see the §1.1 trap below), validated each op bit-exact in isolation BEFORE it entered the graph, and went e2e at the smallest possible slice (1 MoE layer) before scaling. Every shortcut attempted around this ladder cost days (9 debug rounds on a missing AllReduce; an iter-1 hang from re-zeroing a monotonic barrier). Do not reorder the phases.

The pipeline

Phase 0  GRAPH→PLAN      drawio → op inventory → classify {reuse|new-kernel|fused-later}
Phase a  DEMO            builder-first: weights/SHARD_RULES/cache-key/lifetimes; chain of
                         existing v2 tasks; graph-build + test-mode gates (0-GPU first)
Phase b  KERNEL          per-op loop for missing/slow ops; M0→M5 staged bring-up
Phase c  DEBUG           the gate ladder (token-match first, TP-collective blind spot,
                         nondeterminism protocol, hang triage)
Phase d  WORKFLOW        the multi-agent perf loop + box orchestration + history contract

Phase 0 — GRAPH→PLAN (read references/graph-to-plan.md)

1. Parse the draw.io XML → op list. Each op row: name | op kind | input/output shapes AT THE CHOSEN TP/EP | dtype(s) | weight source key(s) | conversion/absorption needs | collective (none / AR / reduce-scatter / EP-dispatch) | grid intuition.
2. Derive PER-RANK shapes: apply the sharding plan (Column→shard N, Row→shard K + AllReduce after, EP→local expert slice). The DSv3 worked table is in the reference.
3. Classify EVERY op:
   • REUSE — an existing v2 task covers it (rmsnorm_v2, silu_mul_v2, linear v2/v3 with its M<=16 contract, embedding_v2, argmax_partial/reduce_v2, nvshmem_tile_allreduce_v2(+residual), tensor_init_v2, mul_sum_add_v2 …). Check runtime_header.h enums 242/243, 326-355 and tasks/blackwell_v2/.
   • NEW KERNEL — no v2 variant exists → Phase-b item. Tag its port kind: leaf (role-split trivial) / collective (v2-safe sync rewrite — HARD) / megakernel-shape Form-2 (num_tasks==num_workers, in-op GMEM barrier).
   • FUSED-BLOCK CANDIDATE — defer; fusion comes only after chain correctness.
4. Output = the model plan doc: the op table, the v2-ABSENT set, the milestone ladder, the risk ranking. Location: .claude/skills/v2-model-support/references/ V2_<MODEL>_MASTER_PLAN.md if it should travel with the repo, scratch/ (git-ignored) for throwaway drafts. Mirror the archived DSv3 instance (references/V2_DSV3_DECODE_MASTER_PLAN.md). This doc is what the phase leads execute against.

Phase (a) — DEMO: builder-first bring-up (read references/demo-stage.md)

Goal: demo/<model>/demo.py + python/mirage/mpk/models/<model>/builder.py that BUILD the graph (no GPU needed yet) and pass test-mode with existing tasks.

• CHAIN-FIRST rule: assemble every layer from existing generic v2 tasks even if slow. Fused megakernels are Phase-b/perf work. A correct chain is your ground truth for every later diff (FUSED_KERNEL_DEBUG_METHODOLOGY.md step 2 depends on it).
• Weight mapping: SHARD_RULES live in TWO places — the builder AND the demo's conversion pass. Update both, always.
• Cache-key contract: the weight cache key hashes config, NOT conversion code. Bump the format-version string on ANY conversion-logic change (silent stale weights is the worst failure). Details + MPK_CONVERT_SEMAPHORE / MPK_BUILD_CACHE_ONLY in the reference.
• v2 wiring is 90% automatic: pass use_v2_runtime=args.use_v2 into PersistentKernel(...); compile() itself runs the v2 queue plan + SMEM plan + the §1.1 deadlock guard (persistent_kernel.py ~:5619-5640). Builder-side work is selecting v2 task names (most *_layer wrappers self-switch on self.use_v2_runtime) and the v2-only allocations (scratch sizing, scale packs).
• Gates before ANY GPU run: (1) graph-build succeeds for the smallest slice; (2) the §1.1 guard passes — v2_unsafe_task_types empty (every graph-used task type has a v2 role variant, else compile() raises instead of wedging the box); (3) test-mode (0-GPU-graph-build + single-pass CPU-launchable subset) green.

Phase (b) — KERNEL: per-op loop on the M0→M5 ladder

For each NEW KERNEL op, dispatch the sibling skill v2-kernel-writing (.claude/skills/v2-kernel-writing/ — the per-kernel inner loop this pipeline plugs into: SPEC→IMPLEMENT→WIRE→VALIDATE→PERF→REVIEW) with the op's spec row from Phase 0 (roles / SMEM regions / sync / correctness reference / validate step — the §3 template in the DSv3 master plan). For pure kernel-PERF rewrites of an op that already passes correctness, ferret-kernel-system/kda-kernel-agent + mpk-faithful-gate are the measurement-honest routes. Every new _v2 task touches: runtime_header.h enum + task_register.cc register_*_v2_task (consumer body MUST begin with emit_dep_wait_consumer_prefix) + graph.cc dispatch + runtime.cc task_type_to_name (+ the task_offset = bid.x block for fused megas) + the .cuh/_spec.h pair in tasks/blackwell_v2/ + the persistent_kernel.py wrapper's "..._v2" if self.use_v2_runtime switch + the builder call site.

Stage the bring-up on the PROVEN ladder (mirror references/V2_DSV3_DECODE_MASTER_PLAN.md — DSv3 worked example; a dense single-GPU model collapses M1 to "none" and M4/M5 shrink):

• M0 — leaves + tail wiring. Trivial role-split tasks (tensor_init-class), tail re-routes (lm_head/argmax path), confirm already-present v2 leaves are reachable from THIS model's builder. Validate in the tests/runtime_python/blackwell_v2/ harness (per-op, deterministically-seeded, vs fp32 torch ref AND vs the v1 twin).
• M1 — collectives at TP2 first. The AllReduce-class ports are the highest-risk items (block-wide __syncthreads()/256-thread bodies vs the 128-thread consumer role → deadlock/half-compute if pasted). Validate on a TP2 micrograph (2 ranks, known vectors, bit-exact sum on both ranks), then TP8. Do this EARLY — it de-risks everything downstream and is the first multi-rank v2 proof.
• M2 — fused blocks. Megakernel-shape (Form-2) ports: check first whether an existing v2 mega can be REUSED via a builder re-route (DSv3's FFN was — kernel already proven, the work was builder-side tensor packing). Each fused mega: bit-match vs its v1 twin in a TP-shaped harness, THEN a small-slice in-MPK smoke (the __align__(1024) extern-smem footgun is only caught in-MPK).
• M3 — FIRST E2E at the smallest real slice (DSv3: --use-v2 --layers 3-3 TP8 EP2 bs=1, --disable-vocab-parallel-lm-head to stay on present tail tasks). TWO hard pre-conditions: (1) reachability diff — build the sliced graph in test-mode and diff the task list vs a full build so no head/tail seed task is silently dropped; (2) §1.1 guard green. Correctness = the Phase-c protocol.
• M4 — scale layers up (+ restore any deferred tail variant). Expect the cold-convert OOM class here at full TP — MPK_CONVERT_SEMAPHORE=K.
• M5 — remaining layer types (DSv3: dense layers 0-2 via one Form-2 task) → full-model e2e. Deliverable: full-layer coherent decode on v2 + tpot vs v1.

Per-kernel gate: test-mode numeric PASS (cos ≥ 0.999, rel_max ≤ 3e-2, no NaN, and bit-exact-vs-v1 for elementwise ops) BEFORE the task enters any e2e graph.