Mpk Lever Cleanup

MPK Lever Cleanup — consolidate env-gated optimizations into one clean path

During exploration, every MPK performance optimization is an env-gated, default-OFF lever (#if MPK_DSV3_XXX + the #else legacy path + os.environ.get(...) → -D injection in persistent_kernel.py). Once the work has settled and is headed for mainline, the levers must be consolidated into a single path: hard-wire the winners as the default, delete the legacy code, remove the control variables. This skill is the complete procedure — plus the pitfalls actually hit — from that wrap-up refactor.

Core mindset: this is a refactor that intentionally changes the default build — the default path flips from "the safe legacy logic" to "the optimized path". The usual "default build must stay byte-identical" commit gate is therefore deliberately waived here; that is the whole point of this commit.

Procedure (7 steps, in order)

1. Enumerate all gates

grep -rhoE "MPK_(DSV3_)?[A-Z0-9_]+" <the megakernel .cuh files> <builder.py> <persistent_kernel.py> \
  | sort -u | grep -vE "MPK_(MAX|PAGE|PROFILING|NUM)"

List both the #ifdef gates in the .cuh files and the os.environ/-D injections in persistent_kernel.py.

2. Classify every gate (VERIFY each one — never assume default-OFF)

⚠️ Classification must be verified, never recalled from memory. Pitfalls actually hit: TOPK_PARALLEL, assumed default-OFF and slated for revert, is in fact ON in every decode build (part of the current winning stack — KEEP); ROUTER_GEMV, assumed obsolete, is in fact the router of the TP<8 fallback (must not be deleted). Use grep -n on the injection condition in persistent_kernel.py and check whether the gate is actually exercised at the production geometry.

3. Codex-vet the classification + refactor plan

Hand the gate inventory + classification + goal to Codex (mcp__codex__codex) for a multi-round discussion: validate the classification, agree on a safe execution order, the correctness risks, how to verify the ABI-revert, and the structural-vs-leaf distinction. Codex will catch conflicts in the classification (see the step-2 pitfalls).

4. Freeze the reference (the correctness comparison baseline)

Before changing anything, run the current winning stack (all levers ON) and record its output: e2e tpot + logits/prose (run it twice to get the A/A nondeterminism envelope). The final "clean default build" is compared against this winning-stack reference, not against the old safe default.

5. Execute in the safe order (build-check after every step)

1. First revert the dead levers that changed the ABI (most dangerous; do it in isolation and verify). If a dead lever touched the (num_in,num_out,TASK_ENUM,variant) tuple in graph.cc/task_register or added a tensor, a single ABI mismatch = "Invalid global read" at runtime. If those changes were never committed, simply git checkout HEAD -- <producer files> to return to a clean ABI, then:

   git diff HEAD -- graph.cc task_register.cc tasks.py multigpu.py allreduce.cuh | wc -l   # expect ≈ 0
   grep -rn "tile_sumsq|<sidecar tokens>|input_ptrs\[N\]" <files>                          # expect 0
   

   Delete the dead lever's consumer half at the same time (otherwise it is a stale-env out-of-bounds landmine).
2. Delete the remaining dead levers + all diagnostic probes → build-check.
3. Hard-wire the LEAF wins (leaf optimizations with a clean #else): remove the gate, delete the #else, keep the geometry guard → build-check.
4. Hard-wire the STRUCTURAL wins last (path-selectors that change the graph shape / large control flow): delete the entire alternate path, keep the TP8/mbt/workers guard → build-check + a --layers 0-3 in-MPK smoke. (Structural gates are more dangerous than leaves — removing one deletes a whole alternate code path. Do it after the tree has already shrunk.)

6. Verify correctness (the default path's math has changed)

• Token-identity cannot be used (DSv3 TP8 decode is FP-nondeterministic — cross-CTA atomicAdd).
• Use instead: the A/A envelope (winning stack compared against itself) + clean default vs the winning-stack reference with the per-step logit-cosine inside the envelope + stable top-k overlap + 512 tokens of coherent prose + no NaN/Inf.
• Perf smoke: e2e tpot should ≈ the winning-stack reference (confirms no win was silently dropped).
• TP8 JIT smoke: --layers 0-3 confirms the hard-wired megakernel really instantiates and runs (the #else-deleted path only instantiates at world_size==8).
• Qwen3 / TP4 regression smokes protect the untouched fallback / non-DSv3 paths.

7. Commit one clean version (for the PR)

• Stage source files only (kernels/.cuh, builder.py, persistent_kernel.py, task_register.cc); exclude .claude/, scratch/, experiment_history/, CSVs/outputs/.pk_compile and other local artifacts.
• Run mpk-commit-reviewer and tell it explicitly that the default-build change is intentional (otherwise it will BLOCK per the standard gate); it still checks staged-path hygiene, the allowed surface, the message, and the correctness story.
• The commit message lists everything: which levers were hard-wired (+ each one's Δ), which dead levers were reverted, which diagnostics were deleted, that the ABI was restored, the verification evidence, and an explicit pre-merge gate (if TP8 runtime validation was blocked by box capacity, write it into the message as a must-run item before merging). Include Co-Authored-By.

Key pitfalls (all hit in practice)

• The default build is intentionally NOT byte-identical — that is the goal, not a bug; waive that one commit gate.
• The ABI-revert is the most dangerous step — do it first and in isolation, grep it clean, diff against the last clean commit, and delete the consumer half together with it.
• Structural wins go last — a path-selector is more dangerous than a leaf-opt (removing it deletes an entire alternate path).
• Classification must be verified — some gates are already unconditionally ON, some are fallbacks; never assume "default-OFF".
• Orphaned legacy functions — after deleting the #else call site, the __device__ function definition may remain (nvcc elides it; harmless but unclean); either delete it or flag it as a known nit in the PR.
• The correctness gate = A/A envelope + coherence, not token-identity (the path is FP-nondeterministic).
• The TP8 runtime gate may be blocked by box capacity — a commit meant for PR review may land with a documented pre-merge gate (a PR is review, not auto-merge); never fabricate numbers.
• Sub-agents can contradict each other about the same fact (e.g. "the lever was deleted" vs "the lever was hard-wired") — verify yourself with grep -c <the win's body symbol> that the win's body is still present (zero refs to the macro ≠ the win was deleted; possibly only the gate was removed and the body became unconditional).