nvfp4-megamoe-kernel/ARCHITECTURE_AND_MEMORY_AUDIT.md

# ARCHITECTURE & MEMORY AUDIT — Post-probe rewrite

**Supersedes:** the prior `ARCHITECTURE_AND_MEMORY.md` (M1 was wrong by 64×
in the bad direction). Incorporates the indexer probe results from
`archived_plans/INDEXER_PROBE_RESULTS_20260602.md`.

**Method.** Every claim verified against `single_shot_inference.py` v16 + the
probe results. Per doctrine.

---

## TL;DR — the picture is much better than the prior audit suggested

**The architecture is faithful to the paper. The 1M-context memory story is
fine on 8×B200. There is no looming OOM crisis.**

That said, the probe surfaced a finding bigger than memory: **the lightning
indexer has never actually run in any production decode to date.** Paris-back
is real, but it ran via dense attention over the full compressed KV history
in CSA layers — the sparse-selection path was silently bypassed because the
indexer's internal compressor never loaded its weights. The system has been
correct because the *fallback* was algebraically correct, not because the
designed CSA path was working.

This is good news. It means:

1. **Fixing the indexer is the next correctness milestone.** It unlocks the
   actual sparse path, which is what makes 1M context tractable at runtime
   (not memory-wise — speed-wise, since dense over 250K compressed entries
   per CSA layer per token is the actual perf wall, not KV storage).
2. **Memory at 1M is dominated by the main compressed KV cache (~10 GB
   total across all CSA+HCA+SWA layers), which is small enough that the
   prior audit's "131 GB" panic was wrong.** No FP4 quantization of the
   indexer cache is needed for memory reasons. (It is still wanted for
   *throughput* per paper §5.2.1, but that's a different fight.)
3. **Three small bugs are blocking the indexer from running correctly.**
   Two are surface (weight-path + buffer-width); one is deeper (the
   scoring einsum's algebra is wrong, treating MQA-on-indexer as full
   multi-head). All three are easy fixes once seen.

---

# PART 1 — WHAT THE PROBE REVEALED

The probe confirmed hypothesis A from the prerequisite doc and surfaced two
collateral findings. The combined picture:

## F1 — Indexer keys are `c_I = 128`-wide, MQA-on-indexer (paper-aligned)

`comp_indexer_kv.shape == (n_comp, 128)`. One vector per compressed block,
**shared across all `n_ih = 64` indexer query heads.** This is the standard
multi-query-attention shape, but applied to the indexer scoring path.

Per-block cost: 128 × 2 bytes = **256 B per compressed block per CSA layer**.
At 1M context (CSA ratio=4 → 250K compressed blocks):

- Per CSA layer: 250K × 256 B = **64 MB**
- × 30 CSA layers = **~1.9 GB total** for indexer KV at 1M context

That's small. ~6× smaller than the main compressed KV cache. The prior
audit's M1 ("indexer KV is 125 GB at 1M, OOM at 250K tokens") was
backwards — the indexer cache is the *smallest* of the three KV streams.

## F2 — The indexer compressor never loaded weights (the real bug)

`Indexer.load:392`:

```python
if f"{pfx}.compressor.kv_proj.weight" in w:
    self.compressor = Compressor(4, self.ihd, 7168, dev)
```

The checkpoint stores the indexer's compressor weights at
`*.indexer.kv_proj.weight`, **not** `*.indexer.compressor.kv_proj.weight`.
So this `if` was always False, `self.compressor` stayed None, and
`Indexer.forward` always returned None at the early-return guard (line
397: `if ... comp_indexer_kv is None or comp_indexer_kv.shape[0] == 0:
return None`).

What this means for every Paris-back run to date:

- CSA layers received `topk_idx = None` from the indexer.
- The gather path at `forward_attention:569–571` checks
  `if ratio == 4 and topk_idx is not None:` → False, so it falls through
  to `elif ratio > 4: all_kv = torch.cat([kv_cache.comp_kv, swa_kv], ...)`.
  Wait — that branch is for `ratio > 4` (HCA), not `ratio == 4` (CSA).
  Need to check what CSA actually did with topk_idx=None.

**The agent should verify which fallback path CSA actually took, and
confirm whether the existing test runs were:**
- (a) attending over **just SWA** (correct only at short context, since
  SWA window is 128 — would explain why Paris works but degrades past
  step 10),
- (b) attending over **the full compressed history** as if it were HCA
  (correct but slow at scale), or
- (c) producing no attention output at all and being saved by a
  downstream operation.

This is a 10-line print insertion at `forward_attention`, not an
investigation campaign. **Add it to the indexer-fix work below, do not
spin up a separate probe.**

## F3 — The scoring einsum has the wrong algebra (MQA vs per-head keys)

The current code at `Indexer.forward:404`:

```python
k_idx = comp_indexer_kv.reshape(n_comp, self.n_ih, self.ihd)
scores = torch.einsum('tnd,cnd->tnc', q_idx.float(), k_idx.float())
```

The reshape requires `comp_indexer_kv` to have `n_ih × ihd = 8192` elements
per block. The probe shows it actually has `ihd = 128` elements. So the
reshape raises today.

**The temptation is to "fix" this by widening `comp_idx_buf` to 8192.**
That would let the reshape succeed and produce numerically plausible
scores. **It would be wrong.** The paper's scoring formula (§2.3.1, eq.
16) is:

