veloxquant_mlx.quantizers.morphkv module

MorphKV KV eviction primitives — recent-window correlation retention.

Inspired by “Dialogue Without Limits: Constant-Sized KV Caches for Extended Responses in LLMs” (Ghadia, Kumar, Jain, Nair, Das, ICML 2025, arXiv:2503.00979). Documented as “MorphKV-adapted (VeloxQuant-MLX implementation)” — not a faithful port. See adaptation limitations below and in cache/morphkv_cache.py.

The paper’s finding: retaining KV by a cumulative attention score (H2O-style) suffers “early-token bias” — tokens that were heavy hitters early dominate the keep set and crowd out context the model is currently attending to. MorphKV keeps a constant-size cache by ranking stored tokens according to their correlation with the attention pattern of a sliding window of recent tokens, so retention tracks what the recent context actually reads and older-but-stale tokens are dropped.

WHERE THIS SITS IN THE REPO

This is the repo’s proxy-attention scorer family (SnapKV / H2O / TOVA / PyramidKV / SqueezeAttention / ChunkKV / CaM / Keyformer). The distinguishing axis: every existing scorer ranks a stored token against either cumulative history (H2O accumulates mass forever) or a single most-recent query (TOVA / SnapKV use the latest position). MorphKV-adapted is the first to rank by correlation with a window of the last window tokens. Two honest reference behaviors bracket it:

  • window = 1 collapses onto a latest-token (TOVA-adapted-style) eviction: the recent-relevance signal is just the newest key’s attention over the keep set. This is the clean, assertable reduction (exercised by a dedicated test), the analogue of Keyformer’s tau = 0 == H2O collapse.

  • a large window approaches averaging over recent context; it never becomes H2O’s cumulative-forever rule (MorphKV recomputes from the live window, it does not accumulate), so we do NOT claim an H2O collapse — only the window = 1 reduction is pinned exactly.

THE HONESTY CRUX (read before trusting any number)

  1. Proxy query. Like H2O / TOVA / SnapKV / Keyformer-adapted, a cache never sees the true query vector, so incoming KEYS are used as proxy queries to estimate the attention each stored key receives. The paper uses the model’s real attention patterns. Documented substitution, not the paper’s math.

  2. Constant-size, recomputed — not accumulated. We keep no cumulative score array. Each step, retention is recomputed from the current keep set and a ring buffer of the last window key rows. That is the mechanism: the cache is a fixed budget refreshed against recent context, not a growing accumulator.

  3. Nothing here is validated on a trained model. The paper’s headline numbers (accuracy / memory savings) are the paper’s, on trained models — NEVER quoted as ours. The mechanism’s benefit is measured only under a constructed “topic-shift” geometry in the benchmark, with a null control where it shows no advantage.

Adaptation limitations (stated plainly):

  • Key-as-query proxy (crux 1).

  • No RoPE position-ID remapping after eviction.

  • Uniform budget / n_sink / window across all heads.

  • The trailing window tokens are protected from eviction (they are the recency context that drives the ranking); leading n_sink tokens are protected as sinks.

Public API (mirrors quantizers/keyformer.py)

MorphKVState — immutable per-head state dataclass init_morphkv_state — construct empty state (validates guards) morphkv_update — absorb S new tokens, evict least recent-relevant if over budget morphkv_get_kv — extract current (keys, values) arrays morphkv_fp16_bytes — bytes stored in current state full_morphkv_fp16_bytes — hypothetical cost without eviction

class veloxquant_mlx.quantizers.morphkv.MorphKVState(keys, values, pos, n_sink, budget, window, head_dim)[source]

Bases: object

Per-head MorphKV 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.

  • pos – Running count of token positions this head has ever inserted (diagnostic; retention itself is stateless-of-history — it is recomputed from keys and the recent window each step).

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

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

  • window – Size of the trailing recent-token window whose aggregate proxy-attention drives retention (>= 1). The last window stored tokens are themselves protected from eviction.

  • head_dim – Head dimension D (for byte accounting before first insert).

Parameters:

  • keys (mx.array | None)

  • values (mx.array | None)

  • pos (int)

  • n_sink (int)

  • budget (int)

  • window (int)

  • head_dim (int)

keys: mx.array | None
values: mx.array | None
pos: int
n_sink: int
budget: int
window: int
head_dim: int
veloxquant_mlx.quantizers.morphkv.init_morphkv_state(n_sink, budget, head_dim, window=8)[source]

Create an empty MorphKVState before any tokens arrive.

Raises:

ValueError – if budget < 1, window < 1, n_sink >= budget, or window > budget (the recent window cannot exceed the cache), or if sinks + window leave no evictable room.

Return type:

MorphKVState

Parameters:
veloxquant_mlx.quantizers.morphkv.morphkv_update(state, new_keys, new_values)[source]

Absorb S new tokens, evicting the least recent-relevant token if over budget.

Return type:

MorphKVState

Parameters:

  • state (MorphKVState)

  • new_keys (mlx.core.array)

  • new_values (mlx.core.array)

For each of the S incoming tokens:

  1. Append the new token to the cache.

  2. If over budget: rank the current keep set by recent-window relevance (_recent_relevance() over the last window stored keys), force the leading n_sink sinks and the trailing window recent tokens to survive (+inf), and evict the non-protected token with the LOWEST recent-relevance. Constant-size: n_kept <= budget after every token.

No cumulative score is carried across steps — the ranking is recomputed fresh each step from the live keep set and recent window (that is the mechanism). Kept tokens are returned in original temporal order.

veloxquant_mlx.quantizers.morphkv.morphkv_update_batched(keys, values, new_keys, new_values, n_sink, budget, window, *, return_indices=False)[source]

Vectorized-over-BH equivalent of calling morphkv_update() once per (batch, head) pair with identical per-row n_sink/ budget/window (true for every real caller: MorphKVKVCache applies one uniform config to every head), so the per-token append/rank/ evict math — otherwise identical for every row — can run as one batched MLX call per step instead of BH separate Python-level calls into morphkv_update(). Same fix, same template, as veloxquant_mlx.quantizers.h2o.h2o_update_batched(). Uses _recent_relevance_batched() internally, which also removes the inner for j in range(window): loop _recent_relevance() runs per call — issue #560’s second, compounding finding.

Numerically identical to the per-head loop it replaces: every op below is the same formula as morphkv_update()’s bootstrap/append/evict branches, applied over a leading BH axis instead of a Python loop.

With return_indices=True, also return indices into prior + new rows.

Return type:

tuple[mx.array, mx.array] | tuple[mx.array, mx.array, mx.array]

Returns:

(keys, values) — [BH, n_kept, D] each.

Parameters:

  • keys (mx.array | None)

  • values (mx.array | None)

  • new_keys (mx.array)

  • new_values (mx.array)

  • n_sink (int)

  • budget (int)

  • window (int)

  • return_indices (bool)

veloxquant_mlx.quantizers.morphkv.morphkv_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 keyformer_get_kv / tova_get_kv).

Return type:

tuple[array, array]

Parameters:

state (MorphKVState)

veloxquant_mlx.quantizers.morphkv.morphkv_fp16_bytes(state)[source]

Bytes currently stored for K + V in fp16.

The recent-window ring is a view into keys (not extra payload), so only K + V are counted — same accounting as H2O / TOVA / Keyformer.

Return type:

int

Parameters:

state (MorphKVState)

veloxquant_mlx.quantizers.morphkv.full_morphkv_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)