17 KiB

Raw Permalink Blame History

TurboQuant KV Cache Compression — Spec

Module: foundry/attention/ Prefix: fa_tq_ (Foundry Attention — TurboQuant)

Based on: TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate (Zandieh et al., Google Research / DeepMind, April 2025)

Overview

TurboQuant compresses the KV cache from FP32 to 2-4 bits per channel with near-zero quality loss. The algorithm is:

Rotate the vector by a fixed random orthogonal matrix (spreads information uniformly)
Scalar quantize each coordinate independently using precomputed Lloyd-Max levels for the Beta distribution
Dequantize by reversing: look up reconstruction values, then rotate back

No codebook training. No preprocessing. No iterative optimization. Just a rotation matrix and a lookup table. The entire quantizer is deterministic given the dimension and bit width.

Why not QJL?

The paper proposes QJL (1-bit quantized Johnson-Lindenstrauss) on residuals for unbiased inner products. Six independent community implementations confirmed that for softmax attention, MSE-only quantization without QJL outperforms the full method. Softmax exponentially amplifies variance, making low MSE more valuable than unbiasedness. We skip QJL entirely.

Asymmetric bit allocation

Keys determine attention routing — errors in keys misroute attention to the wrong tokens. Values just get averaged — errors cancel during the weighted sum. This gives us:

Component	Bits	Levels	Compression vs FP32
Keys	4	16	8x
Values	2	4	16x
Average	3	—	~10.7x

Conservative default: K4/V2. Optional K3/V2 for aggressive compression (~12.8x).

As-Built Notes (implementation vs this spec)

The shipped turboquant.{h,c} follows this design with two deliberate, correctness-driven deviations from the text below:

SRHT signs before the Hadamard transform. The rotation applies the random +-1 sign diagonal D before the FWHT rather than after (H * (D * x)). This spreads the DC component across coordinates and gave measurably better roundtrip error. The transform is still orthogonal and self-inverse up to the 1/sqrt(d) scaling.
L2-norm normalization, not per-head min/max. Step 2 below describes mapping rotated coordinates into [0,1] via per-head min/max. The implementation instead normalizes by the vector L2 norm, which is the true Beta(d/2, d/2) mapping the codebooks are optimized for. The stored per-head range metadata carries the norm-based scale accordingly.

Measured quality: K4/V2 is lossy

The default K4/V2 is not near-lossless in practice. 4-bit keys are near-lossless, but 2-bit values are the error driver. Attention-output cosine similarity (FP vs compressed):

Config	Synthetic (Gaussian)	Realistic (correlated)
K4/V2	~0.94	~0.975
K4/V3	~0.98	-
K4/V4	~0.99	~0.99

Treat the value bit width as a per-model knob: validate on real activations and raise it to V3 or V4 where output fidelity matters. The "attention score preservation >0.999" target in the Testing Strategy below is aspirational for K4/V2 and is reported by validate_tq_attention.c rather than gated. The 11 unit tests in test_turboquant.c (codebook, rotation orthogonality, roundtrip MSE, dot preservation, determinism, packing, cache) all pass.

Algorithm Detail

Step 1: Random Orthogonal Rotation

Generate a fixed d x d orthogonal matrix R where d = head_dim. This is generated once at init from a seed and reused for all tokens and layers.

Implementation: Use a randomized Hadamard transform (fast O(d log d)) instead of dense matrix multiply (O(d^2)):

rotate(x) = D * H * x / sqrt(d)

Where:

H is the Walsh-Hadamard matrix (applied in-place via butterfly ops, no storage)
D is a diagonal of random +1/-1 signs (one d-element array, generated from seed)

This is O(d log d) per vector vs O(d^2) for a full orthogonal matrix. For head_dim=128, that's ~896 FLOPs vs ~16,384.

The rotation spreads information across all coordinates, making each coordinate approximately follow a Beta(d/2, d/2) distribution centered at 0.5 (after shifting to [0,1]). In high dimensions this concentrates sharply around the mean, making simple scalar quantization nearly optimal.

Step 2: Lloyd-Max Scalar Quantization

After rotation, each coordinate is quantized independently using optimal decision boundaries and reconstruction levels for Beta(d/2, d/2).

These are precomputed offline for each (d, bits) pair and stored as constant lookup tables:

