veloxquant_mlx.cache.streaming_llm_cache module

StreamingLLM-adapted KV cache — sink + recency-window structural eviction.

Inspired by “Efficient Streaming Language Models with Attention Sinks” (Xiao et al., ICLR 2024, arXiv:2309.17453). Documented as “StreamingLLM-adapted (VeloxQuant-MLX implementation)” — not a faithful port.

Structural eviction: keep the first stream_n_sink token positions unconditionally (attention sinks) plus the most recent stream_window_size token positions. All other positions are permanently evicted. The cache never grows beyond stream_n_sink + stream_window_size positions, making decode constant-memory.

This is positional eviction — orthogonal to SnapKV-adapted (score-based eviction) and to all quantization methods (which compress all tokens to fewer bits).

Adaptation limitations (stated plainly):

• No attention mask adjustment: the model attends to all returned K/V positions; we only bound what K/V rows exist.

• No RoPE position-ID remapping: original token positions are preserved in the returned rows.

• Fixed sink count (stream_n_sink), not adaptive.

Byte accounting:

stream_kept_bytes — fp16 bytes stored (sink + recent positions, K + V) full_seq_bytes — hypothetical cost if all tokens were kept as fp16 streaming_ratio — full_seq_bytes / stream_kept_bytes (> 1 once window fills) tokens_seen — total positions ever passed to update_and_fetch (all heads avg) tokens_in_window — current sink + recent positions in cache (first head)

class veloxquant_mlx.cache.streaming_llm_cache.StreamingLLMKVCache(config)

Bases: KVCache

KV cache implementing StreamingLLM-adapted sink + recency-window eviction.

Parameters:

config (Any) – KVCacheConfig. Fields consumed: stream_n_sink (int, default 4) — initial positions always kept, stream_window_size (int, default 512) — FIFO recent-token capacity.

Notes

No .bits attribute — stores and returns fp16 K/V directly. The cache never grows beyond stream_n_sink + stream_window_size positions. Both prefill (S > 1) and decode (S == 1) tokens are processed identically — all go through the sink/window logic. This is the key difference from SnapKV- adapted, which evicts only at prefill. StreamingLLM operates continuously. Single-layer (no coordinator); for_model propagates all stream_* fields automatically via dataclasses.replace. 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 sink+window eviction, return full window.

Parameters:

• keys (array) – [B, H, S, D] new key tokens.

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

Returns:

(K_out, V_out) for THIS call’s own attention — the full, un-evicted concatenation of whatever the sink+recent window held before this call plus the S new tokens (see #370 below), NOT capped at stream_n_sink + stream_window_size. What gets stored afterward (the window state, visible to the next call) is capped as before.

mlx_lm builds the attention mask for this call from hidden states — before q/k/v projections exist, let alone this cache’s own update_and_fetch — so it is fixed (as either the “causal” string or an explicit array from make_mask, called with only this call’s query count N) before eviction can possibly run. If this method shrank what it returns to fewer than N keys via window trimming, that already-fixed mask would silently desync from the shape it was built for (VeloxQuant-MLX#370). So trimming is deferred: this call returns the full pre-trim concatenation (matching the mask make_mask already built from the previous call’s true kept positions — see that method), and only the stored window shrinks, for the next call’s make_mask to reflect correctly.

make_mask(N, return_array=False, window_size=None, **_)

Explicit position-based causal mask — see VeloxQuant-MLX#370.

Called BEFORE this step’s own update_and_fetch (and thus before this step’s own window trim, which update_and_fetch defers past this step’s return anyway — see its docstring). self._kept_positions holds the true positions of whatever the sink+recent window already holds from the previous call, which is exactly what update_and_fetch will concatenate its N new tokens onto — so a mask sized [B, 1, N, len(kept) + N] covers this call’s actual returned key count precisely.

Parameters:

• N (int)

• return_array (bool)

• window_size (int | None)

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

Bytes currently stored across all heads (fp16 K + V, sink + recent).

property full_seq_bytes: int

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

property streaming_ratio: float

full_seq_bytes / stream_kept_bytes; > 1 once window fills.

property tokens_seen: int

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

property tokens_in_window: int

Current sink + recent count for head (0, 0).

property tokens_kept: int

Alias for tokens_in_window, matching the tokens_kept name every other eviction cache uses (h2o, tova, pyramidkv, snapkv, squeeze, …) for “tokens currently in the (B=0, H=0) head’s cache”. Telemetry code (e.g. the /v1/kv/stats token-count aggregator) probes for tokens_kept specifically via hasattr/getattr; without this alias it silently reports 0 retained tokens for streaming_llm regardless of actual eviction state, since tokens_seen alone is already enough to pass its “has telemetry” check. Found verifying VeloxQuant-Studio issue #29.

property merge