veloxquant_mlx.cache.rocketkv_cache module

RocketKV-adapted KV cache — two-stage compression (SnapKV eviction + HSA).

Inspired by “RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression” (Behnam, Fu, Zhao, Tsai, Yu, Tumanov; NVIDIA/Georgia Tech; ICML 2025, arXiv:2502.14051). Documented as “RocketKV-adapted (VeloxQuant-MLX implementation)” — not a faithful port. See veloxquant_mlx/quantizers/rocketkv.py for the full list of adaptation decisions and paper/research/surveys/NEW_METHOD_SURVEY_V23.md for the write-up.

Design (paper §3.2-3.6):

Prefill (first call, S > 1):

1. Derive the stage-1/stage-2 split from rocketkv_compression_ratio via split_compression_ratio().

2. Run stage-1 coarse-grain PERMANENT eviction — directly reuses snapkv_compress() (the paper adopts SnapKV verbatim for this stage) — down to a stage-1 token budget derived from that split.

3. Build a paged max/min key summary over the surviving subset (stage 2’s index).

Decode (subsequent calls, S == 1 per step):

1. Append the new key’s contribution to the paged summary incrementally (append_paged_summary_batched()).

2. Using the incoming key as a proxy query (same convention as SnapKV-adapted / A2ATS-adapted / every other query-aware method in this repo — the cache wrapper never sees the true query), run HSA: top-k1 head-dim channels, per-page approximate scores, top-k2 page selection.

3. Gather the exact K/V rows for the selected pages plus the new decode token itself (always kept), and expose that subset as the cache’s return value for this step — the eviction axis (stage 1) determines what is ever STORED; the selection axis (stage 2) determines what is FETCHED per step. Because mlx_lm’s cache protocol has no notion of “fetch a subset, store everything,” this wrapper stores only the stage-1 survivors (bounded, honest storage) and, at decode, narrows the returned K/V to the HSA selection — the base class’s buffer keeps growing with new decode tokens (never evicted, matching RocketKV’s paper: stage 1 only touches the input prompt).

Byte accounting:

stage1_bytes / stage2_aux_bytes — storage after eviction (fp16 kept

tokens) and HSA’s paged max/min auxiliary summary, respectively

full_fp16_bytes — hypothetical cost without any compression compression_ratio — full_fp16_bytes / (stage1_bytes + stage2_aux_bytes) tokens_kept / tokens_total — diagnostic token counters (stage 1 only —

stage 2 never drops a stored token, only narrows what a given decode step attends to)

Limitations (stated plainly):

• Key-as-query proxy at both stage-1 eviction and stage-2 HSA selection — inherited from SnapKV-adapted, not a new approximation.

• No fused kernel: HSA gather/attend happens in eager MLX ops each step.

• No RocketKV-MT (multi-turn) variant — see issue #239.

