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. ## Phase (c) — DEBUG (read `references/debug-gates.md` BEFORE debugging anything) The distilled ladder — full checklists in the reference: 1. Full-layer TOKEN-MATCH first; NEVER judge correctness from few-layer coherence. 2. Broken → diff vs the CHAIN stage-by-stage (clean token position, FULL vectors). 3. Garbage at TP>1 + gate/TP1 fine ⇒ suspect a MISSING CROSS-RANK COLLECTIVE first — a single-rank gate is structurally blind to it (6 gate-fidelity classes). 4. Token-identity only on DETERMINISTIC configs. The TP8 FFN atomicAdd path is nondeterministic → use the 3-part gate: deterministic canary + NaN poison-fill + full-model coherence-in-envelope (with an OFF1-vs-OFF2 control). 5. iter-0-fine / iter-1-hang ⇒ a PERSISTENT state got re-initialized (monotonic barrier + `skip_after_step0` class), not a missing event. 6. Hangs: watchdog (names the hung task) > breadcrumb (crash-only; in-flight counts are base-rate artifacts). Illegal address: compute-sanitizer is ground truth. 7. "Dead task / safe to remove" claims: box token-identity A/B is the ONLY ground truth — static analysis + reviewers have been wrong. Run the `mpk-correctness-gate` agent before trusting any baseline and before every math-changing commit. ## Phase (d) — WORKFLOW orchestration (perf loop, after correctness) **The full v2-updated loop is the sibling skill `v2-perf-iteration` (`.claude/skills/v2-perf-iteration/`) — load it to run this phase; the summary below is orientation only.** The multi-agent loop, unchanged: profiler → (analyzer) → planner → iterator → [ablation-logic-reviewer] → implement → correctness-gate → profiler → commit-reviewer → commit → memory-keeper → decide. Standing disciplines: - EVERY non-trivial conclusion through `ablation-logic-reviewer` + a Codex MCP cross-check before acting on it (the over-claim guard; defaults params only). - Every lever lands env-gated default-OFF; DEFAULT BUILD BYTE-IDENTICAL (the whole v2 wiring itself followed this — `--use-v2` opt-in). - Verdict metric at the PRODUCTION config (DSv3: bs=1 TP8 e2e tpot); smaller TP is triage only. Slowest-CTA per-position, never P50/per-kernel aggregates. - `experiment_history/` contract: journal + INDEX row after every experiment, ESPECIALLY NULL/REGRESS (anti-loop) — via `mpk-memory-keeper`. - GPU-safety: never crash-loop the megakernel (D-state zombies), memory-cap every launch, box sessions per `references/box-orchestration.md`. ## Worked example #2 — Qwen3-8B on v2 (in-tree + upstream; the dense single-GPU shape) A COMPLETE second instance of this pipeline's endpoint already exists for a dense model, and it is the natural starting point for any dense/single-GPU v2 campaign (e.g. Qwen3-8B throughput work): - **In-tree (this branch)**: `demo/qwen3/demo.py` has `--use-v2` (argparse ~:127; `use_v2_runtime=args.use_v2` into `PersistentKernel` ~:352). The graph is built INLINE in the demo (the `python/mirage/mpk/models/qwen3/builder.py` GraphBuilder has NO v2 branches — a wiring-style difference vs DSv3's builder-side gating). The v2 branches swap exactly the GEMM-shaped ops to the Channel-based per-tile linear family, `tiles_per_task=1`: qkv_proj + gate_up → `linear_layer_v3` (~:544, ~:711), o_proj + down_proj → `linear_with_residual_layer_v3` (~:637, ~:744), lm_head → `linear_layer_v3` (~:795) (`TASK_LINEAR_SM100_V3` = 244 / `_WITH_RESIDUAL_` = 245). Everything else keeps its task name and runs as the v2 role variant: rmsnorm (`TASK_RMS_NORM_HOPPER_V2` 326), paged attention (`TASK_ATTN_SM100_V2` 329, consumer-only), silu_mul, embedding, argmax partial/reduce. Task-plan wiring is EXPLICIT at demo level (~:850-854): `task_graph["v2_worker_task_queues"] = build_v2_worker_task_queues(...)` + `add_v2_region_smem_plan(...)` before `mpk.compile()` — the older of the two wiring styles (DSv3 relies on `compile()` doing both internally; see `references/demo-stage.md` §7). - **Single-GPU capable**: yes — the demo runs at `world_size == 1` with a local `--model-path` or the HF default `Qwen/Qwen3-8B`; no NVSHMEM collectives exist at TP1, so no box is needed. The tracked calibration script `tests/runtime_python/blackwell_v2/e2e_qwen3_check.sh` runs v1-vs-`--use-v2`