veloxquant_mlx.integration.mlx_vlm_patch module

Monkey-patches an mlx-vlm vision-language model to use VeloxQuant KV caches.

Usage:

from mlx_vlm import load, generate
from veloxquant_mlx.integration.mlx_vlm_patch import patch_vlm_kv_cache
from veloxquant_mlx.cache import KVCacheConfig

model, processor = load("mlx-community/Qwen2-VL-2B-Instruct-4bit")
config = KVCacheConfig(method="turboquant_rvq", bit_width_inlier=2, seed=42)
patch_vlm_kv_cache(model, config)

# mlx_vlm.generate() now builds a quantized cache automatically
output = generate(model, processor, prompt, image)

How it hooks in (verified against mlx-vlm 0.6.5): mlx_vlm.generate builds its cache via mlx_vlm.models.cache.make_prompt_cache(model.language_model, ...), which defers to model.language_model.make_cache() when defined. This patch overrides exactly that hook — and only on the language model.

The top-level model is deliberately left unpatched: mlx-vlm’s batch / session path (generate/ar.py:_make_cache) calls the top-level model.make_cache() and then converts each returned cache with to_batch_cache(), which raises ValueError on cache types it does not know — including VeloxQuant’s. Batched generation therefore keeps mlx-vlm’s own caches (its built-in kv_bits quantization still works there); single-prompt generation gets the VeloxQuant cache.

Unlike patch_model_kv_cache() (which returns one persistent cache list), make_cache here builds a fresh cache list on every call, so repeated generate() calls never leak KV state between generations.

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

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

Parameters:

  • model (Any) – A loaded mlx-vlm model instance (must expose .language_model).

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