veloxquant_mlx.cache package¶

KV-cache implementations for every quantization/eviction method in VeloxQuant-MLX.

Each sibling module implements one method’s KVCache (TurboQuant, PolarQuant, QJL, KIVI-family, eviction methods like SnapKV/H2O/PyramidKV/ChunkKV, cross-layer methods like XQuant/MiniCache/xKV, etc.); this package re-exports the most commonly used cache classes plus the construction entry points — KVCacheConfig (the shared configuration dataclass), KVCacheFactory (builds a single cache from a config), and KVCacheBuilder (fluent config builder / per-model cache-list construction via for_model). See base.py for the full method registry and the standalone-vs-mlx_lm-protocol distinction.

class veloxquant_mlx.cache.KVCacheBuilder[source]

Bases: object

Fluent builder for KVCache construction with validation.

Example:

cache = (
    KVCacheBuilder()
    .with_method("turboquant_prod")
    .with_head_dim(128)
    .with_bit_width(inlier=2, outlier=3)
    .with_jl_dim(128)
    .with_seed(42)
    .build()
)

with_method(method)[source]

Set the quantisation method.

Parameters:: method (str) – One of ‘turboquant_prod’, ‘turboquant_mse’, ‘polar’, ‘qjl’.
Return type:: KVCacheBuilder

with_head_dim(d)[source]

Set the attention head dimension.

Parameters:: d (int) – Head dimension (must be a power of 2).
Return type:: KVCacheBuilder

with_bit_width(inlier, outlier=None)[source]

Set bit-width(s).

Parameters:

inlier – Bit-width for inlier channels. Either an int (uniform across all layers) or a list[int] of length n_layers for RateQuant-style per-layer allocation. When a list is supplied, this builder must be consumed via KVCacheBuilder.for_model(model, config); direct .build() rejects the list.
outlier (int | None) – Bit-width for outlier channels (defaults to inlier if None).

Return type:

KVCacheBuilder

with_jl_dim(m)[source]

Set the JL projection dimension.

Parameters:: m (int) – Must be <= head_dim.
Return type:: KVCacheBuilder

with_n_outlier_channels(n)[source]

Set the number of outlier channels to detect.

Parameters:: n (int) – Must be < head_dim.
Return type:: KVCacheBuilder

with_n_calib_tokens(n)[source]

Set calibration token count for outlier activation.

Return type:: KVCacheBuilder
Parameters:: n (int)

with_vectorized_attend(enabled=True)[source]

Enable vectorized packed-key unpack in attend().

Return type:: KVCacheBuilder
Parameters:: enabled (bool)

with_outlier_two_stream(enabled=True)[source]

Enable outlier/inlier split cache after calibration.

Return type:: KVCacheBuilder
Parameters:: enabled (bool)

with_fused_query_dot(enabled=True)[source]

Enable fused rotated-query + codebook-dot path.

Return type:: KVCacheBuilder
Parameters:: enabled (bool)

with_seed(seed)[source]

Set the random seed.

Parameters:: seed (int) – Integer seed.
Return type:: KVCacheBuilder

with_precision(dtype)[source]

Set the compute dtype.

Parameters:: dtype (Any) – MLX dtype (e.g. mx.float16).
Return type:: KVCacheBuilder

with_capacity(max_tokens)[source]

Set the maximum number of tokens to store.

Parameters:: max_tokens (int) – Positive integer.
Return type:: KVCacheBuilder

with_artifact_store(store)[source]

Provide an ArtifactStore for loading precomputed artifacts.

Parameters:: store (ArtifactStore) – ArtifactStore instance.
Return type:: KVCacheBuilder

with_observer(observer)[source]

Attach a QuantizationObserver.

Parameters:: observer (QuantizationObserver) – Observer instance.
Return type:: KVCacheBuilder

with_sliding_window(window_size)[source]

Wrap the cache with a sliding-window eviction policy.

Parameters:: window_size (int) – Number of tokens to keep.
Return type:: KVCacheBuilder

build()[source]

Validate the configuration and construct the KVCache.

Return type:: KVCache | _MLXKVCache
Returns:: Configured KVCache instance.
Raises:: QuantizerConfigError – If any validation check fails.

static for_model(model, config)[source]

Build one KVCache per language-model layer, sized per-layer.

Works for text-only and VLM models (Qwen2-VL, Qwen3-VL, Mistral, etc.). Layers without a self_attn attribute (MoE gates, etc.) fall back to a standard fp16 KVCache so the list length always matches model.layers.

Per-layer bit-widths (RateQuant)¶

If config.bit_width_inlier is a list[int], element i is used for layer i. The list length must equal the number of attention-bearing layers (layers without self_attn are skipped from the count). This lets RateQuant-style mixed-precision allocations be passed through the standard API without manual cache wiring.

type model:
param model:: Loaded mlx_lm model instance.
type config:: KVCacheConfig
param config:: KVCacheConfig specifying method, bit_width_inlier, seed, etc. head_dim is overridden per-layer.
rtype:: list
returns:: List of KVCache instances, one per language-model layer.
raises QuantizerConfigError:: If config.method is a standalone method (see STANDALONE_METHODS) that does not implement the mlx_lm KVCache serving contract.

Parameters:: config (KVCacheConfig)
Return type:: list

