# mlx


<!-- WARNING: THIS FILE WAS AUTOGENERATED! DO NOT EDIT! -->

# mlx

> [mlx-lm](https://github.com/ml-explore/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](index.html) and [core](core.html).

``` python
from fastcore.test import test_eq, test_fail, test_close
```

## Loading models

MLX models are HuggingFace repos of quantized weights, most of them
under [mlx-community](https://huggingface.co/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`](https://vedicreader.github.io/rishi/mlx.html#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`.

``` python
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`](https://vedicreader.github.io/rishi/core.html#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](https://github.com/Blaizzy/mlx-vlm), which depends on mlx-lm
and adds the vision/audio towers - `pip install 'rishi[mlx-vlm]'`.

[`MlxVlmChat`](https://vedicreader.github.io/rishi/mlx.html#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.

``` python
chat = MlxChat(qwen3_06b, sp='You are concise.', think=False)
```

``` python
r = chat('Reply with exactly: pong')
assert 'pong' in resp_text(r).lower(), resp_text(r)
r
```

pong

``` python
# 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

``` python
# 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'}

``` python
# 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()
```

``` python
# 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.

``` python
# 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.

``` python
# 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

``` python
# 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()
```

``` python
# 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.

``` python
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.

``` python
# 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.

``` python
for c in (chat, vchat):
    try: c.close()
    except Exception: pass
```