mlx_lm Integration

VeloxQuant-MLX is designed to work as a drop-in extension for mlx_lm. This guide covers the integration patterns: the KVCacheBuilder helper, the mlx_lm_patch monkey-patch, the fused SDPA kernel, and PrefixCache for multi-call prefix reuse.

Pattern 1 — KVCacheBuilder (recommended)

KVCacheBuilder is the primary integration point. It inspects the model's config and constructs one KVCache per transformer layer, matching num_key_value_heads and head_dim automatically.

import mlx_lm
from veloxquant_mlx.cache.base import KVCacheConfig, KVCacheBuilder

model, tokenizer = mlx_lm.load("mlx-community/Llama-3.2-3B-Instruct-4bit")

config = KVCacheConfig(method="turboquant_rvq", bit_width_inlier=1)
caches = KVCacheBuilder.for_model(model, config)

# Pass the per-layer cache list directly to mlx_lm.generate
response = mlx_lm.generate(
    model,
    tokenizer,
    prompt="Hello, world!",
    max_tokens=256,
    kv_cache=caches,
)

KVCacheBuilder.for_model() works with any model that exposes model.layers or model.model.layers with per-layer head_dim/num_key_value_heads — which covers all major mlx_lm model families (text-only and VLM). It returns a list of one KVCache per layer, not a single object.

Pattern 2 — mlx_lm monkey-patch

The monkey-patch approach automatically intercepts the default KV cache creation inside mlx_lm and replaces it with VeloxQuant-MLX caches. This requires zero changes to your generation call:

import mlx_lm
from veloxquant_mlx.integration.mlx_lm_patch import patch_model_kv_cache
from veloxquant_mlx.cache import KVCacheConfig

model, tokenizer = mlx_lm.load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")

config = KVCacheConfig(method="turboquant_rvq", bit_width_inlier=1, seed=42)
patch_model_kv_cache(model, config)  # overrides model.make_cache()

# No cache argument needed — mlx_lm.generate() builds the quantized cache
response = mlx_lm.generate(model, tokenizer, prompt="...", max_tokens=512)

tip

The monkey-patch is useful when integrating with third-party code that calls mlx_lm.generate directly and does not expose a cache argument.

Vision-language models (mlx-vlm)

patch_vlm_kv_cache wires VeloxQuant caches into mlx-vlm models (Qwen2-VL, LLaVA, etc.). mlx-vlm's single-prompt generation builds its cache through model.language_model.make_cache() — the patch overrides exactly that hook (verified against mlx-vlm 0.6.5):

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)

output = generate(model, processor, prompt, image)

Behaviour to know:

• Fresh caches per generation. Unlike the text patch, make_cache rebuilds the cache list on every call, so repeated generate() calls never leak KV state between generations.
• Single-prompt path only. mlx-vlm's batched/session path converts caches with its own to_batch_cache(), which rejects foreign cache types — so batched generation keeps mlx-vlm's built-in caches (its native kv_bits quantization still works there). The top-level model is deliberately left unpatched to keep that path safe.
• Eviction methods warn. Token-dropping methods (snapkv, h2o, pyramidkv, …) may discard image tokens from the prompt prefix; the patch emits a UserWarning and quantization-only methods are recommended for multimodal prompts.

Pattern 3 — Fused SDPA

patch_mlx_lm_for_fused_sdpa replaces the attention computation with a Metal kernel that dequantizes keys/values and computes attention in a single GPU pass:

from veloxquant_mlx.metal.fused_sdpa import patch_mlx_lm_for_fused_sdpa

# Call once, after the model is loaded (it patches the already-imported
# mlx_lm.models.* modules in place — it takes no model argument).
patch_mlx_lm_for_fused_sdpa()

# Subsequent generate calls use the fused kernel automatically
response = mlx_lm.generate(model, tokenizer, prompt="...", max_tokens=1024, kv_cache=caches)

Fused SDPA is most beneficial when:

• The KV cache is large (long sequences, many layers)
• You are using VecInfer (the fused kernel is optimised for its codebook format)
• Throughput is the priority over latency on individual calls

Check compatibility before patching:

from veloxquant_mlx.metal.fused_sdpa import supports_shape

# Verify your VecInfer codebook shape is supported by the kernel's caps
# (n_centroids, n_sub, head_dim) — not a batch/seq_len/heads check.
is_supported = supports_shape(
    n_centroids=2**8,
    n_sub=model.args.head_dim // 8,
    head_dim=model.args.head_dim,
)
print(f"Fused SDPA supported: {is_supported}")

Pattern 4 — PrefixCache (multi-call prefix reuse)

