from fastcore.test import test_eq, test_fail, test_closemlx
mlx
mlx-lm on Apple Silicon — explicit prompt cache, speculative decoding, quantized KV, and mlx-vlm routing for vision/audio.
Linux CI skips real-model cells (skip_exec: true). Shared Chat behavior is in index and core.
Loading models
MLX models are HuggingFace repos of quantized weights, most of them under mlx-community. Unlike litert (one .litertlm file) and llama (one .gguf file), an MLX model is a directory, so mlx_lm.load takes the repo id straight and handles the download and cache itself - there is no file-picking ladder to write here.
read_config fetches just the repo’s config.json, which is how the two things we need to know before loading are decided: the context window, and whether this is a vision/audio model that needs mlx-vlm instead of plain mlx-lm.
test_eq(ctx_len({'max_position_embeddings': 40960}), 40960)
test_eq(ctx_len({'text_config': {'max_position_embeddings': 128000}}), 128000) # VLMs nest it
test_eq(ctx_len({}), 8192) # fallback
assert is_vlm(None, cfg={'vision_config': {}})
assert is_vlm(None, cfg={'audio_config': {}})
assert not is_vlm(None, cfg={'hidden_size': 2048})
test_eq(read_config('/definitely/not/a/model'), {})The prompt cache
MLX exposes prompt-cache ownership directly. LiteRT retains KV state in its Conversation, while llama.cpp re-renders the message list and internally reuses its longest matching KV prefix. Here rishi tracks the rendered token IDs itself, so turn two prefills only what turn one did not already cover.
Keeping that correct means tracking which tokens the cache actually holds. Each turn we render the whole conversation to token ids, compare against _cache_ids with common_prefix_len, and:
- everything up to the first difference is already cached, and is skipped;
- everything after it is stale, and is trimmed off the cache before generating.
The stale part is a real case, not a theoretical one: a chat template does not necessarily re-render an assistant turn as the exact tokens the model generated, and recover_context rewrites history outright. Comparing tokens rather than trusting the cache makes both self-correcting - a mismatch just costs a re-prefill.
trim_prompt_cache removes tokens from the end of the cache, which is what we want here (it is the same thing mlx-lm’s own server does to reuse a cache across requests).
MlxChat
Same callable surface as other backends. MLX-specific: kv_bits, draft_model, adapter_path, save_cache/load_cache, and chat.use.cached_tokens from the explicit prompt cache.
Vision and audio
mlx-lm is text-only. Images and audio go through mlx-vlm, which depends on mlx-lm and adds the vision/audio towers - pip install 'rishi[mlx-vlm]'.
MlxVlmChat is picked automatically: MlxChat(repo) reads the repo’s config.json and routes to it when it finds a vision or audio tower, so Chat('mlx-community/Qwen3-VL-4B-Instruct-4bit') just works. Pass vlm=True/False to decide yourself.
It overrides only the three places mlx-vlm differs - loading, prompt building, and the generate call - and inherits history, callbacks, the tool loop and everything else unchanged. Two honest limitations: mlx-vlm has no per-model tool-call parsers (so tool calls rely on <tool_call> tags, which not every vision model emits), and it manages its own vision-feature cache rather than mlx-lm’s token KV cache, so cross-turn prefix reuse is off on this path.
Audio rides the same path: hand a Path or bytes to a model with an audio tower and mlx-vlm decodes it and resamples it to whatever rate the model’s feature extractor wants. The one thing rishi has to take back from mlx-vlm is the thinking switch - mlx-vlm asks the chat template for enable_thinking=False, which prefills an empty <think></think> block, and a model handed a finished thought can simply end the turn (Qwen3-Omni transcribes nothing at all). rishi passes its own think through instead, so the default is whatever the model’s template does on its own.
Tests
Fake-tokenizer tests run everywhere; #| eval: false cells need Apple Silicon + downloaded weights.
Examples (manual)
Use qwen3_4b (~2.5GB) unless you need vision/audio. Demos below cover cache reuse, thinking, MLX-only knobs, and VLM routing.
chat = MlxChat(qwen3_06b, sp='You are concise.', think=False)r = chat('Reply with exactly: pong')
assert 'pong' in resp_text(r).lower(), resp_text(r)
rpong
# multi-turn: the second turn must actually remember the first
chat('My favourite number is 17. Remember it.')
r = chat('What is my favourite number? Reply with just the number.')
assert '17' in resp_text(r), resp_text(r)
print(chat.use)total=75|in=72|out=3|turns=1|cached=44# Real MLX KV-cache test: capture the exact token tail handed to mlx-lm on each turn.
import time
chat2 = MlxChat(qwen3_06b, sp='You are concise.', think=False, max_output_tokens=16)
fed, orig_generate = [], chat2._generate
def counted_generate(ids, max_output_tokens=None):
fed.append(len(ids))
return orig_generate(ids, max_output_tokens)
chat2._generate = counted_generate
t0 = time.time(); chat2(('cache verification context ' * 96) + '\nReply with exactly: stored'); t1 = time.time() - t0
u1, first_feed = chat2.use, fed[-1]
t0 = time.time(); r = chat2('What exact word did I ask you to reply with? Reply with only that word.'); t2 = time.time() - t0
u2, second_feed = chat2.use, fed[-1]
print(dict(turn1_seconds=round(t1, 2), turn2_seconds=round(t2, 2), first_feed=first_feed,
second_feed=second_feed, cached=u2.cached_tokens, answer=resp_text(r)))
assert u1.cached_tokens == 0, 'first turn has nothing to reuse'
assert u2.cached_tokens > 0, 'second turn should have reused the cache'
assert second_feed < first_feed, 'turn two re-prefilled the full old prompt'
assert chat2.cached_tokens > 0
assert 'stored' in resp_text(r).lower(), resp_text(r)
chat2.close(){'turn1_seconds': 0.25, 'turn2_seconds': 0.08, 'first_feed': 315, 'second_feed': 32, 'cached': 311, 'answer': 'stored'}# and with the cache off, nothing is ever reported as cached
plain = MlxChat(qwen3_4b, prompt_cache=False, think=False)
plain('Say hi.'); plain('Say hi again.')
test_eq(plain.use.cached_tokens, 0)
plain.close()# streaming, with thinking rendered as a blockquote
thinker = MlxChat(qwen3_4b, think=True, max_output_tokens=512)
md = display_stream(thinker('Briefly: why is the sky blue?', stream=True))
assert thought(thinker.turn_res), 'expected a <think> block from a thinking model'
thinker.close()🧠 Thinking
Okay, the user is asking why the sky is blue. I need to explain this in a brief way. Let me recall the scientific explanation. The sky appears blue because of Rayleigh scattering. The atmosphere scatters sunlight, and shorter wavelengths (blue) are scattered more than longer ones (red). So when we look up, the blue light is scattered in all directions, making the sky appear blue. But wait, I should make sure I get the details right. Also, mention that the sun’s light is white, but the scattering makes it look blue. Maybe add that the reason is the way the atmosphere interacts with sunlight. Keep it simple and concise. Avoid jargon. Maybe start with the main point, then the reason, then the effect. That should cover it.
The sky appears blue because sunlight scatters more in the atmosphere when it interacts with the gases and particles in the air. Shorter wavelengths (blue light) are scattered more than longer wavelengths (red light), so when we look up, the scattered blue light dominates, giving the sky its blue color.
# quantized KV cache: same answer, less memory for long contexts
q = MlxChat(qwen3_4b, kv_bits=4, quantized_kv_start=256, think=False)
print(resp_text(q('Name three primes.')))
q.close()Sure! Here are three prime numbers: **2, 3, and 5**.
A prime number is a number greater than 1 that has no positive divisors other than 1 and itself.# speculative decoding: a small draft model proposes, the big one verifies
import time
spec = MlxChat(qwen3_8b, draft_model=qwen3_06b, think=False, max_output_tokens=256)
t0 = time.time(); spec('Write two sentences about monsoons.'); print(f'{time.time()-t0:.1f}s', spec.use)
spec.close()1.9s total=60|in=20|out=40|turns=1# a warmed cache can be saved and reloaded, so a long system prompt is prefilled once, ever
import tempfile
warm = MlxChat(qwen3_4b, sp='You are a terse assistant.', think=False)
warm('Remember the codeword: albatross.')
fn = Path(tempfile.mkdtemp())/'warm.safetensors'
warm.save_cache(fn)
n = warm.cached_tokens
warm.close()
again = MlxChat(qwen3_4b, sp='You are a terse assistant.', think=False)
again.load_cache(fn)
test_eq(again.cached_tokens, n)
again.close()# the point of a shared history format: carry a conversation from one backend to another
from rishi.llama import LlamaChat, qwen3_17b as llama_qwen
lc = LlamaChat(llama_qwen, n_ctx=4096)
lc('My name is Karthik. Remember it.')
hist = lc.hist # canonical rishi history, backend-independent
lc.close()
mc = MlxChat(qwen3_4b, messages=hist, think=False)
r = mc('What is my name? Reply with just the name.')
assert 'karthik' in resp_text(r).lower(), resp_text(r)
mc.close()
resp_text(r)llama_context: n_ctx_seq (4096) < n_ctx_train (40960) -- the full capacity of the model will not be utilized'Karthik'Images and audio
These need pip install 'rishi[mlx-vlm]' and a vision or audio model. Note the class is selected for you.
vchat = MlxChat(qwen3vl_4b, max_output_tokens=256)
test_eq(type(vchat), MlxVlmChat) # routed automatically from the repo's config.json
img = Path('images.jpeg').read_bytes()
r = vchat([img, 'What is in this image? One sentence.'])
print(resp_text(r))
assert resp_text(r).strip()[transformers] Kwargs passed to `processor.__call__` have to be in `processor_kwargs` dict, not in `**kwargs`A happy German Shepherd dog with its tongue out, standing on a gravel path outdoors.# an omni build (audio tower) transcribes; gemma-4-e4b does vision *and* audio in about 5GB
achat = MlxChat(gemma4_e4b, max_output_tokens=256)
r = achat([Path('speech.wav'), 'Transcribe this audio.'])
print(resp_text(r))
assert 'masquerade' in resp_text(r).lower()
# mlx-vlm's own default is `enable_thinking=False`, which prefills an empty `<think></think>` block
# into the reply. rishi passes `think` through instead: with the block, `qwen3omni_30b` ends the turn
# without transcribing anything at all.
assert '</think>' not in achat._prompt([], ['speech.wav'])
achat.close()/Users/71293/code/personal/orgs/ramabana/.venv/lib/python3.13/site-packages/transformers/audio_utils.py:723: UserWarning: At least one mel filter has all zero values. The value for `num_mel_filters` (128) may be set too high. Or, the value for `num_frequency_bins` (257) may be set too low.
warnings.warn(Dancing in the masquerade, idol truth in plain sight jaded. Pop, roll, click, shot. Buhulabi today or not? But such a tide as moving seems asleep, too full for sound and foam. When that which drew from out the boundless deep turns again home, twilight and evening bell, and after that.for c in (chat, vchat):
try: c.close()
except Exception: pass