typedef struct {
    int bits;                    // 2, 3, or 4
    int levels;                  // 2^bits
    float boundaries[15];        // Decision boundaries (levels - 1 values)
    float centroids[16];         // Reconstruction centroids (levels values)
} fa_tq_codebook;

For quantization: binary search into boundaries[] to get the level index. For dequantization: index into centroids[] to get the reconstruction value.

The Beta distribution parameters depend only on head_dim, which is constant per model. So we need one codebook per (head_dim, bits) pair — precomputed at compile time or at fa_tq_init().

Step 3: Packing

Quantized values are packed into bytes:

Bits	Values per byte	Packing
2	4	`[v0:2][v1:2][v2:2][v3:2]`
3	2 (+ 2 waste)	`[v0:3][v1:3][__:2]`
4	2	`[v0:4][v1:4]`

3-bit packing wastes 25% of bits. Alternative: pack 8 values into 3 bytes (no waste, more complex addressing). Start with the simple per-byte packing; optimize later if memory-bound.

Memory Layout

Compressed cache entry (one token, one layer)

key_packed:    [ceil(head_dim * k_bits / 8)] bytes per head, n_heads total
value_packed:  [ceil(head_dim * v_bits / 8)] bytes per head, n_heads total
key_range:     [2 floats per head] — min/max before rotation (for denormalization)
value_range:   [2 floats per head] — min/max before rotation

The min/max per head are needed to map the rotated coordinates into [0, 1] for the Beta quantizer.

Memory per token per layer

Config	Keys	Values	Range metadata	Total	vs FP32 (`2 * n_heads * head_dim * 4`)
K4/V2, 16 heads, dim 128	1024 B	512 B	256 B	1792 B	vs 16,384 B → 9.1x
K4/V2, 32 heads, dim 128	2048 B	1024 B	512 B	3584 B	vs 32,768 B → 9.1x
K3/V2, 16 heads, dim 128	768 B	512 B	256 B	1536 B	vs 16,384 B → 10.7x

Total cache savings

RWKV-X 0.2B (4 attention layers, 16 heads, dim 128, max_cache 2048):

	FP32	TurboQuant K4/V2
Per layer	2 * 2048 * 2048 * 4 = 32 MB	2048 * 1792 = 3.5 MB
Total (4 layers)	128 MB	14 MB

Llama-style 3.6B (24 attention layers, 32 heads, dim 128, max_seq 4096):

	FP32	TurboQuant K4/V2
Per layer	2 * 4096 * 4096 * 4 = 128 MB	4096 * 3584 = 14.3 MB
Total (24 layers)	3 GB	343 MB

API

Types

/* Bit allocation configuration */
typedef struct {
    int key_bits;       /* 2, 3, or 4 (default: 4) */
    int value_bits;     /* 2, 3, or 4 (default: 2) */
} fa_tq_config;

/* Opaque quantizer state (holds rotation signs + codebooks) */
typedef struct fa_tq_state fa_tq_state;

/* Compressed KV cache (replaces fa_context for quantized mode) */
typedef struct fa_tq_cache fa_tq_cache;

Lifecycle

/* Initialize quantizer for a given head_dim and bit config.
 * seed determines the random rotation (deterministic, reproducible).
 * Call once at model load time. */
fa_tq_state* fa_tq_init(int head_dim, fa_tq_config config, uint64_t seed);
void          fa_tq_free(fa_tq_state *state);

/* Default config: K4/V2 */
fa_tq_config fa_tq_config_default(void);

Cache

/* Create compressed KV cache.
 * Same shape semantics as fa_context but stores quantized data. */
fa_tq_cache* fa_tq_cache_new(fa_tq_state *state, int num_layers,
                              int num_heads, int max_seq_len);
void          fa_tq_cache_free(fa_tq_cache *cache);
void          fa_tq_cache_clear(fa_tq_cache *cache);
int           fa_tq_cache_seq_len(fa_tq_cache *cache);

Quantize / Dequantize

/* Quantize a single K or V vector into the cache at position pos.
 * vec: [num_heads * head_dim] float input
 * is_key: true for key, false for value (determines bit width) */
void fa_tq_quantize(fa_tq_state *state, fa_tq_cache *cache,
                    int layer, int pos, const float *vec, bool is_key);

