veloxquant_mlx.cache.polar_cache module

PolarQuant-backed KV cache: polar-coordinate key compression, int8 values.

Wraps PolarQuantizer in the VeloxQuant KVCache ABC (append_key/append_value/attend), storing per-token angle indices and a scalar radius per recursion level for keys, and per-token int8 + fp16 scale for values. This is a “standalone” method (see STANDALONE_METHODS): it does not implement the mlx_lm serving protocol.

class veloxquant_mlx.cache.polar_cache.PolarQuantKVCache(config)

Bases: KVCache

KV cache backed by PolarQuantizer for key compression.

Parameters:

config (Any) – KVCacheConfig instance.

append_key(k)

Encode and cache a single key vector.

Parameters:

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

Return type:

None

append_value(v)

Quantize and cache a value vector.

Parameters:

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

Return type:

None

attend(q)

Compute attention output for a query.

Parameters:

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

Return type:

Any

Returns:

Attention output, shape (d,), fp16.

memory_bytes()

Bytes held for all cached tokens (angle codes, radii, int8 values and scales).

Angle codes are stored one uint8 per angle, not packed to b bits.

Return type:

int

reset()

Clear all stored tokens; the key quantizer (seeded, deterministic polar rotation) is untouched.

Return type:

None