diff --git a/ARCHITECTURE_AND_MEMORY_AUDIT.md b/ARCHITECTURE_AND_MEMORY_AUDIT.md index 213a3928..d0b57910 100644 --- a/ARCHITECTURE_AND_MEMORY_AUDIT.md +++ b/ARCHITECTURE_AND_MEMORY_AUDIT.md @@ -1,434 +1,467 @@ -# ARCHITECTURE & MEMORY AUDIT — 1M-context viability +# ARCHITECTURE & MEMORY AUDIT — Post-probe rewrite -**Method.** Verified against `single_shot_inference.py` v16 and the DSv4 -paper §2.3.1–§2.3.4 (CSA/HCA) and §3.5.1 (heterogeneous KV cache). Every -finding has a line number. Per doctrine. +**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`. -**Framing.** The Paris demo runs ≤ 50 tokens. The model is built for 1M. -That's a **20,000×** gap. Several things in the current single_shot are -fine at 50 tokens, will OOM hard at 4–10K tokens, and don't even resemble -the paper's KV design at 1M. Below: drift first, then memory, in order of -how badly each one blocks the 1M-context goal. +**Method.** Every claim verified against `single_shot_inference.py` v16 + the +probe results. Per doctrine. --- -# PART 1 — ARCHITECTURE DRIFT FROM PAPER +## TL;DR — the picture is much better than the prior audit suggested -## D1 — `comp_idx_buf` shape is wrong (CRITICAL — silent corruption or crash) +**The architecture is faithful to the paper. The 1M-context memory story is +fine on 8×B200. There is no looming OOM crisis.** -`single_shot_inference.py:419`: +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 -self.comp_idx_buf = torch.zeros(max_comp, head_dim, dtype=torch.bfloat16, device=device) +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 — WRONG + 512 — should be 128 ``` -But indexer keys are `n_ih × ihd` wide. From `:1030`: +`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 -n_ih = cfg.get("index_n_heads", 64) -ihd = cfg.get("index_head_dim", 128) +# In __init__: +self.comp_idx_buf = torch.zeros( + max_comp, + indexer_key_dim, # passed from caller, = indexer.ihd = 128 + dtype=torch.bfloat16, device=device, +) ``` -So indexer keys have width `64 × 128 = 8192`, not 512. **The indexer KV -buffer is 16× too narrow.** What happens in practice depends on whether -the assignment broadcasts, raises, or silently truncates — and from the -fact that Paris-back works, the indexer probably isn't being used at all -yet (CSA layers may be producing 0 compressed blocks at 50 tokens — see -D2). At any context where CSA actually compresses, this either crashes -or stores garbage in the top-k selection input. - -**Fix.** Read the actual indexer key width from the indexer's compressor -output (`indexer.compressor.kv_dim = 2 * ihd = 256` for the indexer-side -CSA, since the indexer's compressor takes `(4, ihd, H)`). Then check what -the indexer's compressor actually produces — print its output shape on -first call instead of guessing — and size `comp_idx_buf` to match. - -Verification step: instrument `Compressor.forward` to print -`compressed.shape` on first call from both the **main** compressor (kv_dim -= 512 or 1024) and the **indexer's** compressor (kv_dim = 256). Code to -the observed values. Do not infer them from variable names. - -## D2 — The Compressor is built twice per CSA layer, with different config - -`:394` (inside `Indexer.load`): -```python -self.compressor = Compressor(4, self.ihd, 7168, dev) -``` - -`:1034` (in `main()`): -```python -if ratio > 0: compressors[li] = Compressor(ratio, hd, H, dev) -``` - -For a CSA layer: -- The **layer's** compressor has `(ratio=4, hd=512, H=7168)`, output dim - `kv_dim = 2*hd = 1024`. -- The **indexer's** compressor has `(ratio=4, ihd=128, H=7168)`, output - dim `kv_dim = 2*ihd = 256`. - -Both are constructed, both load weights independently, both have their own -NVFP4 GEMM Nvfp4Linear instances. This matches the paper (§2.3.1: the -indexer has its own compressed key path that's narrower than the main -compressed KV path) — but **it is being done as two completely -independent code paths**, with the indexer's compressor's existence -implicit and easy to miss. That's why D1 was missed. - -It also means the main compressor's `forward` runs twice per layer at -prefill: once for the main KV path, once for the indexer's keys. The -hidden states being projected are *identical*; the GEMM weights are -different. Two separate launches. - -**Architecturally correct, but the code shape hides it.** Two consequences: - -1. The "compressor" abstraction has a different meaning depending on - which compressor instance you're looking at. Rename for clarity: - `Compressor` (main) and `IndexerKeyCompressor` (smaller, for indexer). -2. The Indexer should expose its compressed key width as a property - (`indexer.compressor.kv_dim // 2`) so D1 can compute the buffer width - from data instead of assuming `head_dim`. - -## D3 — KV gather still uses `torch.cat`, undoing P3's pre-allocation - -`:569–571`: - -```python -if ratio == 4 and topk_idx is not None: - 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) -``` - -P3 preallocated `comp_kv_buf`, but the gather **immediately allocates a -fresh `(top_k + ws, hd) = (1024 + 128, 512)` BF16 tensor per layer call** -just to pass to FMHA. For 61 layers × per-token decode: - -- 61 × (1152 × 512 × 2 bytes) ≈ **72 MB allocated and freed per token**. - -That's small in absolute terms, but it's allocator churn on the hot path, -and the entire point of P3 was to remove this pattern. The `comp_kv[tk]` -gather already allocates (it's an advanced-indexing copy); the `cat` then -allocates again to merge with SWA. Two allocs per layer that don't need -to exist. - -**Fix.** Preallocate one more buffer at cache init: - -```python -self.all_kv_buf = torch.zeros(top_k + window_size, head_dim, ...) -``` - -Write the gathered top-k into `all_kv_buf[:top_k]`, the SWA into -`all_kv_buf[top_k:top_k+swa_len]`, pass `all_kv_buf[:top_k+swa_len]` to -FMHA. The gather becomes `torch.index_select(comp_kv_buf, 0, tk, -out=all_kv_buf[:top_k])` — zero allocs. - -## D4 — The HCA path concatenates the ENTIRE compressed history, every layer, every token - -Same line, the `elif` branch: - -```python -elif ratio > 4: all_kv = torch.cat([kv_cache.comp_kv, swa_kv], dim=0) -``` - -HCA layers don't run the indexer. They attend over the **full compressed -history**. At 1M context with HCA ratio=128, that's `1M / 128 = 7813` -compressed entries. Per-token-per-layer FMHA input grows linearly with -prefill length. That's correct math (paper §2.3.2: HCA is dense over -compressed entries), but the *implementation* is allocating a new -contiguous tensor for it every single layer call. - -At 1M context, that's `(7813 + 128) × 512 × 2 bytes ≈ 7.7 MB` per HCA -layer call. Across ~30 HCA layers per token ≈ **230 MB of alloc-and-free -per token, on the decode hot path.** - -**Fix.** Same as D3: preallocate `all_kv_buf` sized for the worst case -(HCA full history + SWA). Use `out=` parameters on the gather. Or skip -the concat entirely — the FMHA kernel can take two K/V tensors and the -mask can encode the boundary. Today it can't, but it should; this is a -real kernel-side ask (see "M5" below). - -## D5 — Indexer top-k attends *only* compressed entries; SWA is appended separately. Confirm against paper - -Paper §2.3.1 figure: CSA's FMHA input is -`Concatenate(selected_compressed_KV, sliding_window_KV)`. Yes, that's -what the code does at `:571`. **Architecturally correct.** - -One subtle thing worth checking: the paper says the sliding window provides -*local* fine-grained context the compressor can't (since compressed entries -each represent m tokens). The current SWA window is **128** (from config). -At 1M context that's still 128 local tokens visible per query — correct. -But the window slide in `KVCache.append_swa` evicts when full -(`self.swa_head = (self.swa_head + T) % self.ws`), which is correct. -✅ no drift. - -## D6 — Attention sink is wired, but per-CSA-layer correctness needs checking - -`_run_production_fmha:489`: - -```python -sinks = w.get(f"{pfx}.sinks") -if sinks is not None: sink_bias = sinks.to(device=dev).float().reshape(n_h) -attn_out = dsv4_attention(q=q, k=k, v=v, scale=scale, n_comp=0, sink_bias=sink_bias) -``` - -`n_comp=0` is passed regardless. Paper §2.3.3 ("Attention Sink"): the sink -is a per-head additive logit to the softmax denominator. The kernel -signature in v15 said `n_comp` is "reserved for future kernel integration" -for the **D5c sink merge** (different softmax over compressed vs SWA). -v16 still passes `n_comp=0`, which means we're using global sink, not -per-segment sink merge. - -The paper isn't explicit about whether sink should be per-segment, but if -the production FMHA was designed for D5c merge and it's being bypassed, -that's an unfinished integration, not necessarily a bug. - -**Action:** confirm against the kernel's actual handling. Print the -sink_bias usage in `dsv4_attention` for one layer. If sink merge is -needed for CSA correctness at long context, that's a real wiring gap. - -## D7 — mHC residual growth (|X|→500–700 at L60) was flagged but not understood - -Your perf audit notes (line 88): "Residual |X| grows to 500–700 at L60 -— mHC bounds it but residual is high." - -Paper §2.2 designed mHC specifically to **bound** the residual via the -doubly-stochastic B matrix (spectral norm ≤ 1). The growth from |X|=1 at -L0 to |X|=700 at L60 suggests B isn't actually norm-bounded at runtime, -or A·C are amplifying. - -This is the **same** issue I flagged in the v14 docs and it's still open. -Not a perf bug, but it's an architecture-fidelity bug, and it's the -**single most likely cause of decode degradation past step 10** (the -"...the" repetition loop noted in your audit). Compounded across decode -steps, a slightly-not-bounded residual becomes a numerically saturated -residual, which makes the final logits less informative. - -**Action.** Print Sinkhorn-Knopp B's row/col sums per layer for one -forward pass. They should be `1.0 ± 1e-6` if Sinkhorn converged. If -they're e.g. `1.02–1.05`, t_max=20 isn't converging at this scale; bump -it or check the dynamic-parameter generation. The single_shot does -`sinkhorn_iters=20` (`:937`) which matches the paper, so the issue is -likely upstream: either A or C is producing values outside [0, 1] or -[0, 2], or the dynamic parameter generation has FP32 noise that breaks -doubly-stochastic. +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 +# PART 2 — MEMORY AT 1M CONTEXT, REVISED -This is the part that should be terrifying. The single_shot was sized for -50 tokens; the model targets 1M. Below is what each KV-cache structure -costs at the four interesting scales, with all numbers worked out, not -estimated. +The numbers below replace the prior audit's. They are conservative and +worst-case. -## Per-layer KV cache sizes — read off the code +## Per-layer KV cache sizes — read off the (corrected) code -Layer setup: -- **CSA layer** (compressor ratio=4): main compressed at `hd=512` BF16, - indexer keys at `n_ih * ihd = 64 * 128 = 8192` BF16, SWA at `ws=128 × hd` -- **HCA layer** (compressor ratio=128): main compressed at `hd=512` BF16, - no indexer, SWA at `ws=128 × hd` -- **SWA-only layer** (first 2 layers of Flash, or HCA per-layer-2 of Pro): - SWA only - -Layer counts per V4-Pro: 61 total, alternating CSA/HCA after layer 1. So -roughly **30 CSA + 30 HCA + 1 SWA-only** (paper §4.2.1). - -### Per-layer KV growth (per token of context): - -| Component | Per token | Bytes / token | × 1M tokens | +| 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, 8192 BF16) | 0.25 × 16384 B | 4096 B | **4.1 GB** | +| **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, fixed 128 × hd × 2) | const | — | 128 KB | +| **SWA per layer** (ring buffer, 128 × hd × 2) | constant | — | 128 KB | -### Total KV cache @ 1M context, all layers, BF16: +## Total KV cache @ 1M context, all layers, BF16: | Layer type | Count | Per-layer @ 1M | Total | |---|---|---|---| -| CSA: main + indexer | 30 | 256 MB + 4.1 GB | **131 GB** | +| 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** | | | **~131 GB** | +| **GRAND TOTAL @ 1M, BF16** | | | **~9.9 GB** | -**The KV cache alone is 131 GB.** That's before model weights (which are -already EP-sharded across 8 GPUs). On 8 × B200 with 192 GB each = 1.5 TB -total, sharding KV across the 8 GPUs gives ~16 GB per GPU — fits, but -it's **15% of HBM dedicated to one request's KV.** And the dominant cost -is the **indexer keys** at 4.1 GB per layer. +**~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. -**Critical observation: the indexer keys are 16× larger than the main -compressed KV per token, and they alone are 86% of total KV.** This is -because `n_ih * ihd = 8192` is much wider than `hd = 512` for the main KV. -The paper does specify this — the indexer is a wide multi-query mechanism -— but if storage is the constraint, the indexer key path is where to -attack. +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. -## M1 — `comp_idx_buf` allocation as written: `(65536, 512)` per layer +## What this changes about priorities -`:419`: -```python -self.comp_idx_buf = torch.zeros(max_comp, head_dim, dtype=torch.bfloat16, ...) -``` - -If this *were* the correct width `(max_comp=65536, 8192)`, that's -`65536 × 8192 × 2 = 1 GB per layer × 30 CSA layers = 30 GB pre-allocated -at startup`, sized for **only 262K tokens of context**, not 1M. - -The current shape `(65536, 512)` allocates 64 MB per layer × 30 = 1.9 GB -— too small by 16× as noted in D1, so either crashes at first compressed -write or silently truncates. Both bad. After fixing D1, you're staring -down 30 GB of pre-allocated KV cache that supports only a quarter of the -target context. And it's 30 GB × 8 GPUs = 240 GB if everything is -replicated, or 30 GB total if cache is sharded. - -**This is the load-bearing memory bug for 1M context.** Two viable fixes: - -1. **Quantize the indexer keys to FP4** (paper §5.2.1: "QK activations - are cached, loaded, and multiplied entirely in FP4"). The indexer keys - are *designed* to be FP4. That's 16× smaller: 4.1 GB/layer → 256 MB at - 1M. Total cache becomes ~10 GB instead of 131 GB. **This is the - correct fix per paper.** -2. **Page the indexer KV.** Only the top-k indices' worth (≤ 1024) need to - be in attention. A paged cache (per the paper's §3.5.1 "heterogeneous - KV cache" with on-disk overflow) only keeps recent + selected pages in - HBM. - -(1) is required regardless of (2). The current BF16 indexer cache is the -single biggest blocker between Paris-demo-works and 1M-context-works. - -## M2 — `max_comp = 65536` hardcoded - -`:411`: -```python -def __init__(self, head_dim, window_size=128, max_comp=65536, device='cuda:0'): -``` - -For CSA (ratio=4), 65536 compressed entries = `262144` tokens of context. -That's the ceiling. At 262K tokens you get `IndexError` on -`comp_kv_buf[self.n_comp:end] = ckv`. There is no graceful behavior, no -on-disk overflow, no error message — it just dies. - -For 1M context, `max_comp` needs to be `1M / 4 = 262144` for CSA layers -and `1M / 128 = 7813` for HCA layers. Currently both layer types share -the same 65536 default. - -**Fix.** Size the buffer per-layer using compress ratio: -```python -max_comp_csa = ceil(target_context / 4) # 262144 for 1M -max_comp_hca = ceil(target_context / 128) # 7813 for 1M -``` -And, critically, make `target_context` a CLI flag with a sensible default -(say 8K) so the script can run small while staying honest about the -ceiling. Hardcoded 65536 with no docstring on what it means is a footgun. - -## M3 — Allocator churn from gather (D3, D4) compounds at 1M - -Repeated for emphasis with numbers: the per-token `torch.cat` in the KV -gather allocates and frees memory proportional to context length. At 1M -context with HCA at all layers, that's **~230 MB of alloc/free per -decoded token**. PyTorch caching allocator handles this but it's still -fragmentation pressure across thousands of decoded tokens. After hours -of decoding, the allocator's cached blocks bloat. - -Already covered in D3/D4. Restating because at 1M scale it's no longer -"small overhead" — it's GB/min of churn. - -## M4 — `get_swa` does `.clone()` every call - -`:457–460`: -```python -def get_swa(self): - if self.swa_len == 0: return torch.zeros(0, self.hd, ...), torch.zeros(0, ...) - if self.swa_len < self.ws: return self.swa[:self.swa_len].clone(), self.swa_pos[:self.swa_len].clone() - idx = torch.arange(self.swa_head, self.swa_head + self.ws) % self.ws - return self.swa[idx].clone(), self.swa_pos[idx].clone() -``` - -Three clones in the second return path (and the third one allocates an -arange too). At `ws=128 hd=512`, the SWA tensor is 128 KB — small — but -**this runs every layer every token**. 61 × decoded tokens × 128 KB ≈ a -few MB/token in allocator pressure. Same fix pattern as D3: return views -into the ring buffer, let the FMHA gather kernel consume them with strides. - -## M5 — KV gather memory could be eliminated entirely with a smarter kernel - -D3 and D4 both pre-allocate a fused `all_kv` buffer because the FMHA -takes one K/V tensor. The deeper fix is to have the FMHA take **two K/V -inputs** (compressed + SWA) and handle them with masking inside. Then no -fused buffer ever has to exist — the kernel reads compressed entries -directly from `comp_kv_buf` (with a gather indices vector for the top-k -selection) and SWA entries directly from `swa_buf`. - -This is the **right** long-term design (and matches how the paper §3.5.1 -envisions the heterogeneous cache). It's a kernel ask, not a script fix, -but it's worth flagging now: the gather → cat → FMHA pattern is *the* -memory inefficiency at long context, and it can be designed out at the -kernel boundary. +- **"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 +# PART 3 — PRIORITY ORDER (REVISED) -These are sequenced by **what blocks 1M context viability**, not by -implementation cost. +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.** -| # | Item | Required for 1M? | Effort | -|---|---|---|---| -| **A1** | **D1: Fix `comp_idx_buf` width to actual indexer key width** | **Critical — silent corruption today** | XS | -| **A2** | **M1: Quantize indexer KV to FP4** (paper §5.2.1) | **Critical — saves 121 GB at 1M** | M-L | -| **A3** | **M2: Make `max_comp` per-layer-type + a CLI flag** | **Critical — current ceiling is 262K** | XS | -| **A4** | D3/D4/M3: Preallocate `all_kv_buf`, eliminate `torch.cat` | High — perf and stability over hours of decode | S | -| **A5** | D7: Investigate mHC residual growth (Sinkhorn convergence print) | High — likely root of decode degradation | S | -| **A6** | M4: Return SWA views, no clone | Medium — small per-call, large in aggregate | XS | -| **A7** | D2: Rename `Compressor` → split into `MainCompressor` and `IndexerKeyCompressor` | Medium — clarifies the duplicate-build pattern | XS | -| **A8** | D6: Verify sink merge semantics with kernel author | Medium — possible silent numerical drift | S | -| **A9** | M5: Two-buffer FMHA kernel (eliminate gather buffer entirely) | Long-term — production design | L | +## Tier 1 — Make the indexer actually work (correctness) -**The "should I run on 1M context tomorrow" answer is no, regardless of -anything else, until A1/A2/A3 are done.** Without A1 you get garbage -top-k. Without A2 you OOM at ~250K tokens even with 8×B200. Without A3 -you crash at 262K. Together those three are the gating set. +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. -**The "is it still architecturally DSv4?" answer is yes — mostly.** The -hot path is faithful to the paper: CSA does overlapped-2m compression -with softmax weights, HCA does heavy non-overlapped compression, indexer -does ReLU(QK)·w_h reduction, attention concatenates selected-compressed -+ SWA, sinks are applied. The drifts above (D1, D2, D6, D7) are -implementation flaws or unfinished wiring, not architectural deviations. +### 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. --- -# DOCTRINE — applies to every priority above +# REVISED PRIORITY TABLE -1. **DSL wall → raw CUDA C++, not Python.** Most of the fixes above are - pre-allocation and shape correctness, not new kernels. The two - exceptions (A2 FP4 indexer KV, A9 two-buffer FMHA) are kernel work and - must follow doctrine: tcgen05/UMMA/TMA, not scalar. +| # | 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 | -2. **Raw CUDA ≠ scalar math.** When A2 lands the FP4 indexer cache, the - dequant on the read side must use `__constant__` LUT (per the original - indexer LUT fix from issue #1), not branch arithmetic. +**Land A1+A2+A3 together as one atomic correctness fix.** That is the +critical path. Everything else is sequential and not gating. -3. **Print, don't guess.** A1's fix is the canonical example: do not - assume the indexer key width is `head_dim` or `n_ih * ihd` — print - `indexer.compressor.forward(...)[0].shape` on first call and code to - that. The current bug exists *because* someone wrote `head_dim` - thinking it was right. +--- -4. **Integration over exploration.** No `KVCache_v2`. Edit `KVCache`. - The 4 fixes (A1/A3/A4/A6) are surgical edits to one class. +# DOCTRINE — applies to every priority -5. **Falsifiable gates.** Numbers to hit: - - A1: `comp_idx_buf.shape[1] == indexer.compressor.kv_dim // 2` (or - whatever the print reveals). Test: run with VERBOSE=2 at 100 tokens - of context; no shape mismatch, no garbage in top-k selection. - - A2: KV cache footprint at 1M context (measured via - `torch.cuda.memory_allocated()` after prefill) drops from ~131 GB to - ≤ 15 GB. Recall@1024 vs FP32 indexer oracle ≥ 99.7% per paper. - - A3: A `--max-context N` flag works; running with `N=1048576` does - not OOM during prefill (it might be slow — that's a separate fight). - - A4: `torch.cuda.memory_reserved()` measured every 10 decode steps is - flat (±50 MB) across 1000 steps. \ No newline at end of file +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. \ No newline at end of file diff --git a/single_shot_inference.py b/single_shot_inference.py index 9b95384f..37e0b499 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -594,7 +594,8 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, _pt('gather_start') swa_kv, swa_pos = kv_cache.get_swa() if kv_cache.comp_kv is not None and kv_cache.n_comp > 0: - if ratio == 4 and topk_idx is not None: + if ratio == 4: + assert topk_idx is not None, f"CSA layer {li}: indexer returned no top-k — 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)