/* Dequantize a range of cached vectors back to float.
 * out: [count * num_heads * head_dim] float output
 * Returns pointer to out for convenience. */
float* fa_tq_dequantize(fa_tq_state *state, fa_tq_cache *cache,
                        int layer, int start, int count,
                        bool is_key, float *out);

Attention with TurboQuant cache

/* Drop-in replacement for fa_attention_cached.
 * Quantizes K, V into cache, dequantizes cached K, V for attention.
 * Q is always full precision (current token only, never cached). */
tm_tensor* fa_tq_attention_cached(tm_tensor *Q, tm_tensor *K, tm_tensor *V,
                                   fa_tq_state *tq, fa_tq_cache *cache,
                                   int layer_idx, int pos, fa_config *config);

Integration Points

RWKV-X (`foundry/rwkvx/rwkvx.c`)

The sparse attention blocks (MOBA layers) maintain per-layer rx_attn_state with flat k_cache/v_cache buffers. Integration:

Add fa_tq_state* and fa_tq_cache* to rx_model or rx_state
In the attention block forward pass (~line 1163), after computing k_vec/v_vec:
- Call fa_tq_quantize() to store into compressed cache
- Call fa_tq_dequantize() to reconstruct K/V for the attention dot products
Remove the FP32 k_cache/v_cache from rx_attn_state when TQ is enabled
Cache eviction (line 1181) works the same — just operates on packed bytes instead of floats, memmove still works

Pantheon / Llama (`Pantheon/src/models/llama.c`)

The per-head [n_kv_head, max_seq, head_dim] layout maps directly:

Add fa_tq_state* to pn_llama_model, fa_tq_cache* to pn_llama_kv_cache
In pn_llama_forward() (~line 1484), replace the memcpy into K/V cache with fa_tq_quantize()
Before the attention head loop (~line 1370), fa_tq_dequantize() the needed range
Circular buffer addressing: cache_pos = pos % max_seq still works — the compressed cache just stores packed bytes at each position

Runtime toggle

Both engines already have a --quantize flag for INT8 weights. Add --tq-cache (or --kv-compress) to enable TurboQuant on the KV cache independently. The two are orthogonal: INT8 weights + TQ cache is the highest compression combo.

File Plan

foundry/attention/
├── attention.h          # Existing — add fa_tq_* declarations
├── attention.c          # Existing — unchanged
├── turboquant.h         # TurboQuant types + API (public header)
├── turboquant.c         # Implementation
├── tq_codebooks.h       # Precomputed Lloyd-Max tables (generated, static const)
├── spec.md              # This file
└── tests/
    └── test_turboquant.c

turboquant.h/c are self-contained. They depend only on <math.h>, <string.h>, <stdint.h>. No tensor library dependency — they operate on raw float* and uint8_t* buffers. This keeps them usable from both the foundry attention module and the hand-rolled attention loops in RWKV-X and Pantheon.

Codebook Generation

Lloyd-Max optimal levels for Beta(d/2, d/2) can be computed offline via iterative Lloyd's algorithm:

Start with uniform boundaries in [0, 1]
Compute centroids as conditional expectations E[X | boundary_i < X < boundary_{i+1}] under Beta(d/2, d/2)
Update boundaries as midpoints between adjacent centroids
Repeat until converged (~20 iterations)

Ship a Python script tools/gen_codebooks.py that generates tq_codebooks.h for all needed (d, bits) pairs. The output is a static C header — no Python dependency at runtime.

Target dimensions: 64, 96, 128, 160 (covers RWKV-X 0.2B through 3.6B and common Llama configs). Bit widths: 2, 3, 4.

That's 4 * 3 = 12 codebooks, each under 200 bytes. Total: ~2 KB of static const data.

Hadamard Transform Detail

The Walsh-Hadamard transform of dimension d (must be power of 2) is:

/* In-place Fast Walsh-Hadamard Transform, O(d log d) */
static void fwht(float *x, int d) {
    for (int len = 1; len < d; len <<= 1) {
        for (int i = 0; i < d; i += len << 1) {
            for (int j = 0; j < len; j++) {
                float u = x[i + j];
                float v = x[i + j + len];
                x[i + j]       = u + v;
                x[i + j + len] = u - v;
            }
        }
    }
}

