veloxquant_mlx.integration package¶

Integration shims that wire VeloxQuant-MLX caches into third-party inference loops.

Monkey-patches mlx_lm and mlx_vlm model instances so their normal generate() entry points build a VeloxQuant KV cache instead of the stock mlx_lm.models.cache.KVCache, without requiring callers to modify the host library. Re-exports patch_model_kv_cache() and patch_vlm_kv_cache().

veloxquant_mlx.integration.patch_model_kv_cache(model, config)[source]

Wire a quantized KV cache into an mlx-lm model’s generation path.

mlx_lm.generate() builds its prompt cache via mlx_lm.models.cache.make_prompt_cache(), which calls model.make_cache() if the model defines one, and falls back to a default fp16 cache otherwise. There is no pre-existing .cache attribute on attention sub-modules to overwrite — caches are only ever constructed lazily, on demand, through that hook.

This function builds the real per-layer cache list via KVCacheBuilder.for_model() (which already handles VLM wrappers, MoE-gate fallback layers, per-layer bit-width lists, and cross-layer methods like XQuant/MiniCache/PyramidKV) and overrides model.make_cache so every call builds a fresh cache list — matching veloxquant_mlx.integration.mlx_vlm_patch.patch_vlm_kv_cache(). A prior version closed over a single cache list built once at patch time, so a second generate() call on the same patched model would reuse the first call’s leftover KV state instead of starting clean.

Parameters:

model (Any) – A loaded mlx_lm model instance.
config (KVCacheConfig) – KVCacheConfig describing the quantization scheme.

Return type:

list[Any]

Returns:

The first freshly built cache list (also validates the config eagerly). Subsequent generate() calls receive their own fresh lists via the installed make_cache hook.

Raises:

QuantizerConfigError – If config.method is a standalone method (does not implement the mlx_lm KVCache serving contract — see veloxquant_mlx.cache.base.STANDALONE_METHODS). Raised immediately by the underlying KVCacheBuilder.for_model() call, before model.make_cache is touched.

Example:

config = KVCacheConfig(
    method="turboquant_rvq",
    bit_width_inlier=1,
    seed=42,
)
patch_model_kv_cache(model, config)

veloxquant_mlx.integration.patch_vlm_kv_cache(model, config)[source]

Wire a quantized KV cache into an mlx-vlm model’s generation path.