veloxquant_mlx.cache.age_tiered_cache module

AgeTieredKV KV cache — position/age-driven three-tier precision (issue #256).

Every token is retained (compression-only, like AMC — never eviction). Each token’s bit-width is set purely by how long ago it was written: age = current_position - token_position. As a token ages past a tier boundary, its stored slice is re-quantized at the coarser tier’s bit-width; this mirrors KIVI’s “flush on boundary crossing” model rather than quantizing once at write time and leaving the choice frozen — a token written under this scheme is never coarser than its current age entitles, which is what issue #256’s “recent -> high precision, older -> low precision” framing calls for.

See veloxquant_mlx.quantizers.age_tiered for the tiering/quantize primitives this class wraps around a per-(batch, head) token buffer, the same storage shape AMC and KIVI both use.

class veloxquant_mlx.cache.age_tiered_cache.AgeTieredKVCache(config)

Bases: KVCache

KV cache implementing position/age-gated three-tier quantization.

Parameters:

config (Any) –

KVCacheConfig. Fields consumed:

age_recent_boundary (int, default 128) — tokens with

age < this stay at age_bits_recent.

age_mid_boundary (int, default 1024) — tokens with age

in [age_recent_boundary, age_mid_boundary) use age_bits_mid; age >= age_mid_boundary uses age_bits_old.

age_bits_recent (int, default 8) age_bits_mid (int, default 4) age_bits_old (int, default 2) age_group_size (int, default 32) — token-axis group size

for the shared min/max quantizer.

Notes

No .bits attribute — stores and returns fp16 K/V directly (same quantize-then-dequantize convention as every other method here). Both prefill (S > 1) and decode (S == 1) go through the identical per-step re-tiering pass. Writes through to the base mlx_lm KVCache’s self.keys / self.values / self.offset on every call, same as AMC (see #83) — without this, .state stays at its __init__ default and mlx_lm’s chunked prefill crashes. is_trimmable() reports False for the same reason AMC’s does: the internal per-token tier state can’t be rolled back by a base-class trim().

Keeps a full fp16 copy of every token’s raw K/V alongside the quantized output (see #397’s fix), so this class’s own runtime memory footprint is closer to fp16 than the compression_ratio implies – the reported compression only describes what would be transmitted/stored if the quantized form were serialized on its own. This raw copy is what lets a group be re-quantized once per tier change without compounding error from its own prior quantize-dequantize output.

update_and_fetch(keys, values)

Absorb new K/V tokens, re-tier the whole buffer by age, return all tokens.

Parameters:

• keys (array) – [B, H, S, D] new key tokens (any dtype; cast to fp16).

• values (array) – [B, H, S, D] new value tokens.

Returns:

(K_out, V_out) both [B, H, n_seen, D] fp16 — AgeTieredKV never evicts, so n_seen equals the total tokens passed so far.

is_trimmable()

False: trim() would only roll back base-class offset bookkeeping, not the internal per-token tier state, silently corrupting future calls (same rationale as AMC’s, see #83).

Return type:

bool

property age_tiered_bytes: int

Actual bytes stored across all heads (fp16-equivalent K + V, tiered).

property full_seq_bytes: int

Hypothetical fp16 K + V cost if AgeTieredKV were never applied.

property compression_ratio: float

full_seq_bytes / age_tiered_bytes; > 1 means savings over fp16.

property tokens_seen: int

Total token positions ever passed to update_and_fetch (all heads summed).

property tokens_kept: int

Tokens currently in the (B=0, H=0) head’s cache — always == tokens per head seen.

property tokens_recent: int

Tokens currently in the RECENT (highest-precision) tier, across all heads.

property tokens_mid: int

Tokens currently in the MID tier, across all heads.

property tokens_old: int

Tokens currently in the OLD (lowest-precision) tier, across all heads.