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: 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.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.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.ArtifactStore[source]

Bases: ABC

DAO interface for loading and saving precomputed quantization artifacts.

abstractmethod load_rotation_matrix(d, seed)[source]

Load a precomputed rotation matrix.

Parameters:
  • d (int) – Matrix dimension.

  • seed (int) – Random seed used to generate the matrix.

Return type:

Any

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.

Parameters:
  • Pi (Any) – Array of shape (d, d).

  • d (int) – Dimension.

  • seed (int) – Seed used to generate Pi.

Return type:

None

abstractmethod load_codebook(distribution, b, d)[source]

Load a precomputed codebook.

Parameters:
  • distribution (str) – Distribution name (e.g. ‘gaussian’, ‘beta’).

  • b (int) – Bit-width.

  • d (int) – Vector dimension.

Return type:

Any

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.

Parameters:
  • cb (Any) – Codebook centroids array.

  • distribution (str) – Distribution name.

  • b (int) – Bit-width.

  • d (int) – Dimension.

Return type:

None

abstractmethod load_jl_matrix(d, m, seed)[source]

Load a precomputed JL projection matrix.

Parameters:
  • d (int) – Input dimension.

  • m (int) – Output (sketch) dimension.

  • seed (int) – Random seed.

Return type:

Any

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.

Parameters:
  • S (Any) – Array of shape (m, d).

  • d (int) – Input dimension.

  • m (int) – Output dimension.

  • seed (int) – Seed.

Return type:

None

abstractmethod exists(artifact_type, **kwargs)[source]

Check whether a specific artifact exists in the store.

Parameters:
  • artifact_type (str) – One of ‘rotation’, ‘codebook’, ‘jl’.

  • **kwargs (Any) – Identifying parameters (d, seed, b, distribution, m).

Return type:

bool

Returns:

True if the artifact is present and loadable.

class veloxquant_mlx.KVCache[source]

Bases: ABC

Abstract base class for KV cache implementations.

abstractmethod append_key(k)[source]

Append a new key vector to the cache.

Parameters:

k (Any) – Key vector, shape (d,), fp16.

Return type:

None

abstractmethod append_value(v)[source]

Append a new value vector to the cache.

Parameters:

v (Any) – Value vector, shape (d,), fp16.

Return type:

None

abstractmethod attend(q)[source]

Compute attention-weighted value for a query.

Parameters:

q (Any) – Query vector, shape (d,), fp16.

Return type:

Any

Returns:

Attention output, shape (d,), fp16.

abstractmethod memory_bytes()[source]

Return current memory footprint of the cache in bytes.

Return type:

int

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 SlidingWindowKVCache rely on reset() to rebuild an emptied cache without re-running (or losing) prior calibration.

Return type:

None

append(k, v)[source]

Append a key-value pair in one call.

Parameters:
  • k (Any) – Key vector, shape (d,).

  • v (Any) – Value vector, shape (d,).

Return type:

None

class veloxquant_mlx.Quantizer[source]

Bases: ABC

Abstract 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:

EncodedVector

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:

Any

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:

Any

Returns:

Estimated inner products, shape (batch_size,), fp16.

class veloxquant_mlx.QuantizationObserver[source]

Bases: ABC

Observer for pipeline events (timing, memory, distortion).

abstractmethod on_event(event)[source]

Handle a quantization pipeline event.

Parameters:

event (Any) – QuantizationEvent dataclass instance.

Return type:

None

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: object

Typed 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
indices: Any | None = None
norm: Any | None = None
signs: Any | None = None
residual_norm: Any | None = None
angles: list[Any] | None = None
final_radius: Any | None = None
outlier_idx: ndarray | None = None
outlier_encoded: EncodedVector | None = None
inlier_encoded: EncodedVector | None = None
memory_bytes()[source]

Compute exact memory footprint of this encoded representation.

Return type:

int

Returns:

Total bytes occupied by all stored arrays.

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: object

Mutable 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)

  • norm (Any | None)

  • rotated (Any | None)

  • indices (Any | None)

  • signs (Any | None)

  • residual_norm (Any | None)

  • angles (list[Any] | None)

  • final_radius (Any | None)

  • outlier_idx (ndarray | None)

  • packed_bits (ndarray | None)

  • metadata (dict[str, Any])

