# ROADMAP — Stage E and the path to a runnable model **Method.** Every state claim below was verified against the source in this revision. Agent status doc and agent assessment were both read, then ignored where they contradict the code. The doctrine rules from prior issue docs apply: DSL wall → raw CUDA C++ (not Python); raw CUDA ≠ scalar math; print, don't guess; integration over exploration; falsifiable gates only. --- ## TL;DR — where we actually are **The FMHA kernel itself is in genuinely good shape.** `fmha_6warp_tma_multirow_multitile.cuh` is one file, properly Blackwell-native (`tcgen05.mma` SS, UMMA descriptors, TMEM, `cp.async.bulk.tensor.2d` + `mbarrier.arrive.expect_tx`, in-kernel multi-tile rescale via SMEM accumulator, multi-row softmax). The TMA hang from P4 was correctly diagnosed (missing `expect_tx`) and resolved. The C-API + ctypes shim works. **P3, P4, P5, P6, P7, P8 numerics are done.** **The agent's "Stage E" assessment is correct in spirit but wrong on specifics.** Where it says "AttentionSubBlock has `NotImplementedError` stubs" — it doesn't. The layer is structurally complete and *imports the right things*. The real problem is one layer down: it calls cache methods like `gather_compressed_kv()`, `gather_swa_kv()`, and properties like `cache.num_query_heads` / `cache.head_dim` **that don't exist on `LayerCacheHandle`**. The handle exposes `read_*_view()` returning paged FP8 buffers; nothing materializes them into the BF16 dense tensors `dsv4_attention` consumes. That's the integration gap, not stubs. **Three structural shortcuts still live in production code.** They are not performance-fatal yet but they will be at scale: 1. Two `torch.cuda.synchronize()` calls remain on the kernel-launch path (`fmha_multitile_op.py:130`, `production.py:279`). The second is in the still-resident "slow path." The first is unconditional, even on the fast path. 2. The slow CuTeDSL+Python-merge path is **not deleted**, only relabeled "slow path." 471 lines of code, 48 unit tests bound to its `FmhaKernel` symbol, still imported at `production.py:53`. 3. The batched call shape (4D q) does a **Python for-loop over batch items** at `production.py:380` instead of folding batch into the grid. **Cleanup is non-trivial.** ~22 dead/stale code paths live alongside the working ones. Some are old CuTeDSL drafts (`ops/decode_sparse.py`, `ops/decode_swa.py`), some are pre-consolidation Stage D probes (`fmha_sm100.cuh`, `fmha_sm100_tc.cuh`, `fmha_sm100_launch.cu`, `fmha_qk_verify.cuh`, `fmha_epilogue_sm100.cuh`), some are misplaced root files (`debug_linear.py`, `test_mapping.py`, `tests/working_softmax_maybe.py`), and 46 of 177 unit tests are pinned to the deprecated `FmhaKernel`. Left in place, this is the primary path by which an agent will drift back into the wrong codebase the next time it context-switches. --- # PART 1 — CLEANUP (do this first, in this order) Cleanup is gated *before* Stage E work because the next stage's failure mode is "agent picks the wrong file to extend." Right now `dsv4/kernels/attention/` has 14 files; only 6 are live. That ratio is the diagnostic. ## C1. Confirm what's live with `grep`, not by reading status docs Before deleting anything: produce `audit_attention_live.md` listing each file under `dsv4/kernels/attention/` and `dsv4/ops/` with its **outbound reference count from non-archive code** (production, layers, model, integration tests). Anything with 0 refs is a deletion candidate; anything with >0 refs gets inspected to confirm the reference is live, not a stale import. **Falsifiable gate:** the audit table goes into `archived_plans/audit_attention_live.md`, file count column adds up to the actual directory listing. Mention deletion candidates by name with the ref count next to them. ## C2. Delete the dead CuTeDSL FMHA scaffolding These have **zero external refs** outside their own `#include` chain and self-references: | File | Status | Action | |---|---|---| | `dsv4/kernels/attention/fmha_sm100.cuh` | Phase-1 scalar reference, header says "WORKING cos 0.999999" | Delete. The 6-warp kernel supersedes it. | | `dsv4/kernels/attention/fmha_sm100_tc.cuh` | Earlier tensor-core fork before consolidation. | Delete. The 6-warp kernel is its successor. | | `dsv4/kernels/attention/fmha_sm100_launch.cu` | Header says "COMPILES but doesn't run via torch.utils.cpp_extension" | Delete. Replaced by `fmha_multitile_capi.cu` + ctypes. | | `dsv4/kernels/attention/fmha_qk_verify.cuh` | Header says "QK GEMM verification" — one test references it. | Move to `tests/unit/qk_verify_kernel.cuh` (it's a test fixture, not a library kernel). | | `dsv4/kernels/attention/fmha_epilogue_sm100.cuh` | Phase-2 epilogue prototype; the multi-tile kernel already does in-kernel rescale via SMEM and doesn't use this. | Delete after confirming it's not `#include`'d by anything live. Audit table from C1 answers this. | | `dsv4/ops/decode_sparse.py` | Old CuTeDSL CSA decode draft, not called from anywhere. | Delete. Path is now `layers/attention.py:_forward_csa` → `kernels/attention.sparse_fmha_with_swa` → `kernels/attention/production.py`. | | `dsv4/ops/decode_swa.py` | Old CuTeDSL SWA decode draft, same story. | Delete. | | `dsv4/kernels/decode/` | Two files: `__init__.py` (empty) and `_NOTES_fp8_bf16.md` (scratch). | Delete the whole directory. | **Falsifiable gate:** after C2, `find dsv4/kernels/attention -name '*.cuh' -o -name '*.cu' | wc -l` returns the value in `audit_attention_live.md`'s live count. No reference warnings on `python -c "import dsv4"`. ## C3. Decide the fate of the CuTeDSL slow path (`fmha.py`) `fmha.py` is 645 lines and is the entire pre-P3-P8 CuTeDSL path. It is currently imported by `production.py:53` (`from dsv4.kernels.attention.fmha import FmhaKernel`) and used by the `else` branch starting at `production.py:410` ("SLOW PATH"). It is also imported by **48 unit tests**. **Two options. Pick one, document it, do not maintain both.** - **Option A (preferred, matches doctrine).** Delete `fmha.py`, delete the SLOW PATH branch from `production.py`, delete the 48 `test_d*.py` unit tests bound to `FmhaKernel`. The 6-warp multi-tile kernel covers `hd∈{64,128,256,512}` × `T=1` decode × any `n_segments`. That's the production surface. Everything else is exploration that already paid off — it does not need to ship. - **Option B (fallback, only if Option A is blocked).** Move `fmha.py` and the 48 tests to `archived_plans/cutedsl_legacy/`. Delete the SLOW PATH from `production.py` regardless. The kernel and tests stay buildable for reference but are no longer in the active import graph. **Failure mode to refuse:** "keep it as fallback in case the 6-warp kernel has a regression." If there's a regression, you fix the 6-warp kernel — that's the production kernel. Carrying a parallel implementation guarantees both rot. **Falsifiable gate:** `grep -r "FmhaKernel" dsv4/ tests/ scripts/ | wc -l` returns 0 (Option A) or only matches inside `archived_plans/` (Option B). The SLOW PATH branch is gone from `production.py` either way. ## C4. Empty the `archive/` subdirectory the same way `dsv4/kernels/attention/archive/` currently holds 5 backups (`fmha_backup_pre_epilog.py`, `fmha_backup_v2.py`, `fmha_smem_acc.py`, `fmha_6warp_tma_driver_api.cuh`, `fmha_tma_driver_api.cuh`). These were correct to *not* delete during P8 consolidation. Now: move them out of `dsv4/` entirely into `archived_plans/code_archive/` and delete `dsv4/kernels/attention/archive/`. Active code directories should contain only active code. `archived_plans/` is the right place for everything that isn't. **Falsifiable gate:** `find dsv4 -name 'archive' -type d | wc -l` returns 0. ## C5. Move the root-level scratch files Three files live at the repo root and don't belong there: - `debug_linear.py` (58 lines, computes expected O for a debug pattern — Stage D probe). → `tests/unit/archive/` or delete if no longer useful. - `test_mapping.py` (49 lines, SMEM-P coordinate mapping check — Stage D probe). → same. - `run_router_tests.py` (hardcoded path `/root/dsv4-nvfp4-workspace/kernel`). → Either generalize the path and move to `scripts/`, or delete. Hardcoded absolute paths in tracked files are bug magnets. - `tests/working_softmax_maybe.py` (496 lines, name speaks for itself). → If it's a draft, archive it. If it's superseded, delete it. The doctrine is "code or archive, never `maybe`." **Falsifiable gate:** repo root has no `.py` files except generated/build artifacts. `tests/` contains no `*_maybe.py` or `working_*.py`. ## C6. Prune the unit test D-probe explosion 46 `test_d*.py` files were Stage D debugging probes. Many are bound to `FmhaKernel` and will be deleted in C3 if you pick Option A. The remainder that test components still in use (`FmhaTmaMultiRowMultiTile`, the indexer, the compressor, mHC, MoE) should either be promoted to named integration tests (`test_fmha_decode.py`, `test_csa_indexer.py`, etc.) or archived. **Falsifiable gate:** `ls tests/unit/test_d*.py | wc -l` returns ≤ 5 (only the ones genuinely testing currently-live behavior, with descriptive names). ## C7. Status doc consolidation After C1–C6, write **one** new top-level `STATUS.md` describing the **post-cleanup** state of the production path, and delete `NEXT_PRIORITIES.md` (this doc replaces it). `archived_plans/` already holds the historical paper trail; do not add to it during cleanup. **Falsifiable gate:** repo root has exactly two living `.md` files: `README.md` and `STATUS.md` (plus this `ROADMAP.md`). `archived_plans/` holds the history, unchanged. --- # PART 2 — STAGE E PRIORITY ORDER (post-cleanup) The agent is right that Stage E = "make the model run end-to-end." The order below is sequenced so each priority unblocks the next, and so the performance-decisive moves come *after* end-to-end correctness, not before. ## E1. Wire `LayerCacheHandle` to expose what `AttentionSubBlock` calls This is the load-bearing missing piece. `dsv4/kernels/attention/__init__.py` calls four things that don't exist on `dsv4/cache/handle.py`: - `cache.gather_compressed_kv(selected_indices) → (k_compressed, v_compressed)` BF16 dense - `cache.gather_all_compressed_kv() → (k_compressed, v_compressed)` BF16 dense (HCA, no top-k) - `cache.gather_swa_kv() → (k_swa, v_swa)` BF16 dense - `cache.num_query_heads`, `cache.head_dim` (properties) The handle already exposes the right *raw* state via `read_classical_view()` (FP8 entries + RoPE BF16 + paged block_table) and `read_swa_view()` (FP8 entries + RoPE BF16 + positions). The work is the gather+dequantize+concat path. **Constraints (doctrine):** - Each gather is a **single fused kernel**, raw CUDA C++ if CuTeDSL can't do block-table-strided FP8→BF16 dequant+gather in one pass. Not a Python loop over selected indices. Not five eager torch ops. - The kernel uses TMA + warp-level dequant, not scalar loads. The MoE GEMM already proves out the FP8 dequant pattern on Blackwell — reuse it. - The dequant decodes E4M3 with the inverse-scale stored alongside. **Print the actual layout of `paged.entries_fp8` + `paged.inv_scale` before writing the kernel** — do not assume the layout matches an MoE convention. **Definition of done:** 1. The four methods exist on `LayerCacheHandle`, implemented in `dsv4/kernels/cache/gather.{cu,py}`. 2. `python -c "from dsv4.layers.attention import AttentionSubBlock; ..."` runs a forward pass on one layer without AttributeError, NotImplementedError, or shape mismatch. 3. Numerical parity vs a pure-Python reference (FP32 dequant + dense FMHA over the gathered tensors): `cos ≥ 0.9995`. FP8 dequant loses a little, that's expected. 4. **Launch count for one CSA layer** measured with Nsight: gather is ≤ 3 launches (compressed gather, SWA gather, concat). Not 30, not 5+5+cat. **Failure modes to refuse:** - A Python gather loop over the block table. The block table is on GPU; the gather is on GPU. No host trips. - BF16 materialization "for now" without FP8 dequant. The paged cache stores FP8 because that's the design — bypassing it loses 2x KV memory and the whole point of the cache structure. ## E2. End-to-end smoke test through one full layer After E1, both halves of `dsv4/model/layer.py` (`mhc_attn` → `attn.forward` → `mhc_ffn` → `ffn.forward`) should execute without error. The test is the forcing function. **Definition of done:** 1. `tests/e2e/test_one_layer.py` runs `DSV4Layer.forward(X, token_ids, cache)` on synthetic inputs with cached weights loaded from `loader/`, returns a tensor of the right shape and dtype, no errors. 2. The same test with `cache` containing actual paged FP8 entries (not zeros) produces output with `cos ≥ 0.99` against a Python reference of the same layer (compressor + indexer + FMHA + mHC, all in eager PyTorch). 3. Test exists for all three attention types (CSA, HCA, SWA) by toggling the layer's `attn_type`. **Failure mode to refuse:** wrapping the test in a try/except to "make it pass." The cos threshold is the gate. ## E3. Top-level `model/dsv4.py` (currently a 2-line TODO) `dsv4/model/dsv4.py` is literally: ```python """Full DSV4 model.""" # TODO: Phase 1 ``` Build the actual model class: embedding → N×DSV4Layer → final norm → prediction head. Use `layer_schedule.py` to alternate CSA/HCA per the V4 spec (Pro: layers 0–1 HCA, then interleaved; Flash: layers 0–1 SWA, then interleaved). **Definition of done:** 1. `DSV4Model.forward(token_ids, cache_manager) → logits` works on a single token decode. 2. Layer schedule matches the V4 paper (verify against `config.py`). 3. MTP module is wired but optional (off by default; doctrine note in code that MTP fine-tuning is a separate workstream). 4. End-to-end test: `tests/e2e/test_decode.py` generates one token from a 3-layer toy config, succeeds, cos ≥ 0.99 vs Python reference. ## E4. Delete the two remaining `torch.cuda.synchronize()` calls on the fast path After C3 the slow-path sync at `production.py:279` is gone with the slow path. The other sync is at `fmha_multitile_op.py:130`: ```python # Synchronize to catch async errors torch.cuda.synchronize() return o, lse ``` This is on the *fast path*. It's there to catch async kernel errors, but the correct way to do that is to check the return code from the C-API launch function (which `fmha_multitile_decode_launch` already returns) plus `cudaPeekAtLastError`, not a full device sync. **Definition of done:** 1. The sync is removed from `fmha_multitile_op.py`. 2. Error checking moves to a `cudaPeekAtLastError` call + the existing return code check. 3. Nsight measurement on a 4-layer decode: zero `cudaDeviceSynchronize` calls between layer 0 entry and final logits. ## E5. Fold the batch loop into the kernel grid `production.py:380–388` is a Python for-loop over batch items. Each iteration calls back into `dsv4_attention` on a 3D slice. For batch=8 and 61 layers, that's 488 launches per decoded step instead of 61. The multi-tile kernel's grid is currently `(1, n_h, 1)`. It needs to become `(1, n_h, batch)` with `blockIdx.z` used to index batch. The kernel already has `_batch_stride` params (read off `FmhaTmaMultiRowMultiTileParams`); they just aren't wired to a grid dim. **Definition of done:** 1. `production.py:380–388` for-loop deleted; the batched path calls the multi-tile kernel once with `gridDim.z = batch`. 2. Numerical parity: `cos ≥ 0.999998` vs the deleted Python loop. 3. Launch count for batch=8 decode: 1 per layer, not 8. ## E6. NVFP4-1.2: FP4 output fusion for FMHA → `wo_a` Right now the FMHA writes BF16 to GMEM, then `wo_a` re-quantizes it to FP4. That's a redundant memory pass on the attention output. The epilogue already has the one-way TMEM→regs→SMEM→GMEM shape — adding an FP4 pack with amax in registers, gated by a template param, is the same pattern the MoE epilogue already uses for NVFP4-1.1. **Constraints (doctrine):** - The FP4 pack is **part of the same single launch**. No second kernel for quantize. Per-128-element block amax reduces via warp shuffle in registers. - The amax → scale → E2M1 LUT path matches `dsv4/ops/quantize.py` byte for byte. Print the LUT from the Python side; assert equality in a unit test. The indexer LUT bug we already fixed is the cautionary tale. - This is gated on E4 (`epilogue_op` slot exists in the multi-tile kernel — let's verify it does before starting E6). **Definition of done:** 1. `fmha_6warp_tma_multirow_multitile_kernel` accepts a template arg `Epilogue` (default = identity); a `FP4Epilogue` specialization packs E2M1 + emits SF blocks in the layout `wo_a` expects. 2. `tests/unit/test_fp4_epilogue_parity.py`: BF16 reference vs FP4 fused, `cos ≥ 0.999` after dequant. 3. End-to-end CSA layer: BF16 output → FP4 output replaces one full BF16 R/W pass through (T, n_h * hd) memory. Measured bandwidth drop with Nsight. ## E7. Stage F: Lightning indexer FP4 tensor-core scoring Currently `indexer_score_topk.cu` uses scalar FP32 cores after the LUT dequant. The paper §5.2.1 is explicit that this path runs in FP4 on tensor cores ("multiplied entirely in FP4 ... 99.7% recall"). The LUT bug fix made the scalar path *correct*; this priority makes it *fast*, the way V4 actually specifies it. **Constraints (doctrine):** - `tcgen05.mma` with `mxf4nvf4` kind. Keys and queries both FP4. Scales E4M3 in shared mem, hardware decodes. No scalar fallback in the live path. - Recall vs the FP32 oracle ≥ 99.7% per paper. The recall test from the LUT fix is the same test; the bar is the same. - The top-k selector goes warp-level (`__shfl_xor` ballot) at the same time — the current serial single-thread heap merge at top_k=1024 is its own latency source on Pro decode. **Definition of done:** 1. `indexer_score_topk.cu` has a `tcgen05`-MMA FP4 path, gated by a config flag, with the scalar path archived (not deleted, since the FP4 path needs the oracle). 2. Recall@k ≥ 99.7% for k ∈ {512, 1024} at compressed_blocks ∈ {2k, 8k, 32k}. 3. End-to-end CSA layer latency drops measurably (Nsight). ## E8. Multi-CTA grid for prefill (Priority 4 from the original ROADMAP) Decode is single-CTA per head and that's correct. Prefill (T > 1) wants multi-CTA to parallelize over M. The one-way epilogue from E6 is the prerequisite — once it lands, the grid expansion is wiring. **Definition of done:** 1. `gridDim.x = ceil_div(M, M_TILE)` for `T > 1` paths. 2. Prefill latency at T=128 with M-tile=64: 2-CTA grid, measured speedup ≥ 1.5× over single-CTA. 3. Decode path unchanged (single CTA, no regression). ## E9. CUDA graph capture (Stage E item from the agent's notes) Once E5 lands and the decode hot path is sync-free, graph capture is straightforward. This is the biggest single perf win on autoregressive decode because it eliminates host launch overhead entirely. **Definition of done:** 1. `DSV4Model.decode_step(token, cache)` is `cudaGraph`-capturable. 2. Replay latency vs uncaptured: ≥ 2× faster on a 3-layer toy at batch=1. 3. Memory: graph capture doesn't `torch.zeros` on the hot path (E1 already forced this for gather; verify it holds for compressor + indexer + FMHA). --- # DOCTRINE — applies to every item above 1. **DSL wall → raw CUDA C++, not Python.** The C-API + ctypes pattern in `fmha_multitile_op.py` is the correct shape for new kernels. 2. **Raw CUDA ≠ scalar math.** Every kernel above is `tcgen05` / UMMA / TMA / warp reductions. No "temporary scalar dot product" without a labeled replacement target. 3. **Print, don't guess.** Two places this round specifically: - E1 gather kernels: print the FP8 cache layout + inv-scale layout before writing the dequant. - E6 FP4 pack: print the LUT from `quantize.py` and assert in a test that the kernel produces byte-identical packed nibbles. 4. **Integration over exploration.** Do not create `fmha_6warp_tma_multirow_multitile_v2.cuh`. Extend the chosen file or stop and replan. Variant proliferation is the agent's primary failure mode in this repo. 5. **Falsifiable gates only.** Every "done" above has a number or a binary check. "Looks fine," "milestone complete," and "Stage X done" without a number are read as "not done." 6. **The slow path is not a fallback.** Once C3 deletes the CuTeDSL+Python-merge path, it stays deleted. If the 6-warp kernel breaks, you fix the 6-warp kernel — not resurrect 471 lines of dead code.