points. The budget-responsive curve confirms the offset fix end-to-end —
pre-174 position drift grows without bound with sequence length and would
flatten every arm regardless of budget. Documented in
`docs-site/docs/algorithms/knorm.md`, replacing its "no model-level
benchmark has been run" note. **Scope**: one model, two budgets — larger
models and a wider sweep remain open under
[181](https://github.com/rajveer43/VeloxQuant-MLX/issues/181) /
[180](https://github.com/rajveer43/VeloxQuant-MLX/issues/180); L2Norm
TTFT/throughput is still unmeasured. Docs only, no code changes.
**Reviewed StreamingLLM-adapted's RoPE position semantics against the paper**
([189](https://github.com/rajveer43/VeloxQuant-MLX/issues/189)) — the
implementation already preserved original absolute token positions after
eviction rather than the paper's cache-slot renumbering (a deliberate
divergence noted briefly in the docs since the 171 offset fix), but this
had never been reviewed against the paper explicitly, and unlike every
other eviction cache in the library, `StreamingLLMKVCache` had **no
regression test** guarding the `_true_offset`/`171` fix. Added three
offset-tracking tests to `test_streaming_llm_cache.py` (mirroring the
KNorm/SnapKV/TOVA/Q-Filters coverage: true position through sustained
eviction, block-size advance on prefill, and correctness across a
prefill-then-decode mix), and expanded `docs-site/docs/algorithms/streaming_llm.md`
with a full review section explaining why absolute-position preservation
was kept over the paper's renumbering (RoPE relativity makes it exact;
consistency with every sibling cache; avoids a per-step re-rotation cost;
the cache-wrapper boundary can't implement the paper's scheme cleanly
anyway). **Not measured**: whether the paper's renumbering scheme would
produce different generation quality — that needs real model inference
and remains open, same as [187](https://github.com/rajveer43/VeloxQuant-MLX/issues/187).
**Reviewed SnapKV-adapted's RoPE position semantics against the paper**
([188](https://github.com/rajveer43/VeloxQuant-MLX/issues/188)) — unlike
StreamingLLM, the SnapKV paper doesn't define a position-renumbering scheme
at all (it compresses the cache once, at the end of prefill, before
generation starts); common HF-`DynamicCache`-based reference
implementations instead fall back to the retained-row-count convention this
repo's `171` fix moved away from. Reviewed why absolute-position
preservation isn't a stylistic choice here but the behavior mathematically
required by how the eviction mechanism works: SnapKV selects from K rows
the model already rotated with RoPE during its own prefill forward pass, so
survivors' rotations are permanently baked in before eviction happens —
reporting anything but their true position would break RoPE's
relative-distance identity. Unlike StreamingLLM (continuous per-step
eviction), a hypothetical renumbering scheme here would only cost a
one-time `O(budget)` re-rotation right after prefill compression — still not
worth implementing, since exact positions already require none. Test
coverage already existed
(`test_offset_tracks_true_position_not_retained_rows` plus the
chunked-prefill offset assertions) — no code or test changes needed, this
was purely a documentation review. Added to
`docs-site/docs/algorithms/snapkv.md`.
Fixed
**RoPE positions after eviction in TOVA-adapted**
([175](https://github.com/rajveer43/VeloxQuant-MLX/issues/175)) — audited
H2O-adapted and TOVA-adapted for the same class of defect fixed for
Q-Filters, L2Norm, SnapKV, and StreamingLLM. Findings:
- **H2O-adapted already fixed** (shipped in v0.44.4, predating this audit):
`self.offset` already tracks the true absolute step count directly, and
because H2O renumbers positions to close gaps when an interior eviction
happens (e.g. with `h2o_n_sink > 0`), it additionally re-rotates the
shifted survivors — a stronger fix than the other caches need, since they
never renumber. No change required.
- **TOVA-adapted carried the defect**: `self.offset` reported the retained
row count rather than the true absolute token position once eviction
pinned the kept set at `tova_budget`. `tova_update` drops exactly the
evicted row and keeps the rest in temporal order (never renumbers), so
the same `_true_offset` counter that sufficed for Q-Filters/L2Norm applies
directly — no `offset` property split like SnapKV's was needed, since
`TOVAKVCache.update_and_fetch` fully resets `self.keys`/`self.values`/
`self.offset` on every call, prefill and decode alike.
Regression tests added to `test_tova_cache.py` mirroring the Q-Filters/L2Norm
coverage: offset tracks true position through sustained eviction, advances
by block size on prefill, and stays correct across a prefill-then-decode mix.
**RoPE positions after eviction in L2Norm-adapted**
([174](https://github.com/rajveer43/VeloxQuant-MLX/issues/174)) — carried
the same defect fixed earlier for Q-Filters, SnapKV, and StreamingLLM:
`self.offset` reported the **retained row count** rather than the true
absolute token position once eviction pinned the kept set at
`knorm_budget`. `mlx_lm` rotates both the query and the incoming key at
`offset=cache.offset` *before* `update_and_fetch` runs, so subsequent
tokens were rotated at a stale, non-advancing position — the offset drift
that excluded L2Norm as a fair baseline in Q-Filters benchmarks.
L2Norm's `knorm_update` already restores kept rows to temporal order after
top-k selection (never renumbers survivors), so the same fix that sufficed
for Q-Filters applies directly: a `_true_offset` counter incremented by the
incoming block size `S`, reported as `self.offset` after every
`update_and_fetch` call. Unlike SnapKV, no `offset` property split was
needed — `L2NormKVCache.update_and_fetch` fully resets
`self.keys`/`self.values`/`self.offset` on *every* call (prefill and decode
alike), so the base class's cursor arithmetic never observes the true
position as a stale row count between calls.
Regression tests added mirroring the Q-Filters/SnapKV coverage: offset
tracks true position through sustained eviction, advances by block size
(not retained rows) on prefill, and stays correct across a prefill-then-decode
mix. L2Norm can now be re-enabled as a fair comparison arm in Q-Filters
benchmarks.
**RoPE positions after eviction in SnapKV-adapted and StreamingLLM-adapted**
([171](https://github.com/rajveer43/VeloxQuant-MLX/issues/171)) — both
carried the same defect fixed earlier for Q-Filters: `self.offset` reported
the **retained row count** rather than the true absolute token position.
`mlx_lm` rotates both the query and the incoming key at
`offset=cache.offset` *before* `update_and_fetch` runs, so once eviction
started every subsequent token was rotated at the wrong position. Measured
over 200 decode steps after a 64-token prefill at budget 32:
- **StreamingLLM** — offset froze at **32** while the true position reached
**263** (drift **+231**, growing without bound once the window saturated).
- **SnapKV** — offset advanced but stayed **exactly 32 behind**, the constant
deficit being the tokens dropped during prefill compression.
Both preserve original positions (`snap_select_indices` returns kept indices
sorted ascending; StreamingLLM's window drops rows without renumbering) and
RoPE is relative, so reporting the true position is sufficient — survivors
need no re-rotation, unlike H2O/Keyformer which renumber.
SnapKV needed more than the `_true_offset` counter that sufficed for
StreamingLLM: its decode path appends deltas, so the base class's cursor
arithmetic and return slice (`self.keys[..., :self.offset, :]`) genuinely
require the row count. `offset` is now a property yielding the row count
while the base class is on the stack and the true position outside it.
Note this diverges from the StreamingLLM paper, which assigns positions by
index *within* the cache; matching that would require re-rotating every
survivor each step. Recorded in the cache docstring.
Two existing tests asserted `offset == retained rows` — the old meaning —
and were updated, with a dedicated regression test added.
**A2ATS-adapted paper-fidelity fixes** ([29](https://github.com/rajveer43/VeloxQuant-MLX/issues/29)) —
four deviations from the source paper (He et al., ACL 2025 Findings), three
of which contradicted its equations rather than knowingly adapting them:
- **Far keys are no longer rotated.** `a2ats_apply_windowed_rope` applied a
fixed `R_window` rotation to out-of-window keys. Paper Eq. (12) leaves them
**unrotated** (`k̃_i = k_i`), with the constant `R_b` belonging on the
*query* (Eq. 11, `u_ij = q_i R_b k_j^T`). The old form computed
`q_i R_w^T k_j^T` — wrong operand, wrong direction — and left far keys in a
rotated frame, defeating the position-decoupling that makes a shared
codebook viable (§3.1, Observation 2). Adds `a2ats_apply_far_query_rope`
for the query-side half, exposed as `A2ATSKVCache.far_query_rope`.
- **`b` is now independent of `w`.** The far-token offset was derived from
`a2ats_window`, hardcoding `b == w`; the paper's §5.1 uses `w=64`,
`b=2048`. New `a2ats_b` config field (default `2048`).
- **Distance gating tracks the decode position.** Rotated keys were written
into the parent `KVCache`, which never revisits them — freezing each
token's near/far class at write time, so a token classified "near" during
prefill kept exact RoPE forever. The cache now stores *pre-RoPE*
reconstructions and re-applies windowed RoPE to the accumulated cache each
step, against the current query position. Costs an `O(total_tokens)` pass
per step, inherent to honest gating under this protocol.
- **The paper's Eq. (14) assignment is now implemented.** Adds
`a2ats_query_second_moment`, `a2ats_cholesky_factor`, and
`a2ats_h_weighted_assignment` — the `H`-weighted quadratic form, computed
exactly via the Eq. (15)–(18) Cholesky identity. Enabled by the new
`a2ats_query_h` config field. The previous cosine blend is retained as a
documented *substitute* for the decode path (no calibrated `H` available),
not relabeled as the paper's estimator.
Benchmark now measures **attention-score** error rather than key-vector
error — under Eq. (12) far keys are intentionally unrotated, so the old
target scored the method's design as error. Windowed RoPE still loses to
always-exact in both geometries (2.9x / 8.3x), and this **survived** the
fixes: the near bucket is now numerically identical to always-exact (max gap
`5e-07`), so the entire penalty is far tokens (~92–96% of the sequence).
Sweeping `b` does not remove it. Tests: 51 → 67, including the strengthened
Eq. (12) assertion that the old `R_window` bug slipped past.
**AdaKV-proxy: the default configuration was not adaptive** ([31](https://github.com/rajveer43/VeloxQuant-MLX/issues/31)) —
`adakv_target_avg_bits` defaulted to `2.0` while `adakv_lo_bit` defaulted to
`2`. Because per-head adaptation requires headroom on both sides of the target
(raising one head must be payable by lowering another), a target sitting exactly
on the floor forced *every* head to `lo_bit` for every possible importance
vector. The shipped default was therefore bit-identical to plain KIVI while the
docs advertised per-head adaptation. The default is now `2.5`, and
`allocate_head_bits` emits a `UserWarning` when the target sits at or outside an
endpoint of the allowed set rather than silently flattening.
**AdaKV-proxy: clamp-before-normalize saturated the allocation** — the
real-valued per-head budget was computed as
`clip(importance_share × H × target, lo, hi)`. Since raw key-norm variances span
orders of magnitude, this pinned nearly every head to `lo` or `hi` and discarded
the interior ordering that the snap and greedy-correction passes exist to act
on: importance vectors differing by 10 000× produced byte-identical allocations.
Replaced with rank-normalisation to a bounded, mean-centred spread placed around
the target, so the mean lands on the budget by construction and no single
extreme head can flatten the result.
**AdaKV-proxy: allocation depended on head ordering** — greedy correction scanned
heads in index order and kept the first strict minimum, so `[10,1,1,1]` and
`[1,1,1,10]` starved different heads. Ties are now broken on importance, making
the allocation permutation-equivariant for distinct importances. (With exact ties
and an indivisible budget some tied head must still lose a bit — an
integer-allocation fact, now documented rather than incidental.)
**AdaKV-proxy: corrected an inverted claim about the importance signal** — the
docstring and docs described inter-token key-norm variance as "a proxy for high
attention entropy". It is the opposite: high norm-variance indicates a few
outlier-magnitude tokens dominating the logits, i.e. an attention-*sparse* head,
which Ada-KV (§3.3, Fig. 1b) would give *less* budget. The signal is sound for
allocating bits — it measures quantization sensitivity — but it is a *different*
criterion from the paper's, not an approximation of it. Now stated plainly and
pinned by a test.
Added
**`compute_head_attention_entropy`** — an AdaKV-proxy importance signal that
carries Ada-KV's own sign: dispersed heads score higher and receive more budget.
Estimates per-head attention entropy over an observation window using the
keys-as-proxy-queries substitution already established by SnapKV-adapted,
normalised by `ln(S)`. Select via `adakv_importance="attention_entropy"`;
window size via `adakv_obs_window` (default 32).
**`KVCacheConfig`** — new fields `adakv_importance`
(default `"norm_variance"`) and `adakv_obs_window` (default `32`).
Changed
- `adakv_target_avg_bits` default `2.0` → `2.5` (see above).
- `benchmark_scripts/benchmark_adakv.py` now sweeps targets inside the adaptive
range and adds an `attention_entropy` arm at matched budget.
<!-- version list -->