veloxquant_mlx.quantizers.keyformer module

Keyformer KV eviction primitives — Gumbel-regularized accumulating scorer.

Inspired by “Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference” (Adnan et al., MLSys 2024, arXiv:2403.09054). Documented as “Keyformer-adapted (VeloxQuant-MLX implementation)” — not a faithful port. See adaptation limitations below and in cache/keyformer_cache.py.

The paper’s finding: naively evicting by an accumulated attention score is unstable — a token that scores low early (before the tokens that will attend to it have arrived) is evicted and can never recover, even if it would have become a heavy hitter. Keyformer regularizes the eviction decision with Gumbel noise on the score logits: a temperature-controlled perturbation that keeps the retained set stochastically “soft” enough that borderline tokens are not deterministically pruned on the first low reading. As decoding proceeds the temperature is annealed toward 0, so late decisions are sharp.

WHERE THIS SITS IN THE REPO

This is the repo’s proxy-attention scorer family (SnapKV / H2O / TOVA / PyramidKV / SqueezeAttention / ChunkKV / CaM). Structurally it is the H2O pair (quantizers/h2o.py): an additive accumulation of proxy attention mass with a protected-sink top-budget eviction. The only new mechanism is the Gumbel-noise regularizer added to the accumulated logits before the keep/evict selection. Setting tau_init = tau_end = 0 removes the noise and recovers H2O-adapted’s deterministic behavior exactly — that is the honest ablation, exercised by the benchmark and a dedicated test.

FIXED, CONFIRMED-BY-PAPER-COMPARISON GAP #1 — no temperature annealing. Section 3.3.1 / Equation 10 of the paper is not a footnote: the annealed temperature schedule tau = tau_init + t * delta_tau (delta_tau = (tau_end - tau_init) / T) is the mechanism that makes Keyformer’s score function behave like a near-standard softmax during the prompt phase (tau approx 1, no tokens discarded yet, nothing to regularize against) and grows deliberately more randomized as more tokens are discarded during decoding (tau -> tau_end, paper default tau_end = 2) — see Section 3.3.1 and Appendix A.8. A single frozen tau for the entire run (the previous state of this module) cannot reproduce that: it is either always too sharp (no regularization when it is needed most, late in generation) or always too soft (unnecessary randomization during the prompt, when every token is still present and nothing has been discarded yet). Fixed by tracking tau_init/tau_end/anneal_steps and computing the current step’s temperature as tau_init + min(pos, anneal_steps) * delta_tau — see KeyformerState. tau_init == tau_end (the default) reproduces a constant temperature exactly, so this is backward compatible with every existing constant-tau caller and test.

FIXED, CONFIRMED-BY-PAPER-COMPARISON GAP #2 — no RoPE position tracking. This module never tracked which absolute position each kept key was rotated at, so an interior eviction (anything evicted that is not the newest token — now routine once sink/recent protection is combined with annealed noise picking a different eviction target than plain recency) silently desynced the survivors’ storage index from the absolute position baked into their rotation, corrupting every subsequent proxy-attention computation and the model’s own downstream attention math. This is the exact bug class H2O-adapted had and fixed (see quantizers/h2o.py’s module docstring, “RoPE position remapping”). Fixed the same way: every kept row’s true original position is now tracked in KeyformerState.positions, and keyformer_update re-rotates the shifted survivors via veloxquant_mlx.quantizers.a2ats_rope.rope_remap_positions() whenever an eviction changes the kept set’s positions.

FUSED METAL KERNEL for the over-budget branch (optional, mirrors H2O-adapted’s kernel): the per-token eviction step — append, Gumbel-regularized argmin, evict, RoPE-remap — can optionally be replaced with two fused Metal dispatches via veloxquant_mlx.metal.keyformer_fused_evict(), structurally identical to veloxquant_mlx.metal.h2o_fused_evict() with one addition: a per-row frozen Gumbel value is threaded through both dispatches and folded into the reduction as score + tau * gumbel. Equivalent to the MLX eviction branch it replaces (keys within 1 fp16 ULP; the kernel rotates in fp32) (see veloxquant_mlx/tests/metal/test_keyformer_evict.py). Used automatically when veloxquant_mlx.metal.metal_available() is true; falls back to the pure-MLX loop otherwise.