```
I[t,s] = Σ_h w^I_{t,h} · ReLU(q^I_{t,h} · K^IComp_s)
```

`K^IComp_s` has no head subscript. It's **one key vector per block, shared
across all `n_ih` indexer query heads.** The score is computed by dotting
each of the 64 query heads against the *same* key, applying ReLU, then
weighting and summing across heads. That's MQA — the same trick used for
the main attention path in DSv4 (§2.3.1 "Shared Key-Value MQA").

The correct einsum:

```python
# q_idx: (T, n_ih, ihd) = (T, 64, 128)
# k_idx: (n_comp, ihd) = (n_comp, 128)        <-- no head dim
# w_h:   (T, n_ih) = (T, 64)
scores = torch.einsum('tnd,cd->tnc', q_idx.float(), k_idx.float())  # 'cd', not 'cnd'
scores = F.relu(scores)
total = (scores * w_h.unsqueeze(-1).float()).sum(1)  # (T, n_comp)
tk = min(self.top_k, n_comp)
_, idx = total.topk(tk, -1)
return idx
```

The `k_idx.reshape(n_comp, self.n_ih, self.ihd)` line goes away entirely —
no reshape needed when keys are MQA-shared.

**Why this matters beyond "the reshape stops crashing":** without this
correction, an agent fixing F2 (load the indexer compressor) and "fixing"
F3 by widening the buffer would produce silently wrong top-k selections.
Same shape as the original indexer LUT bug — code runs, produces plausible
numbers, but the *ranking* of compressed blocks is corrupted because the
math doesn't match the model. Recall@k drops from paper's 99.7% to
something much lower, and we'd be back to debugging "model gets dumber at
long context" by ripping apart the FMHA kernel that isn't broken.

## F4 — The buffer width is wrong but smaller than the prior audit claimed

`KVCache:419`:

```python
self.comp_idx_buf = torch.zeros(max_comp, head_dim, dtype=torch.bfloat16, ...)
                                          ^^^^^^^^
                                          512 — should be 128
```

`head_dim = 512` (main attention head dim). Indexer keys want `c_I = 128`.
The buffer is **4× too wide**, not 16× as the prior audit assumed. Storage
waste at 1M context (CSA only): 30 layers × 250K × (512 - 128) × 2 bytes
= **5.7 GB wasted**. Real, fixable, not catastrophic.

The fix needs a value to use, and that value should come from the indexer
instance, not hard-coded:

```python
# In __init__:
self.comp_idx_buf = torch.zeros(
    max_comp,
    indexer_key_dim,   # passed from caller, = indexer.ihd = 128
    dtype=torch.bfloat16, device=device,
)
```

The construction site at `single_shot_inference.py` (where `KVCache` is
created per layer) needs to pass `indexer.ihd` for CSA layers and skip
the buffer for HCA layers (which have no indexer).

---

# PART 2 — MEMORY AT 1M CONTEXT, REVISED

The numbers below replace the prior audit's. They are conservative and
worst-case.

## Per-layer KV cache sizes — read off the (corrected) code

