|
| 1 | +# ITKV x Shard: Importance-Tiered Structured KV Cache Compression |
| 2 | + |
| 3 | +Status: research roadmap. NOT for V1 launch. Combines two independent |
| 4 | +techniques the kolm codebase already understands: |
| 5 | + |
| 6 | +- **Shard** (`src/kv-cache-shard.js`, `github.com/krish1905/shard`) — |
| 7 | + structural compression: K via PCA + int4 after undoing RoPE; V via |
| 8 | + Hadamard rotation + VQ256. Roughly 10x at parity quality on RoPE-based |
| 9 | + decoder-only models. |
| 10 | +- **ITKV** (`src/itkv-profile.js`, kolm W722 / Invention 4) — token-class |
| 11 | + scorer that labels tokens as sink, policy, schema, retrieved_evidence, |
| 12 | + conversation_recent, boilerplate, or irrelevant_span, then allocates |
| 13 | + precision by class (BF16 sinks + policy, INT8 warm, INT4 cold). |
| 14 | + |
| 15 | +The two are orthogonal: Shard tells you HOW to compress per token; ITKV |
| 16 | +tells you WHICH tokens are worth more bits. This document sketches the |
| 17 | +combination. |
| 18 | + |
| 19 | +## 1. Background |
| 20 | + |
| 21 | +### 1.1 What Shard does |
| 22 | + |
| 23 | +Shard's `ShardCache` is a `transformers.Cache` subclass with three regions |
| 24 | +per layer per attention head: |
| 25 | + |
| 26 | +- **Sink region** (tokens 0..S-1, default S=4) — FP16, never compressed. |
| 27 | +- **Window region** (tokens T-W..T-1, default W=64) — FP16, never compressed. |
| 28 | +- **Compressed tail** (tokens S..T-W-1) — K via PCA + int4 in the principal |
| 29 | + subspace; V via Hadamard rotation + 256-entry codebook index (VQ256). |
| 30 | + |
| 31 | +The compressed tail uniformly spends ~1.5 bits per original element. It does |
| 32 | +not know whether a given token is a system-prompt instruction, a tool |
| 33 | +description, or a throwaway boilerplate line — they all get the same budget. |
| 34 | + |
| 35 | +### 1.2 What ITKV does |
| 36 | + |
| 37 | +ITKV scores each token at insertion time with a class label. The default |
| 38 | +precision policy in `src/itkv-profile.js` is: |
| 39 | + |
| 40 | +``` |
| 41 | +sink -> BF16 (highest) |
| 42 | +policy -> BF16 |
| 43 | +schema -> FP8 or INT8 |
| 44 | +retrieved_evidence -> precision by citation confidence |
| 45 | +conversation_recent -> BF16 |
| 46 | +boilerplate -> INT4 or prefix-cache reference |
| 47 | +irrelevant_span -> compress or evict |
| 48 | +``` |
| 49 | + |
| 50 | +The ITKV profile is a SCHEMA + SCORER today. The runtime tier-dispatch is |
| 51 | +out of scope in W722 — it plugs into vLLM PagedAttention / SGLang radix |
| 52 | +cache when those land. |
| 53 | + |
| 54 | +### 1.3 The gap |
| 55 | + |
| 56 | +Shard underspends on high-importance tokens (sink + policy + tool schemas |
| 57 | +get the same 1.5 bpe as boilerplate). ITKV has no compressor at the byte |
| 58 | +level — only labels. Combining them turns ITKV labels into per-token |
| 59 | +bit-budgets for the Shard codebook. |
| 60 | + |
| 61 | +## 2. Novel combination: per-token VQ codebook size by ITKV importance |
| 62 | + |
| 63 | +The kernel insight: VQ256 is one specific choice on a continuum. |
| 64 | +VQ-{N} uses log2(N) bits per token slot. Per-token codebook size by class: |
| 65 | + |
| 66 | +``` |
| 67 | +sink, policy -> VQ-512 (9 bits/slot) high precision |
| 68 | +schema, evidence -> VQ-256 (8 bits/slot) Shard default |
| 69 | +conv_recent -> VQ-256 (kept in window in practice) |
| 70 | +boilerplate -> VQ-128 (7 bits/slot) |
| 71 | +irrelevant_span -> VQ-64 (6 bits/slot) or evict |
| 72 | +``` |
| 73 | + |
| 74 | +K side: rank of the PCA subspace is also class-aware. High-importance |
| 75 | +tokens get the full Shard rank (e.g. 16); irrelevant_span tokens get a |
| 76 | +truncated rank (e.g. 8) with the residual snapped to zero. |
| 77 | + |
| 78 | +## 3. Architecture (ASCII) |
| 79 | + |
| 80 | +``` |
| 81 | + +---------------------------+ |
| 82 | + | Incoming token batch | |
| 83 | + +-------------+-------------+ |
| 84 | + | |
| 85 | + v |
| 86 | + +---------------------------+ |
| 87 | + | ITKV Scorer | (existing, src/itkv-profile.js) |
| 88 | + | -> class label per | |
| 89 | + | token | |
| 90 | + +-------------+-------------+ |
| 91 | + | |
| 92 | + v |
| 93 | + +---------------------------+ |
| 94 | + | ShardITKVCache (NEW) | |
| 95 | + | - K: PCA rank by class | |
| 96 | + | - V: VQ-{N} by class | |
| 97 | + | - Sink + window untouched| |
| 98 | + +-------------+-------------+ |
| 99 | + | |
| 100 | + v |
| 101 | + +---------------------------+ |
| 102 | + | Per-layer tiered storage | |
| 103 | + | sink/policy : FP16 | |
| 104 | + | schema/ev : VQ256+r16| |
| 105 | + | conv_recent : FP16 | |
| 106 | + | boilerplate : VQ128+r12| |
| 107 | + | irrelevant : VQ64+r8 | |
| 108 | + +---------------------------+ |
| 109 | +``` |
| 110 | + |
| 111 | +A class-label sidecar (uint8 per token per layer) lives alongside the |
| 112 | +compressed K and V tensors so decompression knows which codebook to load. |
| 113 | + |
| 114 | +## 4. Math: bits per token with importance weights |
| 115 | + |
| 116 | +Let: |
| 117 | + |
| 118 | +- `f_c` = expected fraction of tokens in class c |
| 119 | +- `bV_c` = V VQ bits per slot for class c |
| 120 | +- `r_c` = K PCA rank for class c |
| 121 | +- `d` = head_dim (per-head dimension) |
| 122 | +- `H_kv` = num_key_value_heads |
| 123 | +- `L` = num_hidden_layers |
| 124 | + |
| 125 | +Bits per token, averaged across classes (ignoring sink + window for now): |
| 126 | + |
| 127 | +``` |
| 128 | +bits_per_token = 2 * L * H_kv * sum_c f_c * (r_c * 4 / d + bV_c) |
| 129 | +``` |
| 130 | + |
| 131 | +The factor `r_c * 4 / d` reflects K stored as int4 over a rank-`r_c` |
| 132 | +subspace projected from a `d`-dimensional space (so per-token-per-head |
| 133 | +K cost is `r_c * 4` bits, amortized over `d` original slots). |
| 134 | + |
| 135 | +Baseline Shard (uniform): `r = 16, bV = 8` for all c, giving roughly: |
| 136 | + |
| 137 | +``` |
| 138 | +bits_per_token_shard ~= 2 * L * H_kv * (16*4/d + 8) |
| 139 | + ~= 2 * L * H_kv * (0.5 + 8) for d=128 |
| 140 | + ~= 17 * L * H_kv bits/token |
| 141 | +``` |
| 142 | + |
| 143 | +ITKV-weighted (assuming typical agent/RAG workload mix: |
| 144 | +sink/policy=10%, schema/evidence=30%, conv_recent=15%, boilerplate=35%, irrelevant=10%): |
| 145 | + |
| 146 | +``` |
| 147 | +sum_c f_c * (r_c*4/d + bV_c) |
| 148 | + = 0.10 * (16*4/128 + 9) # sink/policy at VQ-512 |
| 149 | + + 0.30 * (16*4/128 + 8) # schema/evidence at VQ-256 |
| 150 | + + 0.15 * (16*4/128 + 8) # conv_recent |
| 151 | + + 0.35 * (12*4/128 + 7) # boilerplate at VQ-128, rank 12 |
| 152 | + + 0.10 * (8*4/128 + 6) # irrelevant at VQ-64, rank 8 |
| 153 | + = 0.10*(0.5+9) + 0.30*(0.5+8) + 0.15*(0.5+8) + 0.35*(0.375+7) + 0.10*(0.25+6) |
| 154 | + = 0.95 + 2.55 + 1.275 + 2.581 + 0.625 |
| 155 | + = 7.98 bits/slot |
| 156 | +``` |
| 157 | + |
| 158 | +vs uniform Shard at 8.5 bits/slot — so the V-side budget alone drops about |
| 159 | +6%. Adding K rank reduction on the low-importance classes (rank 8 vs 16) |
| 160 | +gives another ~3% savings. |
| 161 | + |
| 162 | +Combined: total bits/token drops from ~17 * L * H_kv to ~16 * L * H_kv, |
| 163 | +a ~6% improvement at average quality — BUT crucially, that 6% is taken |
| 164 | +from the tokens that don't need the bits, and the high-importance tokens |
| 165 | +gain a full bit of headroom (9 vs 8 bpe on V, full rank vs truncated on K). |
| 166 | + |
| 167 | +The headline savings comes from the irrelevant_span class being EVICTED |
| 168 | +not compressed. In agent + RAG workloads the irrelevant fraction is |
| 169 | +often 20%-30%, not 10%. At 25% eviction the bits/token drops by an |
| 170 | +additional 25% on top of the 6%, yielding ~12-15x compression vs the |
| 171 | +FP16 baseline (vs Shard's 10x). |
| 172 | + |
| 173 | +## 5. Implementation sketch (algorithm, no code) |
| 174 | + |
| 175 | +1. **Token insertion path** |
| 176 | + - When transformers calls `cache.update(new_K, new_V, layer_idx, ...)`: |
| 177 | + - For each new token, look up its class label from the ITKV scorer. |
| 178 | + - Sink + window tokens: FP16 path (unchanged from Shard). |
| 179 | + - All others: pick the codebook + rank for the class, encode K and V. |
| 180 | + - Append the class label to the per-layer label sidecar. |
| 181 | + |
| 182 | +2. **Eviction path** |
| 183 | + - When the compressed tail reaches a configured max length: |
| 184 | + - First evict irrelevant_span tokens (these are tagged at insertion). |
| 185 | + - If still over budget, evict the oldest boilerplate tokens. |
| 186 | + - Never evict sink or policy tokens within a session. |
| 187 | + |
| 188 | +3. **Attention computation** |
| 189 | + - For each decode step, gather (K, V) for all stored tokens: |
| 190 | + - Sink + window: read FP16 directly. |
| 191 | + - Compressed: per-class decode (rank-`r_c` K project back to full d; |
| 192 | + VQ-`N_c` V index back into the codebook). |
| 193 | + - The class-aware decode path is the only mechanical addition vs |
| 194 | + vanilla Shard. |
| 195 | + |
| 196 | +4. **Codebook lifecycle** |
| 197 | + - Per class, per layer, per head: maintain a codebook of size N_c. |
| 198 | + Online K-means update (Lloyd's algorithm) every `K` insertions to |
| 199 | + adapt to the running token distribution. Codebook state is part of |
| 200 | + the cache's serialized form. |
| 201 | + |
| 202 | +5. **Calibration** |
| 203 | + - Before deployment, run the kolm benchmark suite to find per-class |
| 204 | + {rank, codebook size} that minimize quality_delta subject to a target |
| 205 | + compression ratio. Ship the resulting profile in the runtime passport. |
| 206 | + |
| 207 | +## 6. Expected gains (back-of-envelope) |
| 208 | + |
| 209 | +At workloads dominated by agent + RAG traffic (the kolm target): |
| 210 | + |
| 211 | +| Cache | Compression | Quality delta on agent eval | Notes | |
| 212 | +|---|---|---|---| |
| 213 | +| Default FP16 | 1.0x | 0 (baseline) | reference | |
| 214 | +| Shard alone | ~10x | ~ -0.002 (parity) | uniform 1.5 bpe | |
| 215 | +| ITKV x Shard | ~12-15x | ~ 0 to -0.001 | with 20-25% irrelevant eviction | |
| 216 | + |
| 217 | +The win is biggest where the workload has high boilerplate + irrelevant |
| 218 | +fraction (agents repeating long system prompts, RAG dumping unranked |
| 219 | +chunks). On uniform creative-writing workloads ITKV degenerates to |
| 220 | +"all tokens are conv_recent" and the cache reverts to plain Shard. |
| 221 | + |
| 222 | +## 7. Validation plan |
| 223 | + |
| 224 | +The combination should be validated against these benchmarks before any |
| 225 | +production rollout: |
| 226 | + |
| 227 | +1. **Quality**: kolm benchmark suite (`scripts/bench-quality-calibration.mjs`) |
| 228 | + on the standard model x task matrix, comparing FP16 vs Shard vs ITKVxShard. |
| 229 | + Target: quality_delta within +/- 0.005 of FP16 baseline. |
| 230 | +2. **Compression**: peak VRAM measurement on a long-context agent trace |
| 231 | + (32K + tokens, tool calls + RAG retrievals). Target: >= 1.2x reduction |
| 232 | + vs Shard alone, >= 12x vs FP16. |
| 233 | +3. **Latency**: per-token decode latency. The class-aware decode adds one |
| 234 | + branch + one codebook lookup per token. Target: <= 10% latency overhead |
| 235 | + vs Shard alone. |
| 236 | +4. **Stability**: run a 24-hour rolling agent simulation. The online |
| 237 | + codebook update must not degrade quality monotonically. If it does, |
| 238 | + freeze the codebook after warmup and re-evaluate. |
| 239 | + |
| 240 | +## 8. What's NOT in V1 |
| 241 | + |
| 242 | +This document is research-track. V1 ships Shard alone (uniform 1.5 bpe). |
| 243 | +ITKV remains a separate profile/scorer with no runtime dispatch. The |
| 244 | +combination requires: |
| 245 | + |
| 246 | +- A `ShardITKVCache` Python class extending `ShardCache` with a class |
| 247 | + label sidecar. |
| 248 | +- A class-to-codebook size mapping serialized in the runtime passport. |
| 249 | +- Per-class codebook lifecycle management. |
| 250 | +- Benchmark calibration to find production-ready per-class rank/codebook |
| 251 | + settings. |
| 252 | +- A revised eviction policy with class-aware priority. |
| 253 | + |
| 254 | +Estimated effort: ~6 person-weeks across the Python (cache class + |
| 255 | +codebook lifecycle), JS (passport schema extension, policy selector |
| 256 | +extension), and benchmark (calibration sweep + agent eval harness) |
| 257 | +surfaces. |
| 258 | + |
| 259 | +## 9. Risks |
| 260 | + |
| 261 | +- **Codebook overfitting** — online K-means on a non-stationary token |
| 262 | + stream can drift. Mitigation: freeze the codebook after a warmup |
| 263 | + window and rebuild on session boundaries. |
| 264 | +- **Class label cost** — the ITKV scorer adds latency at token insertion. |
| 265 | + Mitigation: the existing W722 scorer runs at well under 1% of attention |
| 266 | + cost on the reference traces. |
| 267 | +- **Eviction surprises** — if the ITKV scorer mislabels a critical token |
| 268 | + as `irrelevant_span` and the eviction policy throws it away, decode |
| 269 | + quality silently degrades. Mitigation: shadow-mode the eviction policy |
| 270 | + for the first deployment and compare retained-vs-evicted attention |
| 271 | + contribution against ground truth before enabling. |
| 272 | +- **HF Cache contract churn** — both Shard and any subclass live downstream |
| 273 | + of the upstream `transformers.Cache` ABI. Mitigation: pin the |
| 274 | + `transformers` version in the runtime passport. |
| 275 | + |
| 276 | +## 10. Pointers |
| 277 | + |
| 278 | +- `src/kv-cache-shard.js` — Shard module (this codebase) |
| 279 | +- `src/itkv-profile.js` — ITKV profile + token-class scorer (W722) |
| 280 | +- `docs/kv-cache-shard.md` — V1 Shard integration notes |
| 281 | +- `github.com/krish1905/shard` — upstream Shard reference (Apache-2.0) |
| 282 | +- `docs/research/kolm-billion-dollar-distillation-lab-2026-05-24.md` |
| 283 | + lines 1434-1466 — ITKV invention writeup |
0 commit comments