veloxquant_mlx.cache.pyramidkv_cache module

PyramidKV-adapted KV cache — layer-adaptive budget attention-mass eviction.

Inspired by “PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling” (Cai et al., 2024, arXiv:2406.02069). Documented as “PyramidKV-adapted (VeloxQuant-MLX implementation)” — not a faithful port.

PyramidKV is H2O’s cumulative-attention-mass eviction with a per-layer budget instead of a single global one. The observation is pyramidal information funneling: early layers attend broadly (need a big cache), deep layers attend narrowly (need a small one). Allocating a pyramid of budgets — large early, small deep, same average — beats a uniform budget at equal total memory.

Each layer’s budget is resolved once at KVCacheBuilder.for_model() build time by pyramid_budgets(n_layers, avg, n_sink, beta) and baked into that layer’s config as pyramid_resolved_budget. Unlike XQuant / MiniCache, PyramidKV needs no runtime coordinator — layers never exchange data during generation; the only cross-layer information is each layer’s index, consumed at build time.

Single-cache construction (KVCacheFactory.create) has no layer context, so it falls back to pyramid_budget (the average) — behaving as one uniform-budget H2O layer. The pyramid only takes effect via for_model.

This is the fifth 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 per-layer pyramid budget.

Adaptation limitations (stated plainly):

  • Key-as-query proxy (same as H2O-adapted / SnapKV-adapted).

  • Fixed monotone (linear) budget schedule rather than the paper’s prefill-entropy-derived allocation — funneling shape preserved, exact per-layer values not data-driven.

  • No RoPE position-ID remapping after eviction.

  • Uniform budget across heads within a layer.

Byte accounting:

pyramid_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 / pyramid_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 budget (diagnostic)

class veloxquant_mlx.cache.pyramidkv_cache.PyramidKVCache(config)[source]

Bases: KVCache

KV cache implementing PyramidKV-adapted layer-adaptive eviction for one layer.

Parameters:

config (Any) –

KVCacheConfig. Fields consumed: pyramid_resolved_budget (int or None) — this layer’s budget from

the pyramid schedule; set by for_model. If None, falls back to pyramid_budget (single-layer / uniform behaviour).

pyramid_budget (int, default 512) — average/fallback budget. pyramid_n_sink (int, default 4) — leading positions never evicted.

Notes

No .bits attribute — stores and returns fp16 K/V directly. Both prefill (S > 1) and decode (S == 1) tokens go through the same eviction loop. Per-head state is lazily initialised on the first call. 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 PyramidKV eviction, return retained window.

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 per-layer budget (from the pyramid schedule).

property pyramid_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 / pyramid_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