veloxquant_mlx.quantizers.streaming_llm module

StreamingLLM-adapted quantizer primitives — sink + recency-window token 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.

What StreamingLLM adds: structural positional eviction — tokens are kept or dropped purely by their position (first N sink positions + last W recent positions), with no attention scoring, no calibration, and no model weights. This is orthogonal to SnapKV-adapted (score-based eviction) and all quantization methods.

The paper’s key insight: attention sinks (the first few tokens of any sequence) consistently receive high attention weight regardless of content, and most other attention is locally recency-biased. Keeping N_sink + W_recent tokens preserves both effects while bounding the cache to a constant size during streaming decode.

Adaptation decisions (documented, never hidden):

  1. Cache-level implementation. The paper patches the model’s forward pass (attention mask + KV cache together). We implement equivalent semantics inside update_and_fetch by maintaining a frozen sink buffer and a FIFO recent window, concatenating them per call. Functionally equivalent within the cache-wrapper level.

  2. No attention mask adjustment. The paper adjusts position masks so tokens beyond the window are invisible to the query. A cache wrapper cannot inject attention masks; all returned K/V positions will be attended to by the model. Documented as a known limitation — the memory budget is still bounded.

  3. No position-ID remapping. We drop tokens and preserve original positions inside the returned K/V rows. The paper remaps RoPE position IDs; that requires model-level patching. Documented plainly.

  4. Fixed sink count. stream_n_sink is a fixed hyperparameter, not adaptive.

The cache never exceeds stream_n_sink + stream_window_size token positions, making it constant-memory for arbitrarily long generation once the window fills.

This module holds the pure, side-effect-free numerics: concatenation, window trimming, and byte accounting.

class veloxquant_mlx.quantizers.streaming_llm.StreamingWindow(sink_keys, sink_values, recent_keys, recent_values, n_sink, n_recent, tokens_seen)[source]

Bases: NamedTuple

A single-head sink+recent window for StreamingLLM-adapted eviction.

Variables:

  • sink_keys – [n_sink, D] fp16 — frozen attention-sink key rows.

  • sink_values – [n_sink, D] fp16 — frozen attention-sink value rows.

  • recent_keys – [n_recent, D] fp16 — rolling recent-window key rows.

  • recent_values – [n_recent, D] fp16 — rolling recent-window value rows.

  • n_sink – int — sink count (== sink_keys.shape[0]).

  • n_recent – int — current recent-window occupancy (<= window_size).

  • tokens_seen – int — total token positions seen since creation.

Parameters:

  • sink_keys (mlx.core.array)

  • sink_values (mlx.core.array)

  • recent_keys (mlx.core.array)

  • recent_values (mlx.core.array)

  • n_sink (int)

  • n_recent (int)

  • tokens_seen (int)

sink_keys: mlx.core.array

Alias for field number 0

sink_values: mlx.core.array

Alias for field number 1

recent_keys: mlx.core.array

Alias for field number 2

recent_values: mlx.core.array

Alias for field number 3

n_sink: int

Alias for field number 4

n_recent: int

Alias for field number 5

tokens_seen: int

Alias for field number 6

veloxquant_mlx.quantizers.streaming_llm.init_streaming_window(n_sink, D, dtype=mlx.core.float16)[source]

Create an empty streaming window with the given sink count and head dimension.

Return type:

StreamingWindow

Parameters:

  • n_sink (int)

  • D (int)

  • dtype (mlx.core.Dtype)

veloxquant_mlx.quantizers.streaming_llm.stream_update(window, new_keys, new_values, n_sink, window_size)[source]

Absorb [S, D] new K/V tokens into the streaming window.

The first n_sink positions ever seen are unconditionally frozen as sinks. Once frozen, sinks are never evicted. All subsequent tokens occupy the recent window, a FIFO buffer of capacity window_size. When the FIFO exceeds window_size, the oldest tokens are dropped from the front.

Parameters:

  • window (StreamingWindow) – Current StreamingWindow state.

  • new_keys (array) – [S, D] fp16/fp32 key rows to absorb.

  • new_values (array) – [S, D] fp16/fp32 value rows to absorb.

  • n_sink (int) – Number of initial tokens to freeze as sinks.

  • window_size (int) – FIFO capacity for recent tokens.

Return type:

StreamingWindow

Returns:

Updated StreamingWindow.

veloxquant_mlx.quantizers.streaming_llm.stream_get_kv(window)[source]

Return the concatenated [n_sink + n_recent, D] fp16 K and V tensors.

If only sinks exist (no recent tokens yet), returns the sink arrays. If both exist, concatenates sink rows first, then recent rows.

Return type:

tuple[array, array]

Parameters:

window (StreamingWindow)

veloxquant_mlx.quantizers.streaming_llm.stream_fp16_bytes(window)[source]

Bytes stored in the streaming window (fp16 K + V, sink + recent rows).

Return type:

int

Parameters:

window (StreamingWindow)

veloxquant_mlx.quantizers.streaming_llm.full_stream_fp16_bytes(n_tokens, D)[source]

Bytes for uncompressed fp16 K + V at n_tokens positions, dim D.

Return type:

int

Parameters: