veloxquant_mlx.cache.nsnquant_cache module

NSNQuant KV cache wrapper for mlx_lm integration.

Implements “NSNQuant: A Double Normalization Approach for Calibration-Free Low-Bit Vector Quantization of KV Cache” (Son, Choi, Yoo — NeurIPS 2025; arXiv:2505.18231) on top of the standard mlx_lm update_and_fetch protocol. Documented as “NSNQuant-adapted (VeloxQuant-MLX implementation)” — not a faithful port (post-RoPE keys, explicit value Hadamard, k-means-only codebook, fp16 metadata; see the module docstring of veloxquant_mlx/quantizers/nsnquant.py and paper/NEW_METHOD_SURVEY_V11.md).

Chunk-flush residual buffer (the paper’s decode-time story, and the same idiom as KIVI’s fp16 residual window): tokens accumulate at fp16; every time nsn_residual_length tokens age past the quantized frontier, that chunk is round-tripped through NSN + Hadamard + universal-codebook VQ as one self-contained unit — the chunk computes its own channel mean o, so no statistics are frozen, no coordinator exists, and chunk i is forever independent of anything that arrives later. Prefill and decode produce identical chunk boundaries by construction (the flush frontier only ever advances in whole chunks), so the quantized state is path-independent.

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. Both keys and values are quantized (mirroring the paper), unlike the keys-only SVDq/xKV precedent. 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.

Per-token storage at nsn_bits b for head_dim D (per tensor): D * b / 8 bytes of codes (2-bit: uint8 sign mask + uint8 index per 8-dim subvector; 1-bit: index only) + 4 bytes fp16 (s1, s2) per token + 2 * D bytes fp16 (channel mean o) amortized per chunk. All three metadata terms are counted in compressed_*_bytes — not waved away (the paper 4-bit double-quantizes them; we do not).

class veloxquant_mlx.cache.nsnquant_cache.NSNQuantKVCache(config)

Bases: KVCache

KV cache implementing NSNQuant universal-codebook VQ for one layer.

Parameters:

config (Any) – KVCacheConfig. Fields consumed: head_dim, nsn_bits (2 = sign mask + index, 1 = index only), nsn_residual_length (chunk size / fp16 window; the paper recommends 128 for 1-bit), nsn_codebook_size, nsn_subvector_dim, nsn_seed, nsn_max_ctx.

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, so decode-time tokens do get quantized once they age past the chunk frontier (unlike KIVI’s incoming-block-only simplification).

trim(n)

Drop the last n tokens and roll the chunk-flush frontier and byte accounting back with them.

The frontier snaps down to a chunk multiple. Rows between it and the new offset were already round-tripped (their chunk’s channel mean included discarded tokens) and are quantized again on the next flush.

Return type:

int

Parameters:

n (int)

property compressed_key_bytes: int

Realized stored bytes for the compressed key cache (double-normalized VQ codes + fp16 scale metadata, flushed chunks only).

property compressed_value_bytes: int

Realized stored bytes for the compressed value cache (double-normalized VQ codes + fp16 scale metadata, flushed chunks only).

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

Bytes currently held in the fp16 residual window (keys + values) — a snapshot of the un-flushed tail, not a cumulative counter.

property quantized_tokens: int

Tokens behind the chunk-flush frontier (multiple of the chunk size).

property assigned_avg_bits: float

Effective bits/element over the quantized region, including the fp16 s1/s2/o metadata (excludes the fp16 residual window; for an end-to-end ratio use (compressed_*_bytes + residual_fp16_bytes) / fp16_*_bytes).

property merge