feat: Lloyd-Max optimal codebooks for KV cache quantization

[influenist]

Summary

Adds CacheLayer_lloyd_max — drop-in alternative to CacheLayer_quant with three stacked improvements:

1. Lloyd-Max codebooks: k-means optimal centroids for post-WHT distribution
2. Sub-block scales (per-8): 4 scales per 32-element block instead of 1
3. Asymmetric K quantization: scale + zero-point captures per-channel bias

All hardware-validated on RTX 3090 with real TinyLlama-1.1B KV projections.

Results (RTX 3090, TinyLlama-1.1B)

Bits	Stock Uniform	+ Lloyd-Max	+ Sub-8 scales	+ Asymmetric	Total gain
2	8.19 dB	9.51 dB	9.88 dB	12.05 dB	+3.86 dB
3	14.49 dB	15.51 dB	16.80 dB	19.47 dB	+4.98 dB
4	20.63 dB	21.81 dB	23.72 dB	26.09 dB	+5.46 dB
5	26.69 dB	28.05 dB	30.17 dB	32.38 dB	+5.69 dB
6	32.73 dB	34.19 dB	36.37 dB	38.54 dB	+5.81 dB

Latency: identical to stock (all improvements are in codebook values and scale storage, not algorithm complexity).

Per-technique contribution (4-bit)

Technique	SQNR gain	MSE reduction
Lloyd-Max codebooks	+1.19 dB	24%
Sub-block scales (per-8)	+1.90 dB	—
Asymmetric (scale + zero-point)	+2.37 dB	—
Combined	+5.46 dB	~71%

Why asymmetric matters

Analysis of TinyLlama-1.1B K/V projections shows K vectors have fundamentally different statistics than V:

Metric	K	V	K/V ratio
Per-channel bias	0.095	0.024	4x
Channel variance heterogeneity	0.731	0.077	9.5x
Kurtosis (heavy tails)	1.77	0.26	6.8x
Dynamic range	14.08	2.68	5.3x

K has systematic per-channel non-zero mean and outlier channels. Symmetric quantization wastes half the code range. Asymmetric (scale + zero-point) captures this bias directly.

Supported by KIVI (ICML 2024): per-channel K + per-token V is optimal.

Usage

from exllamav3.cache import CacheLayer_lloyd_max

# Full stack: Lloyd-Max + sub-8 + asymmetric
cache = Cache(model, max_tokens,
    layer_type=CacheLayer_lloyd_max,
    k_bits=4, v_bits=4,
    sub_scale_size=8,
    asymmetric=True)

Files (5 commits, clean branch from master)

File	Purpose
`lloyd_max_codebooks.cuh`	k-means optimal codebooks (2-8 bit)
`lm_cache_kernels.cuh`	All kernel variants: standard, sub-block, asymmetric
`lm_cache.cu/.cuh`	Wrappers: cont + paged × {standard, sub, asym}
`cache/lloyd_max.py`	`CacheLayer_lloyd_max(sub_scale_size, asymmetric)`
`compute_kmeans_codebooks.py`	Reproducible codebook computation
`bench_lloyd_max_v2.py`	Benchmark script

References

• TurboQuant (ICLR 2026) — WHT + Lloyd-Max
• TQ3_4S — sub-block scale insight
• KIVI (ICML 2024) — asymmetric K/V
• Lloyd 1957 / Max 1960 — the algorithm

🤖 Generated with Claude Code

[influenist]

Update: validated on real model weights (TinyLlama-1.1B)

Ran the Lloyd-Max kernels on actual KV projection outputs (hidden states projected through real K/V weight matrices, not random data):

RESULTS ON REAL TinyLlama-1.1B KV PROJECTIONS (RTX 3090)
  2-bit: Uniform=8.20dB  LloydMax=9.51dB  delta=+1.32dB  MSE_red=26.2%
  3-bit: Uniform=14.48dB  LloydMax=15.51dB  delta=+1.03dB  MSE_red=21.1%
  4-bit: Uniform=20.62dB  LloydMax=21.81dB  delta=+1.19dB  MSE_red=24.0%
  5-bit: Uniform=26.70dB  LloydMax=28.05dB  delta=+1.34dB  MSE_red=26.6%
  6-bit: Uniform=32.73dB  LloydMax=34.19dB  delta=+1.46dB  MSE_red=28.5%

Consistent with the random-data results. Full model perplexity (end-to-end inference with quantized cache) is still TODO.

[influenist]

Withdrawn — moved to separate project