veloxquant_mlx.quantizers.nestedkv module

NestedKV-adapted quantizer primitives — 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 — this is a one-time, user-directed exception to this repo’s standing rule that every shipped method trace to a verified peer-reviewed venue. See paper/research/surveys/NEW_METHOD_SURVEY_V21.md for the full rationale. Every other method in this repo (38 of 39) has a verified venue; this is the only exception, and the next method survey reverts to requiring one.

What makes this a new axis: every eviction method already in the repo scores a token from ONE importance signal — cumulative attention (H2O), prefill observation-window attention (SnapKV), layer-adaptive budget (PyramidKV), key-norm proxy (Keyformer/MorphKV), reconstruction-reliance (KVzip), or joint key-value leverage score (CurDKV). NestedKV keeps THREE parallel key-only statistics at three time scales (stable/global, episodic/block-local, current/recent-window), scores a token’s anomaly against each independently, and combines the three rankings via a training-free two-axis outer rule: a per-head blend (which scale is most discriminative on this head) and a per-token surprise gate (route to the single strongest scale when the three disagree, rather than average them).

Adaptation limitations (stated plainly):

• Unpublished preprint, no verified venue (see module header above) — the headline exception for this method.

• One-shot prefill compression, unbounded during decode. The paper’s own design (Appendix A) computes scores, blend weights, and surprise gates ONCE at the end of prefill; decoded tokens are appended normally, never rescored or evicted. This is a faithful port of the paper’s actual design, not a shortcut — but it means NestedKV’s cache, unlike every other eviction method here (H2O, CurDKV, StreamingLLM), does not stay bounded through a very long decode run. Mirrors SnapKV-adapted’s prefill-once / decode-append structure, not H2O’s/CurDKV’s per-step loop.

• Episodic block means are computed over fixed prefill-time token positions — a faithful, non-approximated port of the paper’s formula (no eviction-collapses-indices problem exists here, since eviction only happens once, after all three memory scales are already computed).

• Gate/blend constants — beta=3.0, tau=0.60, kappa=10.0, log-prior (0.4, 0.4, 0.2), safeguard_alpha=0.20 — are all taken directly from the paper’s Appendix A (“Hyperparameters”), not guessed.

• Key-only: no query/attention-score access at all, not even a proxy (stronger than H2O/SnapKV/CurDKV’s key-as-query proxy — NestedKV never approximates attention in the first place).

• No RoPE position-ID remapping after eviction; kept tokens preserve their relative order.

• No PyTorch/CUDA reference kept; pure MLX from the start.

Public API

NestedKVState — per-head prefill/decode state dataclass init_nestedkv_state — construct empty state nestedkv_score — one-shot per-head anomaly scoring over full prefill nestedkv_allocate_head_budgets — cross-head budget competition (paper’s component 5) nestedkv_compress_prefill — one-shot per-head scoring + eviction (prefill only) nestedkv_append_decode — plain unscored append (decode only) nestedkv_get_kv nestedkv_fp16_bytes full_nestedkv_fp16_bytes

class veloxquant_mlx.quantizers.nestedkv.NestedKVState(keys, values, n_sink, compressed=False)

Bases: object

Per-head NestedKV state.

Variables:

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

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

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

• compressed – Whether the one-shot prefill compression has already run.

Parameters:

• keys (mx.array | None)

• values (mx.array | None)

• n_sink (int)

• compressed (bool)

keys: mx.array | None

values: mx.array | None

n_sink: int

compressed: bool = False

veloxquant_mlx.quantizers.nestedkv.init_nestedkv_state(n_sink)

Create an empty NestedKVState before prefill.

Return type:

NestedKVState

Parameters:

n_sink (int)

veloxquant_mlx.quantizers.nestedkv.block_size_for(n)

b = clip(floor(n / 32), 128, 256), per the paper’s schedule (Section 2.2).

Return type:

int

Parameters:

n (int)

veloxquant_mlx.quantizers.nestedkv.per_scale_anomaly_scores(k_hat, block_size, window)

Compute (a_s_hat, a_e_hat, a_c_hat) — min-max-normalized per-scale anomaly.

Parameters:

• k_hat (array) – [n, D] L2-normalized keys.

• block_size (int) – episodic block size b.

• window (int) – current-memory trailing window W.

Return type:

tuple[array, array, array]

Returns:

Three [n] float32 arrays, each on [0, 1].

veloxquant_mlx.quantizers.nestedkv.head_adaptive_blend(a_s_hat, a_e_hat, a_c_hat, beta=3.0, prior=(0.4, 0.4, 0.2))

Blend the three per-scale anomaly scores with a head-adaptive softmax weight.

