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:

"""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:

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