veloxquant_mlx.cache.spectral_cache module

SpectralQuantKVCache: KV cache backed by SpectralQuant.

Matches the KVCache interface used by TurboQuantKVCache. Keys are compressed with SpectralQuant (spectral rotation + selective QJL + per-group codebooks). Values are compressed identically.

class veloxquant_mlx.cache.spectral_cache.SpectralQuantKVCache(config)

Bases: KVCache

KV cache backed by SpectralQuant for both key and value compression.

Keys are compressed with a SpectralQuantizer using the key eigenvectors. Values are compressed with a separate SpectralQuantizer using the value eigenvectors (which have d_s ≈ 50 rather than ≈ 4).

Call calibrate(rotations) after construction to inject the per-layer rotation matrices from calibrate_spectral_rotation(). Without calibration the quantizers fall back to random orthogonal rotation (equivalent to TurboQuant without QJL).

Parameters:

config (Any) – KVCacheConfig with method=’spectral’.

calibrate(rotation_entry)

Inject calibrated rotation matrices for this layer.

Parameters:

rotation_entry (tuple) – Tuple (key_U, val_U, key_ev, val_ev, key_ds, val_ds) as returned by calibrate_spectral_rotation() for this layer.

Return type:

None

append_key(k)

Encode and store a key vector.

Parameters:

k (Any) – Key vector, shape (d,) or (1, d), fp16.

Return type:

None

append_value(v)

Encode and store a value vector.

Parameters:

v (Any) – Value vector, shape (d,) or (1, d), fp16.

Return type:

None

attend(q)

Compute softmax attention output for query q.

Parameters:

q (Any) – Query vector, shape (d,), fp16.

Return type:

Any

Returns:

Attention output, shape (d,), fp16.

memory_bytes()

Return compressed memory footprint in bytes.

Return type:

int

compression_ratio()

Key compression ratio vs FP16.

Return type:

float

reset()

Clear all stored tokens; the key/value quantizers are untouched, so a prior calibrate() injection survives the reset.

Return type:

None