veloxquant_mlx.cache.a2ats_cache module

A2ATS-adapted KV cache — windowed RoPE + query-aware retrieval VQ.

Inspired by “A2ATS: Retrieval-Based KV Cache Reduction via Windowed Rotary Position Embedding and Query-Aware Vector Quantization” (He, Xing, Wang, Xu, Wu, Zhou, Liu, Xue, Li — ACL 2025 Findings, aclanthology.org/2025.findings-acl.644). Documented as “A2ATS-adapted (VeloxQuant-MLX implementation)” — not a faithful port.

No query visible at cache level. Like every other query-aware method in this repo (AMC-adapted’s amc_use_query_saliency, H2O-adapted’s key-as-query proxy, SnapKV-adapted’s prefill window), a cache wrapper’s update_and_fetch only ever receives keys and values — the true decode-time query vector is not part of the mlx_lm cache protocol. This port substitutes the incoming key vector itself as a proxy query for both the retrieval-set selection (a2ats_select_retrieval_set()) and the query-aware codebook assignment (a2ats_query_aware_assignment()) — the same category of approximation, not a new one.

Retrieval set gets preferential codebook assignment, not eviction. Every token is quantized and retained; the retrieval-fraction split only changes which centroid a token is matched against (query-aware vs. plain nearest-centroid for the bulk). No token is ever dropped — a compression-only method, not an eviction method (same family framing as AMC-adapted).

Incoming keys are de-rotated before quantization. mlx_lm always rotates keys with self.rope(k, offset=cache.offset) inside the model’s attention module before calling update_and_fetch — this cache never receives truly pre-RoPE keys. update_and_fetch undoes that rotation via rope_remap_positions() before quantizing, so the rest of this cache’s logic (below) can assume a pre-RoPE frame as documented (:issue:`613`).

Windowed RoPE is applied at fetch time, every step. The parent buffer holds the pre-RoPE reconstruction; update_and_fetch re-applies windowed RoPE to the whole accumulated cache against the current decode position on each call, so a token’s near/far class tracks the advancing query as Eq. (11) requires. Baking the rotation in at write time would freeze that class at the moment the token was written (:issue:`29`, finding 3).

Far keys are returned unrotated (Eq. 12, k̃_i = k_i) — that is what makes a shared codebook viable across inputs (§3.1). The constant R_b encoding “far” relative position belongs on the query side (Eq. 11); since this cache never sees the query, it is exposed via A2ATSKVCache.far_query_rope() for callers that do.

Offline codebook calibration required — same footgun class as VecInfer-adapted/CommVQ-adapted/Palu-adapted/SVDq-adapted/AMC-adapted: using a2ats with the default random-init codebook (present only so wiring/shape tests don’t require a calibration pass) degrades to near-random quantization. Supply a calibrated a2ats_codebook via config for real use.

No CUDA kernel fusion reproduced — same MLX/Metal disclaimer as every other VQ-family method in this repo: the benefit on Apple Silicon is memory footprint, not throughput.

Byte accounting:

compressed_key_bytes / compressed_value_bytes — actual stored bytes fp16_key_bytes / fp16_value_bytes — hypothetical full-rank cost compression_ratio — fp16 / compressed (K + V)

class veloxquant_mlx.cache.a2ats_cache.A2ATSKVCache(config)

Bases: KVCache

KV cache implementing A2ATS-adapted windowed RoPE + query-aware VQ.

Parameters:

config (Any) –

KVCacheConfig. Fields consumed:

head_dim (int, required) — must be even (RoPE requirement)

and divisible by a2ats_sub_dim.

a2ats_codebook_bits (int, default 8) — codebook size 2^bits. a2ats_rope_scaling (dict | None) — the model’s rope_scaling

config (model.args.rope_scaling). Required for Llama 3.x: without it the cache de-rotates keys with plain RoPE, which does not invert Llama 3’s scaled rotation. Unsupported scaling types raise rather than silently mis-rotating.

a2ats_sub_dim (int, default 8) — VQ sub-vector width. a2ats_window (int, default 128) — trailing exact-RoPE

window w.

a2ats_b (int, default 2048) — constant far-token relative

position b (Eq. 11), independent of w. The paper’s §5.1 uses w=64, b=2048. Consumed by far_query_rope(), which callers apply to the query.

a2ats_use_query_aware (bool, default True) — paper’s primary

reported path; off degrades to plain nearest-centroid VQ.

a2ats_beta (float, default 0.5) — query/reconstruction blend,

