veloxquant_mlx.cache.squeeze_cache module

SqueezeAttention-adapted KV cache — 2D layer×token budget attention-mass eviction.

Inspired by “SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget” (Wang et al., 2024, arXiv:2404.04793). Documented as “SqueezeAttention-adapted (VeloxQuant-MLX implementation)” — not a faithful port.

SqueezeAttention is H2O’s cumulative-attention-mass eviction with a data-driven per-layer budget. Where PyramidKV assumes a fixed positional taper resolved at build time, SqueezeAttention measures each layer’s attention concentration from the prefill keys and reallocates a fixed total budget toward broad layers and away from concentrated ones — a 2D (layer × token) budget grid.

The re-budget is one-shot at the prefill boundary, mediated by a shared SqueezeCoordinator (injected at KVCacheBuilder.for_model build time):

  1. On the first update_and_fetch (the prompt), the layer computes its concentration score over the incoming keys and reports it to the coordinator.

  2. Once every attention layer has reported, the coordinator computes the per-layer budget schedule (squeeze_budgets) and publishes it.

  3. The layer pulls its resolved budget and stamps it onto every head’s eviction state; decode steps then run against the frozen budget.

Until the coordinator finalises (or when constructed with no coordinator, e.g. KVCacheFactory.create), the layer uses squeeze_budget (the average) as a fallback and behaves as one uniform-budget H2O layer. The 2D reallocation only takes effect via for_model.

This is the sixth distinct eviction configuration in VeloxQuant-MLX:

  • SnapKV-adapted : score-based, once at prefill end.

  • StreamingLLM-adapted : positional (recency + sink), every step.

  • H2O-adapted : cumulative attention mass, uniform budget, every step.

  • TOVA-adapted : current-step attention weight (memoryless), every step.

  • PyramidKV-adapted : H2O scoring with a fixed per-layer pyramid budget.

  • SqueezeAttention-adapted : H2O scoring with a DATA-DRIVEN per-layer budget (measured concentration, one-shot re-budget at prefill).

Adaptation limitations (stated plainly):

  • Key-as-query proxy for both concentration and eviction (same as H2O-adapted / PyramidKV-adapted).

  • Cosine-dispersion proxy for attention entropy (paper uses actual attention maps, not visible at cache level).

  • One-shot re-budget at the prefill boundary; frozen for decode.

  • No RoPE position-ID remapping after eviction.

  • Uniform budget across heads within a layer.

Byte accounting:

squeeze_kept_bytes — fp16 bytes for currently retained K + V tokens full_seq_bytes — hypothetical fp16 cost if all tokens were kept compression_ratio — full_seq_bytes / squeeze_kept_bytes (> 1 = savings) tokens_seen — total token positions ever passed to update_and_fetch tokens_kept — tokens currently in the first (B=0, H=0) head’s cache layer_budget — this layer’s resolved (or fallback) budget (diagnostic) concentration — this layer’s measured concentration score (diagnostic)

class veloxquant_mlx.cache.squeeze_cache.SqueezeAttentionCache(config, layer_id=None, coordinator=None)[source]

Bases: KVCache

KV cache implementing SqueezeAttention-adapted 2D eviction for one layer.

Parameters:

  • config (Any) –

    KVCacheConfig. Fields consumed: squeeze_budget (int, default 512) — average/fallback budget. squeeze_n_sink (int, default 4) — leading positions never evicted. squeeze_resolved_budget (int or None) — an explicit budget override

    (mainly for single-cache / testing); normally None and the coordinator supplies the budget after prefill.

  • layer_id (int | None) – This layer’s index (used to report/query the coordinator). None for single-cache construction (behaves as uniform H2O).

  • coordinator (SqueezeCoordinator | None) – Shared SqueezeCoordinator, or None. When present, the layer reports concentration at prefill and pulls a resolved budget.

Notes

No .bits attribute — stores and returns fp16 K/V directly. Per-head state is lazily initialised on the first call. The one-shot re-budget happens on the first update_and_fetch once the coordinator finalises. 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)[source]

Absorb new K/V tokens, apply SqueezeAttention eviction, return retained window.

On the first call the layer reports its concentration and (once the coordinator finalises) adopts its resolved budget before eviction runs.

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_kept, D] fp16, where n_kept <= layer_budget for all heads.

is_trimmable()[source]

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

This layer’s resolved (or fallback) per-layer budget.

property concentration: float

This layer’s measured concentration score (0.0 before prefill).

property squeeze_kept_bytes: int

Bytes currently stored across all heads (fp16 K + V, kept tokens only).

property full_seq_bytes: int

Hypothetical fp16 K + V cost if all tokens were kept.

property compression_ratio: float

full_seq_bytes / squeeze_kept_bytes; > 1 means memory 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 (diagnostic).

property merge