diff --git a/CLEAN_UP.md b/CLEAN_UP.md new file mode 100644 index 00000000..b2b502da --- /dev/null +++ b/CLEAN_UP.md @@ -0,0 +1,244 @@ +# 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. \ No newline at end of file diff --git a/NEXT_SESSION.md b/NEXT_SESSION.md deleted file mode 100644 index 998673c0..00000000 --- a/NEXT_SESSION.md +++ /dev/null @@ -1,35 +0,0 @@ -Here's the summary for next session: -## Summary: What Was Done & What to Investigate - -### What was done today -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!!!!!) \ No newline at end of file diff --git a/PERFORMANCE_AUDIT.md b/docs/PERFORMANCE_AUDIT.md similarity index 100% rename from PERFORMANCE_AUDIT.md rename to docs/PERFORMANCE_AUDIT.md diff --git a/single_shot_PYTORCH_REFERENCE.py b/dsv4/reference/single_shot_PYTORCH_REFERENCE.py similarity index 100% rename from single_shot_PYTORCH_REFERENCE.py rename to dsv4/reference/single_shot_PYTORCH_REFERENCE.py diff --git a/dump_checkpoint_keys.py b/helpers/dump_checkpoint_keys.py similarity index 100% rename from dump_checkpoint_keys.py rename to helpers/dump_checkpoint_keys.py diff --git a/helpers/import_closure.py b/helpers/import_closure.py new file mode 100644 index 00000000..c05f49fb --- /dev/null +++ b/helpers/import_closure.py @@ -0,0 +1,36 @@ +# 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"] # vLLM has 0 imports of dsv4 (Step 0 confirmed) + +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) diff --git a/probe_hf_indexer.py b/helpers/probe_hf_indexer.py similarity index 100% rename from probe_hf_indexer.py rename to helpers/probe_hf_indexer.py diff --git a/probe_indexer_shapes.py b/helpers/probe_indexer_shapes.py similarity index 100% rename from probe_indexer_shapes.py rename to helpers/probe_indexer_shapes.py diff --git a/probe_keys.py b/helpers/probe_keys.py similarity index 100% rename from probe_keys.py rename to helpers/probe_keys.py diff --git a/probe_shapes.py b/helpers/probe_shapes.py similarity index 100% rename from probe_shapes.py rename to helpers/probe_shapes.py diff --git a/hf_reference_test.py b/tests/integration/hf_reference_test.py similarity index 100% rename from hf_reference_test.py rename to tests/integration/hf_reference_test.py diff --git a/test_gemm_1group.py b/tests/integration/test_gemm_1group.py similarity index 100% rename from test_gemm_1group.py rename to tests/integration/test_gemm_1group.py diff --git a/test_quantize_gpu.py b/tests/integration/test_quantize_gpu.py similarity index 100% rename from test_quantize_gpu.py rename to tests/integration/test_quantize_gpu.py diff --git a/test_se_dequant.py b/tests/integration/test_se_dequant.py similarity index 100% rename from test_se_dequant.py rename to tests/integration/test_se_dequant.py diff --git a/test_se_gpu.py b/tests/integration/test_se_gpu.py similarity index 100% rename from test_se_gpu.py rename to tests/integration/test_se_gpu.py diff --git a/test_se_l1_direct.py b/tests/integration/test_se_l1_direct.py similarity index 100% rename from test_se_l1_direct.py rename to tests/integration/test_se_l1_direct.py diff --git a/test_se_multi_gpu.py b/tests/integration/test_se_multi_gpu.py similarity index 100% rename from test_se_multi_gpu.py rename to tests/integration/test_se_multi_gpu.py diff --git a/tests/unit/test_compressor_position_bias.py b/tests/unit/test_compressor_position_bias.py index 04ab3d3f..cf408327 100644 --- a/tests/unit/test_compressor_position_bias.py +++ b/tests/unit/test_compressor_position_bias.py @@ -35,7 +35,7 @@ def test_csa_position_bias(): # --- CUDA kernel path --- compressed_cuda = csa_compress_production(kv, gate, position_bias, kv_norm_weight, m=m) - # --- PyTorch reference path (matches single_shot_PYTORCH_REFERENCE.py) --- + # --- PyTorch reference path (matches dsv4/reference/single_shot_PYTORCH_REFERENCE.py) --- kv_ref = kv.clone() gate_ref = gate.clone() # Add position_bias cyclic per block diff --git a/tests/unit/test_layer_comparison.py b/tests/unit/test_layer_comparison.py index deeda5fa..4a255f81 100644 --- a/tests/unit/test_layer_comparison.py +++ b/tests/unit/test_layer_comparison.py @@ -31,7 +31,7 @@ def main(): # --- Load PyTorch reference pipeline --- print("\nLoading PyTorch reference pipeline...") - 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 all_w = load_weights(CHECKPOINT_DIR) print("Reference pipeline loaded.") diff --git a/tests/unit/test_mhc_comparison.py b/tests/unit/test_mhc_comparison.py index 7ed825d6..7cf88d2d 100644 --- a/tests/unit/test_mhc_comparison.py +++ b/tests/unit/test_mhc_comparison.py @@ -72,7 +72,7 @@ def main(): # Create a realistic hidden state (simulate running through a few layers) # Use token embedding + a few layers of mHC - 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 ref_all_w = ref_load_weights(CHECKPOINT_DIR) # Build mHC blocks for first 3 layers