x_original: Any
mode: Literal['encode', 'decode']
x_current: Any
norm: Any | None = None
rotated: Any | None = None
indices: Any | None = None
signs: Any | None = None
residual_norm: Any | None = None
angles: list[Any] | None = None
final_radius: Any | None = None
outlier_idx: ndarray | None = None
packed_bits: ndarray | None = None
metadata: dict[str, Any]
class veloxquant_mlx.TransformResult(angles, final_radius, n_levels)[source]

Bases: object

Output 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:
angles: list[Any]
final_radius: Any
n_levels: int
exception veloxquant_mlx.ArtifactNotFoundError[source]

Bases: FileNotFoundError

Raised when a required precomputed artifact is missing from the store.

exception veloxquant_mlx.BlockPoolExhaustedError[source]

Bases: RuntimeError

Raised when a BlockPoolAllocator has no free blocks left to allocate.

exception veloxquant_mlx.CodebookDimensionMismatch[source]

Bases: ValueError

Raised when a codebook’s shape does not match the expected dimension.

exception veloxquant_mlx.CyclicPipelineError[source]

Bases: RuntimeError

Raised when a QuantizationGraph contains a cycle.

exception veloxquant_mlx.QuantizerConfigError[source]

Bases: ValueError

Raised when a quantizer or KV cache is misconfigured.

class veloxquant_mlx.BlockPoolAllocator(config=None)[source]

Bases: object

Fixed-size block allocator for KV-cache storage.

Pre-allocates n_blocks fixed-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 owner as 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 owner already has blocks checked out or was registered and not yet released.

Return type:

None

Parameters:

owner (int)

release_owner(owner)[source]

Clear owner from 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 via register_owner() but never actually allocated anything for it.

Return type:

None

Parameters:

owner (int)

allocate(stream, n_tokens, owner, format='fp16')[source]

Allocate enough blocks to hold n_tokens slots.

Parameters:
  • stream (str) – “k” or “v” (ignored, treated as “kv”, when the pool was built with separate_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, call register_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:

list[Block]

Returns:

List of newly allocated Block instances, 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 or stream is invalid.

free(blocks)[source]

Return blocks to the free list for their stream.

Parameters:

blocks (list[Block]) – Blocks previously returned by allocate().

Return type:

None

free_all(owner)[source]

Free every block currently checked out to owner.

Parameters:

owner (int) – Opaque id passed to a prior allocate() call.

Return type:

None

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 + N list.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.

Return type:

list[Block]

Parameters:

owner (int)

n_free(stream='k')[source]

Number of free blocks remaining for stream.

Return type:

int

Parameters:

stream (str)

shrink(stream, target_free)[source]

Permanently retire fully-free blocks from stream down to target_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:
  • stream (str) – “k” or “v” (ignored, treated as “kv”, when the pool was built with separate_kv=False).

  • target_free (int) – Desired number of free blocks to retain for this stream after shrinking. If the stream already has this many or fewer free blocks, this is a no-op.

Return type:

int

Returns:

Number of blocks actually retired.

Raises:

ValueError – If target_free < 0.

mark_used(block, n_used)[source]

Update block.n_used and 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 writing block.n_used directly – AllocationStats._fragmented_count is 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.

Parameters:
  • block (Block) – A block currently checked out from this pool (owner is not None). Blocks that are free are never fragmented by definition, so calling this on one is a no-op for stats (n_used is still updated).

  • n_used (int)

Return type:

None

class veloxquant_mlx.PoolConfig(block_size=16, n_blocks=1024, separate_kv=True, grow_on_exhaustion=False, max_blocks=None)[source]

Bases: object

Configuration 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)

  • n_blocks (int)

  • separate_kv (bool)

  • grow_on_exhaustion (bool)

  • max_blocks (int | None)

block_size: int = 16
n_blocks: int = 1024
separate_kv: bool = True
grow_on_exhaustion: bool = False
max_blocks: int | None = None
class veloxquant_mlx.PooledKVCache(inner, pool, owner=-1, format='fp16')[source]

Bases: KVCache

Wraps 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 shared BlockPoolAllocator in 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_size tokens: 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’s release() reclaiming the wrong blocks.

  • format (str) – Compression-format tag recorded on each block this cache checks out (purely informational — passed through to pool.allocate(format=...)).

append_key(k)[source]

Check out a K block from the pool if needed, then append to the inner cache.

Parameters:

k (Any) – Key vector, shape (d,), fp16.

Return type:

None

append_value(v)[source]

Check out a V block from the pool if needed, then append to the inner cache.

Parameters:

v (Any) – Value vector, shape (d,), fp16.

Return type:

None

attend(q)[source]

Delegate attention computation to the inner cache.

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 (compressed) cache.

Return type:

int

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:

None

n_blocks_held()[source]

Total K + V blocks currently checked out to this cache.

Return type:

int

reset()[source]

Return all held blocks to the pool and clear the inner cache.

Return type:

None

class veloxquant_mlx.QuantizerFactory[source]

Bases: object

Factory 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:

Quantizer

Returns:

Configured Quantizer instance.

Raises:
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 matches round(target_avg_bits * N).

Parameters:
  • sensitivities – Iterable of per-layer sensitivity weights w_i > 0, typically from calibrate_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:

list[int]

Returns:

list[int] of length len(sensitivities). Pass directly as KVCacheConfig.bit_width_inlier and consume via KVCacheBuilder.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:

list[float]

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:

tuple[float, float]

Returns:

Tuple (alpha, beta). For TurboQuantRVQ at d=128 the paper reports β ≈ 3.5; you will get a similar value from this function. For production use, just pass beta=3.5 to allocate_bits_ratequant() and skip the fit.

Parameters:
class veloxquant_mlx.VecInferKVCache(config)[source]

Bases: KVCache

KV cache implementing VecInfer’s dual-transform product VQ.

Parameters:

config (Any) – KVCacheConfig with VecInfer-specific fields populated by KVCacheFactory.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 checks hasattr(cache, "bits") to route to a different kernel path. We expose .assigned_avg_bits instead.

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:

bool

trim(n)[source]

Drop the last n tokens, keeping every parallel buffer in step.

The inherited trim only lowers offset. In fused mode the index ring buffers are advanced by _stored_S_kv instead, so left alone they keep the trimmed tokens: the next token is written n slots past offset and fused_sdpa attends over the stale tail. The byte accounting is scaled back with the token count (an estimate: it is tracked per call, not per token).

Return type:

int

Parameters:

n (int)

size()[source]

Number of tokens currently stored.

Return type:

int

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 (typically 1/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 of q.

property compressed_key_bytes: int

see nbytes.

Type:

Estimated packed bytes for the compressed key cache (product-VQ codes at bits each, 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 bits each, 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 when smooth is 2D.

  • smooth (array) – Either [head_dim] (broadcast across heads) or [n_heads, head_dim] (per-head); see calibrate_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) @ H so q_tilde @ K_tilde.T == q @ K.T.

Parameters:
  • q (array) – Queries, shape [..., head_dim]. Under GQA, the head axis (-3) may have H_q = H_kv * n_rep heads, a multiple of smooth’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 against apply_dual_transform_keys()’s per-KV-head factors.

  • smooth (array) – Same convention as apply_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.

Parameters:
  • x (array) – Flat [n_samples, sub_dim] MLX array of training sub-vectors.

  • n_centroids (int) – Codebook size 2**b.

  • max_iter (int) – K-means iteration cap.

  • seed (int) – RNG seed.

Return type:

array

Returns:

Codebook of shape [n_centroids, sub_dim] as an MLX float32 array.

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] with H @ H.T == I.

class veloxquant_mlx.KeyNormObserver[source]

Bases: QuantizationObserver

Accumulates 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:

None

Parameters:

event (QuantizationEvent)

report()[source]

Build a KeyNormReport from the norms accumulated so far.

Return type:

KeyNormReport

reset()[source]

Clear all accumulated norm statistics.

Return type:

None

class veloxquant_mlx.KeyNormReport(n_tokens, mean_norm_sq, min_norm_sq, max_norm_sq)[source]

Bases: object

Summary statistics produced by KeyNormObserver.

Parameters:
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: KVCache

Wraps a KVCache instance, 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:
  • cache (KVCache) – The KVCache instance to profile.

  • head_dim (int | None) – Head dimension, used to compute the fp16 compression baseline (2 bytes/element). Defaults to reading cache._d if present, else 0 (disables compression_ratio).

  • layer_id (Any) – Optional label for this layer, used in reports.

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.