THE HONESTY CRUX (read before trusting any number)

  1. Proxy query. Like H2O/SnapKV-adapted, a cache never sees the true query vector, so the incoming KEY is used as a proxy query to estimate the attention each stored key receives. The paper accumulates the model’s real attention logits. This is a documented substitution, not the paper’s math.

  2. Frozen per-position noise, not redrawn sampling. The paper redraws Gumbel noise fresh each decoding step. A cache processes blocks with no global step counter it can trust to redraw against, so we draw ONE deterministic Gumbel value per token position (seeded from a fixed base seed + a per-head running position) and freeze it; the (now-annealed, see gap #1 above) temperature scales that frozen noise. This preserves the mechanism’s intent — borderline tokens are perturbed so a single low reading does not doom them — while staying reproducible and order-diagnosable. It is NOT the paper’s redraw-every-step sampling, and we do not claim it is.

  3. Nothing here is validated on a trained model. The regularizer’s benefit is measured only under constructed “late-riser” geometry in the benchmark, with a control where it shows no advantage.

Adaptation limitations (stated plainly):

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

  • Frozen deterministic per-position Gumbel, not redrawn each step (crux 2).

  • Uniform budget / n_sink / tau schedule across all heads.

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

Public API (mirrors quantizers/h2o.py)

KeyformerState — immutable per-head state dataclass init_keyformer_state — construct empty state (validates guards) keyformer_update — absorb S new tokens, evict if over budget, remap RoPE keyformer_get_kv — extract current (keys, values) arrays keyformer_fp16_bytes — bytes stored in current state full_keyformer_fp16_bytes — hypothetical cost without eviction

class veloxquant_mlx.quantizers.keyformer.KeyformerState(keys, values, scores, gumbel, positions, pos, n_sink, budget, recent, tau_init, tau_end, anneal_steps, rope_base, next_pos, seed)[source]

Bases: object

Per-head Keyformer eviction state.

Variables:

  • keys – [n_kept, D] fp16 stored key rows, RoPE’d at positions (contiguous 0..n_kept-1 after any eviction), or None before first update.

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

  • scores – [n_kept] cumulative proxy-attention mass (float32), or None.

  • gumbel – [n_kept] frozen per-position Gumbel noise (float32), or None. Drawn once when a token is inserted; never redrawn. Added (scaled by the current annealed temperature) to scores only for the keep/evict decision — the stored cumulative mass itself stays clean.

  • positions – [n_kept] int32 absolute positions each kept key is currently rotated at. Reassigned to a contiguous range whenever eviction changes which rows survive, and each key is re-rotated (via rope_remap_positions) to match. Same fix as H2O-adapted’s H2OState.positions — see module docstring gap #2.

  • pos – Running count of token positions this head has ever inserted; seeds the deterministic Gumbel draw AND drives the temperature-annealing schedule, so both are reproducible and independent of block boundaries.

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

  • tau_init – Gumbel-noise temperature (>= 0) at pos == 0. The paper’s default is 1 (Section 3.3.1): during the prompt phase, when no tokens have been discarded yet, this keeps the score function close to a plain softmax.

  • tau_end – Gumbel-noise temperature (>= tau_init) once pos >= anneal_steps. The paper’s default is 2: the more tokens discarded, the more randomized the score function becomes, per Equation 10. tau_end == tau_init reproduces a constant temperature (backward compatible with every existing constant-tau caller).

  • anneal_steps – Number of update steps (T in the paper) over which tau ramps linearly from tau_init to tau_end. Steps beyond this hold at tau_end. Ignored (no annealing) when tau_init == tau_end.

  • rope_base – RoPE frequency base used to de-rotate/re-rotate kept keys. Must match the model’s own attention module base.

  • next_pos – Absolute position the next incoming token will occupy — tracks true sequence position across calls, independent of how many rows have been evicted so far.

  • seed – Base seed for the deterministic per-position Gumbel draw.

Parameters:

  • keys (mx.array | None)

  • values (mx.array | None)

  • scores (mx.array | None)

  • gumbel (mx.array | None)

  • positions (mx.array | None)

  • pos (int)

  • n_sink (int)

  • budget (int)

  • recent (int)

  • tau_init (float)

  • tau_end (float)

  • anneal_steps (int)

  • rope_base (float)

  • next_pos (int)

  • seed (int)

keys: mx.array | None
values: mx.array | None
scores: mx.array | None
gumbel: mx.array | None
positions: mx.array | None
pos: int
n_sink: int
budget: int
recent: int
tau_init: float
tau_end: float
anneal_steps: int
rope_base: float
next_pos: int
seed: int
property tau: float

Backward-compatible alias for tau_init (pre-annealing field name).

veloxquant_mlx.quantizers.keyformer.init_keyformer_state(n_sink, budget, head_dim, recent=0, tau=None, tau_init=1.0, tau_end=1.0, anneal_steps=0, rope_base=10000.0, seed=0)[source]

Create an empty KeyformerState before any tokens arrive.

Parameters:

  • n_sink (int) – Number of leading sink positions to protect from eviction.

  • budget (int) – Maximum total tokens kept (sinks + non-sinks + recent).

  • head_dim (int) – Head dimension D (unused; API symmetry with init_h2o_state).

  • recent (int) – Trailing protected window (extension, off by default).

  • tau (float | None) – Deprecated alias — if given, sets both tau_init and tau_end to this constant value (no annealing), overriding tau_init/tau_end. Kept so existing constant-temperature callers need no changes.

  • tau_init (float) – Gumbel-noise temperature at pos == 0 (paper default 1).

  • tau_end (float) – Gumbel-noise temperature once annealing completes (paper default 2). Equal to tau_init means no annealing.

  • anneal_steps (int) – Update steps (T) over which tau ramps from tau_init to tau_end. 0 disables annealing (temperature is constant at tau_init) regardless of tau_end.

  • rope_base (float) – RoPE frequency base — must match the model’s own attention module, or position remapping after eviction will not cancel out the original rotation correctly.

  • seed (int) – Base seed for the deterministic per-position Gumbel draw.

Raises:

ValueError – if tau_init/tau_end is negative, or the protected positions (sinks + recent) leave no evictable room within budget.

Return type:

KeyformerState

veloxquant_mlx.quantizers.keyformer.keyformer_update(state, new_keys, new_values)[source]

Absorb S new tokens, evicting the lowest Gumbel-regularized token if over budget.

Return type:

KeyformerState

Parameters:

  • state (KeyformerState)

  • new_keys (mlx.core.array)

  • new_values (mlx.core.array)

For each of the S incoming tokens:

  1. Accumulate proxy-attention mass of the new key over all stored keys (H2O-adapted’s additive rule).

  2. Append the new token with cumulative score 0, its true absolute position (state.next_pos), and a frozen per-position Gumbel draw (seeded by the head’s running position).

  3. If over budget: evict the non-protected token with the lowest score + tau * gumbel, where tau is the current annealed temperature (Equation 10) — the Gumbel term is the Keyformer mechanism; at tau == 0 this is exactly H2O-adapted’s argmin on the raw score. Survivors positioned after the evicted token’s position shift down by one and are re-rotated (RoPE remap), same as H2O-adapted.

Uses the fused Metal kernel (_evict_via_metal()) when available, falling back to the pure-MLX reference (_evict_via_mlx()) otherwise — same policy as h2o_update.

Kept tokens are returned in original temporal order.

veloxquant_mlx.quantizers.keyformer.keyformer_update_batched(keys, values, scores, gumbel, positions, new_keys, new_values, n_sink, budget, recent, tau_init, tau_end, anneal_steps, rope_base, next_pos, pos, seeds, return_indices=False)[source]

Vectorized-over-BH equivalent of calling keyformer_update()

With return_indices=True, append chronological indices into the supplied prior + new rows (or the supplied candidate rows for eviction). once per (batch, head) pair with identical per-row config.

All BH rows share n_sink/budget/recent/the tau schedule/ rope_base/next_pos/pos (true for every real caller: KeyformerKVCache applies one uniform config to every head; only each head’s Gumbel noise stream differs, via a distinct seeds[hh] — see _gumbel_at_batched()), so the per-token score/append/evict/ RoPE-remap math — otherwise identical for every row — can run as one batched MLX call per step instead of BH separate Python-level calls into keyformer_update(). Same fix, same template, as veloxquant_mlx.quantizers.h2o.h2o_update_batched() (Keyformer is documented as the “Metal-fused sibling” of H2O-adapted — see this module’s docstring).

Numerically identical to the per-head loop it replaces: every op below is the same formula as keyformer_update()’s bootstrap/score-update/ evict branches, applied over a leading BH axis instead of a Python loop — verified bit-for-bit (fp32-rounding-only) equivalent.

Return type:

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

Returns:

(keys, values, scores, gumbel, positions, next_pos, pos) — the first five [BH, n_kept, D]/[BH, n_kept], next_pos the updated absolute step count and pos the updated per-head running position count (both shared across all rows, as in the loop this replaces).

Parameters:

  • keys (mx.array | None)

  • values (mx.array | None)

  • scores (mx.array | None)

  • gumbel (mx.array | None)

  • positions (mx.array | None)

  • new_keys (mx.array)

  • new_values (mx.array)

  • n_sink (int)

  • budget (int)

  • recent (int)

  • tau_init (float)

  • tau_end (float)

  • anneal_steps (int)

  • rope_base (float)

  • next_pos (int)

  • pos (int)

  • seeds (list[int])

  • return_indices (bool)

veloxquant_mlx.quantizers.keyformer.keyformer_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 (KeyformerState)

veloxquant_mlx.quantizers.keyformer.keyformer_fp16_bytes(state)[source]

Bytes currently stored for K + V in fp16.

Scores/gumbel are transient bookkeeping (float32, n each) — negligible beside K+V and, like H2O’s scores, not counted as cache payload.

Return type:

int

Parameters:

state (KeyformerState)

veloxquant_mlx.quantizers.keyformer.full_keyformer_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)