RFC: Planar3 KV-cache quantization for compressed attention (experimental)

[hexxyan]

Thanks @antirez for ds4 — a remarkably clean and well-structured inference engine.

Summary

This PR adds an opt-in, experimental Planar3 KV-cache quantization based on the PlanarQuant approach. 2D Givens rotations decorrelate 128-dim blocks of KV cache rows, then Lloyd-Max 3-bit centroids quantize the rotated coefficients. For ds4's 512-dim heads, each compressed row is 200 bytes (4 blocks × 50 bytes) — a 5.12× density ratio vs FP16.

What's Included

Core Codec (ds4_planar_quant.c/h)

• CPU quantize/dequantize with inverse Givens rotation
• ds4_row_planar3 struct: 200 bytes per 512-dim row (4 × ds4_block_planar3 at 50 bytes each)
• Subnormal-safe FP16 rounding

Metal GPU Path

• GPU quantize kernel (kernel_planar3_quantize_row) — cooperative 128-thread per row
• Planar3→F16 dequant kernel (kernel_planar3_dequant_to_f16_rows) — pre-stage for FlashAttention
• Indexed attention: cooperative dequant inline in kernel_dsv4_indexed_mixed_attention_heads8 and _rb16
• FlashAttention: Planar3→F16 pre-dequant in all FA encoder paths (gathered_heads, decode_mixed_batch, prefill_static_mixed nonvec+vec)

CPU Hot-Path

• Per-layer planar_staging buffer for dequant (no per-token malloc)
• comp_kv_for_attn() transparent dequant from Planar3

Planar-Only Mode (--planar-kv-cache-only)

• Skips FP16/FP32 compressed cache allocation entirely
• Only Planar3 compressed cache is allocated (actual memory savings)
• Checkpoint save/restore: writes/reads Planar3 bytes directly
• Implies --planar-kv-cache

Checkpoint Compatibility

• GPU quantize on restore: read FP16 comp cache → CPU quantize → upload Planar3
• Planar-only: read/write Planar3 directly, no FP16/FP32 staging

Offline Quality Evaluator (tools/planar_eval)

• Random-normal, worst-case, and file-based test modes
• Cosine similarity, MSE, per-row and batch metrics

Tests (tests/planar_quant_test.c)

• Block size, roundtrip (random/basis/large-norm), batch, compression ratio, block independence, dim mismatch, zero-norm, single-element — 10/10 passing

Status

Experimental / RFC. Implementation and static/unit verification are solid. The codec round-trips random vectors at cosine similarity ~0.985 avg. End-to-end quality validation on a real 80GB+ DS4 model with compressed KV dumps is needed before recommending production use.

Memory Savings

With --planar-kv-cache-only, each compressed KV row is 200 bytes instead of 1024 bytes (FP16). The dual-cache overhead is eliminated — only Planar3 bytes are stored persistently. FA paths use a transient F16 scratch buffer for dequant.

Usage

# Enable Planar3 quantization (keeps FP16/FP32 cache too)
./ds4 -m model.bin --planar-kv-cache -p "Hello"

# Planar-only mode (actual memory savings)
./ds4 -m model.bin --planar-kv-cache-only -p "Hello"

References & Related Work

Planar3 sits within a family of rotation-based KV-cache compression methods that reduce the d×d random orthogonal projection (originally from TurboQuant) to progressively lighter structures:

Core Methods

Method	Rotation	Params/row	FMA/row	PPL (Llama 3.1 8B, 3-bit)	Source
TurboQuant	Dense d×d Walsh-Hadamard	16,384	~16K	7.07	Google, ICLR 2026
RotorQuant	Cl(3,0) Clifford SO(3) rotor	372	~2,400	7.05*	scrya-com/rotorquant
IsoQuant	4D quaternion SO(4) isoclinic	256	1,024	6.91	ParaMind2025/isoquant, arXiv:2603.28430
PlanarQuant (= IsoQuant-2D)	2D Givens rotation	128	256	7.05	experolk/planar-llama

* RotorQuant PPL from production llama.cpp benchmarks on RTX 5090.

Evolution: TurboQuant (dense WHT) → RotorQuant (3D Clifford, by John D. Pope) → IsoQuant (4D quaternion) / PlanarQuant (2D Givens, both by Zhongping Ji / ParaMind2025). RotorQuant's README credits ParaMind2025 for designing PlanarQuant and IsoQuant.

PlanarQuant is the lightest variant — 128 params, 256 FMAs — making it the best fit for fused GPU kernels in production inference engines. Its quality is competitive with TurboQuant at a fraction of the compute.

Community Projects

• Multi-TurboQuant — Unified Python toolkit with 12 methods including all above, plus community Metal/MLX kernels.
• HybridQuant — PlanarQuant + semantic activation indexing for layer-adaptive compression.
• RotorQuant llama.cpp integration — Production reference for TurboQuant/RotorQuant/IsoQuant in llama.cpp with RTX 5090 benchmarks.