Return type:

None

Parameters:

k (Any)

append_value(v)[source]

Append a value to the wrapped cache, timing it as a write call.

Return type:

None

Parameters:

v (Any)

attend(q)[source]

Run attention on the wrapped cache, timing it as a dequantize call.

Return type:

Any

Parameters:

q (Any)

memory_bytes()[source]

Return the wrapped cache’s current memory_bytes() reading.

Return type:

int

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:

LayerProfile

reset()[source]

Clear all accumulated profiling stats (does not affect the wrapped cache).

Return type:

None

class veloxquant_mlx.MLXCacheProfiler(cache, layer_id=0)[source]

Bases: object

Wraps an mlx_lm.models.cache.KVCache-style cache’s update_and_fetch.

Servable methods (the ones cli/serve.py and cli/profile.py drive through real mlx_lm.generate() runs) implement update_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 interface KVCacheProfiler wraps. So this class times update_and_fetch as one unit and records it as quantize_ms_total with is_fused=True; dequantize_ms_total and write_ms_total stay 0 because there is nothing else to attribute them to. Memory comes from the cache’s own nbytes property (the same true packed-byte accounting cli/serve.py’s accounting-only warning already describes — see ACCOUNTING_NOTE in cli/profile.py).

Substituted directly into the list KVCacheBuilder.for_model() returns and passed to mlx_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.

Return type:

Any

Parameters:
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:

LayerProfile

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: object

Aggregated 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 — see is_fused. There is no separate number to split it into.

  • dequantize_ms_total – Cumulative attend wall time, milliseconds. Always 0 when is_fused is True.

  • write_ms_total – Cumulative append_value wall time, milliseconds. Always 0 when is_fused is 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)

  • n_dequantize_calls (int)

  • quantize_ms_total (float)

  • dequantize_ms_total (float)

  • write_ms_total (float)

  • peak_memory_bytes (int)

  • tokens_written (int)

  • fp16_baseline_bytes (int)

  • is_fused (bool)

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: object

Profiling 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 = 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:

str

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:
Return type:

ProfileReport

Returns:

A ProfileReport summarizing all layers.

class veloxquant_mlx.AutoConfigResult(config, reason)[source]

Bases: object

The 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 --explain flag.

Parameters:
config: KVCacheConfig
reason: str
class veloxquant_mlx.HardwareInfo(total_memory_bytes=None, active_memory_bytes=0)[source]

Bases: object

Describes 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:
  • total_memory_bytes (int | None)

  • active_memory_bytes (int)

total_memory_bytes: int | None = None
active_memory_bytes: int = 0
pressure_fraction(additional_bytes)[source]

Fraction of total memory used once additional_bytes is added.

Returns None if total_memory_bytes is unknown.

Return type:

float | None

Parameters:

additional_bytes (int)

class veloxquant_mlx.WorkloadSpec(head_dim=128, seq_len=4096, n_layers=1, batch_size=1, n_kv_heads=1)[source]

Bases: object

Describes 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)

  • seq_len (int)

  • n_layers (int)

  • batch_size (int)

  • n_kv_heads (int)

head_dim: int = 128
seq_len: int = 4096
n_layers: int = 1
batch_size: int = 1
n_kv_heads: int = 1
fp16_kv_bytes()[source]

Estimated fp16 footprint of the full K+V cache for this workload.

Return type:

int

veloxquant_mlx.detect_hardware_info()[source]

Auto-detect HardwareInfo from the running MLX device.

Falls back to an empty (unknown-memory) HardwareInfo if MLX is unavailable or device introspection fails — callers then get sequence-length-only selection instead of an exception.

Return type:

HardwareInfo

veloxquant_mlx.select_kv_cache_config(workload, hardware=None)[source]

Pick a KVCacheConfig for the given workload and hardware.

Selection rules (in priority order):

  1. Memory pressure: if hardware reports the workload’s estimated fp16 footprint would push total usage above MEMORY_PRESSURE_FRACTION of device memory, pick the most aggressive method in the pool (gear, 2-bit) regardless of sequence length.

  2. 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).

  3. 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 via detect_hardware_info().

Return type:

AutoConfigResult

Returns:

AutoConfigResult with the selected KVCacheConfig and a human-readable reason string.

