veloxquant_mlx.quantizers.a2ats module

A2ATS-adapted query-aware VQ assignment + retrieval-set selection.

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.

Retrieval-selection signal is approximated. The paper selects its high-fidelity “retrieval set” using real attention-relevance signals that are not observable inside a cache wrapper (the same category of gap as every other saliency/attention proxy method in this repo: H2O-adapted’s key-as-query proxy, SnapKV-adapted’s prefill window, ZipCache-adapted’s key-norm proxy). This port substitutes a query-similarity proxy: the top retrieval_fraction of tokens by cosine similarity between their key vector and the current query are treated as the retrieval set (kept at a finer codebook / lower compression than the remainder). Every token is still quantized — nothing is evicted; the retrieval set only gets preferential codebook assignment (see a2ats_query_aware_assignment()). Whether the paper’s own retrieval step additionally drops non-retrieved tokens outright is a question for the cache wrapper layer, not this module — this module exposes the split, the cache wrapper decides what to do with each half.

Query-aware codebook assignment comes in two forms here:

1. a2ats_h_weighted_assignment() — the paper’s actual objective (Eq. 13/14): minimize (k̃ − c) H (k̃ − c)ᵀ where H = E[q̃ᵀq̃] is the query second-moment matrix. Computed exactly via the Eq. (15)-(18) Cholesky identity, which turns the H-weighted form into a plain Euclidean distance in the space z = k̃L. Requires a calibrated H (see a2ats_query_second_moment()), so it is available on paths with offline calibration.

2. a2ats_query_aware_assignment() — a cosine blend between reconstruction error and query-centroid alignment, reusing the shape of veloxquant_mlx.quantizers.amc.amc_query_aware_saliency(). This is a substitute for Eq. (14), not an implementation of it: the cosine term is a constant per-centroid bias rather than a per-token coupling. Retained for the decode path, where only a single proxy query is available and no calibrated H exists. See its docstring for the precise differences (:issue:`29`, finding 4).

Public API:

a2ats_query_second_moment — calibrate H = E[qᵀq] (Eq. 10) a2ats_cholesky_factor — L with H = L Lᵀ (Eq. 15) a2ats_h_weighted_assignment — the paper’s Eq. (14) assignment a2ats_query_aware_assignment — cosine-blend substitute (decode path) a2ats_select_retrieval_set — split token indices into retrieval / bulk sets

veloxquant_mlx.quantizers.a2ats.a2ats_query_second_moment(queries, ridge=0.0001)

Second-moment matrix H = E[q^T q] of a set of query states (Eq. 10).

This is the object the paper’s §3.2 is about: the query-aware and plain-VQ objectives coincide only if ``H ∝ I``, and Figure 2 shows it is not — that non-diagonal structure is the paper’s contribution.

Parameters:

• queries (array) – [M, d] post-PE query states, collected offline on a representative dataset (paper §4.2’s calibration step).

• ridge (float) – Added to the diagonal so H is positive definite rather than merely semi-definite, which Cholesky (Eq. 15) requires. Needed whenever M < d or the queries are rank-deficient.

Return type:

array

Returns:

[d, d] float32 positive-definite second-moment matrix.

veloxquant_mlx.quantizers.a2ats.a2ats_cholesky_factor(h)

Lower-triangular Cholesky factor L with H = L Lᵀ (Eq. 15).

Eq. (16)-(18) use L to map the H-weighted objective into a plain Euclidean one on z = k̃L, which is what lets conventional k-means++ build the codebook. Implemented directly (no mx.linalg.cholesky dependency) for the small d values this repo’s sub-vector VQ uses.

Parameters:

h (array) – [d, d] symmetric positive-definite matrix.

Return type:

array

Returns:

[d, d] float32 lower-triangular L.

veloxquant_mlx.quantizers.a2ats.a2ats_h_weighted_assignment(x, codebook, cholesky_factor)

The paper’s query-aware assignment, Eq. (14)/(17).

f'(k̃; C) = argmin_j (k̃ − c_j) H (k̃ − c_j)ᵀ

Implemented via the Eq. (17) identity — with H = L Lᵀ, the H-weighted quadratic form is a plain squared Euclidean distance in the transformed space z = k̃L:

(k̃ − c_j) H (k̃ − c_j)ᵀ = (k̃L − c_jL)(k̃L − c_jL)ᵀ

So this is exact, not an approximation: it computes the paper’s objective, just in the coordinates where it is cheap.

Unlike the cosine blend in a2ats_query_aware_assignment(), the H term couples k̃ and the query distribution per token — it is not a constant per-centroid bias added to every row.

Parameters:

• x (array) – [N, d] sub-vectors to quantize.

• codebook (array) – [K, d] centroids, in the original (untransformed) space.

• cholesky_factor (array) – [d, d] L from a2ats_cholesky_factor().

Return type:

array

Returns:

[N] int32 assigned centroid indices.

veloxquant_mlx.quantizers.a2ats.a2ats_query_aware_assignment(x, codebook, query, beta=0.5)

Cosine-blend query-aware assignment — a substitute for Eq. (14), not an implementation of it.

score(c) = beta * (-||x - c||^2 normalized) + (1 - beta) * cosine_similarity(query, c)

Picks, for each row of x, the centroid maximizing this blended score rather than the plain nearest centroid by reconstruction error alone. beta=1.0 reduces to plain nearest-centroid VQ (equivalent to veloxquant_mlx.allocators.vecinfer.quantize_vq()); beta=0.0 is pure query-direction alignment, ignoring reconstruction quality entirely.

Warning

This is not the paper’s estimator. Paper Eq. (13)/(14) minimizes an H-weighted quadratic form, where H = E[q̃ᵀq̃] is the query second-moment matrix; §3.2’s entire argument is that H ∦ I. Two concrete differences:

• cos_sim here depends only on (codebook, query), so it adds an identical per-centroid bias to every row of ``x``. Eq. (14) couples k̃ and H per token.

• err_term is normalized per-row by max_err, making beta scale-dependent: the same beta means different things for different key magnitudes. H-weighting has no such sensitivity.

Use a2ats_h_weighted_assignment() for the paper’s actual objective. This function is retained for the single-query decode path, where no calibrated H is available (see :issue:`29`, finding 4).

Parameters:

• x (array) – [N, d] sub-vectors to quantize (d == codebook.shape[-1]).

• codebook (array) – [K, d] candidate centroids.

• query (array) – [d] current query sub-vector (same sub-dim slice as x and codebook).

• beta (float) – Blend coefficient in [0, 1].

Return type:

array

Returns:

[N] int32 assigned centroid indices.

veloxquant_mlx.quantizers.a2ats.a2ats_select_retrieval_set(keys, query, retrieval_fraction=0.2)

Split token indices into a high-fidelity retrieval set and the bulk.

Uses veloxquant_mlx.dsa.heap.MaxHeap to select the top ceil(retrieval_fraction * N) tokens by query-cosine-similarity in better-than-full-sort fashion, consistent with this repo’s DSA-first convention (the same heap-based top-k pattern used by veloxquant_mlx.quantizers.amc.amc_assign_tiers()).

Parameters:

• keys (array) – [N, D] key vectors.

• query (array) – [D] current query vector.

• retrieval_fraction (float) – Fraction of tokens routed to the retrieval set, in [0, 1].

Return type:

tuple[array, array]

Returns:

(retrieval_idx, bulk_idx) — int32 index arrays, disjoint, covering every position 0..N-1 exactly once.