# 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 .`"* --- ## 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=`, 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 "" .` shows zero dangling refs, and `single_shot_inference.py` still generates coherent output.