class veloxquant_mlx.AutoOptimizer(options=None)[source]

Bases: object

One-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_n or 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:

HardwareProfile

profile_model(config=None, **overrides)[source]

Profile a model from a HF-style config dict + optional overrides.

Return type:

ModelProfile

Parameters:
estimate_memory(workload, model=None)[source]

Analytical per-method memory estimates for a workload.

Return type:

dict[str, MemoryEstimate]

Parameters:
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:

RecommendationResult

Parameters:
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:

str

Parameters:
fallback_method()[source]

Library-default method name used when nothing survives filtering.

Return type:

str

class veloxquant_mlx.AutoOptimizerOptions(benchmark_db_dir=None, probe_top_n=3)[source]

Bases: object

Constructor options for AutoOptimizer.

Variables:
  • benchmark_db_dir – Path to a BenchmarkDatabase storage 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. 0 skips probing (faster, but a crash-prone standalone method could be recommended).

Parameters:
  • benchmark_db_dir (str | None)

  • probe_top_n (int)

benchmark_db_dir: str | None = None
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: object

Comprehensive 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)

  • chip_generation (int)

  • total_memory_bytes (int | None)

  • available_memory_bytes (int)

  • mlx_version (str | None)

  • macos_version (str | None)

  • metal_available (bool)

  • peak_memory_bandwidth_gbps (float | None)

  • avg_quantize_latency_ms_per_token (float | None)

chip: str = 'unknown'
chip_generation: int = 0
total_memory_bytes: int | None = None
available_memory_bytes: int = 0
mlx_version: str | None = None
macos_version: str | None = None
metal_available: bool = True
peak_memory_bandwidth_gbps: float | None = None
avg_quantize_latency_ms_per_token: float | None = None
classmethod detect()[source]

Build a HardwareProfile for the current machine (see detect_hardware_profile).

Return type:

HardwareProfile

property bandwidth_gbps: float | None

measured, else nominal, else None.

Type:

Likely peak bandwidth

to_dict()[source]

Serialize all fields to a plain dict.

Return type:

dict

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: object

Normalized 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_heads for 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)

  • architecture (str)

  • num_layers (int)

  • num_query_heads (int)

  • num_kv_heads (int)

  • head_dim (int)

  • attention_type (str)

  • dtype (str)

  • parameter_count (int | None)

  • hidden_size (int | None)

  • _extra (dict[str, Any])

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'
parameter_count: int | None = None
hidden_size: int | None = None
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.

Return type:

int

Parameters:
  • num_layers (int | None)

  • batch (int)

to_dict()[source]

Serialize all fields to a plain dict.

Return type:

dict[str, Any]

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: object

Normalized 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 WorkloadObjective names; 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)

  • generation_length (int)

  • batch_size (int)

  • num_concurrent_requests (int)

  • objective (str)

  • max_latency_ms (float | None)

context_length: int = 4096
generation_length: int = 512
batch_size: int = 1
num_concurrent_requests: int = 1
objective: str = 'balanced'
max_latency_ms: float | None = None
to_dict()[source]

Serialize all fields to a plain dict.

Return type:

dict[str, Any]

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: object

Identifier 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_WEIGHTS only assigns weights to the canonical ones below and falls back to balanced for 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: object

Full 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 = False
fallback_reason: str | None = None
property best: ScoredMethod | None

Top-ranked method, or None if nothing survived filtering.

to_dict()[source]

Serialize the full result (ranked list, candidates, inputs) to a plain dict.

Return type:

dict[str, Any]

class veloxquant_mlx.MemoryEstimate(method, baseline_bytes, compressed_bytes, workspace_bytes=0, peak_bytes=0, resident_bytes=0, confidence='high', assumptions=<factory>)[source]

Bases: object

Breakdown 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)

  • peak_bytes (int)

  • resident_bytes (int)

  • confidence (str)

  • assumptions (list[str])

method: str
baseline_bytes: int
compressed_bytes: int
workspace_bytes: int = 0
peak_bytes: int = 0
resident_bytes: int = 0
confidence: str = 'high'
assumptions: list[str]
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).

to_dict()[source]

Serialize all fields (plus derived ratio/percent) to a plain dict.

Return type:

dict[str, Any]

