veloxquant_mlx.cache.skvq_cache module

SKVQ KV cache wrapper for mlx_lm integration.

Implements “SKVQ: Sliding-window Key and Value Cache Quantization for Large Language Models” (Duanmu, Yuan, Li, Duan, Zhang, Lin — COLM 2024; arXiv:2405.06219) on top of the standard mlx_lm update_and_fetch protocol. Documented as “SKVQ-adapted (VeloxQuant-MLX implementation)” — not a faithful port (no offline calibration, no weight-fused permutation, no 1.5-bit values, no FP8 metadata; see the module docstring of veloxquant_mlx/quantizers/skvq.py and paper/NEW_METHOD_SURVEY_V13.md).

Three mechanisms compose:

1. Sliding fp16 window — the NSNQuant chunk-flush idiom: tokens accumulate at fp16; every time skvq_window tokens age past the quantized frontier, that chunk is round-tripped through reorder → clipped group quant → dequant → inverse reorder once and frozen. The frontier only advances in whole chunks, so prefill and token-by-token decode produce identical quantized state by construction (path independence, pinned by test).

2. Channel reordering — per-head permutations for K and V computed from the first flushed chunk (the paper computes them offline from a calibration corpus and fuses them into projection weights — our documented deviation), then frozen for the cache’s lifetime. skvq_reorder=False is the identity-permutation ablation.

3. Clipped dynamic quantization — per-token, per-group asymmetric min/max quant whose window is shrunk by a per-group grid-searched clip factor α (skvq_clip_search=True) or a fixed skvq_clip_alpha. Quant groups run along the channel axis for both K and V; reordering is what makes that viable for keys.

The paper’s attention-sink filter is kept: the first skvq_n_sink tokens are restored to fp16 after chunk 0’s flush (and accounted as fp16, not compressed).

Like every method in this repo, the quantize→dequantize round-trip happens inside update_and_fetch so the downstream SDPA call sees standard fp16 tensors. The paper’s throughput gains come from fused CUDA kernels that do not port to Metal — on Apple Silicon the win is memory, measured honestly by the byte accounting below. Fully deterministic: no RNG anywhere.

Per-token storage at bits b, head_dim D, group size g (per tensor): D * b / 8 bytes of codes + ceil(D/g) * 2 * 2 bytes fp16 (lo, scale) per token. The searched α adds nothing (folded into lo/scale); the frozen permutations add 2 * H * D * 4 bytes per layer total, counted once in perm_bytes.

class veloxquant_mlx.cache.skvq_cache.SKVQKVCache(config)

Bases: KVCache

KV cache implementing SKVQ sliding-window quantization for one layer.

Parameters:

config (Any) – KVCacheConfig. Fields consumed: head_dim, skvq_bits_key, skvq_bits_value, skvq_group_size, skvq_window (fp16 sliding window == flush chunk size), skvq_n_sink (leading tokens restored to fp16, the paper’s filter), skvq_reorder, skvq_clip_search, skvq_clip_alpha (used when search is off), skvq_max_ctx (optional hard cap on context length; None, the default, is unlimited).

Notes

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

update_and_fetch(keys, values)

Append the incoming block, then flush every completed chunk.

The flush overwrites the aged-out tokens’ fp16 storage in place with their dequantized round-trip (quantize once, frozen), so decode-time tokens are quantized when they age past the window — the paper’s sliding-window semantics, unlike KIVI’s incoming-block-only simplification. Chunk 0 additionally restores the first skvq_n_sink rows to their exact fp16 values (the sink filter) and computes the frozen channel permutations.

property compressed_key_bytes: int

Realized stored bytes for the quantized (behind-the-flush-frontier) key region, excluding the fp16 sliding-window tail and sinks.

property compressed_value_bytes: int

Realized stored bytes for the quantized (behind-the-flush-frontier) value region, excluding the fp16 sliding-window tail and sinks.

property fp16_key_bytes: int

Hypothetical fp16 key cost if nothing were compressed.

property fp16_value_bytes: int

Hypothetical fp16 value cost if nothing were compressed.

property residual_fp16_bytes: int

the un-flushed sliding-window tail plus the sink rows that flushing skipped — a snapshot, not a cumulative counter.

Type:

Bytes currently held at fp16 (keys + values)

property perm_bytes: int

One-off storage for the frozen per-head permutations (int32).

property quantized_tokens: int

Tokens behind the flush frontier (multiple of the window), including the fp16-restored sink rows.

property tokens_seen: int

Total tokens processed through update_and_fetch so far.

property key_perms: mx.array | None

Frozen per-head key channel permutations [H, D] (None before the first flush or when skvq_reorder=False).

property value_perms: mx.array | None

Frozen per-head value channel permutations [H, D] (None before the first flush or when skvq_reorder=False).

property assigned_avg_bits: float

Effective bits/element over the quantized key region, including the fp16 (lo, scale) metadata (excludes the fp16 window and sinks; for an end-to-end ratio use (compressed_*_bytes + residual_fp16_bytes) / fp16_*_bytes).

is_trimmable()

False: trim() would only roll back base-class offset bookkeeping, not the internal flush frontier (_q_end) or the frozen per-head channel permutations (_perm_k/_perm_v, computed once from chunk 0 and never revisited) that actually determine what gets quantized and how. Unlike a pure-eviction cache, the information loss here is irreversible: once a chunk is flushed, its rows are overwritten in place with a lossy quantize/dequantize round-trip, so there is no exact fp16 value to restore even if _q_end were rolled back too. Reproduced directly: trimming past _q_end leaves offset < _q_end, and the next update_and_fetch treats already-flushed rows sitting in the buffer at those same indices as already-quantized instead of overwriting them with the new tokens actually meant to land there — silently corrupting future calls rather than crashing (see VeloxQuant-Studio issue #26).

Return type:

bool

property merge