nvfp4-megamoe-kernel/archived_plans/NEXT_PRIORITIES_PART_2.md

# 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.