class veloxquant_mlx.CacheRoutePlanner(n_shards, qcap, warm_slots_per_shard)[source]

Bases: object

Builds a periodic RoutingTable from 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 returned RoutingTable for 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 two RoutingTable instances themselves before switching.

Parameters:
  • n_shards (int) – Number of pool shards to plan across (R in the paper). Use 1 if 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 as SessionRate.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, W in the paper. Total admission budget is n_shards * warm_slots_per_shard, counted in the same units as rate/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:

RoutingTable

Returns:

A RoutingTable valid 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: object

Tracks 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, then rates() to get the SessionRate list a CacheRoutePlanner needs. 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 _counts and one int in _last_seen each), not by how many requests have been recorded — an owner’s whole contribution to memory is released by forget(). That bound is only enforced if the caller calls forget() 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 set max_owners so _counts self-bounds instead of growing without limit.

Parameters:
  • half_life (float) – Number of record() 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.

  • _counts (dict[int, float])

  • _last_seen (dict[int, int])

  • _clock (int)

half_life: float = 20.0
max_owners: int | None = None
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 pow and two dict writes for an already-tracked owner; the eviction scan only runs for a new owner when max_owners is set.

Return type:

None

Parameters:

owner (int)

rate(owner)[source]

Current smoothed rate estimate for owner (0.0 if never seen).

Return type:

float

Parameters:

owner (int)

rates()[source]

Snapshot every tracked owner’s current rate as a SessionRate list.

Return type:

list[SessionRate]

forget(owner)[source]

Drop all tracked state for owner (e.g. on session close).

Return type:

None

Parameters:

owner (int)

class veloxquant_mlx.RoutingTable(shards, expected_load, n_shards, admitted_rate_total, unadmitted_rate_total)[source]

Bases: object

A 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’s kb replication (Eq. 1) — an owner needs more than one shard only when its rate exceeds qcap on its own.

  • expected_load – Per-shard expected load after this plan’s placement, indexed by shard id. Sums admitted owners’ rate / kb over 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 rate over admitted owners, for computing what fraction of traffic the warm set actually covers.

  • unadmitted_rate_total – Sum of rate over owners considered but not admitted (the plan’s cold-tail fallback traffic).

Parameters:
shards: dict[int, tuple[int, ...]]
expected_load: list[float]
n_shards: int
admitted_rate_total: float
unadmitted_rate_total: float
is_admitted(owner)[source]

Whether owner has a warm-set placement in this plan.

Return type:

bool

Parameters:

owner (int)

preferred_shard(owner, fallback=0)[source]

The single best shard for owner, or fallback if not admitted.

For an owner with kb > 1 shards, 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.

Return type:

int

Parameters:
imbalance()[source]

Ratio of the most-loaded shard’s expected load to the mean.

Matches the “imbalance” metric in the paper’s Table 4/5 (1.0x is perfectly balanced; higher is worse). Returns 1.0 for zero total load (nothing to balance).

Return type:

float

class veloxquant_mlx.SessionRate(owner, rate)[source]

Bases: object

A session’s estimated request rate, in requests per unit time.

Variables:
  • owner – Opaque session/owner id, matching BlockPoolAllocator owner ids.

  • rate – Estimated request rate (lambda_b in the paper). Any consistent unit works as long as qcap uses 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.

KVCacheBuilder()

Fluent builder for KVCache construction with validation.

KVCacheConfig([method, head_dim, ...])

Configuration for a KVCache instance.

KVCacheFactory()

Factory for creating KVCache instances from a KVCacheConfig.

ArtifactStore()

DAO interface for loading and saving precomputed quantization artifacts.

KVCache()

Abstract base class for KV cache implementations.

Quantizer()

Abstract base class for all vector quantizers.

QuantizationObserver()

Observer for pipeline events (timing, memory, distortion).

EncodedVector(quantizer_type, batch_size, dim)

Typed output of Quantizer.encode().

QuantizationContext(x_original, mode, x_current)

Mutable payload flowing through a QuantizationHandler chain.

TransformResult(angles, final_radius, n_levels)

Output of RecursivePolarTransform.forward().

ArtifactNotFoundError

Raised when a required precomputed artifact is missing from the store.

BlockPoolExhaustedError

Raised when a BlockPoolAllocator has no free blocks left to allocate.

CodebookDimensionMismatch