class veloxquant_mlx.cache.KVCacheConfig(method='turboquant_rvq', head_dim=128, bit_width_inlier=2, bit_width_outlier=None, jl_dim=None, n_outlier_channels=None, n_calib_tokens=None, enable_vectorized_attend=True, enable_outlier_two_stream=False, enable_fused_query_dot=False, seed=42, dtype=None, capacity=None, sliding_window=None, store=None, observers=<factory>, options=None, key_sub_dim=4, value_sub_dim=8, key_codebook_bits=12, value_codebook_bits=8, residual_length=128, kivi_group_size=32, svdq_rank=None, svdq_energy_threshold=0.95, svdq_bit_schedule=(8, 4, 2, 1, 1, 0, 0, 0), svdq_group_size=32, kitty_hi_fraction=0.25, kitty_hi_bit=4, kitty_lo_bit=2, kitty_group_size=32, adakv_target_avg_bits=2.5, adakv_lo_bit=2, adakv_mid_bit=3, adakv_hi_bit=4, adakv_group_size=32, adakv_update_interval=1, adakv_importance='norm_variance', adakv_obs_window=32, xquant_group_size=2, xquant_base_bits=2, xquant_residual_bits=4, xquant_group_quant_size=32, xquant_max_ctx=8192, kvquant_bits=3, kvquant_outlier_fraction=0.01, kvquant_group_size=32, kvquant_lloyd_iters=8, kvquant_refit_interval=0, kvquant_n_sink=1, palu_rank=None, palu_energy_threshold=0.9, palu_n_head_groups=4, palu_hi_bit=4, palu_lo_bit=2, palu_hi_fraction=0.25, palu_group_size=32, palu_quantize_values=True, cachegen_bits=4, cachegen_group_size=32, cachegen_use_delta=True, cachegen_per_channel=True, cachegen_layer_groups=3, cachegen_resolved_bits=None, minicache_start_frac=0.5, minicache_group_size=2, minicache_retention_threshold=0.9, minicache_slerp_t=0.5, minicache_max_ctx=8192, gear_bits=2, gear_rank=None, gear_energy_threshold=0.9, gear_sparse_fraction=0.01, gear_group_size=32, gear_quantize_values=True, zipcache_hi_bits=4, zipcache_lo_bits=2, zipcache_hi_fraction=0.2, zipcache_group_size=32, zipcache_quantize_values=True, snap_dtype='auto', snap_batched_scoring=True, snap_backend='auto', snap_budget=512, snap_obs_window=32, snap_n_sink=4, stream_n_sink=4, stream_window_size=512, h2o_budget=512, h2o_n_sink=4, h2o_rope_base=10000.0, h2o_grace=16, h2o_decay=0.98, tova_budget=512, tova_n_sink=4, tova_backend='auto', pyramid_budget=512, pyramid_backend='reference', pyramid_n_sink=4, pyramid_beta=2.0, pyramid_resolved_budget=None, squeeze_budget=512, squeeze_n_sink=4, squeeze_strength=1.0, squeeze_resolved_budget=None, chunkkv_budget=512, chunkkv_chunk_size=8, chunkkv_n_sink=4, chunkkv_score='attn_mass', chunkkv_reuse_layers=1, cam_budget=512, cam_n_sink=4, cam_merge='sim_weighted', cam_merge_keys=False, cam_merge_gate=True, xkv_group_size=2, xkv_rank=None, xkv_energy_threshold=0.95, xkv_latent_bits=4, xkv_group_quant_size=32, xkv_max_ctx=8192, nsn_bits=2, nsn_residual_length=64, nsn_codebook_size=256, nsn_subvector_dim=8, nsn_seed=1234, nsn_max_ctx=8192, knorm_budget=512, knorm_n_sink=4, knorm_recent=0, knorm_keep='low', skvq_bits_key=2, skvq_bits_value=2, skvq_group_size=32, skvq_window=128, skvq_n_sink=5, skvq_reorder=True, skvq_clip_search=True, skvq_clip_alpha=1.0, skvq_max_ctx=None, qfilters_budget=512, qfilters_n_sink=4, qfilters_recent=0, qfilters_calib_tokens=128, qfilters_sign=1, qfilters_min_retention=0.9, qfilters_on_low_retention='warn', keyformer_budget=512, keyformer_n_sink=4, keyformer_recent=0, keyformer_tau=None, keyformer_tau_init=1.0, keyformer_tau_end=1.0, keyformer_anneal_steps=0, keyformer_rope_base=10000.0, keyformer_seed=0, morphkv_budget=512, morphkv_n_sink=4, morphkv_window=8, kvzip_budget=512, kvzip_n_sink=4, kvzip_probe='context', kvtc_bit_budget=512, kvtc_bit_choices=(0, 1, 2, 3, 4, 6, 8), kvtc_beta=3.5, curdkv_budget=512, curdkv_n_sink=4, curdkv_rank_cap=16, curdkv_rope_base=10000.0, nestedkv_budget=512, nestedkv_n_sink=4, nestedkv_window=64, nestedkv_beta=3.0, nestedkv_tau=0.6, nestedkv_kappa=10.0, nestedkv_safeguard_alpha=0.2, amc_k_high=0.2, amc_k_mid=0.3, amc_use_query_saliency=False, amc_query_alpha=0.5, amc_adaptive_thresholds=False, amc_threshold_window=64, amc_gamma=0.1, amc_calib_variance=None, amc_group_size=32, a2ats_codebook_bits=8, a2ats_sub_dim=8, a2ats_window=128, a2ats_b=2048, a2ats_use_query_aware=True, a2ats_beta=0.5, a2ats_retrieval_fraction=0.2, a2ats_rope_base=10000.0, a2ats_rope_scaling=None, a2ats_query_h=None, a2ats_codebook=None, anchorkv_theta=0.05, anchorkv_window=32, anchorkv_rho=0.7, anchorkv_anchor_frac=0.0078125, anchorkv_residual_bits=2, anchorkv_seed=42, rocketkv_compression_ratio=8.0, rocketkv_page_size=None, rocketkv_head_topk1=None, rocketkv_obs_window=32, rocketkv_n_sink=4, age_recent_boundary=128, age_mid_boundary=1024, age_bits_recent=8, age_bits_mid=4, age_bits_old=2, age_group_size=32, n_sink_tokens=5, smooth_factors=None, key_codebook=None, value_codebook=None, spectral_key_d_eff=4, spectral_val_d_eff=50, spectral_apply_qjl=True, spectral_model_name='model', use_metal_kernels=None, fused_sdpa=False, fused_sdpa_max_ctx=None, fused_sdpa_memory_bound=False)[source]

Bases: object

Configuration for a KVCache instance.

Variables:

method – Quantisation algorithm.
head_dim – Attention head dimension (d).
bit_width_inlier – Bit-width for inlier channels. Either a single int applied uniformly across all layers, OR a list[int] of length n_layers for per-layer RateQuant-style allocation. When passed as a list, KVCacheBuilder.for_model() consumes element i for layer i; KVCacheFactory.create() requires an int.
bit_width_outlier – Bit-width for outlier channels (None → same as inlier).
jl_dim – JL projection dimension m. None means m = head_dim for QJL, which is cheap but noisy (score correlation ~0.64 at d=128 vs ~0.88 at m=4d); raise it when attention quality matters. See #774.
n_outlier_channels – Number of outlier channels to detect.
seed – Random seed.
dtype – MLX dtype for computations.
capacity – Maximum tokens to store (None → unlimited).
sliding_window – If set, wrap cache with sliding-window eviction.
store – ArtifactStore to load precomputed artifacts from.
observers – List of QuantizationObserver instances.
options – Method-specific option dataclass from veloxquant_mlx.cache.options (e.g. H2OOptions, SnapKVOptions), grouping the flat per-method fields below (h2o_budget, snap_budget, …) into one small object per method. Preferred over setting the flat fields directly (#420): values from options populate the matching flat fields at construction time (see __post_init__), so every existing cache class – which reads the flat fields via getattr(config, "h2o_budget", 512) – keeps working unchanged. Setting a flat per-method field directly still works but emits a DeprecationWarning naming its options replacement.

Parameters:

method (Literal['turboquant_prod', 'turboquant_mse', 'turboquant_rvq', 'polar', 'qjl', 'vecinfer', 'spectral', 'kivi', 'kivi_sink', 'svdq', 'kitty', 'adakv', 'xquant', 'kvquant', 'palu', 'cachegen', 'minicache', 'gear', 'zipcache', 'snapkv', 'streaming_llm', 'h2o', 'tova', 'pyramidkv', 'squeeze', 'chunkkv', 'cam', 'xkv', 'nsnquant', 'knorm', 'skvq', 'qfilters', 'keyformer', 'morphkv', 'kvzip', 'kvtc', 'curdkv', 'nestedkv', 'amc', 'a2ats', 'anchorkv', 'rocketkv', 'age_tiered'])
head_dim (int)
bit_width_inlier (int | list)
bit_width_outlier (int | None)
jl_dim (int | None)
n_outlier_channels (int | None)
n_calib_tokens (int | None)
enable_vectorized_attend (bool)
enable_outlier_two_stream (bool)
enable_fused_query_dot (bool)
seed (int)
dtype (Any)
capacity (int | None)
sliding_window (int | None)
store (ArtifactStore | None)
observers (list)
options (object | None)
key_sub_dim (int)
value_sub_dim (int)
key_codebook_bits (int)
value_codebook_bits (int)
residual_length (int)
kivi_group_size (int)
svdq_rank (int | None)
svdq_energy_threshold (float)
svdq_bit_schedule (tuple[int, ...])
svdq_group_size (int)
kitty_hi_fraction (float)
kitty_hi_bit (int)
kitty_lo_bit (int)
kitty_group_size (int)
adakv_target_avg_bits (float)
adakv_lo_bit (int)
adakv_mid_bit (int)
adakv_hi_bit (int)
adakv_group_size (int)
adakv_update_interval (int)
adakv_importance (str)
adakv_obs_window (int)
xquant_group_size (int)
xquant_base_bits (int)
xquant_residual_bits (int)
xquant_group_quant_size (int)
xquant_max_ctx (int)
kvquant_bits (int)
kvquant_outlier_fraction (float)
kvquant_group_size (int)
kvquant_lloyd_iters (int)
kvquant_refit_interval (int)
kvquant_n_sink (int)
palu_rank (int | None)
palu_energy_threshold (float)
palu_n_head_groups (int)
palu_hi_bit (int)
palu_lo_bit (int)
palu_hi_fraction (float)
palu_group_size (int)
palu_quantize_values (bool)
cachegen_bits (int)
cachegen_group_size (int)
cachegen_use_delta (bool)
cachegen_per_channel (bool)
cachegen_layer_groups (int)
cachegen_resolved_bits (int | None)
minicache_start_frac (float)
minicache_group_size (int)
minicache_retention_threshold (float)
minicache_slerp_t (float)
minicache_max_ctx (int)
gear_bits (int)
gear_rank (int | None)
gear_energy_threshold (float)
gear_sparse_fraction (float)
gear_group_size (int)
gear_quantize_values (bool)
zipcache_hi_bits (int)
zipcache_lo_bits (int)
zipcache_hi_fraction (float)
zipcache_group_size (int)
zipcache_quantize_values (bool)
snap_dtype (str)
snap_batched_scoring (bool)
snap_backend (str)
snap_budget (int)
snap_obs_window (int)
snap_n_sink (int)
stream_n_sink (int)
stream_window_size (int)
h2o_budget (int)
h2o_n_sink (int)
h2o_rope_base (float)
h2o_grace (int)
h2o_decay (float)
tova_budget (int)
tova_n_sink (int)
tova_backend (str)
pyramid_budget (int)
pyramid_backend (str)
pyramid_n_sink (int)
pyramid_beta (float)
pyramid_resolved_budget (int | None)
squeeze_budget (int)
squeeze_n_sink (int)
squeeze_strength (float)
squeeze_resolved_budget (int | None)
chunkkv_budget (int)
chunkkv_chunk_size (int)
chunkkv_n_sink (int)
chunkkv_score (str)
chunkkv_reuse_layers (int)
cam_budget (int)
cam_n_sink (int)
cam_merge (str)
cam_merge_keys (bool)
cam_merge_gate (bool)
xkv_group_size (int)
xkv_rank (int | None)
xkv_energy_threshold (float)
xkv_latent_bits (int)
xkv_group_quant_size (int)
xkv_max_ctx (int)
nsn_bits (int)
nsn_residual_length (int)
nsn_codebook_size (int)
nsn_subvector_dim (int)
nsn_seed (int)
nsn_max_ctx (int)
knorm_budget (int)
knorm_n_sink (int)
knorm_recent (int)
knorm_keep (str)
skvq_bits_key (int)
skvq_bits_value (int)
skvq_group_size (int)
skvq_window (int)
skvq_n_sink (int)
skvq_reorder (bool)
skvq_clip_search (bool)
skvq_clip_alpha (float)
skvq_max_ctx (int | None)
qfilters_budget (int)
qfilters_n_sink (int)
qfilters_recent (int)
qfilters_calib_tokens (int)
qfilters_sign (int)
qfilters_min_retention (float | None)
qfilters_on_low_retention (Literal['warn', 'raise', 'ignore'])
keyformer_budget (int)
keyformer_n_sink (int)
keyformer_recent (int)
keyformer_tau (float | None)
keyformer_tau_init (float)
keyformer_tau_end (float)
keyformer_anneal_steps (int)
keyformer_rope_base (float)
keyformer_seed (int)
morphkv_budget (int)
morphkv_n_sink (int)
morphkv_window (int)
kvzip_budget (int)
kvzip_n_sink (int)
kvzip_probe (str)
kvtc_bit_budget (int)
kvtc_bit_choices (tuple)
kvtc_beta (float)
curdkv_budget (int)
curdkv_n_sink (int)
curdkv_rank_cap (int)
curdkv_rope_base (float)
nestedkv_budget (int)
nestedkv_n_sink (int)
nestedkv_window (int)
nestedkv_beta (float)
nestedkv_tau (float)
nestedkv_kappa (float)
nestedkv_safeguard_alpha (float)
amc_k_high (float)
amc_k_mid (float)
amc_use_query_saliency (bool)
amc_query_alpha (float)
amc_adaptive_thresholds (bool)
amc_threshold_window (int)
amc_gamma (float)
amc_calib_variance (float | None)
amc_group_size (int)
a2ats_codebook_bits (int)
a2ats_sub_dim (int)
a2ats_window (int)
a2ats_b (int)
a2ats_use_query_aware (bool)
a2ats_beta (float)
a2ats_retrieval_fraction (float)
a2ats_rope_base (float)
a2ats_rope_scaling (Any)
a2ats_query_h (Any)
a2ats_codebook (Any)
anchorkv_theta (float)
anchorkv_window (int)
anchorkv_rho (float)
anchorkv_anchor_frac (float)
anchorkv_residual_bits (int)
anchorkv_seed (int)
rocketkv_compression_ratio (float)
rocketkv_page_size (int | None)
rocketkv_head_topk1 (int | None)
rocketkv_obs_window (int)
rocketkv_n_sink (int)
age_recent_boundary (int)
age_mid_boundary (int)
age_bits_recent (int)
age_bits_mid (int)
age_bits_old (int)
age_group_size (int)
n_sink_tokens (int)
smooth_factors (Any)
key_codebook (Any)
value_codebook (Any)
spectral_key_d_eff (int)
spectral_val_d_eff (int)
spectral_apply_qjl (bool)
spectral_model_name (str)
use_metal_kernels (bool | None)
fused_sdpa (bool | None)
fused_sdpa_max_ctx (int | None)
fused_sdpa_memory_bound (bool)

method: Literal['turboquant_prod', 'turboquant_mse', 'turboquant_rvq', 'polar', 'qjl', 'vecinfer', 'spectral', 'kivi', 'kivi_sink', 'svdq', 'kitty', 'adakv', 'xquant', 'kvquant', 'palu', 'cachegen', 'minicache', 'gear', 'zipcache', 'snapkv', 'streaming_llm', 'h2o', 'tova', 'pyramidkv', 'squeeze', 'chunkkv', 'cam', 'xkv', 'nsnquant', 'knorm', 'skvq', 'qfilters', 'keyformer', 'morphkv', 'kvzip', 'kvtc', 'curdkv', 'nestedkv', 'amc', 'a2ats', 'anchorkv', 'rocketkv', 'age_tiered'] = 'turboquant_rvq'

head_dim: int = 128

bit_width_inlier: int | list = 2

bit_width_outlier: int | None = None

jl_dim: int | None = None

n_outlier_channels: int | None = None

n_calib_tokens: int | None = None

enable_vectorized_attend: bool = True

enable_outlier_two_stream: bool = False

enable_fused_query_dot: bool = False

seed: int = 42

dtype: Any = None

capacity: int | None = None

sliding_window: int | None = None

store: ArtifactStore | None = None

observers: list

options: object | None = None

key_sub_dim: int = 4

value_sub_dim: int = 8

key_codebook_bits: int = 12

value_codebook_bits: int = 8

residual_length: int = 128

kivi_group_size: int = 32

svdq_rank: int | None = None

svdq_energy_threshold: float = 0.95

svdq_bit_schedule: tuple[int, ...] = (8, 4, 2, 1, 1, 0, 0, 0)

svdq_group_size: int = 32

kitty_hi_fraction: float = 0.25

kitty_hi_bit: int = 4

kitty_lo_bit: int = 2

kitty_group_size: int = 32

adakv_target_avg_bits: float = 2.5

adakv_lo_bit: int = 2

adakv_mid_bit: int = 3

adakv_hi_bit: int = 4

adakv_group_size: int = 32

adakv_update_interval: int = 1

adakv_importance: str = 'norm_variance'

adakv_obs_window: int = 32

xquant_group_size: int = 2

xquant_base_bits: int = 2

xquant_residual_bits: int = 4

xquant_group_quant_size: int = 32

xquant_max_ctx: int = 8192

kvquant_bits: int = 3

kvquant_outlier_fraction: float = 0.01

kvquant_group_size: int = 32

kvquant_lloyd_iters: int = 8

kvquant_refit_interval: int = 0

kvquant_n_sink: int = 1

palu_rank: int | None = None

palu_energy_threshold: float = 0.9

palu_n_head_groups: int = 4

palu_hi_bit: int = 4

palu_lo_bit: int = 2

palu_hi_fraction: float = 0.25

palu_group_size: int = 32

palu_quantize_values: bool = True

cachegen_bits: int = 4

cachegen_group_size: int = 32

cachegen_use_delta: bool = True

cachegen_per_channel: bool = True

cachegen_layer_groups: int = 3

cachegen_resolved_bits: int | None = None

minicache_start_frac: float = 0.5

minicache_group_size: int = 2

minicache_retention_threshold: float = 0.9

minicache_slerp_t: float = 0.5

minicache_max_ctx: int = 8192

gear_bits: int = 2

gear_rank: int | None = None

gear_energy_threshold: float = 0.9

gear_sparse_fraction: float = 0.01

gear_group_size: int = 32

gear_quantize_values: bool = True

zipcache_hi_bits: int = 4

zipcache_lo_bits: int = 2

zipcache_hi_fraction: float = 0.2

zipcache_group_size: int = 32

zipcache_quantize_values: bool = True

snap_dtype: str = 'auto'

snap_batched_scoring: bool = True

snap_backend: str = 'auto'

snap_budget: int = 512

snap_obs_window: int = 32

snap_n_sink: int = 4

stream_n_sink: int = 4

stream_window_size: int = 512

h2o_budget: int = 512

h2o_n_sink: int = 4

h2o_rope_base: float = 10000.0

h2o_grace: int = 16

h2o_decay: float = 0.98

tova_budget: int = 512

tova_n_sink: int = 4

tova_backend: str = 'auto'

pyramid_budget: int = 512

pyramid_backend: str = 'reference'

pyramid_n_sink: int = 4

pyramid_beta: float = 2.0

pyramid_resolved_budget: int | None = None

squeeze_budget: int = 512

squeeze_n_sink: int = 4

squeeze_strength: float = 1.0

squeeze_resolved_budget: int | None = None

chunkkv_budget: int = 512

chunkkv_chunk_size: int = 8

chunkkv_n_sink: int = 4

chunkkv_score: str = 'attn_mass'

chunkkv_reuse_layers: int = 1

cam_budget: int = 512

cam_n_sink: int = 4

cam_merge: str = 'sim_weighted'

cam_merge_keys: bool = False

cam_merge_gate: bool = True

xkv_group_size: int = 2

xkv_rank: int | None = None

xkv_energy_threshold: float = 0.95

xkv_latent_bits: int = 4

xkv_group_quant_size: int = 32

xkv_max_ctx: int = 8192

nsn_bits: int = 2

nsn_residual_length: int = 64

nsn_codebook_size: int = 256

nsn_subvector_dim: int = 8

nsn_seed: int = 1234

nsn_max_ctx: int = 8192

knorm_budget: int = 512

knorm_n_sink: int = 4

knorm_recent: int = 0

knorm_keep: str = 'low'

skvq_bits_key: int = 2

skvq_bits_value: int = 2

skvq_group_size: int = 32

skvq_window: int = 128

skvq_n_sink: int = 5

skvq_reorder: bool = True

skvq_clip_search: bool = True

skvq_clip_alpha: float = 1.0

skvq_max_ctx: int | None = None

qfilters_budget: int = 512

qfilters_n_sink: int = 4

qfilters_recent: int = 0

qfilters_calib_tokens: int = 128

qfilters_sign: int = 1

qfilters_min_retention: float | None = 0.9

qfilters_on_low_retention: Literal['warn', 'raise', 'ignore'] = 'warn'

keyformer_budget: int = 512

keyformer_n_sink: int = 4

keyformer_recent: int = 0

keyformer_tau: float | None = None

keyformer_tau_init: float = 1.0

keyformer_tau_end: float = 1.0

keyformer_anneal_steps: int = 0

keyformer_rope_base: float = 10000.0

keyformer_seed: int = 0

morphkv_budget: int = 512

morphkv_n_sink: int = 4

morphkv_window: int = 8

kvzip_budget: int = 512

kvzip_n_sink: int = 4

kvzip_probe: str = 'context'

kvtc_bit_budget: int = 512

kvtc_bit_choices: tuple = (0, 1, 2, 3, 4, 6, 8)

kvtc_beta: float = 3.5

curdkv_budget: int = 512

curdkv_n_sink: int = 4

curdkv_rank_cap: int = 16

curdkv_rope_base: float = 10000.0

nestedkv_budget: int = 512

nestedkv_n_sink: int = 4

nestedkv_window: int = 64

nestedkv_beta: float = 3.0

nestedkv_tau: float = 0.6

nestedkv_kappa: float = 10.0

nestedkv_safeguard_alpha: float = 0.2

amc_k_high: float = 0.2

amc_k_mid: float = 0.3

amc_use_query_saliency: bool = False

amc_query_alpha: float = 0.5

amc_adaptive_thresholds: bool = False

amc_threshold_window: int = 64

amc_gamma: float = 0.1

amc_calib_variance: float | None = None

amc_group_size: int = 32

a2ats_codebook_bits: int = 8

a2ats_sub_dim: int = 8

a2ats_window: int = 128

a2ats_b: int = 2048

a2ats_use_query_aware: bool = True

a2ats_beta: float = 0.5

a2ats_retrieval_fraction: float = 0.2

a2ats_rope_base: float = 10000.0

a2ats_rope_scaling: Any = None

a2ats_query_h: Any = None

a2ats_codebook: Any = None

anchorkv_theta: float = 0.05

anchorkv_window: int = 32

anchorkv_rho: float = 0.7

anchorkv_anchor_frac: float = 0.0078125

anchorkv_residual_bits: int = 2

anchorkv_seed: int = 42

rocketkv_compression_ratio: float = 8.0

rocketkv_page_size: int | None = None

rocketkv_head_topk1: int | None = None

rocketkv_obs_window: int = 32

rocketkv_n_sink: int = 4

age_recent_boundary: int = 128

age_mid_boundary: int = 1024

age_bits_recent: int = 8

age_bits_mid: int = 4

age_bits_old: int = 2

age_group_size: int = 32

n_sink_tokens: int = 5

smooth_factors: Any = None

key_codebook: Any = None

value_codebook: Any = None

spectral_key_d_eff: int = 4

spectral_val_d_eff: int = 50

spectral_apply_qjl: bool = True

spectral_model_name: str = 'model'

use_metal_kernels: bool | None = None

fused_sdpa: bool | None = False

fused_sdpa_max_ctx: int | None = None

fused_sdpa_memory_bound: bool = False

class veloxquant_mlx.cache.KVCacheFactory[source]

Bases: object

Factory for creating KVCache instances from a KVCacheConfig.

static create(config)[source]

Instantiate a KVCache from the given configuration.

Parameters:: config (KVCacheConfig) – KVCacheConfig instance.
Return type:: KVCache | _MLXKVCache
Returns:: Configured KVCache. The concrete type depends on config.method – STANDALONE_METHODS produce VeloxQuant’s own KVCache (core. abstractions), every other method produces an mlx_lm.models. cache.KVCache subclass — see the STANDALONE_METHODS comment above.

class veloxquant_mlx.cache.CacheGenKVCache(config)[source]

Bases: KVCache

KV cache implementing CacheGen entropy-coded compression for one layer.

Parameters:

config (Any) –

KVCacheConfig. Fields consumed: cachegen_bits (int, default 4) — base/uniform bit-width. cachegen_resolved_bits (int or None) — per-layer bit-width injected

by KVCacheBuilder.for_model (§5.1.2 layer-wise schedule); falls back to cachegen_bits when None (uniform behaviour).

cachegen_group_size (int, default 32), cachegen_use_delta (bool, default True — token-delta transform), cachegen_per_channel (bool, default True — group entropy estimate

by channel per §5.1.3 instead of pooling all channels).

Notes

No .bits attribute — keeps mlx_lm SDPA on the clean fp16 path.

update_and_fetch(keys, values)[source]

Group-quantize K/V (KIVI-style), reporting entropy-coded byte estimates in addition to the fixed-width baseline; return dequantized K/V.

Parameters:

keys (mlx.core.array)
values (mlx.core.array)

property compressed_key_bytes: int: Entropy-coded (ideal arithmetic-coder) estimate for the key cache, modelling Shannon entropy of the token-delta code stream.

property compressed_value_bytes: int: Entropy-coded (ideal arithmetic-coder) estimate for the value cache, modelling Shannon entropy of the token-delta code stream.

property fixed_width_key_bytes: int: Naive fixed-width packed baseline for the key cache (no entropy coding).

property fixed_width_value_bytes: int: Naive fixed-width packed baseline for the value cache (no entropy coding).

property fp16_key_bytes: int: Hypothetical fp16 key cost if nothing were compressed.

property fp16_value_bytes: int: Hypothetical fp16 value cost if nothing were compressed.

property entropy_savings: float: Fraction of the fixed-width code size saved by entropy coding (key side).

property assigned_avg_bits: float: Effective key bit-width after entropy coding (vs fp16=16).

class veloxquant_mlx.cache.GEARKVCache(config)[source]

Bases: KVCache

KV cache implementing GEAR error-feedback compression for one layer.

Parameters:: config (Any) – KVCacheConfig. Fields consumed: gear_bits (int, default 2), gear_rank (int | None, default None → energy threshold), gear_energy_threshold (float, default 0.90), gear_sparse_fraction (float, default 0.01), gear_group_size (int, default 32), gear_quantize_values (bool, default True — GEAR values too).

Notes

No .bits attribute — keeps mlx_lm SDPA on the clean fp16 path. Single-layer (no coordinator); for_model propagates the gear_* fields automatically via dataclasses.replace.

update_and_fetch(keys, values)[source]

Compress K (and V, unless disabled) with GEAR’s base group quant + low-rank residual + sparse outlier correction; return reconstructed K/V in the input dtype.

Parameters:

keys (mlx.core.array)
values (mlx.core.array)

trim(n)[source]

Drop the n most-recent tokens and scale the byte/error counters to match.

Return type:: int
Parameters:: n (int)

property compressed_key_bytes: int

base codes + low-rank factors + sparse triples, all heads/batches).