patch_model_kv_cache gives every generate() call a fresh cache — correct for a single call, but a program that calls generate() repeatedly with a shared prefix (a system prompt, an agent's growing conversation history) re-prefills that shared prefix from scratch every time. This is the same gap reported against Ollama's MLX engine (ollama/ollama#17829): with no prefix-cache reuse, time-to-first-token scales with total conversation history instead of just the new turn.

veloxquant serve does not have this problem — it hands off to mlx_lm.server, which already reuses prefixes internally via mlx_lm.models.cache.LRUPromptCache. PrefixCache brings that same mechanism to direct mlx_lm.generate()/stream_generate() callers:

import mlx_lm
from veloxquant_mlx.cache.base import KVCacheConfig
from veloxquant_mlx.integration.prefix_cache import PrefixCache

model, tokenizer = mlx_lm.load("mlx-community/Llama-3.2-3B-Instruct-4bit")
config = KVCacheConfig(method="turboquant_rvq", bit_width_inlier=1, seed=42)
prefix_cache = PrefixCache(config)

system_prompt = "You are a careful coding assistant. ..." * 50  # long, shared prefix

# First call: full prefill.
reply_1 = prefix_cache.generate(model, tokenizer, system_prompt + "What does this function do?")

# Second call: the shared prefix is reused from cache — only the new
# suffix is prefilled, so time-to-first-token no longer scales with the
# full accumulated history.
reply_2 = prefix_cache.generate(model, tokenizer, system_prompt + "Now refactor it.")

For callers that want to drive stream_generate themselves (streaming UIs, custom sampling loops), use the lower-level .fetch() / .insert() pair directly — .generate() is a convenience wrapper around exactly this:

tokens = tokenizer.encode(prompt)
cache, rest = prefix_cache.fetch(model, tokens)  # rest = only the uncached suffix

cache_key = list(tokens)
for response in mlx_lm.stream_generate(model, tokenizer, prompt=rest, prompt_cache=cache):
    cache_key.append(response.token)
    print(response.text, end="", flush=True)

prefix_cache.insert(model, cache_key, cache)  # store prompt + generated tokens together

:::warning Only exact-prefix hits for eviction methods Eviction/hybrid methods (h2o, snapkv, streaming_llm, tova, pyramidkv, and 14 others) intentionally report is_trimmable() == False, so a partial-overlap prefix can never be reused — only a byte-identical repeat of a previously seen prompt hits the cache. This is a deliberate safety boundary, not a bug: these methods keep internal importance/eviction bookkeeping (attention-score EMAs, sink/window pointers) that a generic offset-only trim() cannot roll back without risking silent corruption. PrefixCache prints a one-time note at construction when the configured method falls in this group. Compression-only methods (turboquant_rvq, vecinfer, kivi, and most others) support full partial-prefix reuse. :::

:::info Model-identity keying By default the cache is keyed on id(model), valid only for the lifetime of that Python object — reloading the same weights into a new model object starts a cold cache. Pass model_key= explicitly (any hashable, e.g. the model path) if your process reloads model objects across calls but wants cache reuse to survive that:

prefix_cache.generate(model, tokenizer, prompt, model_key="llama-3.2-3b-instruct")

This has an equivalent, not worse, risk profile to mlx_lm.server's own model_key (a plain path tuple, not a content hash) — swapping weights at the same path between calls can serve a stale cache either way. :::

Streaming generation

VeloxQuant-MLX caches work transparently with mlx_lm's streaming API:

for token in mlx_lm.stream_generate(
    model,
    tokenizer,
    prompt="Tell me a very long story.",
    max_tokens=4096,
    kv_cache=cache,
):
    print(token, end="", flush=True)

Multi-turn conversations

For multi-turn chat, reuse the same cache across turns. The cache grows across turns but retains compression:

config = KVCacheConfig(method="turboquant_rvq", bit_width_inlier=1)
caches = KVCacheBuilder.for_model(model, config)

turns = [
    "What is the capital of France?",
    "And what is it known for?",
    "What's the population?",
]

for turn in turns:
    response = mlx_lm.generate(model, tokenizer, prompt=turn, max_tokens=200, kv_cache=caches)
    print(f"User: {turn}\nAssistant: {response}\n")
    # caches now contain compressed K/V for all prior turns

warning

Cache capacity is bounded by max_seq_len. If the conversation exceeds this, use SlidingWindowKVCache to evict old tokens.

Supported models

All mlx_lm model families have been validated:

Model family	Recommended config
Llama 3.1 / 3.2 / 3.3	method="turboquant_rvq", bit_width_inlier=1
Mistral 7B / Mixtral	method="vecinfer", key_codebook_bits=8, value_codebook_bits=8
Qwen 2.5 (7–72B)	method="spectral", spectral_key_d_eff=4
Phi-3 / Phi-3.5 Mini	method="kivi", bit_width_inlier=2
Gemma 2B / 7B	method="turboquant_rvq", bit_width_inlier=2
Falcon 7B	method="turboquant_rvq" + a bit_width_inlier list from RateQuant (per-layer allocation, not a method= value)

See also

• 5-minute quickstart
• Metal kernels guide
• API — KVCacheBuilder