veloxquant_mlx.cache.qjl_cache module

QJL-backed KV cache: pure 1-bit Johnson-Lindenstrauss key compression.

Wraps QJLQuantizer in the VeloxQuant KVCache ABC (append_key/append_value/attend), storing per-token JL sign bits and a scalar norm 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.qjl_cache.QJLKVCache(config)

Bases: KVCache

Minimal KV cache using pure 1-bit QJL for key compression.

Parameters:

config (Any) – KVCacheConfig instance. jl_dim sets the sketch size m; when unset it defaults to head_dim. The estimator is unbiased but its variance falls as 1/m, so m = head_dim gives noisy scores (correlation with exact q.k ~0.64 at d=128, ~0.88 at m = 4 * head_dim). Set jl_dim higher when attention quality matters.

append_key(k)

Encode and store a key vector.

Parameters:

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

Return type:

None

append_value(v)

Quantize and store a value vector.

Parameters:

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

Return type:

None

attend(q)

Compute attention output.

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.

Signs are stored one per int8 (m bytes per key), not bit-packed, so the key side is 8x what a true 1-bit store would need.

Return type:

int

reset()

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

Return type:

None