Type:: Realized stored bytes for the compressed key cache (GEAR three-part

property compressed_value_bytes: int

base codes + low-rank factors + sparse triples, all heads/batches).

Type:: Realized stored bytes for the compressed value cache (GEAR three-part

property base_only_key_bytes: int: Base-layer-only key bytes (group quant codes alone, no error-feedback), the baseline GEAR compares against.

property base_only_value_bytes: int: Base-layer-only value bytes (group quant codes alone, no error-feedback), the baseline GEAR compares against.

property fp16_key_bytes: int: Hypothetical fp16 key cost if nothing were compressed.

property fp16_value_bytes: int: Hypothetical fp16 value cost if nothing were compressed.

property assigned_avg_bits: float: Effective key bit-width after error-feedback overhead (vs fp16=16).

property error_recovery_ratio: float

Fraction of the base quantization error removed by GEAR (key side).

1 - ||X - GEAR(X)||^2 / ||X - base(X)||^2. 0 = no recovery, →1 = near-lossless. The core GEAR claim, measured not asserted.

class veloxquant_mlx.cache.MiniCacheKVCache(config, role='primary', group_id=0, coordinator=None, n_readers=1)[source]

Bases: KVCache

KV cache implementing MiniCache cross-layer merging for one layer.

Parameters:

config (Any) – KVCacheConfig. Fields consumed: minicache_retention_threshold (float, default 0.9), minicache_slerp_t (float, default 0.5).
role (str) – "primary" or "merge" (default "primary").
group_id (int) – Cross-layer merge group this layer belongs to.
coordinator (MiniCacheCoordinator | None) – Shared MiniCacheCoordinator (None → degenerate).
n_readers (int) – Number of merge layers in this primary’s group (group_size - 1). Only meaningful for role="primary"; tells the coordinator how many fetches to expect before a published entry can be reclaimed.

property merge

update_and_fetch(keys, values)[source]

Primary: publish true KV to the coordinator and pass through losslessly. Merge: SLERP-merge against the paired primary and reconstruct from the merge.

Parameters:

keys (mlx.core.array)
values (mlx.core.array)

property role: str

"primary" or "merge".

Type:: This layer’s cross-layer merge role

property group_id: int: Cross-layer merge group this layer belongs to.

property compressed_key_bytes: int: Realized stored bytes for the compressed key cache (merged-token magnitudes + retained fp16 vectors + shared direction where applicable).

property compressed_value_bytes: int: Realized stored bytes for the compressed value cache (merged-token magnitudes + retained fp16 vectors + shared direction where applicable).

property fp16_key_bytes: int: Hypothetical fp16 key cost if nothing were compressed.

property fp16_value_bytes: int: Hypothetical fp16 value cost if nothing were compressed.

property n_retained: int: Number of tokens retained as full fp16 vectors (above the cosine-similarity retention threshold) instead of merged.

property n_merged: int: Number of tokens SLERP-merged with the paired primary layer.

property retention_rate: float: Fraction of tokens retained as full fp16 vectors rather than merged.

class veloxquant_mlx.cache.PALUKVCache(config)[source]

Bases: KVCache

KV cache implementing PALU true-latent low-rank K/V compression.

Parameters:: config (Any) – KVCacheConfig. Fields consumed: head_dim (D), palu_rank (int | None — explicit rank; None → energy threshold), palu_energy_threshold (float, default 0.90), palu_n_head_groups (int, default 4), palu_hi_bit / palu_lo_bit / palu_hi_fraction (mixed-bit), palu_group_size (int, default 32), palu_quantize_values (bool, default True).

Notes

Does not expose .bits — mlx_lm’s SDPA checks hasattr(cache, "bits") to route to a quantized kernel; we keep that path clean and hand it reconstructed fp16. Exposes .assigned_avg_bits (max of key/value effective bit-widths).

update_and_fetch(keys, values)[source]

Fit group-head SVD projections on prefill, else project + mixed-bit quantize into the latent buffers; return reconstructed fp16 K/V.

Parameters:

keys (mlx.core.array)
values (mlx.core.array)

property offset: int: Number of tokens accumulated so far (own counter — parent ring buffer is bypassed).

size()[source]

Number of tokens currently stored.

Return type:: int

property state: Last reconstructed fp16 (keys, values), since the true-latent buffers bypass the parent ring buffer mlx_lm reads directly.

is_trimmable()[source]

Whether this cache supports trimming.

Return type:: bool

trim(n)[source]

Drop the n most-recently-added tokens from the latent buffers; returns the number actually dropped.

Return type:: int
Parameters:: n (int)

property nbytes: int: Total realized stored bytes (compressed key + value latents).

property compressed_key_bytes: int: Realized stored bytes for the compressed key cache (mixed-bit latents + projection basis, all heads/batches).

property compressed_value_bytes: int: Realized stored bytes for the compressed value cache (mixed-bit latents + projection basis, all heads/batches).

property fp16_key_bytes: int: Hypothetical fp16 key cost if nothing were compressed.

property fp16_value_bytes: int: Hypothetical fp16 value cost if nothing were compressed.

property projection_bytes: int: Per-group SVD projection basis (V + mu, fp32) overhead, shared across heads in a group.

property rank: int: Latent rank actually stored (uniform across head groups).

property assigned_avg_bits: float: Effective bit-width — max of the key and value latent rates.

property merge

class veloxquant_mlx.cache.PolarQuantKVCache(config)[source]

Bases: KVCache

KV cache backed by PolarQuantizer for key compression.

Parameters:: config (Any) – KVCacheConfig instance.

append_key(k)[source]

Encode and cache a single key vector.

Parameters:: k (Any) – Key vector, shape (d,), fp16.
Return type:: None

append_value(v)[source]

Quantize and cache a value vector.

Parameters:: v (Any) – Value vector, shape (d,), fp16.
Return type:: None

attend(q)[source]

Compute attention output for a query.

Parameters:: q (Any) – Query vector, shape (d,), fp16.
Return type:: Any
Returns:: Attention output, shape (d,), fp16.

memory_bytes()[source]

Bytes held for all cached tokens (angle codes, radii, int8 values and scales).

Angle codes are stored one uint8 per angle, not packed to b bits.

Return type:: int

reset()[source]

Clear all stored tokens; the key quantizer (seeded, deterministic polar rotation) is untouched.

Return type:: None

class veloxquant_mlx.cache.QJLKVCache(config)[source]

Bases: KVCache

Minimal KV cache using pure 1-bit QJL for key compression.

Parameters:: config (Any) – KVCacheConfig instance. jl_dim sets the sketch size m; when unset it defaults to head_dim. The estimator is unbiased but its variance falls as 1/m, so m = head_dim gives noisy scores (correlation with exact q.k ~0.64 at d=128, ~0.88 at m = 4 * head_dim). Set jl_dim higher when attention quality matters.

append_key(k)[source]

Encode and store a key vector.

Parameters:: k (Any) – Key vector, shape (d,), fp16.
Return type:: None

append_value(v)[source]

Quantize and store a value vector.

Parameters:: v (Any) – Value vector, shape (d,), fp16.
Return type:: None

attend(q)[source]

Compute attention output.

Parameters:: q (Any) – Query vector, shape (d,), fp16.
Return type:: Any
Returns:: Attention output, shape (d,), fp16.

memory_bytes()[source]

Bytes held for all cached tokens.

Signs are stored one per int8 (m bytes per key), not bit-packed, so the key side is 8x what a true 1-bit store would need.

Return type:: int

reset()[source]

Clear all stored tokens; the key quantizer (seeded, deterministic JL projection) is untouched.

Return type:: None

class veloxquant_mlx.cache.SlidingWindowKVCache(inner, window_size)[source]

Bases: KVCache

Wraps any KVCache with sliding-window token eviction.

Maintains a RingBuffer of the most recent window_size token slots. When the window is full, the oldest token is evicted.

Because the inner cache does not support random deletion, we maintain the window by resetting the inner cache to empty and re-feeding it the current window’s tokens each time one is evicted. This simple strategy is correct but rebuilds the cache every window_size tokens; suitable for inference not training.

We cache all key-value vectors in our own ring buffer so they’re available to re-feed on each window advance.

Parameters:

inner (KVCache) – The underlying KVCache to wrap.
window_size (int) – Number of tokens to keep.

append_key(k)[source]

Buffer a key vector for paired insertion with the next value.

Parameters:: k (Any) – Key vector, shape (d,), fp16.
Return type:: None

append_value(v)[source]

Insert a key-value pair and evict the oldest token if full.

Parameters:: v (Any) – Value vector, shape (d,), fp16.
Return type:: None

attend(q)[source]

Delegate to the inner cache for attention computation.

Parameters:: q (Any) – Query vector, shape (d,), fp16.
Return type:: Any
Returns:: Attention output, shape (d,), fp16.

memory_bytes()[source]

Return memory usage of the inner (windowed) cache.

Return type:: int

reset()[source]

Clear the window and the inner cache, returning both to empty.

Return type:: None

class veloxquant_mlx.cache.TurboQuantKVCache(config)[source]

Bases: KVCache

KV cache backed by TurboQuantProd (or TurboQuantMSE) for key compression.

Keys are stored bit-packed: indices use b_mse bits per coordinate and QJL signs use 1 bit each — matching the paper’s memory model. Values use per-token int8 with fp16 scale.

Parameters:: config (Any) – KVCacheConfig instance.

append_key(k)[source]

Encode and cache a single key vector (bit-packed storage).

Parameters:: k (Any) – Key vector, shape (d,), fp16.
Return type:: None

append_value(v)[source]

Quantize and cache a single value vector (int8 per-token).

Parameters:: v (Any) – Value vector, shape (d,), fp16.
Return type:: None

attend(q)[source]

Compute attention output for a query vector.

Unpacks bit-packed keys, estimates attention scores, and returns the weighted sum of decoded values.

Parameters:: q (Any) – Query vector, shape (d,), fp16.
Return type:: Any
Returns:: Attention output, shape (d,), fp16.

memory_bytes()[source]

Return actual memory of bit-packed key-value storage.

Return type:: int
Returns:: Total bytes occupied by all cached data.

reset()[source]

Clear all stored tokens; the key quantizer (seeded, deterministic rotation/codebook) is untouched. The outlier detector, if enabled, is re-created rather than preserved: its calibrated channels reflect the token stream observed so far, and after a reset that stream starts over (e.g. a fresh sliding-window rebuild) — reusing stale channel picks from evicted tokens would be wrong, not just stale.

Return type:: None

class veloxquant_mlx.cache.ZipCacheKVCache(config)[source]

Bases: KVCache

KV cache implementing ZipCache-adapted per-token mixed-precision for one layer.

Parameters:: config (Any) – KVCacheConfig. Fields consumed: zipcache_hi_bits (int, default 4) — bit-width for salient tokens, zipcache_lo_bits (int, default 2) — bit-width for non-salient tokens, zipcache_hi_fraction (float, default 0.20) — fraction of tokens at hi_bits, zipcache_group_size (int, default 32) — token group size for quant, zipcache_quantize_values (bool, default True) — apply to values too.

Notes

No .bits attribute — keeps mlx_lm SDPA on the clean fp16 path. Single-layer (no coordinator); for_model propagates the zipcache_* fields automatically via dataclasses.replace. Values are quantized uniformly at hi_bits (saliency routing is key-driven; values follow the hi-bit path as the safer default).

update_and_fetch(keys, values)[source]

Sort key tokens by L2-norm saliency, quantize the top hi_fraction at hi_bits and the rest at lo_bits (values uniformly at hi_bits); return reconstructed fp16 K/V.

Parameters:

keys (mlx.core.array)
values (mlx.core.array)

property compressed_key_bytes: int: Realized stored bytes for the compressed key cache (mixed hi/lo-bit saliency-routed codes, all heads/batches).

property compressed_value_bytes: int: Realized stored bytes for the compressed value cache (uniform hi_bits codes, all heads/batches).

property baseline_key_bytes: int: Uniform lo-bit baseline for comparison (no saliency routing).

property baseline_value_bytes: int: Uniform lo-bit baseline for comparison (no saliency routing).

property fp16_key_bytes: int: Hypothetical fp16 key cost if nothing were compressed.

property fp16_value_bytes: int: Hypothetical fp16 value cost if nothing were compressed.

property compression_ratio: float: fp16 bytes / compressed bytes (key side). >1 means storage savings.

property effective_avg_bits: float: Average key bits/element implied by the stored mixed-bit rate.

property merge

class veloxquant_mlx.cache.SnapKVKVCache(config)[source]

Bases: KVCache

KV cache implementing SnapKV-adapted prefill eviction for one layer.

Parameters:: config (Any) – KVCacheConfig. Fields consumed: snap_budget (int, default 512) — max tokens retained after prefill, snap_obs_window (int, default 32) — trailing keys used as proxy queries, snap_n_sink (int, default 4) — initial positions always kept.

Notes

No .bits attribute — stores and returns fp16 K/V directly. Single-layer (no coordinator); for_model propagates all snap_* fields automatically via dataclasses.replace. Eviction happens at every prefill chunk (S > 1) — mlx_lm’s chunked prefill (prefill_step_size) calls update_and_fetch once per chunk for prompts longer than one chunk, so budget enforcement must re-run over the accumulated kept set, not just each chunk in isolation (see #84). Decode tokens (S == 1) are always kept. Accumulated decode tokens grow the cache beyond the initial budget in the decode phase — consistent with the paper’s design (the budget constrains prefill history, not the decode stream).

property offset: int

True absolute token position (NOT the retained row count).

While the base class’s update_and_fetch is on the stack this yields the retained row count instead, so its cursor arithmetic and return slice stay correct.

update_and_fetch(keys, values)[source]

Prefill: score and evict down to the retention budget (SnapKV window-attention proxy), deferred to after this step’s own attention. Decode: pass through unevicted.

Parameters:

keys (mlx.core.array)
values (mlx.core.array)

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

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

The inherited KVCache.make_mask returns the string "causal" whenever offset == 0 (the very first prefill call, before any eviction has run) or otherwise builds a mask keyed off row index rather than true position — both wrong once eviction has made the kept keys a non-contiguous subset. Before the first prefill call (_kept_positions is still None, nothing evicted yet) this falls back to the base class’s behavior exactly, since a plain trailing-window mask is correct there.

Called before this step’s update_and_fetch, whose attention runs over the stored rows plus this step’s N new ones (eviction is deferred past the return), so the key axis is len(kept) + N.

Parameters:

N (int)
return_array (bool)
window_size (int | None)

property evicted_key_bytes: int: Bytes stored for kept key rows (fp16).

property evicted_value_bytes: int: Bytes stored for kept value rows (fp16).

property full_key_bytes: int: Hypothetical fp16 key cost without any eviction.

property full_value_bytes: int: Hypothetical fp16 value cost without any eviction.

property tokens_kept: int: Total token positions retained across all heads and steps.

property tokens_total: int: Total token positions seen (before eviction) across all heads and steps.

property eviction_ratio: float: full_fp16_bytes / kept_fp16_bytes; > 1 means storage savings.

property keep_rate: float: Fraction of tokens retained (tokens_kept / tokens_total).

is_trimmable()[source]

False: trim() reads a corrupted self.offset and both writes and returns the wrong thing, worse than a mere bookkeeping rollback.

The base KVCache.trim(n) does n = min(self.offset, n); self.offset -= n. Outside update_and_fetch (_in_base is False), this class’s offset property returns _true_offset — the absolute token position used for RoPE — not the retained row count (_row_offset) the base class’s buffer actually holds. So n is clamped against the wrong (larger) number, and the setter then writes _row_offset -= n, which can leave _row_offset larger than the number of rows ever written whenever eviction has dropped any tokens (the normal case once prefill has run). The next update_and_fetch then returns a slice reaching past the real data into the pre-allocated buffer’s stale/uninitialized rows, silently feeding garbage into attention rather than crashing. Reproduced directly: budget=4, 20 real prefill tokens (so _row_offset == 4, _true_offset == 20); trim(3) leaves _row_offset == 17 though only 4 rows are real, and the following decode step returns an 18-row slice (see VeloxQuant-Studio issue #27).

Return type:: bool

property merge

class veloxquant_mlx.cache.StreamingLLMKVCache(config)[source]

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)[source]

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, **_)[source]

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()[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 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

class veloxquant_mlx.cache.H2OKVCache(config)[source]

Bases: KVCache

KV cache implementing H2O-adapted heavy-hitter oracle eviction for one layer.

Parameters:: config (Any) – KVCacheConfig. Fields consumed: h2o_budget (int, default 512) — maximum tokens retained at any time, h2o_n_sink (int, default 4) — leading positions never evicted, h2o_rope_base (float, default 10000.0) — RoPE frequency base, must match the model’s own attention RoPE base for the defensive _fix_incoming_rope fallback (see that method) to correct an offset/position desync correctly, should one ever occur; no longer used for post-eviction remapping, which was removed (#609), h2o_grace (int, default 16) — most-recently-arrived tokens protected from eviction, giving each new token this many update steps to accumulate real attention mass before it becomes eviction-eligible. Fixes a real, confirmed-on-real-models problem (see module docstring): without this, a brand-new token’s starting score of 0.0 makes it almost always the eviction target the instant the cache is full, freezing the kept set on whichever tokens filled the budget first. h2o_decay (float, default 0.98) — multiplicative per-step decay on existing scores, applied before new attention mass is added. Fixes a second real, confirmed-on-real-models problem (see module docstring): without this, scores only ever grow, so old tokens permanently outrank newer ones regardless of current relevance, thinning the kept window into a sparse, gapped trickle instead of a contiguous local context even with grace fixing the freeze. Must be in (0, 1]; 1.0 disables decay.

Notes

No .bits attribute — stores and returns fp16 K/V directly. Both prefill (S > 1) and decode (S == 1) tokens go through the same H2O update loop — unlike SnapKV-adapted, there is no prefill-only phase. Single-layer (no coordinator); KVCacheBuilder.for_model() propagates all h2o_* fields automatically via dataclasses.replace. The per-head state is lazily initialised on the first call to update_and_fetch when shapes are known. Does NOT delegate to the base mlx_lm KVCache.update_and_fetch: that implementation assumes self.offset always equals the number of rows physically stored, which eviction breaks. Instead self.keys/self.values hold exactly the n_kept retained rows and self.offset is kept equal to the true absolute step count, so mlx_lm’s self.rope(x, offset=cache.offset) call — made on both the query and the next incoming key, entirely inside the model’s attention module, before this cache ever sees them — rotates at the correct position without this cache needing to intercept the query at all. .state/.size()/is_trimmable() are overridden to keep this contract explicit rather than relying on the base class’s coincidentally-compatible slicing behavior (see #83 for the historical reason .state must stay valid across chunked prefill).

update_and_fetch(keys, values)[source]

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

Manages self.keys / self.values / self.offset directly instead of delegating to the base class’s append-only buffer, because that buffer assumes offset always equals the stored row count — false here the moment eviction drops a row. Keeping self.offset equal to the true absolute step count (not the kept-row count) is what keeps mlx_lm’s attention module rotating the next query and incoming key at the correct position — see the class/module docstrings for why getting this wrong silently corrupted generation.

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) for THIS call’s own attention — the full, un-evicted concatenation of whatever was stored before this call plus the S new tokens (see #370 below), NOT capped at h2o_budget. What gets stored afterward (self.keys / self.values, visible to the next call) is capped at h2o_budget 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 eviction, that already-fixed mask would silently desync from the shape it was built for (VeloxQuant-MLX#370). So eviction is deferred: this call returns the full pre-eviction concatenation (matching the mask make_mask already built from the previous call’s true kept positions — see that method), and only self.keys/self.values shrink, for the next call’s make_mask to reflect correctly.

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

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 eviction, which update_and_fetch defers past this step’s return anyway — see its docstring). self._kept_positions holds the true positions of whatever self.keys already stores from the previous call, which is exactly update_and_fetch’s prev_keys_flat this call 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()[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

size()[source]

Rows actually stored (n_kept), NOT self.offset (the true absolute step count) — the base class conflates the two, which is exactly the assumption this cache must not make. Overridden explicitly rather than relying on the base class’s coincidentally harmless behavior here.

Return type:: int

property state

(keys, values) — always exactly the n_kept stored rows.

Not delegated to the base class’s getter: it slices self.keys[..., :self.offset, :], and since self.offset is the true absolute step count (usually > n_kept once eviction has happened), that would silently clamp to all stored rows anyway via MLX’s slice semantics — correct by coincidence, not by contract. Made explicit here instead.

property h2o_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 / h2o_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).

class veloxquant_mlx.cache.TOVAKVCache(config)[source]

Bases: KVCache

KV cache implementing TOVA-adapted current-step attention eviction for one layer.

Parameters:: config (Any) – KVCacheConfig. Fields consumed: tova_budget (int, default 512) — maximum tokens retained at any time, tova_n_sink (int, default 4) — leading positions never evicted. tova_backend (str, default “auto”) — auto, mlx, metal, or reference.

Notes

No .bits attribute — stores and returns fp16 K/V directly. Both prefill (S > 1) and decode (S == 1) tokens go through the same TOVA update loop — unlike SnapKV-adapted, there is no prefill-only phase. Single-layer (no coordinator); KVCacheBuilder.for_model() propagates all tova_* fields automatically via dataclasses.replace. The per-head state is lazily initialised on the first call to update_and_fetch when shapes are known. 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 TOVA 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) for THIS call’s own attention — the full, un-evicted concatenation of whatever was stored before this call plus the S new tokens (see #370 below), NOT capped at tova_budget. What gets stored afterward (self.keys / self.values, visible to the next call) is capped at tova_budget 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 eviction, that already-fixed mask would silently desync from the shape it was built for (VeloxQuant-MLX#370). So eviction is deferred: this call returns the full pre-eviction concatenation (matching the mask make_mask already built from the previous call’s true kept positions — see that method), and only self.keys/self.values (and self._bh_positions/self._kept_positions) shrink, for the next call’s make_mask to reflect correctly.

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

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 eviction, which update_and_fetch defers past this step’s return anyway — see its docstring). self._kept_positions holds the true positions of whatever self.keys already stores from the previous call, which is exactly update_and_fetch’s previous_k this call 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)

size()[source]

Stored rows; offset separately tracks the absolute RoPE position.

Return type:: int

property state: Stored (keys, values); evicted positions are not separately encoded.

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 tova_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 / tova_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

class veloxquant_mlx.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

class veloxquant_mlx.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

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

Bases: KVCache

KV cache implementing ChunkKV-adapted chunk-level eviction for one layer.

Parameters:

config (Any) –
KVCacheConfig. Fields consumed: chunkkv_budget (int, default 512) — max tokens kept (sinks incl.). chunkkv_chunk_size (int, default 8) — eviction granularity C;

1 reduces to H2O-adapted exactly.

chunkkv_n_sink (int, default 4) — leading positions never evicted. chunkkv_score (str, default “attn_mass”) — “attn_mass” | “key_norm”.
layer_id (int | None) – This layer’s index (used to resolve leader/follower role and report/query the coordinator). None for single-cache construction — the layer then always runs its own eviction (as if it were its own leader).
coordinator (ChunkKVIndexReuseCoordinator | None) – Shared ChunkKVIndexReuseCoordinator, or None. When present and layer_id is not the block’s leader, this layer reuses the leader’s kept-token positions instead of evicting itself.

Notes

update_and_fetch(keys, values)[source]

Absorb new K/V tokens, apply ChunkKV 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) for THIS call’s own attention — the full, un-evicted concatenation of whatever was stored before this call plus the S new tokens (see #370 below), NOT capped at chunkkv_budget or chunk-aligned. What gets stored afterward (visible to the next call) is capped/chunk-aligned 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 chunk eviction, that already-fixed mask would silently desync from the shape it was built for (VeloxQuant-MLX#370). So eviction is deferred: this call returns the full pre-eviction concatenation (matching the mask make_mask already built from the previous call’s true kept positions — see that method), and only the stored per-head states shrink, for the next call’s make_mask to reflect correctly.

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

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 eviction, which update_and_fetch defers past this step’s return anyway — see its docstring). self._kept_positions holds the true positions of whatever every head’s state already stores 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()[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 cache’s per-layer token budget.

property chunk_size: int: This cache’s eviction granularity C (diagnostic).

property is_index_reuse_leader: bool: True if this layer runs its own eviction (leader, or reuse disabled).

property chunkkv_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 / chunkkv_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).

