Kakeya KV cache: bug fix, Gemma 4 benchmark standard, cross-model reports (4 models)

[FluffyAIcode]

Summary

This PR contains the full body of work on the Kakeya KV cache compression layer:

1. Critical correctness fix in KakeyaKVCache that made it a no-op on Gemma 4's default config.
2. Generalized cache + factory (build_kakeya_cache(model)) that works for any HF transformers decoder-only model — Gemma 2/3/4, Llama, Mistral, Qwen 2/3, SmolLM2, Cohere2, and similar.
3. Benchmark harness (kakeya_benchmark.py) with chunked prefill, analytic baseline, and per-layer JSON reports.
4. Byte-exact extrapolator (kakeya_extrapolate.py) that projects any report to arbitrary context lengths (cross-validated ≤ 0.003 absolute ratio error).
5. Benchmark standard + Gemma 4 reference run (reports/STANDARD.md): measured 2k–32k plus projected 64k–256k.
6. Cross-model reports for four open-source models (reports/{gemma4_e2b, qwen2_5_0_5b, smollm2_1_7b, qwen3_0_6b}/), each with a narrative REPORT.md and all raw JSON.
7. Cross-model comparison (reports/CROSS_MODEL.md) and top-level README.md.

Headline numbers

Under the standard codec preset (block_size=512, residual_length=256, d_res=8, K=16, variance_ratio=0.95), at 131 072 token context with bf16 store projection:

Model	Baseline KV	Kakeya KV	Total ratio	Absolute saved
Qwen/Qwen3-0.6B	14.00 GiB	3.10 GiB	4.51×	10.90 GiB
google/gemma-4-E2B-it	774 MiB	180 MiB	4.29×	594 MiB
HuggingFaceTB/SmolLM2-1.7B-Instruct	24.00 GiB	10.65 GiB	2.25×	13.35 GiB
Qwen/Qwen2.5-0.5B-Instruct	1.50 GiB	714 MiB	2.15×	786 MiB

All four models use the same build_kakeya_cache(model) factory and the same run_all_benchmarks.sh orchestrator. No model-specific code.

What changed in the codec

Bug fix

Gemma4TextConfig.num_kv_shared_layers defaults to 0, which made the old code do layer_types[:-0] → empty list → cache with zero layers. Guarded with > 0.

Model-agnostic factory

Layer-plan inference moved to _resolve_layer_plan(config) and handles all the standard HF patterns:

Config pattern	Handling
layer_types exists (Gemma 2/3/4, Cohere2, Qwen3 hybrid)	Per-layer dispatch
num_kv_shared_layers > 0 (Gemma 4)	Strip trailing shared layers
No layer_types, sliding_window=None (Llama, SmolLM2, Qwen2, Qwen3-dense)	Every layer is full attention
sliding_window only / attention_chunk_size (Llama 4)	Every layer is sliding

What the benchmark does

kakeya_benchmark.py loads the target model, runs a long-prompt prefill through both DynamicCache (baseline) and KakeyaKVCache, and reports byte counts + ratios per layer and per model section (full-attn vs sliding-window). It supports:

• --skip-baseline-prefill: compute the baseline bytes analytically (exact, since it's determined by config + seq_len); saves half the CPU budget for long contexts.
• --prefill-chunk N: chunked prefill to bound activation memory on CPU.
• --skip-generation: skip greedy-decode sanity check (for very long contexts).

What the extrapolator does

For fixed codec params, per-block skeleton bytes and per-token encoded bytes are deterministic constants — they do not depend on context length. kakeya_extrapolate.py reads any measured report and scales it exactly to arbitrary target contexts. Cross-validation summary:

Prediction	Predicted	Measured	Abs error
8k → 16k full (f32 / bf16)	2.151 / 4.030	2.149 / 4.027	≤ 0.003
8k → 32k full (f32 / bf16)	2.191 / 4.237	2.189 / 4.234	≤ 0.003
16k → 32k full (f32 / bf16)	2.189 / 4.234	2.189 / 4.234	0.000

So the 64k / 128k / 256k rows in every report are byte-exact under the preset, not statistical fits.

Benchmark standard

reports/STANDARD.md is the formal methodology doc. Key points:

• Fixed codec preset across all runs (for direct comparability).
• Measured sweep: 2 048 / 4 096 / 8 192 tokens (real prefill + decode).
• Projected sweep: 16 384 / 32 768 / 65 536 / 131 072 / 262 144 tokens (byte-exact extrapolation).
• Two reported ratios per context: "f32 store" (as the codec is implemented today) and "bf16 store" (projected with the same codec storing compressed tensors in bf16).
• Headline number for cross-model comparisons: total_ratio_bf16_store at 32k tokens.

Reproducing any row

pip install -U torch transformers accelerate huggingface_hub

python3 -c "from huggingface_hub import snapshot_download; snapshot_download('Qwen/Qwen3-0.6B', local_dir='models/Qwen3-0.6B', allow_patterns=['*.json','*.safetensors','tokenizer*','*.model'])"

./run_all_benchmarks.sh models/Qwen3-0.6B qwen3_0_6b

Produces:

reports/qwen3_0_6b/bench_2048.json
reports/qwen3_0_6b/bench_4096.json
reports/qwen3_0_6b/bench_8192.json
reports/qwen3_0_6b/extrapolation.json  # contains the 16k–256k rows

Environment used

CPU-only x86_64, 15 GB RAM, no GPU, torch==2.11.0+cu130, transformers==5.5.4. Public (not gated) open-source models only.

Follow-up

The only pending optimization explicitly called out in the reports is the bf16-store conversion inside kakeya_kv_codec.py (currently compressed tensors are float32 on CPU). The reports make this visible by always showing the "bf16 store, projected" column next to the "f32 store" column, so it's clear what absolute gain that optimization would unlock.