|
| 1 | +# Laguna-S-2.1 (`LagunaForCausalLM` / `laguna`) — W6: KV cache + incremental decode → 5× decode |
| 2 | + |
| 3 | +**Row:** `MODEL-TEXT-laguna-laguna-for-causal-lm` — stays **RUNNABLE / ACTIVE** (speed step). |
| 4 | +**Claim context:** `CLAIM-LAGUNA-W6`. **Date:** 2026-07-31. **Base:** W5 `6e75928a`. |
| 5 | +**HW:** dgx.casa GB10 (sm_121a), 119 GiB unified. **Branch:** `laguna-s21-w6-kvcache` (NOT pushed). |
| 6 | + |
| 7 | +W6 kills the O(n²) full-recompute. W5 made Laguna RUNNABLE but `LagunaForwardGguf` |
| 8 | +is a STATELESS whole-sequence recompute (re-runs the entire prompt+generated |
| 9 | +context for EVERY token). W6 adds a per-layer K/V cache + a single-token |
| 10 | +incremental-decode forward that is **TOKEN-IDENTICAL** to the W5 recompute (a pure |
| 11 | +FLOP-equivalence, bit-exact by construction) and ~5× faster per decoded token. |
| 12 | + |
| 13 | +## Result (the honest finish — same binary, A/B on the real GGUF) |
| 14 | + |
| 15 | +`examples/laguna-gen`, real 3-shard `unsloth/Laguna-S-2.1-GGUF UD-Q4_K_XL` (GB10, |
| 16 | +`--gpu`, keep-quant), greedy, prompt "The capital of France is", 24 tokens: |
| 17 | + |
| 18 | +``` |
| 19 | +STATELESS (W5, --stateless): ids 22345 83 350 785 989 395 13259 330 4159 9431 377 340 4328 377 444 136 22029 9626 71 493 6396 565 7760 10291 TPOT 3.33 s/tok (decode 76.60s/23, prefill 1.37s) |
| 20 | +KV-CACHE (W6, default): ids 22345 83 350 785 989 395 13259 330 4159 9431 377 340 4328 377 444 136 22029 9626 71 493 6396 565 7760 10291 TPOT 0.66 s/tok (decode 15.10s/23, prefill 1.47s) |
| 21 | +``` |
| 22 | + |
| 23 | +- **TOKEN-IDENTICAL GATE: PASS.** The two paths' `generated ids` lines are |
| 24 | + BYTE-EQUAL (identical md5 `754728c6…`), and both equal the recorded W5 golden |
| 25 | + verbatim. A KV cache changes nothing numerically when correct — this is the |
| 26 | + bit-exact-by-construction proof, and it held on the FIRST run (no cache bug). |
| 27 | +- **SPEED: 3.33 → 0.66 s/tok decode = 5.05× faster.** The stateless path's |
| 28 | + per-step time GREW with context (step 1 = 0.68s … step 23 = 5.08s: the O(n²) |
| 29 | + signature); the cached decode is FLAT ~0.66–0.69 s/tok regardless of context. |
| 30 | +- Load 48.1s (cold, drop_caches), peak resident 71.09 GiB (119 GiB pool). Worker |
| 31 | + down, `flock $HOME/gpu.lock`, one model resident. |
| 32 | + |
| 33 | +## Design — the KV cache + incremental decode |
| 34 | + |
| 35 | +`LagunaKvCache` (`laguna.h`) mirrors `DeepseekV4KvCache`, extended from MLA's |
| 36 | +single `deck` latent to GQA multi-head K/V. Per layer it stores the POST-QK-RMSNorm |
| 37 | +/ POST-RoPE keys and the RAW values at f32 (bit-exact to the W5 recompute); a |
| 38 | +`len` counter is the global position of the next token; `first_pos[layer]` is the |
| 39 | +global position of the layer's first cached row (advanced by sliding eviction). |
| 40 | + |
| 41 | +`LagunaForwardGgufCached` (`laguna.cpp`) is the same keep-quant composition as |
| 42 | +`LagunaForwardGguf` but binds the cache: |
| 43 | +- **Prefill** (first call, `cache.len==0`): all prompt tokens, positions 0..P-1. |
| 44 | + Projects q/k/v for all P rows, qk-norms + RoPEs them, APPENDS the K/V to each |
| 45 | + layer's cache, then attends (the existing full O(P²) masked attention). |
| 46 | +- **Decode** (later calls, T=1): ONE new token at position `cache.len`. Projects |
| 47 | + its q/k/v, qk-norms + RoPEs, APPENDS the single K/V row, attends the new query |
| 48 | + over the layer's cached K/V. Global positions `0..len`. |
| 49 | + |
| 50 | +**Why it is bit-exact:** RoPE and QK-RMSNorm depend only on a token's OWN absolute |
| 51 | +position, and Laguna attention is causal, so a token's hidden state (and thus its |
| 52 | +K/V) depends only on tokens `0..j` — identical whether recomputed at step j+1 or |
| 53 | +at any later step. Caching the post-RoPE K + raw V therefore reproduces exactly |
| 54 | +what the recompute would recompute. To guarantee identical FLOAT ops (not just |
| 55 | +"mathematically equal"), the attention inner loop and the FFN block are EXTRACTED |
| 56 | +into shared helpers (`LagunaAttention`, `LagunaFfnBlock`, `LagunaEmbed`, |
| 57 | +`LagunaFinalLogits`) that BOTH forwards call — the two paths run the same code, |
| 58 | +differing only in the K/V source (fresh projections vs cache) and the query set |
| 59 | +(all T rows vs the one new row). The gate confirms the recompute path is unchanged |
| 60 | +(its ids still match the W5 golden after the refactor). |
| 61 | + |
| 62 | +### The NEW bit — MIXED attention: global vs sliding-window eviction |
| 63 | + |
| 64 | +Laguna is 12 GLOBAL layers (full causal) + 36 SLIDING-WINDOW-512 layers. Cache |
| 65 | +handling is per-layer (grounded in gemma2/3 `is_sliding`): |
| 66 | +- **Global layers** (`window==0`): cache grows unbounded — the whole history is |
| 67 | + kept; `first_pos` stays 0. |
| 68 | +- **Sliding layers** (`window==512`): after appending the new row, if the row |
| 69 | + count exceeds 512 the oldest rows are EVICTED from the front of the flat K/V |
| 70 | + vectors and `first_pos` advances by the drop count. A query at global position |
| 71 | + `pi` scores only kv with `pi - pj < window`, so once >512 rows are cached the |
| 72 | + oldest can never be scored again — evicting them is exact, and it caps the |
| 73 | + sliding-layer K/V at the 512-token window (the memory win). The attention still |
| 74 | + applies the `pi - pj >= window` mask, so correctness does not depend on the |
| 75 | + eviction firing (for the 6+24-token gate the window is never exceeded, so |
| 76 | + eviction is a no-op here and the sliding cache holds all rows — the code path is |
| 77 | + present and unit-safe; a >512-ctx run is the future exercise of the eviction). |
| 78 | + |
| 79 | +Per-layer variable Q-head count (48 global / 72 sliding) and GQA group flow |
| 80 | +through unchanged — only KV heads (8) and head_dim (128) are cached, which are |
| 81 | +uniform across layers. |
| 82 | + |
| 83 | +## Driver |
| 84 | + |
| 85 | +`examples/laguna_gen` gained `--stateless` (force the W5 O(n²) recompute for the |
| 86 | +A/B gate). Default is the W6 KV-cache path: step 0 prefills the whole prompt |
| 87 | +(local logits index = last row), later steps feed ONE new token at `positions = |
| 88 | +{kv.len}` with local logits index 0. Env: no new flags for correctness; the |
| 89 | +existing `--gpu` routes the keep-quant GEMMs to the GB10. |
| 90 | + |
| 91 | +## Gate + build |
| 92 | + |
| 93 | +- **Build (disk-aware):** CUDA `RelWithDebInfo -DVLLM_CPP_CUDA=ON |
| 94 | + -DVLLM_CPP_CUDA_ARCHITECTURES=121a -DVLLM_CPP_CUTLASS_DIR=$HOME/cutlass-4.5.0 |
| 95 | + -DVLLM_CPP_TRITON=ON -DVLLM_CPP_BUILD_TESTS=OFF`, target `laguna-gen` only |
| 96 | + (cutlass found, FA2 enabled). Engine lib + example built RC=0; build tree |
| 97 | + cleaned after (2.7 GiB, box restored). |
| 98 | +- **Correctness gate:** TOKEN-IDENTICAL PASS (byte-equal ids, md5 match, both == |
| 99 | + W5 golden). No cache bug found — the shared-helper design was bit-exact first try. |
| 100 | +- **Speed:** decode 3.33 → 0.66 s/tok (5.05×), prefill unchanged (~1.4s). |
| 101 | + |
| 102 | +## Residuals (future speed) |
| 103 | + |
| 104 | +0.66 s/tok is still host-orchestrated per-token per-expert GEMV over the keep-quant |
| 105 | +tower. The next throughput levers (both already in-tree from ds4): |
| 106 | +`MatmulBTQuantGrouped` grouped-expert GEMM + a device-resident decode chain |
| 107 | +(`ForwardResidentDecodeGguf`/decode CUDA-graph analog). Token-exact-vs-llama.cpp on |
| 108 | +the shared quant is the remaining correctness tightening. The O(n²)→O(n) recompute |
| 109 | +kill (this row) is the dominant decode win and is landed + gated. |
| 110 | + |
| 111 | +## Decision |
| 112 | + |
| 113 | +Row stays **RUNNABLE / ACTIVE**. The KV-cache incremental decode is token-identical |
| 114 | +to the W5 recompute (byte-equal, proven on the real bytes) and 5× faster per token. |
| 115 | +Not fabricated: both id streams are recorded verbatim and their md5s match. |
0 commit comments