class veloxquant_mlx.cache.CaMKVCache(config)[source]

Bases: KVCache

KV cache implementing CaM-adapted cache-merging eviction for one layer.

Parameters:

config (Any) –

KVCacheConfig. Fields consumed: cam_budget (int, default 512) — maximum tokens retained. cam_n_sink (int, default 4) — leading positions never evicted. cam_merge (str, default “sim_weighted”) — “sim_weighted” |

”mean” | “drop”; “drop” reduces bit-for-bit to H2O-adapted.

cam_merge_keys (bool, default False) — merge keys too (values are: always merged).
cam_merge_gate (bool, default True) — apply the paper’s Eq. 14: Bernoulli merge gate (whether to merge at all, not just how strongly). False unconditionally merges every over-budget loser (the paper’s ablated “w.o. Merge Mask” configuration).

seed (int) — base seed for the gate’s deterministic draws.

Notes

No .bits attribute — stores and returns fp16 K/V directly. Both prefill (S > 1) and decode (S == 1) tokens go through the same merge loop. Per-head state is lazily initialised on the first call. Because CaM merges (not drops) it always trims to exactly budget — the output is rectangular [B, H, budget, D] once past budget, so no cross-head alignment is needed. 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 merge 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 CaM merge-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) for THIS call’s own attention — the full, un-evicted concatenation of whatever was stored before this call plus the S new tokens (see #370 below), NOT capped at cam_budget. What gets stored afterward (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 merge- eviction, that already-fixed mask would silently desync from the shape it was built for (VeloxQuant-MLX#370). So eviction is deferred: this call returns the full pre-eviction concatenation (matching the mask make_mask already built from the previous call’s true kept positions — see that method), and only the stored per-head states shrink, for the next call’s make_mask to reflect correctly.

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

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 eviction, which update_and_fetch defers past this step’s return anyway — see its docstring). self._kept_positions holds the true positions of whatever every head’s state already stores 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()[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 cache’s per-layer token budget.

property merge_mode: str: This cache’s merge disposition (diagnostic).

property merge_gate: bool: Whether the Eq. 14 Bernoulli merge gate is active (diagnostic).

property cam_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 / cam_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).