must be in [0, 1].

a2ats_retrieval_fraction (float, default 0.20) — fraction of

tokens routed to the query-aware assignment path, must be in [0, 1].

a2ats_rope_base (float, default 10000.0). a2ats_codebook (mx.array | np.ndarray | None) — pre-trained

codebook; random-init fallback if absent (wiring/tests only, see module docstring’s calibration warning).

Notes

Never exposes .bits — keeps mlx_lm SDPA on the clean fp16 path, same convention as every other VQ-family cache in this repo.

update_and_fetch(keys, values)

Quantize, store pre-RoPE, and return keys rotated against the current decode position.

mlx_lm rotates keys at offset=cache.offset inside the model’s attention module before this method is ever called (e.g. mlx_lm/models/llama.py:90-92), so the incoming keys arrive already post-RoPE. They are de-rotated back to the pre-RoPE frame this cache’s windowed-RoPE logic requires before quantization, and the windowed rotation is re-applied on return — see rope_remap_positions() and :issue:`613`.

The parent mlx_lm.models.cache.KVCache concatenates whatever it is handed and never revisits it. Writing rotated keys into it therefore freezes each token’s near/far class at the moment it was written — a token classified “near” during prefill keeps its exact rotation forever, even thousands of steps later, which silently defeats the distance gating the method is named for (:issue:`29`, finding 3).

So the parent buffer holds the pre-RoPE reconstruction, and windowed RoPE is re-applied to the full accumulated slice on every call, against query_position = self._next_position - 1. Tokens cross from “near” to “far” as the decode position advances, matching Eq. (11), where i is the position of the decoding query.

Cost: the rotation is O(total_tokens) per step rather than O(S). That is inherent to honest distance gating under this protocol — the near/far split is a function of the current query, so it cannot be precomputed. Far tokens are returned unrotated (Eq. 12), so only the window-sized near slice does nontrivial work.

Parameters:

• keys (mlx.core.array)

• values (mlx.core.array)

trim(n)

Roll back n trailing tokens — used by speculative decoding (rejected draft tokens) and prefix-cache reuse.

The inherited base mlx_lm.models.cache.KVCache.trim() only rewinds self.offset; it does not touch this cache’s own _next_position counter, which drives both the windowed-RoPE query position (update_and_fetch()) and the stored absolute position passed to _quantize_head. Left stale, the next update_and_fetch call computes the near/far split against a query position that is n steps ahead of where the trimmed cache actually is, and writes new tokens at the wrong absolute position — confirmed via a drift of up to 5.1 in the returned keys (:issue:`614`).

Byte-accounting counters are rolled back by the same per-token formula _account_bytes() used to add them, so a trim+rewrite reports identical bytes/ratios to a cache that never saw the trimmed tokens.

Return type:

int

Parameters:

n (int)

far_query_rope(query)

Apply the paper’s constant far-token rotation R_b to a query.

WRoPE splits into two halves: far keys stay unrotated (Eq. 12, handled in update_and_fetch()), and the constant R_b that encodes “far” relative position rides on the query instead (Eq. 11, u_ij = q_i R_b k_j^T). This method is the query-side half, using this cache’s configured a2ats_b.

It is exposed rather than applied internally because update_and_fetch never receives the query — the mlx_lm cache protocol passes only keys and values (see the module docstring). A caller with real query access applies this before scoring far tokens; the proxy-query paths inside this cache cannot.

Return type:

array

Parameters:

query (mlx.core.array)

property compressed_key_bytes: int

Realized stored bytes for the compressed key cache (retrieval-aware VQ codes, all heads/batches; excludes amortized codebook).

property fp16_key_bytes: int

Hypothetical fp16 key cost if nothing were compressed.

property compressed_value_bytes: int

Realized stored bytes for the compressed value cache (retrieval-aware VQ codes, all heads/batches; excludes amortized codebook).

property fp16_value_bytes: int

Hypothetical fp16 value cost if nothing were compressed.

property codebook_bytes: int

Static codebook overhead in bytes (fp16 storage), amortized across all tokens.

property compression_ratio: float

fp16_bytes / compressed_bytes (K + V combined); > 1 means memory savings over fp16.

property assigned_avg_bits: float

Effective bits/element from the sub-vector codebook, scaled by sub_dim/head_dim.

property tokens_seen: int

Total tokens processed through update_and_fetch so far.

property tokens_retrieved: int

Total tokens assigned to the query-aware retrieval-set codebook path rather than plain nearest-centroid.

property merge