Performance: _process_prefill’s single-shot path originally looped for b in range(B): for h in range(H):, calling snapkv_compress once per head — the same unbatched-per-head-loop cost pattern that dominated ChunkKV/CurDKV/GEAR before their own fixes (#525/#526/#527). Here it reused _snapkv_compress_batched, the batched primitive SnapKVKVCache already relies on in production (stage 1 is SnapKV verbatim, per this module’s own docstring — this was reinventing its per-head loop rather than reusing the batched call). Verified bit-identical to the old loop. Isolated stage-1 benchmark (B=1, H=32, S=1024, D=128, budget=128): ~4.53ms (old per-head loop) -> ~1.50ms (batched), ~3x faster and markedly more stable run-to-run.

_process_decode had the same unbatched-per-(b,h)-loop pattern, and was the hotter of the two (runs once per generated token, not once per prefill): it called append_paged_summary() once per head, itself an inner Python loop over pages. Batched via append_paged_summary_batched(), legal here because every (b, h) row’s paged summary provably shares the same n_tokens/page count at every point in _process_decode’s call pattern (every row is seeded with the same n_kept at prefill via _snapkv_compress_batched’s one shared budget, and every decode step appends the same S new tokens to every row — verified by grepping every write site of the paged-summary state, not assumed). Verified bit-for-bit identical to the old per-(b,h) loop against 1536 randomized single-step trials plus 18 five-step decode sequences (varying page_size, prior token count, new-token count, head_dim, batch size). Isolated decode-step benchmark (D=128, page_size=8, n_kept=256, 50 steps): H=8: ~2.58ms -> ~0.33ms/step (~7.9x); H=32: ~11.54ms -> ~0.41ms/step (~28x) — speedup grows with B*H, as expected for removing B*H separate Python-dispatched op graphs (each itself looping over pages).

The chunked-prefill re-run path remains an explicit Python loop — unlike _process_decode, per-head kept-token counts here are REDERIVED per call via each head’s own _stage1_head (SnapKV) call, so this path was verified (not merely assumed) to keep every row’s n_kept — and hence page count — in sync via the shared budget/S formula, but the loop itself was left unbatched since only chunked prompts hit it (rare, unlike every-decode-step). See append_paged_summary_batched()’s precondition note for why it isn’t reused there.

class veloxquant_mlx.cache.rocketkv_cache.RocketKVKVCache(config)

Bases: DeferredEvictionMixin, KVCache

KV cache implementing RocketKV-adapted two-stage compression for one layer.

Parameters:

config (Any) –

KVCacheConfig. Fields consumed: rocketkv_compression_ratio (float, default 8.0) — overall

target compression ratio c; drives the adaptive stage-1/ stage-2 split (paper §3.6). Must be > 1.

rocketkv_page_size (Optional[int], default None) — HSA page

size; None derives it from the adaptive split instead of a fixed value.

rocketkv_head_topk1 (Optional[int], default None) — head-dim

channels kept per HSA step; None derives it from the adaptive split (head_dim / head_dim_ratio).

rocketkv_obs_window (int, default 32) — stage-1 SnapKV

observation window.

rocketkv_n_sink (int, default 4) — stage-1 SnapKV sink

tokens always kept.

Notes

No .bits attribute — stores and returns fp16 K/V directly. Single-layer (no coordinator); for_model propagates all rocketkv_* fields automatically via dataclasses.replace. Stage 1 (eviction) runs once at prefill, exactly like SnapKVKVCache. Stage 2 (HSA) runs every decode step and narrows the returned K/V subset without shrinking the stored buffer — decode tokens accumulate in storage (paper’s stage 1 only compresses the input prompt), but each attention call only touches the HSA- selected pages plus the running decode tail.

property offset: int

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

See SnapKVKVCache.offset for why the two must not be conflated — stage 1 here drops tokens exactly the same way.

property state

Expose retained rows, excluding the base cache’s capacity padding.

update_and_fetch(keys, values)

Prefill: run stage-1 SnapKV eviction per head and build paged HSA summaries. Decode: append tokens exactly and update summaries incrementally.

Parameters:

• keys (mlx.core.array)

• values (mlx.core.array)

select_indices(query, b, h, keep_recent=0)

HSA stage-2: approximate top-k page selection for one (b, h) query.

Exposed for callers that want to actually exercise sparse attention with this cache (the base update_and_fetch return value stays dense — see the module docstring’s storage-vs-fetch distinction).

Parameters:

• query (array) – [D] proxy query vector (key-as-query convention).

• b (int) – Batch and head index into the stored paged summaries.

• h (int) – Batch and head index into the stored paged summaries.

• keep_recent (int) – Trailing token count always included regardless of HSA score (e.g. the newest decode tokens).

Return type:

array

Returns:

[n_selected] int32 token indices, ascending, deduplicated, covering the HSA-selected pages and the trailing keep_recent window.

is_trimmable()

False: trim() would only roll back base-class offset bookkeeping, not the internal stage-1 kept-token accounting or stage-2 paged max/min summaries, silently desynchronizing both from the trimmed buffer on the next call.

Return type:

bool

property stage1_bytes: int

Bytes stored for kept fp16 K + V rows after stage-1 eviction.

property stage2_aux_bytes: int

Auxiliary bytes for HSA’s paged max/min summaries (fp16-packed).

property full_fp16_bytes_total: int

Hypothetical fp16 K + V cost without any compression.

property compression_ratio: float

full_fp16_bytes / (stage1_bytes + stage2_aux_bytes); > 1 means savings.

property tokens_kept: int

Total tokens retained after stage-1 eviction (plus all decode tokens, which are always kept).

property tokens_total: int

Total tokens seen across prefill and decode, before stage-1 eviction.

property keep_rate: float

Fraction of tokens retained after stage-1 eviction (tokens_kept / tokens_total).

property stage1_ratio: float

Configured stage-1 (SnapKV) eviction ratio — prefill budget is S / stage1_ratio.

property page_size: int

Configured HSA page size used to build the paged max/min summaries.

property head_topk1: int

Resolved stage-2 per-head top-k1 count used by HSA approximate scoring.