Delta_k measures how discriminative scale k is on this head (gap between its top-10% and bottom-10% scores). Larger gap -> more weight, anchored by a fixed log-prior so a tied/uninformative head falls back to (0.4, 0.4, 0.2).

Return type:

array

Returns:

[n] float32 blended score a_blend.

Parameters:

• a_s_hat (mlx.core.array)

• a_e_hat (mlx.core.array)

• a_c_hat (mlx.core.array)

• beta (float)

• prior (tuple[float, float, float])

veloxquant_mlx.quantizers.nestedkv.surprise_gated_score(a_s_hat, a_e_hat, a_c_hat, a_blend, tau=0.6, kappa=10.0)

Combine a_blend with the strongest individual scale via a surprise gate.

s(i) = std(a_s_hat[i], a_e_hat[i], a_c_hat[i]), then min-max normalized within the head and mean-centered (per the paper’s Appendix A: “surprise scores are min-max normalized within each head and mean-centered before applying the gate”). a_win(i) = max of the three. alpha(i) = sigmoid(kappa * (s(i) - tau)). a_star(i) = (1 - alpha(i)) * a_blend(i) + alpha(i) * a_win(i).

Return type:

array

Returns:

[n] float32 final NestedKV score a_star.

Parameters:

• a_s_hat (mlx.core.array)

• a_e_hat (mlx.core.array)

• a_c_hat (mlx.core.array)

• a_blend (mlx.core.array)

• tau (float)

• kappa (float)

veloxquant_mlx.quantizers.nestedkv.nestedkv_score(keys, window=64, beta=3.0, tau=0.6, kappa=10.0, prior=(0.4, 0.4, 0.2))

One-shot per-head NestedKV score over all N prefill keys.

Parameters:

• keys (array) – [n, D] fp16/fp32 key matrix for one head (full prefill).

• window (int) – current-memory trailing window W.

• beta (float) – gate/blend constants (paper Appendix A defaults).

• tau (float) – gate/blend constants (paper Appendix A defaults).

• kappa (float) – gate/blend constants (paper Appendix A defaults).

• prior (tuple[float, float, float]) – gate/blend constants (paper Appendix A defaults).

Return type:

array

Returns:

[n] float32 final score a_star (higher = more anomalous = keep-worthy).

veloxquant_mlx.quantizers.nestedkv.nestedkv_allocate_head_budgets(head_scores, total_budget, safeguard_alpha=0.2)

Allocate a total layer budget across H heads by global top-B score competition.

Two-step rule (paper Appendix A, “head-wise safeguard”):

1. Each head h keeps a guaranteed floor of floor(safeguard_alpha * (1-r) * n_h) of its own highest-scoring tokens, where r is the implied eviction ratio (derived from total_budget vs sum of n_h, not passed separately).

2. The remaining budget is filled by pooling all NOT-yet-guaranteed (head, position, score) triples across heads and taking the global top-K by score until the remaining budget is exhausted.

Parameters:

• head_scores (list[array]) – length-H list of [n_h] float32 score arrays (one per head).

• total_budget (int) – total tokens to keep across all H heads this layer.

• safeguard_alpha (float) – per-head guaranteed-floor fraction.

Return type:

list[int]

Returns:

length-H list of ints B_h, sum(B_h) == min(total_budget, sum(n_h)).

veloxquant_mlx.quantizers.nestedkv.nestedkv_compress_prefill(state, keys, values, budget, window=64, beta=3.0, tau=0.6, kappa=10.0)

Run the one-shot NestedKV prefill compression for one head.

Computes a_star over all S prefill tokens, pins the first n_sink positions, keeps the top-budget scoring tokens (sinks always included), in ascending original-position order. Runs exactly once per sequence.

Return type:

NestedKVState

Parameters:

• state (NestedKVState)

• keys (mlx.core.array)

• values (mlx.core.array)

• budget (int)

• window (int)

• beta (float)

• tau (float)

• kappa (float)

veloxquant_mlx.quantizers.nestedkv.nestedkv_append_decode(state, keys, values)

Plain unscored append for decode-phase tokens. Never evicts.

If called before prefill compression has run (state.keys is None), this also serves as the bootstrap path.

Return type:

NestedKVState

Parameters:

• state (NestedKVState)

• keys (mlx.core.array)

• values (mlx.core.array)

veloxquant_mlx.quantizers.nestedkv.nestedkv_get_kv(state)

Return (keys, values) arrays from state.

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

Return type:

tuple[array, array]

Parameters:

state (NestedKVState)

veloxquant_mlx.quantizers.nestedkv.nestedkv_fp16_bytes(state)

Bytes currently stored for K + V in fp16.

Return type:

int

Parameters:

state (NestedKVState)

