veloxquant_mlx.cache.knorm_cache module

L2Norm-adapted KV cache — intrinsic key-norm eviction.

Inspired by “A Simple and Effective L2 Norm-Based Strategy for KV Cache Compression” (Devoto, Zhao, Scardapane, Minervini — EMNLP 2024; arXiv:2406.11430). Documented as “L2Norm-adapted (VeloxQuant-MLX implementation)” — not a faithful port.

The repo’s first intrinsic-signal eviction cache: token importance is read directly off the stored key vector’s L2 norm — the paper’s finding is that in trained decoder LMs a low key norm predicts high future attention, so the cache keeps the lowest-norm tokens and evicts the highest-norm ones. No attention scores, no key-as-query proxy (the approximation H2O/SnapKV/TOVA need), no structure-only recency rule (StreamingLLM): the paper’s actual signal is fully observable at the cache level, making this the cleanest adaptation in the eviction family.

Because the score is intrinsic (computed once at insertion, never updated):

  • eviction vectorizes as one protected top-k per incoming block — no per-token softmax-over-cache loop like H2O;

  • with knorm_recent=0 the kept set is path-independent: prefill in one block and token-by-token decode yield bit-for-bit identical caches (the “keep k best with a heap” invariant — see quantizers/knorm.py).

update_and_fetch batches every head’s update in one shot via knorm_update_batched (reshape to [B*H, ...], one batched argsort, reshape back) instead of a Python for b: for h: loop — safe because budget/n_sink/recent/keep are uniform across all heads in one cache instance, so every head’s state has identical shape at every step (same precondition RocketKV’s decode-summary batching relied on). Measured on M4: H=8 1.53ms -> 0.76ms/step (2.0x), H=32 4.07ms -> 0.77ms/step (5.3x).

Adaptation limitations (stated plainly):

  • The low-norm ⇒ high-attention correlation is the paper’s empirical claim about trained models — not validated here on synthetic data (the benchmark’s isotropic control shows no advantage, honestly reported).

  • No RoPE position-ID renumbering after eviction. Surviving tokens keep their original absolute positions, so self.offset reports the true token position and RoPE stays correct without re-rotating survivors (see update_and_fetch and :issue:`171`, :issue:`174`). Positions do become non-contiguous where tokens were dropped.

  • Uniform budget and n_sink across all heads.

  • knorm_recent (trailing protected window) is an extension, off by default; enabling it breaks the path-independence property.

Byte accounting (same names as H2OKVCache):

knorm_kept_bytes — fp16 bytes for currently retained K + V tokens full_seq_bytes — hypothetical fp16 cost if all tokens were kept compression_ratio — full_seq_bytes / knorm_kept_bytes (> 1 = savings) tokens_seen — total token positions ever passed to update_and_fetch tokens_kept — tokens currently in the first (B=0, H=0) head’s cache

class veloxquant_mlx.cache.knorm_cache.L2NormKVCache(config)[source]

Bases: DeferredEvictionMixin, KVCache

KV cache implementing L2Norm-adapted intrinsic key-norm eviction for one layer.

Parameters:

config (Any) – KVCacheConfig. Fields consumed: knorm_budget (int, default 512) — max tokens retained (incl. sinks), knorm_n_sink (int, default 4) — leading positions never evicted, knorm_recent (int, default 0) — trailing protected window (extension), knorm_keep (str, default “low”) — “low” = paper finding; “high” = inverted.

Notes

No .bits attribute — stores and returns fp16 K/V directly. Single-layer (no coordinator); the default KVCacheBuilder.for_model() path returns one L2NormKVCache per attention layer. Per-head state is lazily initialised on the first update_and_fetch. Validation (keep mode, sink/recent-vs-budget guard) happens at construction. Writes through to the base mlx_lm KVCache’s self.keys / self.values / self.offset on every call so .state stays valid (mlx_lm’s generate() reads it unconditionally during chunked prefill); is_trimmable() reports False since the internal per-token state can’t be rolled back by a base-class trim() (see #83).

update_and_fetch(keys, values)[source]

Absorb new K/V tokens, apply key-norm eviction, return retained window.

Parameters:

  • keys (array) – [B, H, S, D] new key tokens (any dtype; cast to fp16).

  • values (array) – [B, H, S, D] new value tokens.

Returns:

Full prior kept rows plus all new rows, before eviction (#610). Only the state stored for the next call is compressed.

is_trimmable()[source]

False: trim() would only roll back base-class offset bookkeeping, not the internal per-token eviction/compression state that actually determines what gets returned, silently corrupting future calls.

Return type:

bool

property merge
property knorm_kept_bytes: int

Bytes currently stored across all heads (fp16 K + V, kept tokens only).

property full_seq_bytes: int

Hypothetical fp16 K + V cost if all tokens were kept.

property compression_ratio: float

full_seq_bytes / knorm_kept_bytes; > 1 means memory savings over fp16.

property tokens_seen: int

Total token positions ever passed to update_and_fetch (all heads summed).

property tokens_kept: int

Tokens currently in the (B=0, H=0) head’s cache (diagnostic).