class veloxquant_mlx.cache.RocketKVKVCache(config)[source]

Bases: DeferredEvictionMixin, KVCache

KV cache implementing RocketKV-adapted two-stage compression for one layer.

Parameters:

config (Any) –

KVCacheConfig. Fields consumed: rocketkv_compression_ratio (float, default 8.0) — overall

target compression ratio c; drives the adaptive stage-1/ stage-2 split (paper §3.6). Must be > 1.

rocketkv_page_size (Optional[int], default None) — HSA page: size; None derives it from the adaptive split instead of a fixed value.
rocketkv_head_topk1 (Optional[int], default None) — head-dim: channels kept per HSA step; None derives it from the adaptive split (head_dim / head_dim_ratio).
rocketkv_obs_window (int, default 32) — stage-1 SnapKV: observation window.
rocketkv_n_sink (int, default 4) — stage-1 SnapKV sink: tokens always kept.

Notes

No .bits attribute — stores and returns fp16 K/V directly. Single-layer (no coordinator); for_model propagates all rocketkv_* fields automatically via dataclasses.replace. Stage 1 (eviction) runs once at prefill, exactly like SnapKVKVCache. Stage 2 (HSA) runs every decode step and narrows the returned K/V subset without shrinking the stored buffer — decode tokens accumulate in storage (paper’s stage 1 only compresses the input prompt), but each attention call only touches the HSA- selected pages plus the running decode tail.

