veloxquant_mlx.cache.nestedkv_cache module

NestedKV-adapted KV cache — multi-scale ensembled prefill eviction.

Inspired by “NestedKV: Nested Memory Routing for Long-Context KV Cache Compression” (Chen, Liu, Gao, Fan, Wang, Chu, Lin, Hu; arXiv:2605.26678). Documented as “NestedKV-adapted (VeloxQuant-MLX implementation)” — not a faithful port. No verified peer-reviewed venue as of 2026-07-14 — a one-time, user-directed exception to this repo’s venue-verification rule. See paper/research/surveys/NEW_METHOD_SURVEY_V21.md.

Multi-scale ensembled eviction: at the end of prefill, each head’s tokens are scored by three parallel continuum-memory anomaly signals (stable/global, episodic/block-local, current/recent-window), combined via a head-adaptive blend and a per-token surprise gate (see quantizers/nestedkv.py). Each head then independently keeps its own top-nestedkv_budget scoring tokens — UNIFORM per-head budget, like H2O/CurDKV (see “Uniform per-head budget” below, issue #21, for why this class does not use the paper’s cross-head budget competition even though quantizers/nestedkv.py still implements it as a tested primitive). Decode tokens are appended unscored, never evicted — this is a one-shot prefill compressor, not a per-step recurring eviction loop.

This is the 15th eviction-family method in VeloxQuant-MLX and the first that ensembles multiple independent importance signals rather than committing to one:

  • H2O / CurDKV / KVzip / Keyformer / MorphKV : one signal, scored every step.

  • SnapKV : one signal (obs-window attention), scored once at prefill.

  • NestedKV : THREE signals, scored once at prefill, combined by a head-adaptive blend + per-token surprise gate.

Adaptation limitations (stated plainly — see quantizers/nestedkv.py for the full crux):

  • Unpublished preprint, no verified venue.

  • One-shot prefill compression; cache is NOT bounded during decode (grows with every decoded token, same as SnapKV-adapted’s decode-phase design).

  • Key-only scoring, no query/attention access at all (not even a proxy).

  • Gate/blend constants (beta=3.0, tau=0.60, kappa=10.0, prior=(0.4,0.4,0.2), safeguard_alpha=0.20) taken directly from the paper’s Appendix A.

  • Uniform per-head budget (issue #21). The paper’s cross-head budget competition (component 5, nestedkv_allocate_head_budgets) can legitimately give different heads different token counts. An earlier version of this wrapper let each head keep its ragged, cross-head- competed count and zero-padded shorter heads (at the front) to stack them into one tensor. That was a real correctness bug, not just a shape accommodation: the padded rows became real (unmasked) cache entries that the downstream attention computation attended to — and mlx_lm’s attention forward computes ONE shared mask from the model’s first layer’s cache and reuses it for every layer, so a per-cache mask override cannot correctly express each NestedKV layer’s independent padding pattern. Rather than build mask machinery this architecture cannot actually support end-to-end, this class now keeps every head’s kept length uniform at nestedkv_budget — same convention as every other eviction method here (H2O, CurDKV, PyramidKV). nestedkv_score (the cross-scale anomaly ranking) still runs per head exactly as before; only the cross-head reallocation of how many tokens each head keeps is dropped, in favor of each head independently keeping its own top-nestedkv_budget scoring tokens (sinks always included) — the same per-head-independent budget model H2O and CurDKV already use. No padding is ever needed, so there is nothing left to mask.

Byte accounting:

nestedkv_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 / nestedkv_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.nestedkv_cache.NestedKVKVCache(config)[source]

Bases: DeferredEvictionMixin, KVCache

KV cache implementing NestedKV-adapted multi-scale ensembled eviction.

Parameters:

config (Any) –

KVCacheConfig. Fields consumed: nestedkv_budget (int, default 512) — per-head budget

(uniform across heads, see #21 — NOT per-head-equivalent summed into a cross-head-competed total as in the paper),

nestedkv_n_sink (int, default 4) — leading sink positions, nestedkv_window (int, default 64) — current-memory window W, nestedkv_beta (float, default 3.0) — head-adaptive blend temperature, nestedkv_tau (float, default 0.60)— surprise gate threshold, nestedkv_kappa (float, default 10.0)— surprise gate sharpness, nestedkv_safeguard_alpha (float, default 0.20) — parsed and

stored for API stability but no longer consumed by this class (see #21): it configured the cross-head budget floor in nestedkv_allocate_head_budgets, which this class no longer calls. Still a real, tested parameter of that quantizer-level function for direct callers.

Notes

No .bits attribute — stores and returns fp16 K/V directly. Eviction happens ONCE at prefill (S > 1), independently per head at a UNIFORM budget (see #21 — NOT the paper’s cross-head budget competition, which produced ragged per-head lengths that could not be safely stacked into one tensor without either corrupting attention with fake padding rows or requiring per-layer attention masks this architecture doesn’t support). Decode tokens (S == 1) are always appended, never rescored or evicted — same convention as SnapKV-adapted, NOT H2O’s/CurDKV’s per-step loop. Single-layer (no coordinator); KVCacheBuilder.for_model() propagates all nestedkv_* fields automatically via dataclasses.replace. 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; prefill compresses once, decode always appends.

Parameters:

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

  • 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 nestedkv_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 / nestedkv_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).