veloxquant_mlx.quantizers.nestedkv.full_nestedkv_fp16_bytes(tokens_seen, head_dim)

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

Return type:

int

Parameters:

• tokens_seen (int)

• head_dim (int)

veloxquant_mlx.quantizers.nestedkv.per_scale_anomaly_scores_batched(k_hat, block_size, window)

Batched-[BH,n,D] equivalent of per_scale_anomaly_scores().

Returns three [BH, n] float32 arrays, each min-max normalized to [0, 1] independently per row (matching the per-head function’s per-head normalization).

Return type:

tuple[array, array, array]

Parameters:

• k_hat (mlx.core.array)

• block_size (int)

• window (int)

veloxquant_mlx.quantizers.nestedkv.head_adaptive_blend_batched(a_s_hat, a_e_hat, a_c_hat, beta=3.0, prior=(0.4, 0.4, 0.2))

Batched-[BH,n] equivalent of head_adaptive_blend().

Each row’s blend weights are computed from that row’s own scores only (no cross-row coupling) — same per-row-independence as every other NestedKV batched primitive. n<=1 rows fall back to the fixed prior exactly as the per-head function does; since n is shared across all rows in one call (same S), the fallback is an all-rows-or-no-rows branch.

Return type:

array

Returns:

[BH, n] float32 blended score a_blend.

Parameters:

• a_s_hat (mlx.core.array)

• a_e_hat (mlx.core.array)

• a_c_hat (mlx.core.array)

• beta (float)

• prior (tuple[float, float, float])

veloxquant_mlx.quantizers.nestedkv.surprise_gated_score_batched(a_s_hat, a_e_hat, a_c_hat, a_blend, tau=0.6, kappa=10.0)

Batched-[BH,n] equivalent of surprise_gated_score().

Return type:

array

Returns:

[BH, n] float32 final NestedKV score a_star.

Parameters:

• a_s_hat (mlx.core.array)

• a_e_hat (mlx.core.array)

• a_c_hat (mlx.core.array)

• a_blend (mlx.core.array)

• tau (float)

• kappa (float)

veloxquant_mlx.quantizers.nestedkv.nestedkv_score_batched(keys, window=64, beta=3.0, tau=0.6, kappa=10.0, prior=(0.4, 0.4, 0.2))

Batched-[BH,n,D] equivalent of nestedkv_score().

Parameters:

• keys (array) – [BH, n, D] fp16/fp32 key matrix (full prefill), one row per head.

• window (int) – gate/blend constants (shared across rows — every real caller applies one uniform config per layer).

• beta (float) – gate/blend constants (shared across rows — every real caller applies one uniform config per layer).

• tau (float) – gate/blend constants (shared across rows — every real caller applies one uniform config per layer).

• kappa (float) – gate/blend constants (shared across rows — every real caller applies one uniform config per layer).

• prior (tuple[float, float, float]) – gate/blend constants (shared across rows — every real caller applies one uniform config per layer).

Return type:

array

Returns:

[BH, n] float32 final score a_star (higher = more anomalous = keep-worthy).

veloxquant_mlx.quantizers.nestedkv.nestedkv_compress_prefill_batched(keys, values, n_sink, budget, window=64, beta=3.0, tau=0.6, kappa=10.0, *, return_indices=False)

Batched-[BH,S,D] equivalent of nestedkv_compress_prefill().

Runs the one-shot NestedKV scoring once across all BH rows, then keeps each row’s own top-budget_eff scoring tokens (sinks always included), in ascending original-position order — same per-row- independent selection as the per-head loop, batched via mx.argsort/ take_along_axis instead of Python .tolist()/sorted(). budget_eff depends only on n_sink/budget/S (all shared across rows in one call), so every row keeps exactly the same count — stacks with no padding, same invariant the per-head loop relies on.

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

Return type:

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

Returns:

(kept_keys, kept_values), both [BH, budget_eff, D] fp16.

Parameters:

• keys (mlx.core.array)

• values (mlx.core.array)

• n_sink (int)

• budget (int)

• window (int)

• beta (float)

• tau (float)

• kappa (float)

• return_indices (bool)

veloxquant_mlx.quantizers.nestedkv.nestedkv_append_decode_batched(prev_keys, prev_values, keys, values)

Batched-[BH,S,D] equivalent of nestedkv_append_decode().

Plain unscored concatenation along the token axis — inherently independent per row, so this is a single vectorized concat over all BH rows instead of a Python loop issuing BH separate concats. Also serves as the bootstrap path when prev_keys is None.

Return type:

tuple[mx.array, mx.array]

Parameters:

• prev_keys (mx.array | None)

• prev_values (mx.array | None)

• keys (mx.array)

• values (mx.array)