**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) |
| 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:
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)
print("REACHABLE:",len(seen)," | DEAD CANDIDATES:",len(dead))
fordindead: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.
`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.
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/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.
# 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`.
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:
- **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:
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.
1.**P4 — Fused RMSNorm + NVFP4 quantize**: New CUDA kernel (2 launches vs 6+). Integrated into `single_shot_inference.py` for the attention path. Unit test passes (cos=0.996-0.999 vs unfused). Decode speed improved 0.53→0.43s/token.
2.**P5 — Fused mHC + RMSNorm + NVFP4 quantize**: New CUDA kernel (2 launches vs 7+). Unit test passes on B200 (cos=0.996-0.999). NOT yet integrated into single_shot — needs moe_forward refactoring.
3.**CRITICAL FIX — Indexer deadlock**: `__syncthreads()` inside a strided loop caused deadlock at production context lengths (always deadlocked when `num_valid % 128 != 0`). Fixed with per-thread local top-k + block-level merge. Both copies updated.
### The regression
After P4 integration, model output went from coherent English to garbled nonsense. This is a NEW regression — output was correct last night after P0-P3 + KV work.
### What to investigate (ordered by likelihood)
1.**`run_from_quantized` gsa shape bug (MOST LIKELY)**: The fused kernel produces per-row gsa (shape `(M,)`). `run_from_quantized` passes this to the CuTeDSL NVFP4 GEMM as `global_scale_a`. But the GEMM expects a **single scalar** — it's one scale for the entire A matrix. For M=1 decode, shape `(1,)` may work as a scalar, BUT the gsa VALUE differs from the unfused path's scalar gsa because:
- Unfused: quantize computes gsa from the full (M, N) tensor — single scalar
- Fused: computes gsa per-row, then `[:1].reshape(1)` takes only the first row's gsa
- These differ when rows have different magnitudes
2.**Dequant→requant noise for compressor**: The compressor gets `x_normed` by dequantizing the fused kernel's FP4 output. This introduces ~0.5% quantization error (cos=0.994) that wasn't present before. This noise propagates into compression scores and indexer queries.
3.**A/B test first**: Run with `_use_fused_rmsnorm_quantize=False` to confirm P4 is the cause. If output is correct with unfused, the bug is in the P4 path.
### Key code paths
- **`dsv4/layers/linear.py`**: `run_from_quantized()` — the gsa handling is suspicious
- **`dsv4/ops/gemm_runner.py`**: `run_nvfp4_grouped_gemm()` — how global_scale_a is consumed
- **`dsv4/kernels/gemm/grouped.py`**: CuTeDSL GEMM — global_scale_a is scalar
- **`single_shot_inference.py:855-870`**: P4 integration point
- **`dsv4/kernels/cuda/fused_rmsnorm_quantize.cu`**: the fused kernel
### Quick fix ideas
- Use **scalar gsa** (reduce per-row gsa to a single max) for the GEMM — keeps the fusion but makes gsa compatible
- Or: don't use `run_from_quantized` — just use the fused kernel for the BF16 output (dequant) and let each linear re-quantize with its own scalar gsa (saves rmsnorm launches but not quantize launches)
- Or: fix the CuTeDSL GEMM to support per-row global_scale_a (THIS IS COMPLEX, BUT IF IT IS THE CORRECT WAY OF DOING THINGS. I WANT IT DONE THAT WAY!!!!!)
Blocking a user prevents them from interacting with repositories, such as opening or commenting on pull requests or issues. Learn more about blocking a user.
