veloxquant_mlx.cache.amc_cache module

AMC-adapted KV cache — per-token saliency-driven tiered rank + precision.

Inspired by “Adaptive Model Compression (AMC)” (Hu, Yuan, Hu, Yin, Li, Suchter — Apple; arXiv:2607.10109, no verified peer-reviewed venue as of 2026-07-14). Documented as “AMC-adapted (VeloxQuant-MLX implementation)” — not a faithful port; see veloxquant_mlx.quantizers.amc module docstring for the full honesty crux (venue status, hardware/RTL scope cut, compression-only-not-eviction framing).

Every token — prefill or decode — is scored, tiered, rank-masked, and quantized on every call. No token is ever dropped: this is the fortieth method in this repo and the first whose family is “adaptive rank+precision” rather than “eviction.” The cache’s stored sequence length always equals the number of tokens seen; only the effective rank/bit-width per token varies.

Byte accounting:

amc_kept_bytes — actual bytes for K + V given each token’s tier full_seq_bytes — hypothetical fp16 full-rank cost if AMC were never

applied

compression_ratio — full_seq_bytes / amc_kept_bytes (> 1 = savings) tokens_seen — total token positions ever passed to update_and_fetch tokens_high/mid/low — cumulative per-tier token counts (observability)

class veloxquant_mlx.cache.amc_cache.AMCKVCache(config)

Bases: KVCache

KV cache implementing AMC-adapted saliency-driven tiered compression.

Parameters:

config (Any) –

KVCacheConfig. Fields consumed:

amc_k_high (float, default 0.20) — top percentile -> High tier amc_k_mid (float, default 0.30) — next percentile -> Mid tier amc_use_query_saliency (bool, default False) amc_query_alpha (float, default 0.5) — Eq. 3 balance coefficient amc_adaptive_thresholds (bool, default False) — Eq. 4-5 closed-loop amc_threshold_window (int, default 64) — trailing window size amc_gamma (float, default 0.1) — threshold attenuation factor amc_calib_variance (float | None) — required if adaptive thresholds on amc_group_size (int, default 32) — quantization group size

Notes

No .bits attribute — stores and returns fp16 K/V directly (the rank mask + quantize round-trip is simulated, matching every other method’s “quantize-then-dequantize to fp16” convention). Both prefill (S > 1) and decode (S == 1) go through the identical per-token tiering loop. Query-aware saliency (Eq. 3) requires a query vector, which is not available at the cache-wrapper level (same limitation as H2O/SnapKV/ CurDKV’s key-as-query-proxy) — when enabled, the mean of the current step’s keys is used as a proxy query, disclosed here as the same category of approximation as the eviction methods’ proxy queries. Writes through to the base mlx_lm KVCache’s self.keys / self.values / self.offset on every call so .state stays valid (mlx_lm’s generate() reads it unconditionally during chunked prefill); is_trimmable() reports False since the internal per-token state can’t be rolled back by a base-class trim() (see #83).

update_and_fetch(keys, values)

Absorb new K/V tokens, apply AMC tiered compression, 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 — AMC 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 eviction/compression state that actually determines what gets returned, silently corrupting future calls.

Return type:

bool

property amc_kept_bytes: int

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

property full_seq_bytes: int

Hypothetical fp16 full-rank K + V cost if AMC were never applied.

property compression_ratio: float

full_seq_bytes / amc_kept_bytes; > 1 means savings over fp16 full-rank.

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_high: int

Cumulative tokens (across all heads/steps) assigned the High tier.

property tokens_mid: int

Cumulative tokens (across all heads/steps) assigned the Mid tier.

property tokens_low: int

Cumulative tokens (across all heads/steps) assigned the Low tier.