KVarN variance-normalized KV-cache quantization: CPU + CUDA (#180)

[pekkah]

Implements KVarN (arXiv:2606.03458) end-to-end per the phased plan in docs/kvarn-feasibility-research.md and #180: calibration-free 4-bit-key / 2-bit-value KV-cache quantization as a selectable quantizer inside the existing TurboQuant machinery, on the CPU (scalar + AVX2) and CUDA paths. Clean-room from the paper — no code from the reference vLLM fork.

Commits (each implemented + adversarially reviewed before landing)

• 82efc4f Core quantizer — Hadamard rotation (reuses WalshHadamard, K and V), log-space dual-axis Sinkhorn variance normalization, asymmetric RTN K4/V2 in 128-token tiles (channel-major K / token-major V packing), 2-bit BitPacking.
• 6040076 Cache + attention integration — TqQuantizer enum in TurboQuantKvCache (KVarN mode: the FP32 window is the staging; whole-tile promotion), TqAttention dispatch, CLI --tq-mode <lloydmax|kvarn>.
• f95fb3e Accuracy gate — new perplexity CLI command (teacher-forced NLL, position buckets) + scripts/kvarn-gate/ (wikitext fetcher, 20-problem math eval). P0 gate: strong pass.
• 44a5e1a P1 AVX2 — fused KeyScores/AggregateValues (9.7×/8.4× per-kernel), scalar fallback + ForceScalar parity hook.
• 0d7cd28 P2 CUDA correctness — llm_kvarn_rotate_query / llm_kvarn_compress_tile / llm_kvarn_attention; packed codes byte-identical to the CPU compressor; EstimateMaxContextKvarn; --tq-mode kvarn -g -1; perplexity -g -1.
• e43f3f0 P2 CUDA decode perf — CUDA-graph re-entry (two topologies, tracked params, promotion hoisted), warp-per-channel / flat-index attention walks.
• 99c3773 Chunked batched prefill — llm_kvarn_prefill_attention (single streaming softmax over compressed tiles + FP32 window, deferred inverse WHT), schedule-independent tiles (re-chunking is byte-exact).

Measured (RTX 4070 Laptop 8 GB · Qwen3-0.6B-Q8_0 / Qwen3-8B-Q4_K_M · wikitext-2)

Claim	Result
Accuracy at FP16 parity	PPL +1.26% (0.6B c=3072) / +0.98% (8B c=4096) vs fp32; math eval 18/20 = fp32; depth penalty flattens at ~+2.2% (no drift accumulation)
~4× KV capacity	8B auto-context 2376 → 18304 (7.7×) in 8 GB
Throughput ≥ FP16	CPU decode 1.57–2.20× faster than fp32; CUDA 8B decode ≥ fp32 (+13% at 4k depth); prefill 7.7–13× vs per-token (15.5k prompt + 300 tokens: 98 s, was ~610 s)

Gates that must keep holding: perplexity (0.6B c=3072 kvarn ≈ 15.60, 8B c=4096 ≈ 8.99), compress byte-exactness tests, graph-vs-direct decode parity (bit-identical over promotion boundaries).

Scope notes

• KVarN CUDA: full offload of dense models, head_dim pow2 ≤ 256 (CPU up to 1024); SnapKV / narrowed-KV / batching / MoE combos rejected at config time (same exclusivity as Lloyd-Max TQ).
• Lloyd-Max path byte-for-byte unchanged (--tq alone). Related: the gate found Lloyd-Max 3-bit collapses on Qwen3 — TurboQuant Lloyd-Max 3-bit collapses on Qwen3 (PPL 61x, degenerate decode) — pre-existing on master #432.
• Follow-ups with calibrated baselines: KVarN follow-ups: int8-TC prefill GEMM, split-tile decode kernel, minors (#180 successor) #433 (int8-TC prefill GEMM, split-tile decode kernel, minors).

Closes #180 P0–P2 (P3 Vulkan optional, tracked in #433).

Tests: 32 KVarN (ForwardPass) + 7 graph + 78 TurboQuant + 80 CLI green; zero-warning Release build (TreatWarningsAsErrors, AOT/trim analyzers on).

🤖 Generated with Claude Code

https://claude.ai/code/session_018EWo1niax1g8A5E79RSALJ

[gemini-code-assist]

Code Review

This pull request implements the KVarN KV-cache quantizer (4-bit keys / 2-bit values) as an alternative to Lloyd-Max codebooks, adding CPU and CUDA support (including CUDA graph capture/replay and chunked prefill), a CLI perplexity evaluation command, and comprehensive tests. Feedback is provided regarding a critical bug in the CUDA prefill attention kernel (llm_kvarn_prefill_attention) where uninitialized shared memory s_r can contain garbage values (such as NaN or Infinity) and poison the accumulator s_racc during the first tile loop iteration; initializing s_r to 1.0f is recommended to prevent this.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

In the first iteration of the tile loop (tile = 0), s_racc[idx] *= s_r[idx / head_dim] is executed. At this point, s_racc is 0.0f, but s_r is uninitialized and contains garbage values from shared memory. If s_r contains NaN or Infinity, the multiplication 0.0f * NaN or 0.0f * Infinity will result in NaN, which will poison s_racc and propagate to the output. To prevent this, initialize s_r to 1.0f at the start of the kernel.

    if (tid < KPF_MQ) { s_m[tid] = sharpi_neg_inf(); s_l[tid] = 0.0f; s_r[tid] = 1.0f; }