Raised when a codebook's shape does not match the expected dimension.

CyclicPipelineError

Raised when a QuantizationGraph contains a cycle.

QuantizerConfigError

Raised when a quantizer or KV cache is misconfigured.

BlockPoolAllocator([config])

Fixed-size block allocator for KV-cache storage.

PoolConfig([block_size, n_blocks, ...])

Configuration for a BlockPoolAllocator.

PooledKVCache(inner, pool[, owner, format])

Wraps any VeloxQuant KVCache with block-pool-backed memory accounting.

QuantizerFactory()

Factory for creating Quantizer instances via the QuantizerRegistry.

allocate_bits_ratequant(sensitivities, ...)

Allocate per-layer bit-widths via RateQuant Theorem 2.

calibrate_layer_sensitivities(model, tokenizer)

Run a calibration forward pass and return per-layer sensitivity.

fit_distortion_curve(head_dim[, ...])

Fit D(b) = α·β^(-b) on synthetic unit-norm Gaussian keys.

VecInferKVCache(config)

KV cache implementing VecInfer's dual-transform product VQ.

apply_dual_transform_keys(K, smooth, H)

Apply K_tilde = (K / lambda) @ H (smooth then Hadamard).

apply_dual_transform_queries(q, smooth, H)

Apply q_tilde = (q * lambda) @ H so q_tilde @ K_tilde.T == q @ K.T.

calibrate_smooth_factors(keys_calib[, eps])

Compute per-(head, channel) smooth scaling factor from a key sample.

train_codebook(x, n_centroids[, max_iter, seed])

Train a VQ codebook on flat sub-vector samples.

walsh_hadamard_matrix(d[, dtype])

Construct an orthonormal Walsh-Hadamard matrix.

KeyNormObserver()

Accumulates per-token key L2 norm² for later RateQuant allocation.

KeyNormReport(n_tokens, mean_norm_sq, ...)

Summary statistics produced by KeyNormObserver.

KVCacheProfiler(cache[, head_dim, layer_id])

Wraps a KVCache instance, recording per-call timing and memory.

MLXCacheProfiler(cache[, layer_id])

Wraps an mlx_lm.models.cache.KVCache-style cache's update_and_fetch.

LayerProfile(layer_id[, n_quantize_calls, ...])

Aggregated profiling stats for a single KV-cache layer.

ProfileReport([layers, elapsed_s])

Profiling results across one or more layers.

format_profile_table(report)

Render a ProfileReport as a human-readable fixed-width table.

profile_layers(profilers[, elapsed_s])

Aggregate a list of KVCacheProfiler instances into one ProfileReport.

AutoConfigResult(config, reason)

The selected configuration plus the reasoning behind it.

HardwareInfo([total_memory_bytes, ...])

Describes the machine the KV cache will run on.

WorkloadSpec([head_dim, seq_len, n_layers, ...])

Describes the job the KV cache needs to serve.

detect_hardware_info()

Auto-detect HardwareInfo from the running MLX device.

select_kv_cache_config(workload[, hardware])

Pick a KVCacheConfig for the given workload and hardware.

AutoOptimizer([options])

One-stop automatic KV-cache strategy recommender with caching and live validation.

AutoOptimizerOptions([benchmark_db_dir, ...])

Constructor options for AutoOptimizer.

HardwareProfile([chip, chip_generation, ...])

Comprehensive description of the Apple-Silicon environment.

ModelProfile([model_id, architecture, ...])

Normalized model-architecture description used by the planner.

WorkloadProfile([context_length, ...])

Normalized description of a serving workload.

WorkloadObjective()

Identifier constants for the planner's supported objectives.

RecommendationResult(objective, ranked, ...)

Full outcome of a planning run, ready for the explainer / CLI.

MemoryEstimate(method, baseline_bytes, ...)

Breakdown of one method's predicted KV-cache memory footprint.

CacheRoutePlanner(n_shards, qcap, ...)

Builds a periodic RoutingTable from measured session rates.

RateEstimator([half_life, max_owners, ...])

Tracks a per-owner exponentially-weighted request rate estimate.

RoutingTable(shards, expected_load, ...)

A fixed admission-and-placement plan for one control interval.

SessionRate(owner, rate)

A session's estimated request rate, in requests per unit time.

Subpackages