veloxquant_mlx.cache.vecinfer_cache module

VecInfer KV cache wrapper for mlx_lm integration.

Wraps veloxquant_mlx.allocators.vecinfer primitives in the standard update_and_fetch cache protocol expected by mlx_lm. The cache:

• applies a per-(head, channel) smooth scaling + Walsh-Hadamard rotation to keys, suppressing outliers before product VQ;

• encodes transformed keys against a pre-trained codebook, immediately dequantizes (then inverse-transforms) so the downstream SDPA call sees fp16 keys;

• tracks compressed vs fp16 byte counts so benchmarks can report a realized compression ratio.

The paper’s CUDA kernel fusion (Section 3.3) is NOT portable to MLX/Metal; the win on Apple Silicon is memory compression, not speedup over fp16.

Per-token storage (keys) at codebook bit-width b_k and sub-vector dimension d_k: (D / d_k) * b_k / 8 bytes — plus an amortized codebook cost of 2**b_k * d_k * 2 bytes shared across all tokens.

class veloxquant_mlx.cache.vecinfer_cache.VecInferKVCache(config)

Bases: KVCache

KV cache implementing VecInfer’s dual-transform product VQ.

Parameters:

config (Any) – KVCacheConfig with VecInfer-specific fields populated by KVCacheFactory.create. Required fields: head_dim, key_codebook_bits, value_codebook_bits, key_sub_dim, value_sub_dim. Optional: smooth_factors (numpy or mx array; identity if absent), key_codebook, value_codebook (random init if absent — for tests only).

Notes

Storage layout deliberately delegates concatenation to mlx_lm’s base _MLXKVCache. We quantize + dequantize on the way in so that subsequent SDPA calls see a standard fp16 key tensor.

Never exposes .bits — mlx_lm’s SDPA checks hasattr(cache, "bits") to route to a different kernel path. We expose .assigned_avg_bits instead.

update_and_fetch(keys, values)

Apply smooth-scale + Hadamard rotation to keys, product-VQ-quantize K/V, and dequantize for downstream SDPA (or stash indices only when memory-bound fused mode is active).

property state

parent fp16 (keys, values). Fused mode: raw (key_indices, value_indices) ring buffers, since no fp16 K_hat/V_hat is materialized.

Type:

Standard mode

empty()

Whether the cache holds no tokens yet (checks the index ring buffer in fused mode, the parent buffer otherwise).

Return type:

bool

trim(n)

Drop the last n tokens, keeping every parallel buffer in step.

The inherited trim only lowers offset. In fused mode the index ring buffers are advanced by _stored_S_kv instead, so left alone they keep the trimmed tokens: the next token is written n slots past offset and fused_sdpa attends over the stale tail. The byte accounting is scaled back with the token count (an estimate: it is tracked per call, not per token).

Return type:

int

Parameters:

n (int)

size()

Number of tokens currently stored.

Return type:

int

property nbytes: int

the parent fp16 K_hat/V_hat buffers plus, in fused mode, the pre-allocated uint32 index buffers (stash mode holds both, so it uses more memory than plain fp16).

Type:

Bytes actually held

fused_sdpa(q, scale, *, causal=True, sliding_window=0)

Compute attention directly from compressed K/V indices.

Parameters:

• q (array) – [B, H_q, S_q, D] uncompressed queries — not yet smooth or Hadamard transformed. This method applies the dual transform internally so callers (mlx_lm SDPA dispatcher) pass exactly the same tensor they’d pass to standard SDPA.

• scale (float) – Attention scale (typically 1/sqrt(D)).

• causal (bool) – Apply causal mask.

• sliding_window (int) – If > 0, sliding-window attention with the given window width.

Return type:

array

Returns:

[B, H_q, S_q, D] attention output in the dtype of q.

property compressed_key_bytes: int

see nbytes.

Type:

Estimated packed bytes for the compressed key cache (product-VQ codes at bits each, all heads/batches; excludes codebook). Not what is held in memory

property fp16_key_bytes: int

Hypothetical fp16 key cost if nothing were compressed.

property compressed_value_bytes: int

see nbytes.

Type:

Estimated packed bytes for the compressed value cache (product-VQ codes at bits each, all heads/batches; excludes codebook). Not what is held in memory

property fp16_value_bytes: int

Hypothetical fp16 value cost if nothing were compressed.

property codebook_bytes: int

Static codebook overhead in bytes (fp16 storage).

property assigned_avg_bits: float

Effective bits/element averaged over keys and values.

Excludes codebook overhead (amortized across many tokens); for an end-to-end byte ratio use compressed_*_bytes / fp16_*_bytes.

property merge