veloxquant_mlx.quantizers.tova module¶
TOVA (Token Omission Via Attention) KV eviction primitives.
Inspired by “Transformers are Multi-State RNNs” (Oren et al., 2024, arXiv:2401.06104), whose TOVA policy retains a fixed-size cache by dropping, at each step, the single token with the lowest attention weight in the current step. Documented as “TOVA-adapted (VeloxQuant-MLX implementation)” — not a faithful port. See adaptation limitations below and in cache/tova_cache.py.
The distinguishing feature versus H2O-adapted: TOVA is memoryless. It scores tokens by the attention weight received at the current step only, with no running accumulation. H2O keeps a cumulative sum (inertial — a token that was a heavy hitter long ago survives); TOVA reacts instantly to the present context (a token that stops being attended to is evicted even if it dominated earlier).
- Adaptation limitations (stated plainly):
Key-as-query proxy: the paper reads the actual attention distribution of the most recent query row from the forward pass. At cache level the query is not visible, so we use the incoming key vector as a proxy query to approximate the current-step attention distribution. Same approximation as SnapKV-adapted and H2O-adapted.
No RoPE position-ID renumbering after eviction — the evicted row is dropped and the rest kept in temporal order, so surviving tokens keep their original absolute positions and the cache layer reports the true token position for RoPE rather than the retained row count (see
cache/tova_cache.pyand :issue:`171`, :issue:`175`).Uniform budget across all heads.
When a multi-token chunk (S > 1) arrives, tokens are absorbed one at a time and the last incoming key of each step is the proxy query. This differs from the paper’s single-step decode framing but keeps budget semantics identical.
Public API¶
TovaState — immutable per-head state dataclass init_tova_state — construct empty state tova_update — absorb S new tokens, evict lowest current-step-weight token tova_get_kv — extract current (keys, values) arrays tova_fp16_bytes — bytes stored in current state full_tova_fp16_bytes — hypothetical cost without eviction
- class veloxquant_mlx.quantizers.tova.TovaState(keys, values, n_sink, budget)[source]¶
Bases:
objectPer-head sliding TOVA state.
- Variables:
keys – [n_kept, D] fp16 stored key rows, or None before first update.
values – [n_kept, D] fp16 stored value rows, or None before first update.
n_sink – Number of leading sink positions — never evicted.
budget – Maximum tokens to keep at any time (including sinks).
- Parameters:
Note
Unlike H2OState there is no
scoresfield — TOVA is memoryless and recomputes the current-step attention weights fresh on every update, discarding them after the eviction decision.
- veloxquant_mlx.quantizers.tova.init_tova_state(n_sink, budget, head_dim)[source]¶
Create an empty TovaState before any tokens arrive.
- Parameters:
- Raises:
ValueError – if there are sink positions to protect but they leave no evictable room within
budget(n_sink=0, budget=0remains a valid “disabled cache” configuration).- Return type:
- veloxquant_mlx.quantizers.tova.tova_update(state, new_keys, new_values, *, backend='auto')[source]¶
Absorb [S,D] K/V using auto, GPU-only mlx, metal, or reference eviction.
Surviving rows retain their original positions and FP16 values. The reference backend retains the original Python-driven algorithm for parity and benchmarking. Auto uses the measured default (see TOVA_METAL_FINDINGS).
- veloxquant_mlx.quantizers.tova.tova_get_kv(state)[source]¶
Return
(keys, values)arrays from state.Returns
([0, 1], [0, 1])zero-row placeholders before the first update.