| Component | Per token (compressed) | Bytes / token | × 1M tokens |
|---|---|---|---|
| **CSA main compressed** (1 entry / 4 tokens, hd=512 BF16) | 0.25 × 1024 B | 256 B | **256 MB** |
| **CSA indexer keys** (1 entry / 4 tokens, c_I=128 BF16) | 0.25 × 256 B | 64 B | **64 MB** |
| **HCA compressed** (1 entry / 128 tokens, hd=512 BF16) | 0.0078 × 1024 B | 8 B | **8 MB** |
| **SWA per layer** (ring buffer, 128 × hd × 2) | constant | — | 128 KB |

## Total KV cache @ 1M context, all layers, BF16:

| Layer type | Count | Per-layer @ 1M | Total |
|---|---|---|---|
| CSA: main + indexer | 30 | 256 MB + 64 MB | **9.6 GB** |
| HCA: main | 30 | 8 MB | 240 MB |
| SWA | 61 | 128 KB | 8 MB |
| **GRAND TOTAL @ 1M, BF16** | | | **~9.9 GB** |

**~10 GB of KV state for a 1M-token context.** On 8×B200 (192 GB each, 1.5 TB
total) that's negligible — about 0.7% of total HBM, or ~1.25 GB per GPU if
sharded EP-style alongside the experts. The system has plenty of memory
headroom for the design target.

For comparison, DeepSeek-V3.2's KV cache at 1M context is ~92 GB (per V4
paper Figure 1). V4 at ~10 GB is a 9× reduction — which is **exactly the
"~10% of V3.2's KV cache" claim from the paper.** The implementation hits
the design memory budget; the prior audit was wrong about how it gets there.

## What this changes about priorities

- **"Quantize indexer KV to FP4 to save 121 GB" is gone.** It was based on
  a wrong width. The indexer cache is 2 GB at 1M; FP4 would shrink it to
  500 MB. Nice; not urgent.
- **"max_comp = 65536 is the ceiling at 262K tokens" is still real.** That
  hardcoded buffer size hasn't changed. At 1M context CSA needs
  `max_comp_csa = 262144`. Still a config fix, just not paired with a
  quantization fight.
- **"Allocator churn from `torch.cat` in the gather" is still real and
  still gets worse with context length.** Pre-allocation still matters at
  long context for perf and stability over hours of decoding. Just not
  urgent for "does it fit in memory."

---

# PART 3 — PRIORITY ORDER (REVISED)

Sequenced by what unblocks correctness first, then performance, then
memory. The big shift from the prior audit: **the indexer fix is the
gating correctness work; memory is no longer the crisis it was framed as.**

## Tier 1 — Make the indexer actually work (correctness)

These are all small edits but they have to land together. The agent
should treat this as one atomic landing, not three independent fixes,
because individually each one either does nothing or makes things worse.

### A1 — Fix the indexer compressor weight path

`Indexer.load:392`. Change the check and the load prefix to match the
checkpoint:

```python
# Was:
if f"{pfx}.compressor.kv_proj.weight" in w:
    self.compressor = Compressor(4, self.ihd, 7168, dev)
    self.compressor.load(w, f"{pfx}.compressor", dev)
# Should be (read the actual key from the checkpoint, not assumed):
if f"{pfx}.kv_proj.weight" in w:
    self.compressor = Compressor(4, self.ihd, 7168, dev)
    self.compressor.load(w, f"{pfx}", dev)
```

The agent's probe already identified this — verify the fix is in v17 by
running a checkpoint-loaded forward and confirming `self.compressor is
not None` for at least one CSA layer.

### A2 — Fix `comp_idx_buf` width to `c_I = 128`

`KVCache:419`. Plumb `indexer_key_dim` through `KVCache.__init__` (or
better: derive it from a probe of the indexer's compressor on first
call). Default for non-CSA layers: skip the buffer.

### A3 — Fix the scoring einsum to MQA-on-indexer

`Indexer.forward:404`. Drop the head-axis reshape and use `'tnd,cd->tnc'`
as shown in F3 above. This is the deeper correctness fix and the easiest
one to get wrong if A1+A2 land first and an agent "fixes" the reshape by
widening the buffer.

**Gate for Tier 1:**
1. `Indexer.forward` returns a non-None `idx` tensor for every CSA layer
   on a prompt of ≥ 4 tokens. Verify with a print on layer 0.
