nvfp4-megamoe-kernel/CLEAN_UP.md

DSV4 Repo Cleanup & Comment Audit — Agent Working Spec

Audience: the LLM agent doing the cleanup. Prime directive: the running code is the source of truth. Docs, .md files, and comments are not. When they disagree, the code wins and the prose gets corrected — never the reverse.

Two hard rules that exist because of past pain:

1. Never delete. Only move/archive. Especially .md files — they contain lessons we still reference.
2. Every time you move a file, update the references in the same commit, then grep the moved basename repo-wide to confirm zero dangling references. The recurring failure mode here is: a file is moved, a reference is missed, the next agent thinks the file is gone, and recreates a divergent copy. That is how this repo got two of everything. Do not let it happen again.

---

Background the agent must internalize first: this repo has TWO lineages

There are two parallel implementations of the model, and the docs describe the wrong one.

	Lineage M (LIVE)	Lineage P (parallel / maybe-serving)
Entry point	single_shot_inference.py (monolith)	dsv4/model/dsv4.py (nn.Module assembly)
Orchestration	manual, inside the script	dsv4/model/layer.py + dsv4/layers/*
Indexer	inline PyTorch einsum in the script's Indexer.forward	dsv4/kernels/indexer/* package
Compressor / KV cache	the script's own Compressor / KVCache classes	dsv4/cache/*, dsv4/kernels/cache/*
Produces coherent output?	Yes — this is what runs	Unconfirmed; dsv4/model/dsv4.py has 0 in-repo importers

single_shot_inference.py is the live path. It imports a subset of dsv4/ primitives and reimplements the rest itself. Lineage P (dsv4/model/dsv4.py + the dsv4/layers/{attention,ffn,embedding,norm} nn.Modules + dsv4/kernels/{indexer,router,cache}) is either the vLLM/sglang integration surface or dead. You cannot tell from inside the repo.

→ Step 0 below resolves this. Do not archive anything in Lineage P until Step 0 is done.

---

PART 1 — Repo Cleanup

Step 0 — Establish the canonical entry points (do this FIRST, before moving anything in dsv4/)

The cleanup is only safe once you know what's reachable. There are (at most) two roots:

• Standalone: single_shot_inference.py.
• Serving: whatever the modified vLLM at /root/dsv4-nvfp4-workspace/vllm imports from dsv4. Find it:

grep -rn "import dsv4\|from dsv4" /root/dsv4-nvfp4-workspace/vllm 2>/dev/null

If that comes back empty, then dsv4/model/dsv4.py and all of Lineage P are not used by serving either → they are archive candidates (Step 2). If it imports dsv4.model.dsv4 (or anything in Lineage P), then Lineage P is live for serving and must be kept, not archived.

Build a reusable "is this file dead?" tool (the durable fix for the recreate problem)

Drop this in helpers/import_closure.py. It computes the import closure from the entry points and prints every dsv4/*.py not reachable. Run it before archiving anything, and any time an agent claims a file is unused.

# helpers/import_closure.py — list dsv4 modules NOT reachable from the entry points.
# Usage: python helpers/import_closure.py  (run from repo root, PYTHONPATH=repo root)
import ast, pathlib, sys
ROOT = pathlib.Path(__file__).resolve().parent.parent
ENTRYPOINTS = ["single_shot_inference.py"]  # + add the vLLM glue module if Step 0 found one

def module_to_path(mod):
    p = ROOT / (mod.replace(".", "/") + ".py")
    if p.exists(): return p
    p = ROOT / mod.replace(".", "/") / "__init__.py"
    return p if p.exists() else None

def imports_of(path):
    tree = ast.parse(path.read_text())
    out = set()
    for n in ast.walk(tree):
        if isinstance(n, ast.Import):
            out |= {a.name for a in n.names}
        elif isinstance(n, ast.ImportFrom) and n.module:
            out.add(n.module)
    return {m for m in out if m.startswith("dsv4")}

seen, stack = set(), list(ENTRYPOINTS)
stack = [ (ROOT / e) for e in stack ]
while stack:
    f = stack.pop()
    if f in seen or f is None or not f.exists(): continue
    seen.add(f)
    for m in imports_of(f):
        mp = module_to_path(m)
        if mp and mp not in seen: stack.append(mp)

all_py = set((ROOT / "dsv4").rglob("*.py"))
dead = sorted(p.relative_to(ROOT) for p in all_py - seen if "__pycache__" not in str(p))
print("REACHABLE:", len(seen), " | DEAD CANDIDATES:", len(dead))
for d in dead: print("  ", d)

This is the anti-recreate safeguard. Wire it into the agent's pre-commit habit: "before deleting/archiving a module, prove it's dead with import_closure.py; before creating a 'missing' module, prove it doesn't already exist with grep -rn <basename> ."

---

Step 1 — Root-level files

Only single_shot_inference.py stays in root (plus standard project files). Verified: all the test/probe/dump scripts below have 0 inbound imports, so moving them needs no code changes — they are run directly with PYTHONPATH=<repo root>, which still resolves their from dsv4 ... imports from any location. Their hardcoded /root/nvidia-meeting/... checkpoint paths are runtime data paths, unaffected by the move.

File	Action	Destination	Code changes needed
single_shot_inference.py	keep	root	—
README.md	keep	root	(but see Part 2 — its package-structure section is stale)
pyproject.toml, Dockerfile, docker-compose.yml, build_and_run.sh, .gitignore, .dockerignore	keep	root	—
PERFORMANCE_AUDIT.md	move	docs/	none (doc)
test_se_dequant.py	move	tests/integration/	none (0 importers)
test_se_gpu.py	move	tests/integration/	none
test_se_l1_direct.py	move	tests/integration/	none
test_se_multi_gpu.py	move	tests/integration/	none
test_gemm_1group.py	move	tests/integration/	none
test_quantize_gpu.py	move	tests/integration/	none
hf_reference_test.py	move	tests/integration/	none
probe_hf_indexer.py	move	helpers/ (new)	none
probe_indexer_shapes.py	move	helpers/	none
probe_keys.py	move	helpers/	none
probe_shapes.py	move	helpers/	none
dump_checkpoint_keys.py	move	helpers/	none
single_shot_PYTORCH_REFERENCE.py	move	dsv4/reference/	YES — 3 edits, see Step 3

mkdir -p helpers (no __init__.py needed; these run as scripts). tests/integration/ and dsv4/reference/ already exist.

The tests/integration/ items load the real checkpoint — keep them if they still pass, send them to tests/archive/ if superseded. That's a judgment call for the human, not an auto-archive.

---

Step 2 — dsv4/ internals

2a. .cu duplication — the loader only ever looks in kernels/cuda/

dsv4/kernels/cuda/loader.py resolves every .cu relative to dsv4/kernels/cuda/, regardless of which Python file calls get_cuda_module. So any .cu sitting in a semantic subfolder (indexer/, etc.) is never compiled — it's dead. Confirmed dead duplicates:

Dead copy (never compiled)	Live copy (what actually compiles)	Status
dsv4/kernels/indexer/indexer_score_topk.cu (292 lines)	dsv4/kernels/cuda/indexer_score_topk.cu (166 lines)	DIFFER — do not blind-delete
dsv4/kernels/indexer/gather_kv.cu (106 lines)	dsv4/kernels/cuda/gather_kv.cu (121 lines)	DIFFER — do not blind-delete

Procedure (because they differ): diff each pair. Decide which is the intended version. The subfolder copy may actually be a newer improvement that's silently dead because the loader can't reach it. If the subfolder copy is the better one, copy it into kernels/cuda/ first (so the live path gets the fix), verify, then delete the subfolder copy. Do not assume "live == canonical."

Decision to make (human): either (a) keep the flat convention — all .cu live in kernels/cuda/, delete subfolder .cu after reconciling — which matches the loader and needs no Python changes; or (b) teach loader.py to accept subdir-qualified source paths and move .cu into semantic folders. (a) is lower risk. Pick one and make loader.py's docstring say which.

2b. Dead-code / orphan modules (archive candidates, gated on Step 0)

From the import-graph scan, these dsv4/ modules have 0 in-repo importers. Confirm with import_closure.py and the Step 0 vLLM check, then move to a new dsv4/_archive/ (mirror the subpath) rather than deleting:

• dsv4/model/dsv4.py ← 0 in-repo importers. This is the "full model." If Step 0 shows vLLM imports it, it is LIVE — keep. Otherwise archive.
• dsv4/model/mtp.py
• dsv4/layers/embedding.py
• dsv4/kernels/indexer/csa_indexer.py (the live indexer is inline in single_shot_inference.py; this is Lineage P)
• dsv4/kernels/router/nvfp4_fused_router_kernel.py
• dsv4/ops/topk.py, dsv4/ops/topk_select.py, dsv4/ops/router.py
• dsv4/loader/hf_checkpoint.py
• dsv4/reference/attention.py, dsv4/reference/csa_attention.py ← keep regardless; they're cheap oracles you run by hand for validation.

Imported by Lineage P only (not by single_shot): dsv4/model/{layer,layer_schedule}.py, dsv4/layers/{attention,ffn,norm}.py, dsv4/cache/*, dsv4/kernels/cache/*, dsv4/kernels/indexer/score_topk.py, dsv4/kernels/router/dense_router_decode.py, dsv4/ops/{rope.py,custom_ops.py}. Keep all of these if Step 0 says Lineage P is the serving path. Archive only if Lineage P is confirmed dead.

Note the ops duplication for the human: ops/rope.py (Lineage P) vs ops/rope_cuda.py (live, used by single_shot); ops/topk.py/topk_select.py (orphan) vs the live topk inside single_shot. Don't merge these blindly — pick the canonical one per lineage decision.

2c. preload_all() is dead and references a non-existent file

dsv4/kernels/cuda/loader.py:preload_all() has no callers and asks for compressor_reduce_quant.cu, which does not exist (the file is compressor_reduce.cu). Either delete preload_all() or fix the filename — see Part 2 #1.

---

Step 3 — Reference-update cheatsheet (the only moves that need code edits)

Everything in Step 1 is zero-edit except single_shot_PYTORCH_REFERENCE.py, which is imported by 3 unit tests via a bare top-level import that only resolves because the file is in repo root.

Pre-move check: open single_shot_PYTORCH_REFERENCE.py and confirm its own imports are absolute (from dsv4. ...) or stdlib. If it bare-imports any sibling root module, fix those first or the move breaks it.

Move: single_shot_PYTORCH_REFERENCE.py → dsv4/reference/single_shot_PYTORCH_REFERENCE.py

Edit 1 — tests/unit/test_layer_comparison.py:34

-    from single_shot_PYTORCH_REFERENCE import mHCBlock, load_weights, forward_layer, rmsnorm
+    from dsv4.reference.single_shot_PYTORCH_REFERENCE import mHCBlock, load_weights, forward_layer, rmsnorm

Edit 2 — tests/unit/test_mhc_comparison.py:75

-    from single_shot_PYTORCH_REFERENCE import mHCBlock, load_weights as ref_load_weights, forward_layer
+    from dsv4.reference.single_shot_PYTORCH_REFERENCE import mHCBlock, load_weights as ref_load_weights, forward_layer

Edit 3 — tests/unit/test_compressor_position_bias.py:38 — this is a comment reference, not an import. Update the text only:

-    # --- PyTorch reference path (matches single_shot_PYTORCH_REFERENCE.py) ---
+    # --- PyTorch reference path (matches dsv4/reference/single_shot_PYTORCH_REFERENCE.py) ---

Verify after the move:

grep -rn "single_shot_PYTORCH_REFERENCE" . | grep -v "dsv4/reference/single_shot_PYTORCH_REFERENCE.py"
# every remaining hit must be one of the three updated lines above

---

PART 2 — Comment / Doc Audit (code is the source of truth)

These are verified mismatches where the prose describes a previous version of the code. Fix the prose to match the code. Listed highest-confidence first.

1. dsv4/kernels/cuda/loader.py — preload_all() names a file that doesn't exist

The code refers to compressor_reduce_quant.cu; the actual file is compressor_reduce.cu. The function also has no callers.

• Fix: delete preload_all() (it's dead), or change "compressor_reduce_quant.cu" → "compressor_reduce.cu" and verify the module's pybind function name matches what callers expect.
• Also re-check the module docstring's usage example (mod.fused_amax_quantize_nvfp4(x, divisor)) against the actual exported symbol in fused_amax_quantize.cu.

2. README.md "Package structure" + ROADMAP.md reference attention files that don't exist

The docs describe the attention kernel as dsv4/kernels/attention/fmha.py (the "592-line main production kernel") and fmha_smem_acc.py, and mention a dsv4/kernels/decode/ directory. None of these exist. The real live attention path is:

production.py  →  fmha_multitile_op.py  →  fmha_multitile_capi.cu  →  fmha_6warp_tma_multirow_multitile.cuh

• Fix: regenerate the README "Package structure" block from the actual tree (find dsv4 -type f | sort), and purge fmha.py / fmha_smem_acc.py / kernels/decode/ references from README and ROADMAP. Keep the lessons prose; correct the file map.

3. dsv4/kernels/attention/production.py docstring contradicts the ROADMAP about the production path

production.py (which single_shot_inference.py imports — i.e., the live attention entry) says, verbatim: "No CuTeDSL runtime dependency. No Python KV merge." But README.md / ROADMAP.md / the status docs describe "Python KV merge ships today" as the production path, and frame Priorities 1/2/4/8 around the CuTeDSL fmha.py + epilogue_tma_store kernel.

• Implication (flag to the human, don't silently rewrite): the live attention path appears to have moved to the C-API multitile kernel (fmha_multitile_* + the .cuh), which would make the entire "D1/D1.5/Python KV merge" framing and several roadmap priorities stale — planning fixes for a kernel you no longer run. Confirm which kernel dsv4_attention actually dispatches, then reconcile: the code (production.py → multitile C-API) wins; rewrite the ROADMAP's "Current status / blockers" to match.

4. dsv4/kernels/indexer/score_topk.py docstring has the wrong scoring formula

Line ~43 writes I[t,s] = Σ_h w_h[t,h] · ReLU(q_I[t,h] · K^IComp[s,h]) — the [s,h] implies a per-head key. The key is shared across heads (MQA, paper c_I=128). The sibling csa_indexer.py docstring and the live single_shot einsum both use the correct shared-key form.

• Fix: K^IComp[s,h] → K^IComp[s]. (If Step 2b archives this module, fix-or-archive — either way don't leave the wrong formula to mislead a future resurrection.)

---

A repeatable comment-audit method (because no one can eyeball 75k lines)

I verified the four above by reading the live path. The rest of the audit should be systematic, not heroic. Run this on the live closure (from import_closure.py), not the whole repo, and prioritize:

1. Top-of-file docstrings and # eq. / formula comments — highest mislead-risk. For each live module, read only the module docstring + any comment containing eq, shape, →, FP4/FP8/BF16, or a hardcoded number, and check it against the code immediately below.
2. Grep for known-stale tokens and review each hit on the live path:

   grep -rn "Python KV merge\|fmha\.py\|fmha_smem_acc\|MLA\|split-KV\|TODO\|FIXME\|XXX\|for now\|Phase 1\|will swap\|deferred" dsv4/ single_shot_inference.py
   

   Each "for now / will swap / Phase 1" comment is a promise that may already be broken — verify against current code.
3. Dtype claims: any comment asserting a tensor is FP8/FP4/BF16/FP32 — confirm against the actual .dtype / cast in code. (The KVCache docstring in single_shot_inference.py is a good example of a correct, valuable one — FP8 nope + BF16 rope — so don't strip long comments reflexively; only fix the wrong ones.)
4. One rule for the agent going forward: when you change code, the diff is not done until the surrounding comment/docstring describes the new code. Treat a stale comment as a build break.

---

Suggested commit sequence

1. helpers/import_closure.py + run Step 0 (record the vLLM finding in this file).
2. Root file moves (Step 1) — zero-edit batch first, then the single_shot_PYTORCH_REFERENCE.py move + 3 edits (Step 3), with the grep verification.
3. .cu dedup (Step 2a) — diff, reconcile into cuda/, delete dead subfolder copies.
4. Lineage-P archive decision (Step 2b) — only after Step 0; move to dsv4/_archive/, never delete.
5. Comment fixes #1–#4 (Part 2), then the grep-driven sweep.

After each step: grep -rn "<moved basename>" . shows zero dangling refs, and single_shot_inference.py still generates coherent output.