veloxquant_mlx.quantizers.knorm module

L2-norm KV eviction primitives — intrinsic key-norm token importance.

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

The paper’s finding (in trained decoder LMs): a low L2 norm of a key embedding usually leads to a high attention score during decoding — a KV pair’s influence is largely determined by the key itself, before it is ever queried. Eviction follows directly: rank cached tokens by key norm, keep the lowest-norm ones. This is the repo’s first intrinsic eviction scorer: every other eviction method scores with attention / a key-as-query proxy (SnapKV, H2O, TOVA, PyramidKV, SqueezeAttention, ChunkKV, CaM) or pure structure (StreamingLLM, sink, sliding-window). Note the sign inversion vs ChunkKV’s key_norm option and ZipCache’s saliency proxy, which treat high norm as important — the inversion is the paper’s empirical content.

Path-independence invariant (recent=0): the norm is computed once at insertion and never updated, so evicting the worst-scoring non-sink token whenever over budget is the classic “keep k best with a heap” algorithm — the final kept set equals the global budget-best over all tokens seen, regardless of how arrivals were grouped. Prefill-in-one-block and token-by-token decode produce bit-for-bit identical kept sets. No accumulating-score method (H2O, TOVA) has this property. (recent > 0 breaks it: the protected window moves with time.)

Adaptation limitations (stated plainly):

  • The low-norm ⇒ high-attention correlation is the paper’s empirical claim about trained models; nothing here validates it on synthetic data (see the benchmark’s isotropic control, where the method shows no advantage).

  • No RoPE position-ID renumbering after eviction — surviving tokens keep their original absolute positions (knorm_update restores temporal order after top-k selection), so the cache layer reports the true token position for RoPE rather than the retained row count (see cache/knorm_cache.py and :issue:`171`, :issue:`174`).

  • Uniform budget and n_sink across all heads.

  • recent (trailing protected window) is an extension, off by default.

Public API (mirrors quantizers/h2o.py)

KnormState — immutable per-head state dataclass init_knorm_state — construct empty state (validates budget vs guards) knorm_update — absorb a whole [S, D] block, evict if over budget knorm_get_kv — extract current (keys, values) arrays knorm_fp16_bytes — bytes stored in current state full_knorm_fp16_bytes — hypothetical cost without eviction

class veloxquant_mlx.quantizers.knorm.KnormState(keys, values, norms, n_sink, budget, recent, keep)[source]

Bases: object

Per-head L2-norm eviction state.

Variables:

  • keys – [n_kept, D] fp16 stored key rows, or None before first update.

  • values – [n_kept, D] fp16 stored value rows, or None before first update.

  • norms – [n_kept] float32 L2 norm of each kept key row — computed once at insertion, never updated (the intrinsic-score property).

  • n_sink – Number of leading sink positions — never evicted.

  • budget – Maximum tokens kept at any time (including sinks).

  • recent – Trailing protected window (0 = off, paper-faithful).

  • keep – “low” = keep lowest-norm tokens (the paper’s finding); “high” = inverted selection (the benchmark’s ablation arm).

Parameters:

  • keys (mx.array | None)

  • values (mx.array | None)

  • norms (mx.array | None)

  • n_sink (int)

  • budget (int)

  • recent (int)

  • keep (str)

keys: mx.array | None
values: mx.array | None
norms: mx.array | None
n_sink: int
budget: int
recent: int
keep: str
veloxquant_mlx.quantizers.knorm.init_knorm_state(n_sink, budget, head_dim, recent=0, keep='low')[source]

Create an empty KnormState before any tokens arrive.

Raises:

ValueError – if keep is invalid or the protected positions (sinks + recent) leave no evictable room within the budget.

Return type:

KnormState

Parameters:
veloxquant_mlx.quantizers.knorm.knorm_update(state, new_keys, new_values)[source]

Absorb a whole block of S tokens, then evict down to budget in one shot.

Vectorized — no per-token loop: norms are intrinsic (never updated), so the over-budget case is a single protected top-k selection. Kept tokens are returned in original temporal order.

Return type:

KnormState

Parameters:

  • state (KnormState)

  • new_keys (mlx.core.array)

  • new_values (mlx.core.array)

veloxquant_mlx.quantizers.knorm.knorm_update_batched(prior_keys, prior_values, prior_norms, new_keys, new_values, budget, n_sink, recent, keep, *, return_indices=False)[source]

Batched knorm_update across N = B*H independent rows.

Requires budget/n_sink/recent/keep to be uniform across all N rows (true for every L2NormKVCache instance — one config per cache, not per head), so every row’s n_total after concatenation is identical and the top-k selection is a single batched argsort along the token axis instead of N independent Python-level calls.

prior_keys/prior_values/prior_norms must be pre-materialized with a uniform n_prior (0 before any tokens have arrived — the caller passes empty [N, 0, D]/[N, 0] arrays rather than None, since batching needs uniform shapes; unlike append_paged_summary there is no partial-page carry to fold in, just concatenation followed by an ordinary top-k).

With return_indices=True, append the selected input-row indices.

Returns (keys, values, norms), each [N, min(n_total, budget), D] / [N, min(n_total, budget)], kept rows in original temporal order — bit-for-bit identical to running knorm_update per row in a loop (exact reindex, no lossy floating-point reduction involved).

Return type:

tuple[array, array, array] | tuple[array, array, array, array]

Parameters:

  • prior_keys (mlx.core.array)

  • prior_values (mlx.core.array)

  • prior_norms (mlx.core.array)

  • new_keys (mlx.core.array)

  • new_values (mlx.core.array)

  • budget (int)

  • n_sink (int)

  • recent (int)

  • keep (str)

  • return_indices (bool)

veloxquant_mlx.quantizers.knorm.knorm_get_kv(state)[source]

Return (keys, values) arrays from state.

Returns ([0, 1], [0, 1]) zero-row placeholders before the first update (same contract as h2o_get_kv).

Return type:

tuple[array, array]

Parameters:

state (KnormState)

veloxquant_mlx.quantizers.knorm.knorm_fp16_bytes(state)[source]

Bytes currently stored for K + V in fp16.

Return type:

int

Parameters:

state (KnormState)

veloxquant_mlx.quantizers.knorm.full_knorm_fp16_bytes(tokens_seen, head_dim)[source]

Hypothetical fp16 K + V bytes if all tokens_seen were stored.

Return type:

int

Parameters:

  • tokens_seen (int)

  • head_dim (int)