veloxquant_mlx.cache.morphkv_cache module

MorphKV-adapted KV cache — 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.

Keeps a constant-size cache by ranking stored tokens according to their correlation with the aggregate proxy-attention of a sliding window of recent tokens (paper’s contribution: eliminate the “early-token bias” of cumulative H2O-style scoring by tracking what the recent context is actually attending to). Setting morphkv_window = 1 reduces this to a latest-token (TOVA-adapted- style) eviction — the honest reference behavior, checked by a dedicated test.

Where it sits: the repo’s proxy-attention scorer family (SnapKV / H2O / TOVA / PyramidKV / SqueezeAttention / ChunkKV / CaM / Keyformer). The distinguishing axis is recent-window correlation, versus cumulative (H2O) or latest-only (TOVA) scoring.

THE HONESTY CRUX:

  1. Proxy query — the incoming KEY stands in for the unseen query (as H2O / TOVA / SnapKV / Keyformer-adapted).

  2. Constant-size, recomputed each step from the live keep set + recent window — NOT a cumulative accumulator. Only window = 1 collapse is pinned exactly (to the latest-token ranking); no H2O collapse is claimed.

  3. Not validated on a trained model; the paper’s accuracy/memory numbers are the paper’s, on trained models — never reproduced or claimed here. The mechanism’s benefit is measured only under a constructed topic-shift geometry, with a null control.

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.

  • Trailing morphkv_window tokens protected (recency context that drives the ranking); leading morphkv_n_sink tokens protected as sinks.

Byte accounting (same names as H2OKVCache / KeyformerKVCache):

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

class veloxquant_mlx.cache.morphkv_cache.MorphKVKVCache(config)[source]

Bases: DeferredEvictionMixin, KVCache

KV cache implementing MorphKV-adapted recent-window retention for one layer.

Parameters:

config (Any) – KVCacheConfig. Fields consumed: morphkv_budget (int, default 512) — max tokens kept (incl. sinks), morphkv_n_sink (int, default 4) — leading positions never evicted, morphkv_window (int, default 8) — trailing recent-attention window.

Notes

No .bits attribute — stores and returns fp16 K/V directly. Both prefill (S > 1) and decode (S == 1) go through the same update loop. Single-layer (no coordinator); KVCacheBuilder.for_model() propagates all morphkv_* fields via dataclasses.replace. Per-head state is lazily initialised on the first update_and_fetch. MorphKV is deterministic (no RNG). Validation (budget/window/sink bounds) 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).

property merge
update_and_fetch(keys, values)[source]

Absorb new K/V tokens, apply recent-window eviction, return 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 morphkv_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 / morphkv_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).