2. `forward_attention` at CSA layers takes the
   `if ratio == 4 and topk_idx is not None` branch, not the fallback.
3. Paris-back still works. Output is identical-or-better than v16's
   Paris-back (since v16 was running the dense fallback, which is a
   correctness *superset* of CSA — it attends over more keys, not fewer).
4. **Recall test:** compare the top-k indices from the indexer against
   an FP32 oracle (just compute the scoring in FP32 outside the kernel
   and topk on that). Recall ≥ 99% at top_k=512 with n_comp ≥ 1024.

## Tier 2 — Verify what the fallback was actually doing (cleanup)

### B1 — Find and document the v16 CSA fallback path

`forward_attention:569–571`: when `topk_idx` was always None, what
actually happened in CSA layers? The branches as read:

```python
if ratio == 4 and topk_idx is not None:        # never taken
    all_kv = torch.cat([kv_cache.comp_kv[tk], swa_kv], dim=0)
elif ratio > 4:                                 # only HCA layers
    all_kv = torch.cat([kv_cache.comp_kv, swa_kv], dim=0)
```

For CSA with `topk_idx=None` and `ratio == 4`, **neither branch fires.**
What `all_kv` is at that point depends on what came before. The agent
should run a 5-line probe in v16 (or look at the bisected behavior) to
confirm whether v16 CSA layers:
- attended over just SWA (would explain decode degradation past step 10),
- attended over the full compressed history (would explain decode
  working but being slower than necessary),
- crashed at this point and something downstream rescued the run (most
  likely if Paris-back still happened).

This is *informational* — it doesn't gate Tier 1 — but it answers "what
exactly did 'Paris-back' validate?" and it tells you whether decode
quality should jump (if v16 was on SWA-only) or stay flat (if v16 was on
full compressed) when Tier 1 lands.

### B2 — Once Tier 1 lands, add explicit error on `topk_idx is None` in CSA

The fact that the CSA fallback was silent for this long is the meta-bug.
After Tier 1, the CSA path should *require* `topk_idx is not None`:

```python
if ratio == 4:
    assert topk_idx is not None, f"CSA layer {li} got no top-k from indexer — indexer is broken"
    tk = topk_idx[0].clamp(0, kv_cache.n_comp - 1)
    all_kv = torch.cat([kv_cache.comp_kv[tk], swa_kv], dim=0)
elif ratio > 4:
    all_kv = torch.cat([kv_cache.comp_kv, swa_kv], dim=0)
```

This is a tripwire for future regressions of the same shape.

## Tier 3 — Memory & allocator hygiene (still real, just not urgent)

### C1 — `max_comp` per-layer-type + CLI flag

`KVCache.__init__:411`. Make `max_comp` a function of context length and
compress ratio:

```python
def __init__(self, head_dim, indexer_key_dim, compress_ratio,
             window_size=128, target_context=8192, device='cuda:0'):
    self.max_comp = (target_context + compress_ratio - 1) // compress_ratio
    ...
```

And expose `target_context` as a CLI arg (`--max-context`). Default
small (8192) so the script stays runnable.

### C2 — Pre-allocate `all_kv_buf`, eliminate `torch.cat` in gather

Same fix as D3/D4 in the prior audit — still valid:

```python
# Once at init:
self.all_kv_buf = torch.zeros(max_top_k + window_size, head_dim, ...)
```

Gather writes into views of this buffer with `out=` arguments. FMHA
consumes the prefix. Zero allocs on hot path.

### C3 — `KVCache.get_swa` returns views, not clones

`KVCache:457–460`. Drop the `.clone()` calls. Return slices.

### C4 — Optional: Quantize indexer KV to FP4 (paper §5.2.1)

For throughput, not memory. Defer until E7 (Stage F indexer FP4 tensor-core
scoring) lands — at that point the FP4 storage and FP4 MMA path are paired,
which is the right shape. **Don't quantize the cache without also
upgrading the scoring kernel** — that would be storage savings paid for
with a dequant kernel that doesn't exist yet.

## Tier 4 — Architecture fidelity nice-to-haves

### D1 — Split `Compressor` class into `MainCompressor` and `IndexerKeyCompressor`

`single_shot_inference.py:272`. Same class is instantiated with totally
different config in two places. Splitting documents the difference and
prevents the "I assumed it was the same thing" bug class (which is how
the buffer width bug happened in the first place).

