veloxquant_mlx.cache.snapkv_cache module

SnapKV-adapted KV cache — prefill observation-window token eviction.

Inspired by “SnapKV: LLM Knows What You are Looking for Before Generation” (Yuan et al., ICLR 2025, arXiv:2404.14469). Documented as “SnapKV-adapted (VeloxQuant-MLX implementation)” — not a faithful port.

Token eviction: during prefill (S > 1), the last snap_obs_window key rows act as proxy queries; their softmax attention over all prefix tokens produces a per-token importance score. Only the top-snap_budget tokens (plus the first snap_n_sink sink positions) are retained as fp16. All subsequent decode tokens (S == 1) are always appended — never evicted.

mlx_lm’s chunked prefill calls update_and_fetch once per prefill_step_size chunk for prompts longer than one chunk. Every S > 1 call re-enforces the budget against the accumulated kept set (prior kept tokens concatenated with the new chunk), not just that chunk in isolation — otherwise the retained count would grow by up to snap_budget per chunk instead of staying capped (see #84).

This is the repo’s first eviction method. Every other method compresses all tokens to fewer bits; SnapKV-adapted stores fewer tokens at full fp16 precision. The two axes compose: wrap a quantizer cache around the selected subset for combined eviction + compression.

Adaptation limitations (stated plainly):

• Key-as-query proxy: obs-window uses key vectors, not true prompt query vectors (not visible at update_and_fetch).

• No max-pool smoothing (paper’s kernel_size > 1 not implemented).

• Uniform snap_budget across all heads.

Byte accounting:

evicted_key_bytes / evicted_value_bytes — fp16 bytes for the kept subset full_key_bytes / full_value_bytes — hypothetical cost without eviction eviction_ratio — full_fp16 / kept_fp16 (> 1 = savings) tokens_kept / tokens_total — diagnostic token counters

class veloxquant_mlx.cache.snapkv_cache.SnapKVKVCache(config)

Bases: KVCache

KV cache implementing SnapKV-adapted prefill eviction for one layer.

Parameters:

config (Any) – KVCacheConfig. Fields consumed: snap_budget (int, default 512) — max tokens retained after prefill, snap_obs_window (int, default 32) — trailing keys used as proxy queries, snap_n_sink (int, default 4) — initial positions always kept.

Notes

No .bits attribute — stores and returns fp16 K/V directly. Single-layer (no coordinator); for_model propagates all snap_* fields automatically via dataclasses.replace. Eviction happens at every prefill chunk (S > 1) — mlx_lm’s chunked prefill (prefill_step_size) calls update_and_fetch once per chunk for prompts longer than one chunk, so budget enforcement must re-run over the accumulated kept set, not just each chunk in isolation (see #84). Decode tokens (S == 1) are always kept. Accumulated decode tokens grow the cache beyond the initial budget in the decode phase — consistent with the paper’s design (the budget constrains prefill history, not the decode stream).

property offset: int

True absolute token position (NOT the retained row count).

While the base class’s update_and_fetch is on the stack this yields the retained row count instead, so its cursor arithmetic and return slice stay correct.

update_and_fetch(keys, values)

Prefill: score and evict down to the retention budget (SnapKV window-attention proxy), deferred to after this step’s own attention. Decode: pass through unevicted.

Parameters:

• keys (mlx.core.array)

• values (mlx.core.array)

make_mask(N, return_array=False, window_size=None, **_)

Explicit position-based causal mask — see VeloxQuant-MLX#370.

The inherited KVCache.make_mask returns the string "causal" whenever offset == 0 (the very first prefill call, before any eviction has run) or otherwise builds a mask keyed off row index rather than true position — both wrong once eviction has made the kept keys a non-contiguous subset. Before the first prefill call (_kept_positions is still None, nothing evicted yet) this falls back to the base class’s behavior exactly, since a plain trailing-window mask is correct there.

Called before this step’s update_and_fetch, whose attention runs over the stored rows plus this step’s N new ones (eviction is deferred past the return), so the key axis is len(kept) + N.

Parameters:

• N (int)

• return_array (bool)

• window_size (int | None)

property evicted_key_bytes: int

Bytes stored for kept key rows (fp16).

property evicted_value_bytes: int

Bytes stored for kept value rows (fp16).

property full_key_bytes: int

Hypothetical fp16 key cost without any eviction.

property full_value_bytes: int

Hypothetical fp16 value cost without any eviction.

property tokens_kept: int

Total token positions retained across all heads and steps.

property tokens_total: int

Total token positions seen (before eviction) across all heads and steps.

property eviction_ratio: float

full_fp16_bytes / kept_fp16_bytes; > 1 means storage savings.

property keep_rate: float

Fraction of tokens retained (tokens_kept / tokens_total).

is_trimmable()

False: trim() reads a corrupted self.offset and both writes and returns the wrong thing, worse than a mere bookkeeping rollback.

The base KVCache.trim(n) does n = min(self.offset, n); self.offset -= n. Outside update_and_fetch (_in_base is False), this class’s offset property returns _true_offset — the absolute token position used for RoPE — not the retained row count (_row_offset) the base class’s buffer actually holds. So n is clamped against the wrong (larger) number, and the setter then writes _row_offset -= n, which can leave _row_offset larger than the number of rows ever written whenever eviction has dropped any tokens (the normal case once prefill has run). The next update_and_fetch then returns a slice reaching past the real data into the pre-allocated buffer’s stale/uninitialized rows, silently feeding garbage into attention rather than crashing. Reproduced directly: budget=4, 20 real prefill tokens (so _row_offset == 4, _true_offset == 20); trim(3) leaves _row_offset == 17 though only 4 rows are real, and the following decode step returns an 18-row slice (see VeloxQuant-Studio issue #27).

Return type:

bool

property merge