Full rotation: multiply each element by the random sign (D), apply FWHT, scale by 1/sqrt(d). Inverse: same operation (FWHT is its own inverse up to scaling, and D is its own inverse).

For head_dim not a power of 2: pad to the next power of 2, apply FWHT, truncate. The extra coordinates are discarded — they don't affect the quantized output. Alternatively, use a dense random orthogonal matrix for non-power-of-2 dims (rare in practice).

Performance Budget

Target: TurboQuant overhead < 5% of total inference time.

Operation	FLOPs per token per layer	Notes
Hadamard rotate K	`d * log2(d)` = 896	d=128
Hadamard rotate V	896
Quantize K (binary search)	`d * log2(levels)` = 512	4-bit, 4 comparisons
Quantize V	`d * log2(levels)` = 256	2-bit, 2 comparisons
Pack K into bytes	`d` = 128	Shift + OR
Pack V into bytes	`d` = 128
Total quantize	~2,816	Per head

Dequantize K (all cached)	`seq * d` = 262,144	seq=2048, lookup + unpack
Hadamard inverse K	`seq * d * log2(d)` = 1,835,008
Dequantize V (all cached)	`seq * d` = 262,144
Hadamard inverse V	`seq * d * log2(d)` = 1,835,008
Total dequantize	~4.2M	Dominates

Attention Q@K^T + softmax + @V	`~2 * seq * d` = 524,288	Standard attention

Dequantize is ~8x the attention cost per head, which is too expensive if done naively for the full cache every token.

Optimization: Fused dequant-dot

Instead of materializing the full dequantized K/V matrix, fuse dequantization into the attention loop:

For each cached position t:
    k_packed = cache[t]
    k_quant = unpack(k_packed)           // byte ops
    k_rotated = centroids[k_quant]       // lookup, no multiply
    score[t] = dot(Q_rotated, k_rotated) // dot in rotated space

Key insight: dot(Q, R^T * k_rotated) = dot(R * Q, k_rotated). Pre-rotate Q once (O(d log d)), then compute dots directly in the rotated domain. No inverse Hadamard on cached vectors at all.

Same for values: accumulate sum(attn[t] * v_rotated[t]) in rotated space, then inverse-rotate the final output once.

This reduces dequant to just unpack + table lookup per element — essentially free compared to the dot product itself.

Revised cost with fused dequant-dot:

Operation	FLOPs	Notes
Rotate Q once	896	Hadamard
Quantize K, V (new token)	~2,816	Per head
Attention loop (unpack + lookup + dot)	`~2 * seq * d`	Same as FP32 attention
Inverse rotate output	896	Hadamard
Overhead vs FP32 attention	~0.7%	Just two Hadamard transforms

Testing Strategy

Unit tests (`test_turboquant.c`)

Roundtrip accuracy: Quantize random vectors, dequantize, measure MSE. Assert MSE < threshold for each bit width.
Rotation orthogonality: Verify R * R^T = I (within float tolerance).
Dot product preservation: |dot(q, k) - dot(q_dequant, k_dequant)| < epsilon for random q, k.
Attention score preservation: Full attention with FP32 cache vs TQ cache, compare output cosine similarity. Target: >0.999 for K4/V2.
Determinism: Same seed produces identical quantization across runs.
Edge cases: All-zero vectors, all-same vectors, very large/small magnitudes.
Cache eviction: Verify memmove on packed data preserves correct values after dequant.
Bit packing: Verify pack/unpack roundtrip for all bit widths (2, 3, 4) with all possible level values.

Integration test

Run RWKV-X 0.2B inference with and without TQ cache on the same prompt. Compare:

Output token agreement (should match for short sequences)
Perplexity difference on a held-out text (target: <0.5% increase)

Implementation Order

tq_codebooks.h — Generate Lloyd-Max tables with Python script, ship as static header
turboquant.c — Core: Hadamard, quantize, dequantize, pack/unpack, cache struct
test_turboquant.c — Unit tests for roundtrip, accuracy, dot preservation
RWKV-X integration — Wire into rwkvx.c attention blocks behind --tq-cache flag
Pantheon integration — Wire into llama.c behind same flag
Benchmarks — tok/s with and without, memory usage comparison

17 KiB Raw Permalink Blame History