veloxquant_mlx package¶
veloxquant_mlx — KV cache quantization for Apple Silicon MLX.
Implements TurboQuant, TurboQuantRVQ, PolarQuant, and QJL plus the RateQuant per-layer bit allocator for production LLM inference.
Preview APIs: AutoOptimizer (automatic
hardware-aware strategy selection, RFC #469) and
CacheRoutePlanner (rate-aware session
routing) are exported here for convenience but are subject to change
without a major version bump — see their module docstrings
(veloxquant_mlx.planning, veloxquant_mlx.routing.cacheroute)
for what “preview” means for each.
- class veloxquant_mlx.KVCacheBuilder[source]
Bases:
objectFluent 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:
- with_head_dim(d)[source]
Set the attention head dimension.
- Parameters:
d (
int) – Head dimension (must be a power of 2).- Return type:
- 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:
- with_jl_dim(m)[source]
Set the JL projection dimension.
- Parameters:
m (
int) – Must be <= head_dim.- Return type:
- with_n_outlier_channels(n)[source]
Set the number of outlier channels to detect.
- Parameters:
n (
int) – Must be < head_dim.- Return type:
- with_n_calib_tokens(n)[source]
Set calibration token count for outlier activation.
- Return type:
- Parameters:
n (int)
- with_vectorized_attend(enabled=True)[source]
Enable vectorized packed-key unpack in attend().
- Return type:
- Parameters:
enabled (bool)
- with_outlier_two_stream(enabled=True)[source]
Enable outlier/inlier split cache after calibration.
- Return type:
- Parameters:
enabled (bool)
- with_fused_query_dot(enabled=True)[source]
Enable fused rotated-query + codebook-dot path.
- Return type:
- Parameters:
enabled (bool)
- with_precision(dtype)[source]
Set the compute dtype.
- Parameters:
dtype (
Any) – MLX dtype (e.g. mx.float16).- Return type:
- with_capacity(max_tokens)[source]
Set the maximum number of tokens to store.
- Parameters:
max_tokens (
int) – Positive integer.- Return type:
- with_artifact_store(store)[source]
Provide an ArtifactStore for loading precomputed artifacts.
- Parameters:
store (
ArtifactStore) – ArtifactStore instance.- Return type:
- with_observer(observer)[source]
Attach a QuantizationObserver.
- Parameters:
observer (
QuantizationObserver) – Observer instance.- Return type:
- 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:
- 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_inlieris alist[int], elementiis used for layeri. 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:
- param config:
KVCacheConfig specifying method, bit_width_inlier, seed, etc. head_dim is overridden per-layer.
- rtype:
- returns:
List of KVCache instances, one per language-model layer.
- raises QuantizerConfigError:
If
config.methodis a standalone method (seeSTANDALONE_METHODS) that does not implement the mlx_lm KVCache serving contract.
- Parameters:
config (KVCacheConfig)
- Return type:
- class veloxquant_mlx.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:
objectConfiguration 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 lengthn_layersfor per-layer RateQuant-style allocation. When passed as a list,KVCacheBuilder.for_model()consumes elementifor layeri;KVCacheFactory.create()requires an int.bit_width_outlier – Bit-width for outlier channels (None → same as inlier).
jl_dim – JL projection dimension m.
Nonemeansm = head_dimfor 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 fromoptionspopulate the matching flat fields at construction time (see__post_init__), so every existing cache class – which reads the flat fields viagetattr(config, "h2o_budget", 512)– keeps working unchanged. Setting a flat per-method field directly still works but emits aDeprecationWarningnaming itsoptionsreplacement.
- 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_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_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
- enable_vectorized_attend: bool = True
- enable_outlier_two_stream: bool = False
- enable_fused_query_dot: bool = False
- seed: int = 42
- dtype: Any = None
- store: ArtifactStore | None = None
- observers: list
- 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_energy_threshold: float = 0.95
- 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_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
- 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_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
- squeeze_budget: int = 512
- squeeze_n_sink: int = 4
- squeeze_strength: float = 1.0
- 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_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
- qfilters_budget: int = 512
- qfilters_n_sink: int = 4
- qfilters_recent: int = 0
- qfilters_calib_tokens: int = 128
- qfilters_sign: int = 1
- qfilters_on_low_retention: Literal['warn', 'raise', 'ignore'] = 'warn'
- keyformer_budget: int = 512
- keyformer_n_sink: int = 4
- keyformer_recent: int = 0
- 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_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_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'
- fused_sdpa_memory_bound: bool = False
- class veloxquant_mlx.KVCacheFactory[source]
Bases:
objectFactory 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.ArtifactStore[source]
Bases:
ABCDAO interface for loading and saving precomputed quantization artifacts.
- abstractmethod load_rotation_matrix(d, seed)[source]
Load a precomputed rotation matrix.
- Parameters:
- Return type:
- Returns:
MLX array of shape (d, d), fp16.
- Raises:
ArtifactNotFoundError – If the artifact does not exist.
- abstractmethod save_rotation_matrix(Pi, d, seed)[source]
Persist a rotation matrix.
- abstractmethod load_codebook(distribution, b, d)[source]
Load a precomputed codebook.
- Parameters:
- Return type:
- Returns:
MLX array of shape (2^b,), fp16.
- Raises:
ArtifactNotFoundError – If the artifact does not exist.
- abstractmethod save_codebook(cb, distribution, b, d)[source]
Persist a codebook.
- abstractmethod load_jl_matrix(d, m, seed)[source]
Load a precomputed JL projection matrix.
- Parameters:
- Return type:
- Returns:
MLX array of shape (m, d), fp16.
- Raises:
ArtifactNotFoundError – If the artifact does not exist.
- abstractmethod save_jl_matrix(S, d, m, seed)[source]
Persist a JL projection matrix.
- abstractmethod exists(artifact_type, **kwargs)[source]
Check whether a specific artifact exists in the store.
- class veloxquant_mlx.KVCache[source]
Bases:
ABCAbstract base class for KV cache implementations.
- abstractmethod append_key(k)[source]
Append a new key vector to the cache.
- abstractmethod append_value(v)[source]
Append a new value vector to the cache.
- abstractmethod attend(q)[source]
Compute attention-weighted value for a query.
- abstractmethod memory_bytes()[source]
Return current memory footprint of the cache in bytes.
- Return type:
- reset()[source]
Clear all stored tokens, returning the cache to its empty state.
Not abstract: only concrete storage-owning caches (TurboQuant, PolarQuant, QJL, Spectral, and decorators that delegate to one — SlidingWindowKVCache, PooledKVCache) implement this. A decorator with an unrelated, already-established meaning for “reset” (e.g. KVCacheProfiler.reset(), which clears only its own accumulated profiling stats and deliberately leaves the wrapped cache untouched) is not required to satisfy this contract, so this default raises rather than being abstract.
Implementers must clear only token storage (ring buffers / lists / counters) and leave any quantizer, calibration, or configuration state untouched — callers such as
SlidingWindowKVCacherely onreset()to rebuild an emptied cache without re-running (or losing) prior calibration.- Return type:
- class veloxquant_mlx.Quantizer[source]
Bases:
ABCAbstract base class for all vector quantizers.
All public-facing quantizer objects must satisfy this interface.
- abstractmethod encode(x)[source]
Encode a batch of vectors into a compact representation.
- Parameters:
x (
Any) – Input array of shape (batch, d), fp16.- Return type:
- Returns:
EncodedVector containing the compressed representation.
- abstractmethod decode(ev)[source]
Reconstruct approximate vectors from an EncodedVector.
- Parameters:
ev (
EncodedVector) – Encoded representation produced by encode().- Return type:
- Returns:
Reconstructed array of shape (batch, d), fp16.
- abstractmethod estimate_inner_product(q, ev)[source]
Estimate inner products between a query and encoded keys.
- Parameters:
q (
Any) – Query vector, shape (d,) or (1, d), fp16.ev (
EncodedVector) – Encoded key cache, batch_size == n_keys.
- Return type:
- Returns:
Estimated inner products, shape (batch_size,), fp16.
- class veloxquant_mlx.QuantizationObserver[source]
Bases:
ABCObserver for pipeline events (timing, memory, distortion).
- class veloxquant_mlx.EncodedVector(quantizer_type, batch_size, dim, indices=None, norm=None, signs=None, residual_norm=None, angles=None, final_radius=None, outlier_idx=None, outlier_encoded=None, inlier_encoded=None)[source]
Bases:
objectTyped output of Quantizer.encode().
Different quantizer types populate different subsets of fields.
- Variables:
quantizer_type – Registry key of the producing quantizer.
batch_size – Number of vectors encoded.
dim – Original vector dimensionality.
indices – Scalar codebook indices, shape (batch, d) uint8.
norm – Per-vector L2 norm, shape (batch,) fp16.
signs – QJL sign bits, shape (batch, m) int8.
residual_norm – QJL residual norm, shape (batch,) fp16.
angles – PolarQuant level angles, list of (batch, d/2^ℓ) arrays.
final_radius – PolarQuant scalar radius, shape (batch,) fp16.
outlier_idx – Outlier channel indices for CompositeQuantizer.
outlier_encoded – Nested EncodedVector for outlier channels.
inlier_encoded – Nested EncodedVector for inlier channels.
- Parameters:
- quantizer_type: str
- batch_size: int
- dim: int
- outlier_encoded: EncodedVector | None = None
- inlier_encoded: EncodedVector | None = None
- class veloxquant_mlx.QuantizationContext(x_original, mode, x_current, norm=None, rotated=None, indices=None, signs=None, residual_norm=None, angles=None, final_radius=None, outlier_idx=None, packed_bits=None, metadata=<factory>)[source]
Bases:
objectMutable payload flowing through a QuantizationHandler chain.
- Variables:
x_original – Original input vectors, shape (batch, d).
mode – Whether the chain is encoding or decoding.
x_current – Working copy mutated by each handler.
norm – L2 norm stored by NormalizationHandler, shape (batch,).
rotated – Vector after RotationHandler, shape (batch, d).
indices – Codebook indices, shape (batch, d) uint8.
signs – QJL sign bits, shape (batch, m) int8.
residual_norm – Residual L2 norm for QJL stage, shape (batch,).
angles – Per-level polar angles from PolarTransformHandler.
final_radius – Scalar radius after all polar levels, shape (batch,).
outlier_idx – Channel positions of outlier channels.
packed_bits – Bit-packed index array from BitPackingHandler.
metadata – Arbitrary stage-specific metadata.
- Parameters:
- x_original: Any
- mode: Literal['encode', 'decode']
- x_current: Any
- class veloxquant_mlx.TransformResult(angles, final_radius, n_levels)[source]
Bases:
objectOutput of RecursivePolarTransform.forward().
- Variables:
angles – List of angle arrays, one per polar level. angles[0] has shape (batch, d/2), angles[ℓ] has shape (batch, d/2^(ℓ+1)).
final_radius – Scalar radius at the end of recursion, shape (batch,).
n_levels – Number of polar recursion levels applied.
- Parameters:
- final_radius: Any
- n_levels: int
- exception veloxquant_mlx.ArtifactNotFoundError[source]
Bases:
FileNotFoundErrorRaised when a required precomputed artifact is missing from the store.
- exception veloxquant_mlx.BlockPoolExhaustedError[source]
Bases:
RuntimeErrorRaised when a BlockPoolAllocator has no free blocks left to allocate.
- exception veloxquant_mlx.CodebookDimensionMismatch[source]
Bases:
ValueErrorRaised when a codebook’s shape does not match the expected dimension.
- exception veloxquant_mlx.CyclicPipelineError[source]
Bases:
RuntimeErrorRaised when a QuantizationGraph contains a cycle.
- exception veloxquant_mlx.QuantizerConfigError[source]
Bases:
ValueErrorRaised when a quantizer or KV cache is misconfigured.
- class veloxquant_mlx.BlockPoolAllocator(config=None)[source]
Bases:
objectFixed-size block allocator for KV-cache storage.
Pre-allocates
n_blocksfixed-size blocks up front and hands them out on request, avoiding per-token malloc/free during generation. Freed blocks return to a free list and are reused by later allocations (LIFO — a just-freed block tends to be cache-hot), so steady-state generation performs zero new-memory allocation once the pool is warm.Keys and values can be tracked on independent free lists (
PoolConfig.separate_kv=True, the default) so K-heavy and V-heavy workloads (e.g. asymmetric bit-widths) don’t compete for the same blocks, or share a single free list for tighter packing when K/V pressure is uneven across requests.Thread-safe: every mutating and read method acquires an internal lock, so one pool can be shared across threads/concurrent request handlers without external synchronization. The lock is a plain (non-reentrant)
threading.Lock— do not call back into the pool from inside a callback invoked while holding it.This class only manages block bookkeeping (which block ids are free, who owns which, how full each is, fragmentation/latency stats). It does not itself own array storage — pair it with a backing store (e.g.
veloxquant_mlx.memory.mlx_storage.MLXBlockStorage) that allocates the actual mx.array buffers indexed by block id.- Parameters:
config (
PoolConfig|None) – Pool sizing/behavior configuration.
Example:
pool = BlockPoolAllocator(PoolConfig(block_size=16, n_blocks=256)) blocks = pool.allocate(stream="k", n_tokens=40, owner=request_id, format="int2") ... pool.free_all(owner=request_id)
- register_owner(owner)[source]
Mark
owneras active, guarding against accidental id reuse.allocate()calls this automatically for any owner not yet seen, but a second allocate() for an owner that is already active is treated as that same request continuing to grow, not a collision (see allocate()’s docstring) — it can’t otherwise tell the two cases apart. Call register_owner() explicitly yourself (once, before any allocate() calls for this owner) if you want a genuine collision — e.g. a request-id counter that wrapped around, or two callers racing on the same id — to raise immediately, including on a second register_owner() call for a still-active id.- Raises:
OwnerAlreadyActiveError – If
owneralready has blocks checked out or was registered and not yet released.- Return type:
- Parameters:
owner (int)
- release_owner(owner)[source]
Clear
ownerfrom the active-owner set without freeing blocks.Rarely needed directly —
free_all()calls this once an owner’s last block is freed. Useful if you registered an owner viaregister_owner()but never actually allocated anything for it.
- allocate(stream, n_tokens, owner, format='fp16')[source]
Allocate enough blocks to hold
n_tokensslots.- Parameters:
stream (
str) – “k” or “v” (ignored, treated as “kv”, when the pool was built withseparate_kv=False).n_tokens (
int) – Number of token slots needed.owner (
int) – Opaque id (e.g. request id) the blocks are checked out to. Repeated allocate() calls for an owner that is already active on this pool are treated as that same request growing (the normal case — see PoolBackedKVCache) and never raise. To catch a different caller accidentally reusing an id that is still checked out, callregister_owner()yourself before the first allocate() — a second register_owner() for an id that is still active raises OwnerAlreadyActiveError.format (
str) – Compression-format tag to record on each block.
- Return type:
- Returns:
List of newly allocated
Blockinstances, in order.- Raises:
BlockPoolExhaustedError – If fewer than the required number of free blocks remain and the pool cannot grow enough to cover the gap (see
PoolConfig.grow_on_exhaustion/max_blocks). No blocks are allocated on failure (all-or-nothing).ValueError – If
n_tokens<= 0 orstreamis invalid.
- free(blocks)[source]
Return blocks to the free list for their stream.
- free_all(owner)[source]
Free every block currently checked out to
owner.Unlike
free(), which frees an arbitrary caller-supplied subset and must locate+remove each block id from the owner’s list, this already has the owner’s complete, correct block-id list – so it pops that list once (O(1)) instead of doing N membership checks + Nlist.remove()calls (VeloxQuant-MLX#510; was 13.4ms at 256K tokens / 16K blocks, confirmed superlinear: 2x blocks -> 2.0-3.4x time).
- blocks_for(owner)[source]
Return the blocks currently checked out to
owner, in allocation order.
- n_free(stream='k')[source]
Number of free blocks remaining for
stream.
- shrink(stream, target_free)[source]
Permanently retire fully-free blocks from
streamdown totarget_free.Use this to give memory back when a pool was sized for peak load that turned out much higher than steady-state usage. Only blocks currently on the free list are eligible — in-use blocks are never touched, so this never breaks a live request. Retired blocks are removed from the pool’s bookkeeping entirely (not just marked free/unfree), so any external backing storage indexed by block id (e.g.
MLXBlockStorage) may drop the corresponding buffer.- Parameters:
- Return type:
- Returns:
Number of blocks actually retired.
- Raises:
ValueError – If
target_free< 0.
- mark_used(block, n_used)[source]
Update
block.n_usedand keep fragmentation stats consistent.Callers that track a block’s fill level themselves (e.g.
PooledKVCache, which appends token-by-token to an already-checked-out block) must go through this instead of writingblock.n_useddirectly –AllocationStats._fragmented_countis maintained incrementally (not recomputed by a full scan every allocate()/free()/shrink() call, which would be O(n_blocks) on every single-block growth step in the decode hot path), so it only stays correct if every write to a block’s fill level is observed here.
- class veloxquant_mlx.PoolConfig(block_size=16, n_blocks=1024, separate_kv=True, grow_on_exhaustion=False, max_blocks=None)[source]
Bases:
objectConfiguration for a
BlockPoolAllocator.- Variables:
block_size – Number of token slots stored per block.
n_blocks – Total number of blocks the pool manages initially.
separate_kv – If True, keys and values are allocated from independent free lists (K and V blocks are never interchanged). If False, a single free list backs both keys and values, which packs tighter when K and V pressure is uneven across requests.
grow_on_exhaustion – If True, an allocate() call that would otherwise raise BlockPoolExhaustedError instead grows the exhausted stream’s block count just enough to satisfy the request (see max_blocks). If False (default), exhaustion always raises.
max_blocks – Upper bound on total blocks across all streams once grown. None means unbounded growth. Ignored when grow_on_exhaustion is False. Growth that would exceed this cap still raises BlockPoolExhaustedError.
- Parameters:
- block_size: int = 16
- n_blocks: int = 1024
- separate_kv: bool = True
- grow_on_exhaustion: bool = False
- class veloxquant_mlx.PooledKVCache(inner, pool, owner=-1, format='fp16')[source]
Bases:
KVCacheWraps any VeloxQuant KVCache with block-pool-backed memory accounting.
The wrapped cache still owns and drives its own storage arrays (its compression format may be anything registered with
KVCacheFactory); this wrapper checks out and returns fixed-size blocks from a sharedBlockPoolAllocatorin lock-step with token appends, so a multi-request server can track allocation counts, reuse, and fragmentation across every active cache from one pool instead of each cache doing its own ad hoc growth.Block granularity is
pool.config.block_sizetokens: a new block is requested from the pool only when the previous one fills up, and every block owned by this cache is returned to the pool in one call when the request completes (release()). No block is ever allocated or freed on a per-token basis, matching the “avoid malloc/free during generation” goal in issue #249.- Parameters:
inner (
KVCache) – The underlying KVCache that performs actual compression.pool (
BlockPoolAllocator) – Shared BlockPoolAllocator to draw block accounting from.owner (
int) – Opaque id identifying this cache’s request/sequence. Two PooledKVCache instances must use different owners to avoid each other’srelease()reclaiming the wrong blocks.format (
str) – Compression-format tag recorded on each block this cache checks out (purely informational — passed through topool.allocate(format=...)).
- append_key(k)[source]
Check out a K block from the pool if needed, then append to the inner cache.
- append_value(v)[source]
Check out a V block from the pool if needed, then append to the inner cache.
- attend(q)[source]
Delegate attention computation to the inner cache.
- release()[source]
Return every block this cache holds to the pool.
Call this once the request finishes (or is evicted) so its blocks become available for reuse by later requests.
- Return type:
- class veloxquant_mlx.QuantizerFactory[source]
Bases:
objectFactory for creating Quantizer instances via the QuantizerRegistry.
All quantizer construction should go through this factory.
Example:
q = QuantizerFactory.create("qjl", d=128, m=128, seed=42) q = QuantizerFactory.create("turboquant_mse", d=128, b=2, seed=42) q = QuantizerFactory.create("turboquant_prod", d=128, b=3, m=128, seed=42) q = QuantizerFactory.create("polar", d=128, b=2, seed=42)
- static create(method, d, b=2, m=None, seed=42, store=None, **kwargs)[source]
Instantiate a Quantizer by method name.
- Parameters:
method (
Literal['qjl','turboquant_mse','turboquant_prod','polar']) – Algorithm name registered in QuantizerRegistry.d (
int) – Vector dimension.b (
int) – Bit-width (default 2).m (
int|None) – JL projection dimension (default = d for QJL methods).seed (
int) – Random seed.store (
ArtifactStore|None) – ArtifactStore for loading precomputed artifacts. If None, artifacts are generated on-the-fly.**kwargs (
Any) – Additional keyword arguments forwarded to the quantizer.
- Return type:
- Returns:
Configured Quantizer instance.
- Raises:
QuantizerConfigError – If parameters are invalid.
KeyError – If method is not registered.
- veloxquant_mlx.allocate_bits_ratequant(sensitivities, target_avg_bits, beta=3.5, bit_choices=(1, 2, 3))[source]
Allocate per-layer bit-widths via RateQuant Theorem 2.
Continuous solution (paper, eq. 2):
b_i = b̄ + (ln w_i − ln_w_bar) / ln(β)
rounded to the nearest integer in
bit_choices, then re-balanced via greedy +1/−1 adjustments so the integer total exactly matchesround(target_avg_bits * N).- Parameters:
sensitivities – Iterable of per-layer sensitivity weights
w_i > 0, typically fromcalibrate_layer_sensitivities().target_avg_bits (
float) – Desired mean bits/dim across layers. May be fractional (e.g. 1.5) — the integer allocation will straddle it.beta (
float) – Distortion-rate decay constant for the underlying quantizer. Paper-reported values: 3.5 for TurboQuant, 5.0 for KIVI/QuaRot. Mismatched β can invert allocation ordering — see paper Section 4.3.bit_choices (
tuple) – Allowed integer bit-widths. RVQ supports {1, 2, 3, …}.
- Return type:
- Returns:
list[int] of length
len(sensitivities). Pass directly asKVCacheConfig.bit_width_inlierand consume viaKVCacheBuilder.for_model().
- veloxquant_mlx.calibrate_layer_sensitivities(model, tokenizer, prompts=None, seq_len=256, verbose=False)[source]
Run a calibration forward pass and return per-layer sensitivity.
- Parameters:
model – Loaded mlx_lm model (e.g. from
mlx_lm.load()).tokenizer – Matching tokenizer.
prompts (
list|None) – Optional list of calibration strings. Defaults to 8 general-domain prompts spanning history, science, and CS.seq_len (
int) – Max tokens per prompt (truncated if longer).verbose (
bool) – Whether to print per-sequence progress.
- Return type:
- Returns:
list[float] of length
n_attention_layers, each strictly > 0. Higher values indicate layers whose key cache is more error-prone at fixed bit-width and should get more bits.
- veloxquant_mlx.fit_distortion_curve(head_dim, bit_choices=(1, 2, 3), seed=0, n_samples=64)[source]
Fit
D(b) = α·β^(-b)on synthetic unit-norm Gaussian keys.- Return type:
- Returns:
Tuple
(alpha, beta). ForTurboQuantRVQat d=128 the paper reports β ≈ 3.5; you will get a similar value from this function. For production use, just passbeta=3.5toallocate_bits_ratequant()and skip the fit.- Parameters:
- class veloxquant_mlx.VecInferKVCache(config)[source]
Bases:
KVCacheKV cache implementing VecInfer’s dual-transform product VQ.
- Parameters:
config (
Any) –KVCacheConfigwith VecInfer-specific fields populated byKVCacheFactory.create. Required fields:head_dim,key_codebook_bits,value_codebook_bits,key_sub_dim,value_sub_dim. Optional:smooth_factors(numpy or mx array; identity if absent),key_codebook,value_codebook(random init if absent — for tests only).
Notes
Storage layout deliberately delegates concatenation to mlx_lm’s base
_MLXKVCache. We quantize + dequantize on the way in so that subsequent SDPA calls see a standard fp16 key tensor.Never exposes
.bits— mlx_lm’s SDPA checkshasattr(cache, "bits")to route to a different kernel path. We expose.assigned_avg_bitsinstead.- update_and_fetch(keys, values)[source]
Apply smooth-scale + Hadamard rotation to keys, product-VQ-quantize K/V, and dequantize for downstream SDPA (or stash indices only when memory-bound fused mode is active).
- property state
parent fp16 (keys, values). Fused mode: raw (key_indices, value_indices) ring buffers, since no fp16 K_hat/V_hat is materialized.
- Type:
Standard mode
- empty()[source]
Whether the cache holds no tokens yet (checks the index ring buffer in fused mode, the parent buffer otherwise).
- Return type:
- trim(n)[source]
Drop the last
ntokens, keeping every parallel buffer in step.The inherited
trimonly lowersoffset. In fused mode the index ring buffers are advanced by_stored_S_kvinstead, so left alone they keep the trimmed tokens: the next token is writtennslots pastoffsetandfused_sdpaattends over the stale tail. The byte accounting is scaled back with the token count (an estimate: it is tracked per call, not per token).
- property nbytes: int
the parent fp16 K_hat/V_hat buffers plus, in fused mode, the pre-allocated uint32 index buffers (stash mode holds both, so it uses more memory than plain fp16).
- Type:
Bytes actually held
- fused_sdpa(q, scale, *, causal=True, sliding_window=0)[source]
Compute attention directly from compressed K/V indices.
- Parameters:
q (
array) –[B, H_q, S_q, D]uncompressed queries — not yet smooth or Hadamard transformed. This method applies the dual transform internally so callers (mlx_lm SDPA dispatcher) pass exactly the same tensor they’d pass to standard SDPA.scale (
float) – Attention scale (typically1/sqrt(D)).causal (
bool) – Apply causal mask.sliding_window (
int) – If > 0, sliding-window attention with the given window width.
- Return type:
array- Returns:
[B, H_q, S_q, D]attention output in the dtype ofq.
- property compressed_key_bytes: int
see
nbytes.- Type:
Estimated packed bytes for the compressed key cache (product-VQ codes at
bitseach, all heads/batches; excludes codebook). Not what is held in memory
- property fp16_key_bytes: int
Hypothetical fp16 key cost if nothing were compressed.
- property compressed_value_bytes: int
see
nbytes.- Type:
Estimated packed bytes for the compressed value cache (product-VQ codes at
bitseach, all heads/batches; excludes codebook). Not what is held in memory
- property fp16_value_bytes: int
Hypothetical fp16 value cost if nothing were compressed.
- property codebook_bytes: int
Static codebook overhead in bytes (fp16 storage).
- property assigned_avg_bits: float
Effective bits/element averaged over keys and values.
Excludes codebook overhead (amortized across many tokens); for an end-to-end byte ratio use compressed_*_bytes / fp16_*_bytes.
- property merge
- veloxquant_mlx.apply_dual_transform_keys(K, smooth, H)[source]
Apply
K_tilde = (K / lambda) @ H(smooth then Hadamard).- Parameters:
K (
array) – Keys, shape[..., head_dim]. Per-head smoothing assumes the second-to-last axis indexes heads whensmoothis 2D.smooth (
array) – Either[head_dim](broadcast across heads) or[n_heads, head_dim](per-head); seecalibrate_smooth_factors().H (
array) – Walsh-Hadamard matrix[head_dim, head_dim].
- Return type:
array- Returns:
Transformed keys, same shape as input.
- veloxquant_mlx.apply_dual_transform_queries(q, smooth, H)[source]
Apply
q_tilde = (q * lambda) @ Hso q_tilde @ K_tilde.T == q @ K.T.- Parameters:
q (
array) – Queries, shape[..., head_dim]. Under GQA, the head axis (-3) may haveH_q = H_kv * n_repheads, a multiple ofsmooth’s per-KV-head row count — each KV head’s factor is repeated across its query group (see :issue:`616`), not averaged, so the dual transform still cancels againstapply_dual_transform_keys()’s per-KV-head factors.smooth (
array) – Same convention asapply_dual_transform_keys().H (
array) – Walsh-Hadamard matrix[head_dim, head_dim].
- Return type:
array- Returns:
Transformed queries, same shape as input.
- veloxquant_mlx.calibrate_smooth_factors(keys_calib, eps=0.0001)[source]
Compute per-(head, channel) smooth scaling factor from a key sample.
- Parameters:
keys_calib (
array) – Calibration keys with shape[n_tokens, n_heads, head_dim]or[n_tokens, head_dim](single-head).eps (
float) – Floor for the per-channel max to avoid divide-by-zero.
- Return type:
array- Returns:
Smooth factors shaped to match the input head layout (
[n_heads, head_dim]or[head_dim]). Formula:lambda_i = sqrt(max_t |K[t, ..., i]|)
- veloxquant_mlx.train_codebook(x, n_centroids, max_iter=30, seed=42)[source]
Train a VQ codebook on flat sub-vector samples.
- veloxquant_mlx.walsh_hadamard_matrix(d, dtype=mlx.core.float32)[source]
Construct an orthonormal Walsh-Hadamard matrix.
Recursive form from VecInfer Eq. 5:
H_1 = [[1]] H_{2k} = (1/sqrt(2)) * [[H_k, H_k], [H_k, -H_k]]
- Parameters:
d (
int) – Output dimension; must be a power of 2.dtype – MLX dtype for the returned matrix.
- Return type:
array- Returns:
Matrix of shape
[d, d]withH @ H.T == I.
- class veloxquant_mlx.KeyNormObserver[source]
Bases:
QuantizationObserverAccumulates per-token key L2 norm² for later RateQuant allocation.
Usage:
obs = KeyNormObserver() config.observers.append(obs) # ... run calibration prompts ... report = obs.report() # report.heterogeneity_ratio indicates RateQuant benefit
- on_event(event)[source]
Accumulate key_l2_norm_sq from the event’s metadata, if present.
- Return type:
- Parameters:
event (QuantizationEvent)
- report()[source]
Build a KeyNormReport from the norms accumulated so far.
- Return type:
- class veloxquant_mlx.KeyNormReport(n_tokens, mean_norm_sq, min_norm_sq, max_norm_sq)[source]
Bases:
objectSummary statistics produced by
KeyNormObserver.- n_tokens: int
- mean_norm_sq: float
- min_norm_sq: float
- max_norm_sq: float
- property heterogeneity_ratio: float
max / min— predictor of when mixed-precision allocation helps.Per RateQuant Theorem 3, when this ratio (geometric variant of AM/GM) is far above 1, per-layer bit allocation produces measurable gains over uniform quantization. Below ~2.0, savings are marginal.
- class veloxquant_mlx.KVCacheProfiler(cache, head_dim=None, layer_id=0)[source]
Bases:
KVCacheWraps a
KVCacheinstance, recording per-call timing and memory.Transparent decorator: implements the same append_key/append_value/attend/ memory_bytes interface as the wrapped cache, so it can be substituted anywhere a KVCache is expected (including inside a list built by
KVCacheBuilder.for_model()).- Parameters:
Example:
cache = KVCacheFactory.create(config) profiled = KVCacheProfiler(cache, head_dim=config.head_dim, layer_id=0) profiled.append_key(k) profiled.append_value(v) out = profiled.attend(q) print(format_profile_table([profiled.profile()]))
- append_key(k)[source]
Append a key to the wrapped cache, timing it as a quantize call.
- append_value(v)[source]
Append a value to the wrapped cache, timing it as a write call.
- attend(q)[source]
Run attention on the wrapped cache, timing it as a dequantize call.
- profile()[source]
Return the accumulated LayerProfile for this wrapped cache.
Reads the wrapped cache’s memory_bytes() here rather than after every append_key/append_value call: these caches are append-only (memory is monotonically non-decreasing, no eviction), so the latest reading is already the peak, and there is no need to pay for a recompute on every token.
- Return type:
- class veloxquant_mlx.MLXCacheProfiler(cache, layer_id=0)[source]
Bases:
objectWraps an
mlx_lm.models.cache.KVCache-style cache’s update_and_fetch.Servable methods (the ones
cli/serve.pyandcli/profile.pydrive through realmlx_lm.generate()runs) implementupdate_and_fetch(keys, values) -> (keys, values), a single fused call that quantizes, stores, and immediately dequantizes for the caller’s attention step — there is no separate append_key/append_value/attend triad to hook into, unlike the standalone interfaceKVCacheProfilerwraps. So this class times update_and_fetch as one unit and records it asquantize_ms_totalwithis_fused=True;dequantize_ms_totalandwrite_ms_totalstay 0 because there is nothing else to attribute them to. Memory comes from the cache’s ownnbytesproperty (the same true packed-byte accountingcli/serve.py’s accounting-only warning already describes — seeACCOUNTING_NOTEincli/profile.py).Substituted directly into the list
KVCacheBuilder.for_model()returns and passed tomlx_lm.generate(..., prompt_cache=profilers).- Parameters:
cache (Any)
layer_id (Any)
- update_and_fetch(keys, values)[source]
Forward to the wrapped cache’s update_and_fetch, timing it as one fused call.
- profile()[source]
Return the accumulated LayerProfile for this wrapped cache.
Reads the wrapped cache’s nbytes here rather than after every update_and_fetch call: nbytes is monotonically non-decreasing across calls, so the latest reading is already the peak, and there is no need to pay for a recompute on every token.
- Return type:
- class veloxquant_mlx.LayerProfile(layer_id, n_quantize_calls=0, n_dequantize_calls=0, quantize_ms_total=0.0, dequantize_ms_total=0.0, write_ms_total=0.0, peak_memory_bytes=0, tokens_written=0, fp16_baseline_bytes=0, is_fused=False)[source]
Bases:
objectAggregated profiling stats for a single KV-cache layer.
- Variables:
layer_id – Index or label identifying the layer.
n_quantize_calls – Number of append_key calls (quantization events).
n_dequantize_calls – Number of attend calls (dequantization events).
quantize_ms_total – Cumulative append_key wall time, milliseconds. For an
MLXCacheProfiler(fused update_and_fetch interface), this holds the combined quantize+dequantize+write time instead — seeis_fused. There is no separate number to split it into.dequantize_ms_total – Cumulative attend wall time, milliseconds. Always 0 when
is_fusedis True.write_ms_total – Cumulative append_value wall time, milliseconds. Always 0 when
is_fusedis True.peak_memory_bytes – memory_bytes()/nbytes of the wrapped cache as of the last profile() call. These caches are append-only (monotonically non-decreasing memory, no eviction), so the latest reading is the peak — read lazily at report time rather than after every call.
tokens_written – Number of append_key (or update_and_fetch) calls, a proxy for tokens stored.
fp16_baseline_bytes – What the incoming K+V would cost uncompressed (4 * head_dim bytes per token standalone; the actual K+V tensor bytes for MLXCacheProfiler), used to derive compression_ratio.
is_fused – True when quantize_ms_total is a combined measurement (MLXCacheProfiler) rather than quantize-only (KVCacheProfiler).
- Parameters:
- layer_id: Any
- n_quantize_calls: int = 0
- n_dequantize_calls: int = 0
- quantize_ms_total: float = 0.0
- dequantize_ms_total: float = 0.0
- write_ms_total: float = 0.0
- peak_memory_bytes: int = 0
- tokens_written: int = 0
- fp16_baseline_bytes: int = 0
- is_fused: bool = False
- property quantize_ms_mean: float
Average append_key (or update_and_fetch) latency in milliseconds.
- property dequantize_ms_mean: float
Average attend latency in milliseconds (always 0 when is_fused is True).
- property compression_ratio: float
fp16 baseline bytes divided by actual bytes (>1 means compressed).
- class veloxquant_mlx.ProfileReport(layers=<factory>, elapsed_s=0.0)[source]
Bases:
objectProfiling results across one or more layers.
- Variables:
layers – Per-layer profiles, in layer order.
elapsed_s – Total wall time covered by the profiling session.
- Parameters:
layers (list[LayerProfile])
elapsed_s (float)
- layers: list[LayerProfile]
- elapsed_s: float = 0.0
- property total_bytes_written: int
Sum of peak memory across all profiled layers.
- property total_tokens: int
Tokens processed by the model (not summed across layers).
Every layer sees every token, so summing per-layer counts would overstate the total by the layer count (#629); take the max instead.
- property tokens_per_sec: float
Total tokens divided by elapsed_s (0.0 if elapsed_s is unset).
- property overall_compression_ratio: float
Combined fp16-baseline-bytes over actual bytes across all layers.
- veloxquant_mlx.format_profile_table(report)[source]
Render a ProfileReport as a human-readable fixed-width table.
Matches the format proposed in issue #252:
Layer Quantize Dequantize Memory ------------------------------------------------ Layer 0 12.3 µs 8.1 µs 1.2 MB Layer 1 11.8 µs 7.9 µs 1.2 MB
- Parameters:
report (
ProfileReport) – A ProfileReport (or a bare list[LayerProfile]).- Return type:
- Returns:
Multi-line formatted string.
- veloxquant_mlx.profile_layers(profilers, elapsed_s=0.0)[source]
Aggregate a list of KVCacheProfiler instances into one ProfileReport.
- Parameters:
profilers (
list[KVCacheProfiler|MLXCacheProfiler]) – One profiler per model layer, in layer order.elapsed_s (
float) – Total wall-clock time of the profiled run, for tokens/sec.
- Return type:
- Returns:
A ProfileReport summarizing all layers.
- class veloxquant_mlx.AutoConfigResult(config, reason)[source]
Bases:
objectThe selected configuration plus the reasoning behind it.
- Variables:
config – The resulting KVCacheConfig, ready for KVCacheFactory.create() or KVCacheBuilder.for_model().
reason – Short human-readable explanation of why this method/bit-width was chosen, for logging or a CLI’s
--explainflag.
- Parameters:
config (KVCacheConfig)
reason (str)
- config: KVCacheConfig
- reason: str
- class veloxquant_mlx.HardwareInfo(total_memory_bytes=None, active_memory_bytes=0)[source]
Bases:
objectDescribes the machine the KV cache will run on.
- Variables:
total_memory_bytes – Total device memory (unified memory on Apple Silicon). None disables memory-pressure-based selection.
active_memory_bytes – Memory already in use before this cache is allocated (e.g. model weights, other requests’ caches).
- Parameters:
- active_memory_bytes: int = 0
- class veloxquant_mlx.WorkloadSpec(head_dim=128, seq_len=4096, n_layers=1, batch_size=1, n_kv_heads=1)[source]
Bases:
objectDescribes the job the KV cache needs to serve.
- Variables:
head_dim – Attention head dimension. Must be a positive power of 2, matching
KVCacheConfig’s own constraint.seq_len – Expected (or worst-case) sequence length in tokens. Drives the short-context / long-context precision tradeoff.
n_layers – Number of attention layers whose caches will share this configuration. Used only to size the memory-pressure estimate; defaults to 1 (single-layer estimate).
batch_size – Number of concurrent sequences. Used only to size the memory-pressure estimate; defaults to 1.
n_kv_heads – Number of KV heads per layer. Multiplies the fp16 K+V footprint used by the memory-pressure rule, so set it to the model’s real value (e.g. 8 for Llama-3-8B); the default of 1 keeps the older single-head estimate and undercounts by this factor for multi-head models (#631).
- Parameters:
- head_dim: int = 128
- seq_len: int = 4096
- n_layers: int = 1
- batch_size: int = 1
- n_kv_heads: int = 1
- veloxquant_mlx.detect_hardware_info()[source]
Auto-detect
HardwareInfofrom the running MLX device.Falls back to an empty (unknown-memory)
HardwareInfoif MLX is unavailable or device introspection fails — callers then get sequence-length-only selection instead of an exception.- Return type:
- veloxquant_mlx.select_kv_cache_config(workload, hardware=None)[source]
Pick a KVCacheConfig for the given workload and hardware.
Selection rules (in priority order):
Memory pressure: if
hardwarereports the workload’s estimated fp16 footprint would push total usage aboveMEMORY_PRESSURE_FRACTIONof device memory, pick the most aggressive method in the pool (gear, 2-bit) regardless of sequence length.Sequence length: short contexts (<
SHORT_CONTEXT_TOKENS) favor higher precision (turboquant_rvq, 4-bit); long contexts (>=LONG_CONTEXT_TOKENS) favor aggressive compression (kvquant, 3-bit with outlier isolation); everything in between uses a balanced default (kivi, 2-bit).Head dimension: if
head_dim >= LARGE_HEAD_DIM, the group size is doubled (min(head_dim, 64)) so each group still holds enough elements to amortize per-group scale/zero-point overhead.
- Parameters:
workload (
WorkloadSpec) – Description of the job (head_dim, seq_len, n_layers, batch_size, n_kv_heads).hardware (
HardwareInfo|None) – Description of the target machine. If None, auto-detects viadetect_hardware_info().
- Return type:
- Returns:
AutoConfigResult with the selected KVCacheConfig and a human-readable reason string.
- class veloxquant_mlx.AutoOptimizer(options=None)[source]
Bases:
objectOne-stop automatic KV-cache strategy recommender with caching and live validation.
Preview API: see the module-level warning above — recommendations are an analytical proxy (±20% memory error band, table-estimated hardware bandwidth), not a measured guarantee. Prefer a low
probe_top_nor a real benchmark database before production use.Purpose: Answers “which of the 43 methods should I use?” by profiling hardware/model, filtering incompatible methods, ranking by objective, and validating top choices against the live serve-tier to prevent crashes.
Deliberately stateful: Hardware detection and benchmark database are cached on the instance, so repeated calls don’t re-detect or re-probe. This makes it efficient for interactive CLI use and Jupyter notebooks.
Design: 1. Hardware caching: First
detect_hardware()call introspects once 2. Model profiling: Each call profiles the given config (no caching; configs vary) 3. Strategy planning: Filters, estimates memory, scores by objective 4. Serve-tier validation: Probes top N candidates to catch crash-tier methods 5. Fallback: If nothing survives, recommends library default (turboquant_rvq)- Usage:
>>> optimizer = AutoOptimizer() >>> result = optimizer.recommend_strategy(model_config=hf_config, objective="latency") >>> print(optimizer.explain(result))
- With benchmarks (optional speedup):
>>> from pathlib import Path >>> optimizer = AutoOptimizer(options=AutoOptimizerOptions( ... benchmark_db_dir=str(Path.home() / ".cache" / "veloxquant" / "benchmarks") ... ))
- Testing (no probing):
>>> optimizer = AutoOptimizer(options=AutoOptimizerOptions(probe_top_n=0)) >>> result = optimizer.recommend_strategy(model_config=...) # Fast, but may recommend crash-prone methods
- Attributes:
options: AutoOptimizerOptions (benchmark_db_dir, probe_top_n)
- Parameters:
options (AutoOptimizerOptions | None)
- detect_hardware()[source]
Detect (and cache) this machine’s HardwareProfile.
- Return type:
- profile_model(config=None, **overrides)[source]
Profile a model from a HF-style config dict + optional overrides.
- Return type:
- Parameters:
- estimate_memory(workload, model=None)[source]
Analytical per-method memory estimates for a workload.
- Return type:
- Parameters:
workload (WorkloadProfile)
model (ModelProfile | None)
- recommend_strategy(model_config=None, *, model=None, workload=None, objective=None, **kwargs)[source]
Recommend + validate the best KV-cache strategy for a model + workload.
Pipeline: 1. Profile model (from config or prebuilt ModelProfile) 2. Detect hardware (cached after first call) 3. Look up benchmarks in local DB (if enabled) 4. Filter incompatible methods (attention type, memory budget) 5. Score survivors on 4 axes (memory, latency, throughput, quality) 6. Rank by objective weights (memory vs latency vs quality) 7. Probe top N candidates with serve-tier (catch crashes) 8. Return recommendation + ranked alternatives + explanation
- Return type:
- Parameters:
model (ModelProfile | None)
workload (WorkloadProfile | None)
objective (str | None)
kwargs (Any)
- Arguments:
- model_config: HuggingFace-style config dict for architecture extraction
(num_layers, hidden_size, num_attention_heads, etc.). One of model_config or model must be provided.
model: Prebuilt ModelProfile; skips config parsing if given. workload: WorkloadProfile with context_length, batch_size, objective.
Defaults to balanced 4K context if not given.
- objective: Convenience override for workload.objective
(“memory”, “latency”, “quality”, or “balanced”).
- **kwargs: Additional PlanningOptions forwarded to the planner:
memory_budget_bytes: Hard cap on KV footprint
prefer_no_calibration: Exclude methods needing setup
require_metal: Exclude CPU-fallback implementations
- Returns:
RecommendationResult with: - recommendation: Top choice (ScoredMethod) - ranked: Full ranked list (1–3 alternatives) - fallback_used: True if no viable method survived - candidates: Filtered set (viable + excluded) - model/hardware/workload: The inputs (for explanation)
Examples:
Basic recommendation: >>> result = optimizer.recommend_strategy( … model_config={“num_layers”: 32, “hidden_size”: 4096, “num_attention_heads”: 32} … )
With objective override: >>> result = optimizer.recommend_strategy(model_config=config, objective=”latency”)
With memory budget: >>> result = optimizer.recommend_strategy( … model_config=config, … memory_budget_bytes=8 * 1024**3 # 8 GiB max … )
Get explanation: >>> print(optimizer.explain(result))
- explain(result, **kwargs)[source]
Render
explain()text.- Return type:
- Parameters:
result (RecommendationResult)
kwargs (Any)
- class veloxquant_mlx.AutoOptimizerOptions(benchmark_db_dir=None, probe_top_n=3)[source]
Bases:
objectConstructor options for
AutoOptimizer.- Variables:
benchmark_db_dir – Path to a
BenchmarkDatabasestorage directory; enabled only when explicitly given, so a stray dir never changes recommendations behind the user’s back.probe_top_n – How many of the ranked candidates are validated against the live serve-tier probe before committing.
0skips probing (faster, but a crash-prone standalone method could be recommended).
- Parameters:
- probe_top_n: int = 3
- class veloxquant_mlx.HardwareProfile(chip='unknown', chip_generation=0, total_memory_bytes=None, available_memory_bytes=0, mlx_version=None, macos_version=None, metal_available=True, peak_memory_bandwidth_gbps=None, avg_quantize_latency_ms_per_token=None)[source]
Bases:
objectComprehensive description of the Apple-Silicon environment.
Unlike the minimal
HardwareInfo, this is the object the auto-selection pipeline passes around: every field is a fact the planner can cite in an explanation.- Parameters:
- chip: str = 'unknown'
- chip_generation: int = 0
- available_memory_bytes: int = 0
- metal_available: bool = True
- classmethod detect()[source]
Build a HardwareProfile for the current machine (see detect_hardware_profile).
- Return type:
- class veloxquant_mlx.ModelProfile(model_id='unknown', architecture='unknown', num_layers=0, num_query_heads=0, num_kv_heads=0, head_dim=0, attention_type='gqa', dtype='float16', parameter_count=None, hidden_size=None, _extra=<factory>)[source]
Bases:
objectNormalized model-architecture description used by the planner.
- Variables:
model_id – Model identifier (any string; the caller’s
model_id).architecture – Short architecture slug (
"llama","qwen", …).num_layers – Number of attention-bearing layers.
num_query_heads – Attention heads per layer.
num_kv_heads – Key/value heads per layer (
== num_query_headsfor MHA).head_dim – Attention head dimension.
attention_type –
"mha","gqa", or"mqa".dtype – Expected compute dtype (
"float16", …).parameter_count – Total weight count, when the config reports one.
hidden_size – Model hidden width, when known (informational).
- Parameters:
- model_id: str = 'unknown'
- architecture: str = 'unknown'
- num_layers: int = 0
- num_query_heads: int = 0
- num_kv_heads: int = 0
- head_dim: int = 0
- attention_type: str = 'gqa'
- dtype: str = 'float16'
- property baseline_kv_bytes_per_token: int
Full K+V footprint in bytes for one token, one layer (fp16).
- baseline_kv_bytes(num_layers=None, batch=1)[source]
Total fp16 K+V footprint across layers/batch for one token.
- class veloxquant_mlx.WorkloadProfile(context_length=4096, generation_length=512, batch_size=1, num_concurrent_requests=1, objective='balanced', max_latency_ms=None)[source]
Bases:
objectNormalized description of a serving workload.
- Variables:
context_length – Maximum prompt context the served cache must hold.
generation_length – Expected number of decode tokens per request.
batch_size – Parallel requests sharing the KV cache.
num_concurrent_requests – Requests served concurrently (drives how many caches the memory budget must cover).
objective – One of the
WorkloadObjectivenames; the planner uses this to weight the memory/latency/throughput/quality axes.max_latency_ms – Hard latency ceiling, when the caller knows one (used to exclude strategies whose analytical proxy exceeds it).
- Parameters:
- context_length: int = 4096
- generation_length: int = 512
- batch_size: int = 1
- num_concurrent_requests: int = 1
- objective: str = 'balanced'
- property total_tokens_per_request: int
Context + generation, the cache lifetime of one request.
- property effective_batch: int
Attention batch the memory estimator must budget for.
- class veloxquant_mlx.WorkloadObjective[source]
Bases:
objectIdentifier constants for the planner’s supported objectives.
Kept as a plain object (not an Enum) so downstream code can pass arbitrary objective names; the planner’s
DEFAULT_OBJECTIVE_WEIGHTSonly assigns weights to the canonical ones below and falls back tobalancedfor anything else.- MEMORY = 'memory'
- LATENCY = 'latency'
- THROUGHPUT = 'throughput'
- QUALITY = 'quality'
- BALANCED = 'balanced'
- class veloxquant_mlx.RecommendationResult(objective, ranked, candidates, model, hardware, workload, fallback_used=False, fallback_reason=None)[source]
Bases:
objectFull outcome of a planning run, ready for the explainer / CLI.
- Parameters:
objective (str)
ranked (list[ScoredMethod])
candidates (CandidateFilterResult)
model (ModelProfile)
hardware (HardwareProfile)
workload (WorkloadProfile)
fallback_used (bool)
fallback_reason (str | None)
- objective: str
- ranked: list[ScoredMethod]
- candidates: CandidateFilterResult
- model: ModelProfile
- hardware: HardwareProfile
- workload: WorkloadProfile
- fallback_used: bool = False
- property best: ScoredMethod | None
Top-ranked method, or None if nothing survived filtering.
- class veloxquant_mlx.MemoryEstimate(method, baseline_bytes, compressed_bytes, workspace_bytes=0, peak_bytes=0, resident_bytes=0, confidence='high', assumptions=<factory>)[source]
Bases:
objectBreakdown of one method’s predicted KV-cache memory footprint.
- Variables:
method – Method name this estimate describes.
baseline_bytes – fp16 K+V for the full context x batch (the no-compression starting point).
compressed_bytes – Projected resident K+V after the method’s compression/eviction model.
workspace_bytes – Predicted non-tensor overhead (codebooks, allocator).
peak_bytes – Largest allocation during prefill (usually baseline).
resident_bytes – Footprint that stays resident during decode.
confidence –
"high"for curated-model methods,"medium"for a generic compression default,"low"for unmodeled methods.assumptions – Human-readable list of what the model assumed.
- Parameters:
- method: str
- baseline_bytes: int
- compressed_bytes: int
- workspace_bytes: int = 0
- peak_bytes: int = 0
- resident_bytes: int = 0
- confidence: str = 'high'
- property reduction_ratio: float
Fraction of baseline the method is projected to keep (0..1).
- property savings_percent: float
Percent of baseline bytes saved (0..100).
- class veloxquant_mlx.CacheRoutePlanner(n_shards, qcap, warm_slots_per_shard)[source]
Bases:
objectBuilds a periodic
RoutingTablefrom measured session rates.Implements Algorithm 1 from the paper: size each admitted session’s shard count from its rate relative to
qcap(Eq. 1), admit sessions in decreasing-rate order up to the warm-slot budget, then place admitted sessions with longest-processing-time-first (LPT) list scheduling so no shard’s expected load is far from the mean.This planner is stateless between calls — call
plan()once per control interval with fresh rates, and hold onto the returnedRoutingTablefor the whole interval rather than recomputing per request. Recomputing every interval, versus reusing a stale table, is a real tradeoff the paper measures directly (Appendix E): a fresh plan changes most owners’ shard assignments (high “cache churn”) even when only a few owners’ rates actually moved, which costs a rewarming dip; a stale plan drifts and loses some hit rate/latency instead. This class always recomputes fresh — callers wanting hysteresis should diff twoRoutingTableinstances themselves before switching.- Parameters:
n_shards (
int) – Number of pool shards to plan across (Rin the paper). Use1if the pool isn’t partitioned — every admitted owner then simply gets warm-set priority with no placement decision to make.qcap (
float) – Load cap per shard, in the same rate unit asSessionRate.rate(Section 3.1). Calibrate from a single shard’s measured latency/load knee, as the paper recommends — this module does not calibrate it for you.warm_slots_per_shard (
float) – Warm-set capacity per shard,Win the paper. Total admission budget isn_shards * warm_slots_per_shard, counted in the same units asrate/qcap(equal-slot admission model; see the paper’s Section 3.1 and 6 for the heterogeneous-prefix-size limitation this inherits).
Example:
planner = CacheRoutePlanner(n_shards=4, qcap=50.0, warm_slots_per_shard=200.0) table = planner.plan(estimator.rates()) shard = table.preferred_shard(owner, fallback=hash(owner) % planner.n_shards)
- plan(rates)[source]
Compute a fresh routing table from the given session rates.
- Parameters:
rates (
list[SessionRate]) – Estimated per-session rates for this control interval. Order does not matter; sessions are re-sorted internally by decreasing rate as the paper’s Algorithm 1 requires.- Return type:
- Returns:
A
RoutingTablevalid for the caller’s next control interval.
- class veloxquant_mlx.RateEstimator(half_life=20.0, max_owners=None, _counts=<factory>, _last_seen=<factory>, _clock=0)[source]
Bases:
objectTracks a per-owner exponentially-weighted request rate estimate.
A minimal helper for callers who don’t already track per-session rates: call
record()on every incoming request, thenrates()to get theSessionRatelist aCacheRoutePlannerneeds. Not required — any source of{owner: rate}works, including offline aggregate statistics as in the paper’s telemetry-derived workload.Memory is bounded by the number of distinct owners tracked (one float in
_countsand one int in_last_seeneach), not by how many requests have been recorded — an owner’s whole contribution to memory is released byforget(). That bound is only enforced if the caller callsforget()on session close; nothing here does it automatically. A caller that can’t guarantee that — e.g. a long-running server seeing many short-lived session IDs — should setmax_ownersso_countsself-bounds instead of growing without limit.- Parameters:
half_life (
float) – Number ofrecord()calls (from any owner) for an owner’s rate contribution to decay by half. Smaller values track bursts faster but are noisier; larger values are stable but slow to react to a session going cold or hot.max_owners (
int|None) – Maximum distinct owners to track at once.None(the default) preserves the original unbounded behavior. When set, recording a request from a not-yet-tracked owner that would push the owner count past this cap first evicts the single currently-lowest-rate owner (an approximation of “the coldest session,” not a full LRU) to make room. Does not add a branch to the already-tracked-owner path, so the hot case (an owner already in_counts) is unaffected._clock (int)
- half_life: float = 20.0
- record(owner)[source]
Record one request arrival for
owner.The estimate is an exponentially weighted count of the owner’s arrivals on a global clock that ticks once per recorded request (from any owner), so it is proportional to the owner’s arrival rate and decays while the owner is silent. Cost is one dict lookup, one
powand two dict writes for an already-tracked owner; the eviction scan only runs for a new owner whenmax_ownersis set.
- rate(owner)[source]
Current smoothed rate estimate for
owner(0.0 if never seen).
- rates()[source]
Snapshot every tracked owner’s current rate as a
SessionRatelist.- Return type:
- class veloxquant_mlx.RoutingTable(shards, expected_load, n_shards, admitted_rate_total, unadmitted_rate_total)[source]
Bases:
objectA fixed admission-and-placement plan for one control interval.
- Variables:
shards – Maps each admitted owner to the tuple of shard ids it may use, ordered by preference (most-loaded-eligible-first is not guaranteed; callers needing a single preferred shard should use
shards[owner][0]). An owner with more than one shard should split its load evenly across them, per the paper’skbreplication (Eq. 1) — an owner needs more than one shard only when its rate exceedsqcapon its own.expected_load – Per-shard expected load after this plan’s placement, indexed by shard id. Sums admitted owners’
rate / kbover every shard each owner was placed on, plus each shard’s cold-tail share of unadmitted traffic. For diagnostics and shadow-replay comparisons only — not enforced at dispatch time.n_shards – Total number of shards this plan was computed for.
admitted_rate_total – Sum of
rateover admitted owners, for computing what fraction of traffic the warm set actually covers.unadmitted_rate_total – Sum of
rateover owners considered but not admitted (the plan’s cold-tail fallback traffic).
- Parameters:
- n_shards: int
- admitted_rate_total: float
- unadmitted_rate_total: float
- is_admitted(owner)[source]
Whether
ownerhas a warm-set placement in this plan.
- preferred_shard(owner, fallback=0)[source]
The single best shard for
owner, orfallbackif not admitted.For an owner with
kb > 1shards, returns the currently least-loaded of its assigned shards (ties broken by lowest shard id), so repeated calls track load drift within the same interval even though the set of eligible shards stays fixed.
- class veloxquant_mlx.SessionRate(owner, rate)[source]
Bases:
objectA session’s estimated request rate, in requests per unit time.
- Variables:
owner – Opaque session/owner id, matching
BlockPoolAllocatorowner ids.rate – Estimated request rate (
lambda_bin the paper). Any consistent unit works as long asqcapuses the same unit.
- Parameters:
- owner: int
- rate: float
Public API¶
The names below are re-exported at the top level for convenience; each is documented once, at its defining module (linked from the subpackage pages below), to avoid duplicate entries.
|
Fluent builder for KVCache construction with validation. |
|
Configuration for a KVCache instance. |
|
Factory for creating KVCache instances from a KVCacheConfig. |
|
DAO interface for loading and saving precomputed quantization artifacts. |
|
Abstract base class for KV cache implementations. |
|
Abstract base class for all vector quantizers. |
|
Observer for pipeline events (timing, memory, distortion). |
|
Typed output of Quantizer.encode(). |
|
Mutable payload flowing through a QuantizationHandler chain. |
|
Output of RecursivePolarTransform.forward(). |
|
Raised when a required precomputed artifact is missing from the store. |
|
Raised when a BlockPoolAllocator has no free blocks left to allocate. |
|
Raised when a codebook's shape does not match the expected dimension. |
|
Raised when a QuantizationGraph contains a cycle. |
|
Raised when a quantizer or KV cache is misconfigured. |
|
Fixed-size block allocator for KV-cache storage. |
|
Configuration for a |
|
Wraps any VeloxQuant KVCache with block-pool-backed memory accounting. |
|
Factory for creating Quantizer instances via the QuantizerRegistry. |
|
Allocate per-layer bit-widths via RateQuant Theorem 2. |
|
Run a calibration forward pass and return per-layer sensitivity. |
|
Fit |
|
KV cache implementing VecInfer's dual-transform product VQ. |
|
Apply |
|
Apply |
|
Compute per-(head, channel) smooth scaling factor from a key sample. |
|
Train a VQ codebook on flat sub-vector samples. |
|
Construct an orthonormal Walsh-Hadamard matrix. |
|
Accumulates per-token key L2 norm² for later RateQuant allocation. |
|
Summary statistics produced by |
|
Wraps a |
|
Wraps an |
|
Aggregated profiling stats for a single KV-cache layer. |
|
Profiling results across one or more layers. |
|
Render a ProfileReport as a human-readable fixed-width table. |
|
Aggregate a list of KVCacheProfiler instances into one ProfileReport. |
|
The selected configuration plus the reasoning behind it. |
|
Describes the machine the KV cache will run on. |
|
Describes the job the KV cache needs to serve. |
|
Auto-detect |
|
Pick a KVCacheConfig for the given workload and hardware. |
|
One-stop automatic KV-cache strategy recommender with caching and live validation. |
|
Constructor options for |
|
Comprehensive description of the Apple-Silicon environment. |
|
Normalized model-architecture description used by the planner. |
|
Normalized description of a serving workload. |
|
Identifier constants for the planner's supported objectives. |
|
Full outcome of a planning run, ready for the explainer / CLI. |
|
Breakdown of one method's predicted KV-cache memory footprint. |
|
Builds a periodic |
|
Tracks a per-owner exponentially-weighted request rate estimate. |
|
A fixed admission-and-placement plan for one control interval. |
|
A session's estimated request rate, in requests per unit time. |
Subpackages¶
- veloxquant_mlx.allocators package
- veloxquant_mlx.artifacts package
- veloxquant_mlx.benchmarks package
- Submodules
- veloxquant_mlx.benchmarks.attend_benchmark module
- veloxquant_mlx.benchmarks.benchmark_db module
- veloxquant_mlx.benchmarks.cacheroute_benchmark module
- veloxquant_mlx.benchmarks.comm_vq_benchmark module
- veloxquant_mlx.benchmarks.falcon_rabitq_benchmark module
- veloxquant_mlx.benchmarks.metal_kernel_benchmark module
- veloxquant_mlx.benchmarks.model_kv_benchmark module
- veloxquant_mlx.benchmarks.rabitq_benchmark module
- veloxquant_mlx.benchmarks.workload_replay_benchmark module
- Submodules
- veloxquant_mlx.cache package
- Submodules
- veloxquant_mlx.cache.a2ats_cache module
- veloxquant_mlx.cache.adakv_cache module
- veloxquant_mlx.cache.age_tiered_cache module
- veloxquant_mlx.cache.amc_cache module
- veloxquant_mlx.cache.anchorkv_cache module
- veloxquant_mlx.cache.base module
- veloxquant_mlx.cache.cachegen_cache module
- veloxquant_mlx.cache.cam_cache module
- veloxquant_mlx.cache.chunkkv_cache module
- veloxquant_mlx.cache.chunkkv_coordinator module
- veloxquant_mlx.cache.curdkv_cache module
- veloxquant_mlx.cache.gear_cache module
- veloxquant_mlx.cache.h2o_cache module
- veloxquant_mlx.cache.keyformer_cache module
- veloxquant_mlx.cache.kitty_cache module
- veloxquant_mlx.cache.kivi_cache module
- veloxquant_mlx.cache.knorm_cache module
- veloxquant_mlx.cache.kvquant_cache module
- veloxquant_mlx.cache.kvtc_cache module
- veloxquant_mlx.cache.kvzip_cache module
- veloxquant_mlx.cache.minicache_cache module
- veloxquant_mlx.cache.minicache_coordinator module
- veloxquant_mlx.cache.morphkv_cache module
- veloxquant_mlx.cache.nestedkv_cache module
- veloxquant_mlx.cache.nsnquant_cache module
- veloxquant_mlx.cache.options module
VecInferOptionsKIVIOptionsSVDqOptionsKittyOptionsAdaKVOptionsXQuantOptionsKVQuantOptionsPALUOptionsCacheGenOptionsMiniCacheOptionsGEAROptionsZipCacheOptionsSnapKVOptionsStreamingLLMOptionsH2OOptionsTOVAOptionsPyramidKVOptionsSqueezeAttentionOptionsChunkKVOptionsCaMOptionsXKVOptionsNSNQuantOptionsL2NormOptionsSKVQOptionsQFiltersOptionsKeyformerOptionsMorphKVOptionsKVzipOptionsKVTCOptionsCurDKVOptionsNestedKVOptionsAMCOptionsA2ATSOptionsAnchorKVOptionsRocketKVOptionsAgeTieredKVOptionsSpectralQuantOptions
- veloxquant_mlx.cache.palu_cache module
- veloxquant_mlx.cache.polar_cache module
- veloxquant_mlx.cache.pyramidkv_cache module
- veloxquant_mlx.cache.qfilters_cache module
- veloxquant_mlx.cache.qjl_cache module
- veloxquant_mlx.cache.registry module
- veloxquant_mlx.cache.rocketkv_cache module
- veloxquant_mlx.cache.sink_cache module
- veloxquant_mlx.cache.skvq_cache module
- veloxquant_mlx.cache.sliding_window_cache module
- veloxquant_mlx.cache.snapkv_cache module
- veloxquant_mlx.cache.spectral_cache module
- veloxquant_mlx.cache.squeeze_cache module
- veloxquant_mlx.cache.squeeze_coordinator module
- veloxquant_mlx.cache.streaming_llm_cache module
- veloxquant_mlx.cache.svdq_cache module
- veloxquant_mlx.cache.tova_cache module
- veloxquant_mlx.cache.turboquant_cache module
- veloxquant_mlx.cache.turboquant_rvq_cache module
- veloxquant_mlx.cache.vecinfer_cache module
- veloxquant_mlx.cache.xkv_cache module
- veloxquant_mlx.cache.xkv_coordinator module
- veloxquant_mlx.cache.xquant_cache module
- veloxquant_mlx.cache.xquant_coordinator module
- veloxquant_mlx.cache.zipcache_cache module
- Submodules
- veloxquant_mlx.cli package
- Submodules
- veloxquant_mlx.cli.auto_config module
- veloxquant_mlx.cli.benchmark module
- veloxquant_mlx.cli.estimate_memory module
- veloxquant_mlx.cli.methods module
- veloxquant_mlx.cli.panel module
- veloxquant_mlx.cli.precompute module
- veloxquant_mlx.cli.profile module
- veloxquant_mlx.cli.profile_hardware module
- veloxquant_mlx.cli.recommend module
- veloxquant_mlx.cli.serve module
- veloxquant_mlx.cli.worker module
- Submodules
- veloxquant_mlx.codebooks package
- veloxquant_mlx.config package
- veloxquant_mlx.core package
- veloxquant_mlx.dsa package
- veloxquant_mlx.handlers package
- Submodules
- veloxquant_mlx.handlers.base module
- veloxquant_mlx.handlers.bit_pack_handler module
- veloxquant_mlx.handlers.normalization module
- veloxquant_mlx.handlers.outlier_split module
- veloxquant_mlx.handlers.polar_handler module
- veloxquant_mlx.handlers.qjl_residual_handler module
- veloxquant_mlx.handlers.rotation_handler module
- veloxquant_mlx.handlers.scalar_quant_handler module
- veloxquant_mlx.handlers.value_quant_handler module
- Submodules
- veloxquant_mlx.integration package
- veloxquant_mlx.math package
- veloxquant_mlx.memory package
- veloxquant_mlx.metal package
- Submodules
- veloxquant_mlx.metal.fused_sdpa module
- veloxquant_mlx.metal.kernels module
tova_fused_evict()tova_fused_evict_virtual_values()tova_fused_evict_indices()crosskv_rope_recode()vecinfer_dequant_metal()vecinfer_quantize_metal()vecinfer_encode_decode_metal()vecinfer_encode_decode_simple_metal()turboquant_bit_pack()turboquant_bit_unpack()turboquant_scalar_quantize()turboquant_scalar_dequantize()turboquant_hadamard_quantize()scalar_fused_decode_attend()scalar_fused_decode_attend_batched()kivi_group_quant_dequant()qjl_encode()qjl_inner_product()turboquant_fused_rvq_decode_attend()rvq_quant_pack()rvq_unpack_decode()comm_vq_decode_metal()rabitq_hamming_score()rabitq_fused_attend()rabitq_encode()rabitq_pack_values()rabitq_prefill_attend()flash_prefill_attend()streaming_prefill_attend()h2o_fused_evict()pyramidkv_fused_evict()keyformer_fused_evict()qfilters_fused_evict()qfilters_score()
- Submodules
- veloxquant_mlx.observers package
- veloxquant_mlx.outlier package
- veloxquant_mlx.planning package
- veloxquant_mlx.preconditioners package
- veloxquant_mlx.profiling package
- veloxquant_mlx.quantizers package
- Submodules
- veloxquant_mlx.quantizers.a2ats module
- veloxquant_mlx.quantizers.a2ats_rope module
- veloxquant_mlx.quantizers.adakv module
- veloxquant_mlx.quantizers.age_tiered module
- veloxquant_mlx.quantizers.amc module
AMCTierConfigAMCThresholdStateinit_amc_threshold_state()amc_saliency()amc_query_aware_saliency()amc_query_aware_saliency_batched()amc_assign_tiers()amc_assign_tiers_batched()amc_adaptive_thresholds()amc_apply_rank_mask()amc_quantize_tier()amc_compress_tokens_batched()amc_pack_low_tier()amc_fp16_bytes()full_amc_fp16_bytes()
- veloxquant_mlx.quantizers.amc_calibration module
- veloxquant_mlx.quantizers.anchorkv module
- veloxquant_mlx.quantizers.base module
- veloxquant_mlx.quantizers.cachegen module
CodeStreamquantize_to_codes()dequant_codes()token_delta()symbol_entropy_bits()entropy_coded_bytes()fixed_width_bytes()layer_group_bits()cachegen_quant_dequant()quantize_to_codes_batched()dequant_codes_batched()token_delta_batched()entropy_coded_bytes_batched()fixed_width_bytes_batched()
- veloxquant_mlx.quantizers.cam module
- veloxquant_mlx.quantizers.chunkkv module
- Public API
chunk_partition()chunk_scores()chunkkv_keep_mask()ChunkKVStateinit_chunkkv_state()chunkkv_update()chunkkv_update_batched()chunkkv_apply_reuse_indices()chunkkv_apply_reuse_indices_batched()chunkkv_trim_to()chunkkv_trim_batched()chunkkv_get_kv()chunkkv_fp16_bytes()full_chunkkv_fp16_bytes()
- veloxquant_mlx.quantizers.comm_vq module
- veloxquant_mlx.quantizers.composite module
- veloxquant_mlx.quantizers.curdkv module
- veloxquant_mlx.quantizers.gear module
- veloxquant_mlx.quantizers.h2o module
- veloxquant_mlx.quantizers.keyformer module
- veloxquant_mlx.quantizers.kitty module
- veloxquant_mlx.quantizers.kivi module
- veloxquant_mlx.quantizers.knorm module
- veloxquant_mlx.quantizers.kvquant module
- veloxquant_mlx.quantizers.kvtc module
- veloxquant_mlx.quantizers.kvzip module
- veloxquant_mlx.quantizers.minicache module
- veloxquant_mlx.quantizers.morphkv module
- veloxquant_mlx.quantizers.nestedkv module
- Public API
NestedKVStateinit_nestedkv_state()block_size_for()per_scale_anomaly_scores()head_adaptive_blend()surprise_gated_score()nestedkv_score()nestedkv_allocate_head_budgets()nestedkv_compress_prefill()nestedkv_append_decode()nestedkv_get_kv()nestedkv_fp16_bytes()full_nestedkv_fp16_bytes()per_scale_anomaly_scores_batched()head_adaptive_blend_batched()surprise_gated_score_batched()nestedkv_score_batched()nestedkv_compress_prefill_batched()nestedkv_append_decode_batched()
- veloxquant_mlx.quantizers.nsnquant module
- veloxquant_mlx.quantizers.palu module
- veloxquant_mlx.quantizers.polarquant module
- veloxquant_mlx.quantizers.pyramidkv module
- veloxquant_mlx.quantizers.qfilters module
- TWO FILTER SOURCES (which one you get changes the guarantees)
- Path-dependence applies to the FALLBACK ONLY
- Public API (mirrors quantizers/knorm.py)
QFiltersStateinit_qfilters_state()estimate_filter_dir()qfilters_update()qfilters_update_batched()qfilters_get_kv()qfilters_fp16_bytes()full_qfilters_fp16_bytes()
- veloxquant_mlx.quantizers.qfilters_calibration module
- veloxquant_mlx.quantizers.qjl module
- veloxquant_mlx.quantizers.rabitq module
- veloxquant_mlx.quantizers.rocketkv module
- veloxquant_mlx.quantizers.skvq module
- veloxquant_mlx.quantizers.snapkv module
- veloxquant_mlx.quantizers.squeeze module
- veloxquant_mlx.quantizers.streaming_llm module
- veloxquant_mlx.quantizers.svdq module
- veloxquant_mlx.quantizers.tova module
- veloxquant_mlx.quantizers.turboquant_mse module
- veloxquant_mlx.quantizers.turboquant_prod module
- veloxquant_mlx.quantizers.turboquant_rvq module
- veloxquant_mlx.quantizers.xkv module
- veloxquant_mlx.quantizers.xquant module
- veloxquant_mlx.quantizers.zipcache module
- Submodules
- veloxquant_mlx.routing package
- veloxquant_mlx.spectral package
- veloxquant_mlx.tools package
- veloxquant_mlx.transfer package
- veloxquant_mlx.transforms package
- veloxquant_mlx.ui package
- veloxquant_mlx.weight package