### D2 — Verify sink merge semantics (D6 from prior audit, unchanged)

`_run_production_fmha:489` passes `n_comp=0` always. The kernel may
expect `n_comp = len(compressed_kv)` for the D5c sink merge. Print the
kernel's actual handling, confirm or fix.

### D3 — Understand mHC residual growth (D7 from prior audit, unchanged)

|X| → 500-700 at L60 still indicates Sinkhorn B isn't doubly-stochastic
at runtime. Print B row/col sums, expect 1.0 ± 1e-6. This may also
partly explain the decode degradation past step 10 (compounding
non-bounded residuals → saturated logits → low-information argmax).
Tier 1 fixing the indexer may improve decode behavior enough that this
stops mattering — but worth still checking once the indexer is correct.

---

# REVISED PRIORITY TABLE

| # | Item | What it unblocks | Effort | Blocks 1M? |
|---|---|---|---|---|
| **A1** | Fix indexer compressor weight path | Indexer runs at all | XS | Yes — correctness |
| **A2** | `comp_idx_buf` width = 128 (not 512) | Indexer can store keys | XS | Yes — correctness |
| **A3** | Scoring einsum `'tnd,cd->tnc'` | Top-k is correct | XS | Yes — correctness |
| **B1** | Document the v16 CSA fallback | Knowing what Paris validated | XS | No |
| **B2** | Assert `topk_idx is not None` in CSA | Future regression tripwire | XS | No |
| **C1** | Per-layer `max_comp` + `--max-context` | Long context doesn't crash at 262K | XS | Yes — but trivial |
| **C2** | Pre-alloc `all_kv_buf`, kill cat | Stable decode over hours | S | No, but real perf |
| **C3** | `get_swa` returns views | Small but everywhere | XS | No |
| **C4** | FP4 indexer cache (paired with E7) | Throughput, paper compliance | M-L | No |
| **D1** | Split Compressor classes for clarity | Prevents the same-class-confusion bug | XS | No |
| **D2** | Sink merge semantics check | Subtle numerics | S | No |
| **D3** | mHC Sinkhorn convergence check | Decode degradation | S | No |

**Land A1+A2+A3 together as one atomic correctness fix.** That is the
critical path. Everything else is sequential and not gating.

---

# DOCTRINE — applies to every priority

1. **DSL wall → raw CUDA C++, not Python.** Doesn't apply much in this
   round — most fixes are 3-line edits to Python orchestration. The
   exception is C4 (FP4 indexer cache) which is a kernel fight and must
   follow doctrine: tcgen05/UMMA/TMA on the read side, `__constant__`
   LUT for any dequant, paired with the E7 scoring kernel.

2. **Raw CUDA ≠ scalar math.** Same — when C4 lands, the indexer's
   `tcgen05.mma` FP4 path replaces the scoring einsum. The current FP32
   einsum (post-fix) is a correctness oracle, not a perf target.

3. **Print, don't guess.** This entire round exists because of a probe
   that printed instead of assuming. **The pattern works.** Use it
   again for:
   - B1: probe what the v16 CSA fallback actually returned.
   - C2: print `all_kv` shape and dtype to verify the pre-allocated
     buffer is being sliced correctly.
   - D3: print Sinkhorn row/col sums per layer.
   Stop running new code until the probes have written their output to
   a `.md` next to this one.

4. **Integration over exploration.** No `Indexer_v2`, no `KVCache_v2`.
   Edit the existing classes. Tier 1 is ~10 line-edits total across
   3 functions.

5. **Falsifiable gates.** Already listed per priority above. The
   meta-gate for the whole audit: after Tier 1, **the indexer's top-k
   recall vs an FP32 oracle is ≥ 99% on a prompt with n_comp ≥ 1024.**
   Until that number is measured and recorded, "the indexer works" is
   an assertion, not a fact.

6. **Don't optimize for a problem you don't have.** The prior audit's
   biggest mistake was framing memory as a 1M-context crisis based on
   a wrong width. The real picture is: V4 hit its KV cache memory
   targets, the implementation is faithful, the actual blocker is a
   handful of small bugs in the sparse-selection path. Fix those first
   and re-measure before adding new infrastructure.