nvfp4-megamoe-kernel/archived_plans/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:

```bash
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.

```python
# 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`**
```diff
-    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`**
```diff
-    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:
```diff
-    # --- PyTorch reference path (matches single_shot_PYTORCH_REFERENCE.py) ---
+    # --- PyTorch reference path (matches dsv4/reference/single_shot_PYTORCH_REFERENCE.py) ---
```

**Verify after the move:**
```bash
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:
   ```bash
   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.