property offset: int

True absolute token position (NOT the retained row count).

See SnapKVKVCache.offset for why the two must not be conflated — stage 1 here drops tokens exactly the same way.

property state: Expose retained rows, excluding the base cache’s capacity padding.

update_and_fetch(keys, values)[source]

Prefill: run stage-1 SnapKV eviction per head and build paged HSA summaries. Decode: append tokens exactly and update summaries incrementally.

Parameters:

keys (mlx.core.array)
values (mlx.core.array)

select_indices(query, b, h, keep_recent=0)[source]

HSA stage-2: approximate top-k page selection for one (b, h) query.

Exposed for callers that want to actually exercise sparse attention with this cache (the base update_and_fetch return value stays dense — see the module docstring’s storage-vs-fetch distinction).

Parameters:

query (array) – [D] proxy query vector (key-as-query convention).
b (int) – Batch and head index into the stored paged summaries.
h (int) – Batch and head index into the stored paged summaries.
keep_recent (int) – Trailing token count always included regardless of HSA score (e.g. the newest decode tokens).

Return type:

array

Returns:

[n_selected] int32 token indices, ascending, deduplicated, covering the HSA-selected pages and the trailing keep_recent window.

is_trimmable()[source]

False: trim() would only roll back base-class offset bookkeeping, not the internal stage-1 kept-token accounting or stage-2 paged max/min summaries, silently desynchronizing both from the trimmed buffer on the next call.

Return type:: bool

property stage1_bytes: int: Bytes stored for kept fp16 K + V rows after stage-1 eviction.

property stage2_aux_bytes: int: Auxiliary bytes for HSA’s paged max/min summaries (fp16-packed).

property full_fp16_bytes_total: int: Hypothetical fp16 K + V cost without any compression.

property compression_ratio: float: full_fp16_bytes / (stage1_bytes + stage2_aux_bytes); > 1 means savings.

property tokens_kept: int: Total tokens retained after stage-1 eviction (plus all decode tokens, which are always kept).

property tokens_total: int: Total tokens seen across prefill and decode, before stage-1 eviction.

property keep_rate: float: Fraction of tokens retained after stage-1 eviction (tokens_kept / tokens_total).

property stage1_ratio: float: Configured stage-1 (SnapKV) eviction ratio — prefill budget is S / stage1_ratio.

property page_size: int: Configured HSA page size used to build the paged max/min summaries.

property head_topk1: int: Resolved stage-2 per-head top-k1 count used by HSA approximate scoring.

veloxquant_mlx.cache package¶

Per-layer bit-widths (RateQuant)¶

Submodules¶