fix: assert topk_idx is not None in CSA layers — no silent fallback to SWA-only
The indexer silently returning None caused CSA layers to attend over only the SWA window (128 tokens), not the compressed sparse KV. This went undetected because the model still produced plausible output at short context. The assert makes any future indexer regression immediately visible.
This commit is contained in:
@@ -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.
|
||||
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.
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user