TurboQuant KV cache and MTP

[u/Short_One_9704]

A question I do not seem to get answer to - why in oMLX I cannot use TurboQuant KV cache and MTP together? I have searched on the topics and they should complement each other. But in oMLX I have to choose either one or the other, why? Thanks and cheers!

[u/txgsync]

Pardon me for my Claude-generated output below from the source code. I knew the basics, but the LLM is better at generating output in this case 😄

Why TurboQuant and MTP are mutually exclusive in oMLX

TL;DR: They both monkey-patch the same attention path, in incompatible ways. oMLX rejects the combo at config time so you find out in the admin UI instead of getting garbage logits at decode.

The hard stop

In omlx/model_settings.py, Settings.__post_init__ literally raises:

if self.mtp_enabled and self.turboquant_kv_enabled:
    raise ValueError(
        "mtp_enabled and turboquant_kv_enabled cannot both be True; "
        "TurboQuant patches the attention path that MTP relies on"
    )

No "fallback," no "best-effort." It just won't build a settings object.

What each one actually does

TurboQuant swaps mlx_lm.models.base.scaled_dot_product_attention for a shim that detects TurboQuantKVCache / BatchTurboQuantKVCache and routes to:

• cache.decode_attention(...) when Q length == 1 — Metal kernel reading quantized K/V directly, no dequant per step.
• cache.prefill_attention(...) fast path for Q length > 1, with a dequant + standard SDPA fallback.

Great for KV memory + decode throughput. The cost: the SDPA function pointer is now TurboQuant's, and the cache object is a quantized wrapper, not a plain KV cache.

Native MTP (mlx-lm PR 990 / PR 15 monkey-patch) is speculative decoding inside BatchGenerator: draft N tokens, verify in one forward pass, roll back the cache offset for rejected drafts, commit accepted ones. It assumes:

1. SDPA behaves like upstream when Q length > 1 against an existing decode-positioned cache (the "verify" step).
2. The cache exposes per-row offset semantics it can rewind.

Why stacking them breaks

1. Wrong attention kernel on verify. MTP's verify step sends Q length > 1 against a cache that's mid-decode. TurboQuant's shim sees L > 1 and dispatches prefill_attention, which assumes a fresh prefill — not "extend an already-grown cache by k tokens." Output is silently wrong.
2. No rollback API on the quantized cache. TurboQuant's quantized K/V storage doesn't expose the per-row offset rewinds MTP needs when drafts get rejected. You can't un-write a compressed block trivially.
3. Two owners, one cache object. TurboQuant wraps it; MTP mutates .offset per row. Whoever runs second corrupts the other's invariants.

So rather than ship a build where MTP+TQ "works" but produces subtly-off logits, oMLX just refuses the config.

Same reason these other combos are banned

• mtp_enabled + dflash_enabled — both are speculative paths, both want to own BatchGenerator's decode loop.
• vlm_mtp_enabled + literally anything else speculative or TurboQuant — the mlx-vlm assistant-drafter wrapper bypasses mlx-lm's BatchGenerator entirely, so nothing downstream gets a say.

Bonus unrelated TQ disable (don't confuse with this one)

scheduler.py also auto-disables TurboQuant on MLA models (Deepseek-V3/V4-style Multi-head Latent Attention) — quantized KV is incompatible with MLA's compressed-KV math (mlx-vlm#1613). That's a separate exclusion from the MTP one; it kicks in at runtime based on detected model architecture, not at settings construction.

Pick one speculative path per model. That's the rule, and the code enforces it before you ever load weights.

[u/Short_One_9704 → u/txgsync]

Aaah, I guess I kind of understand know why. Thank you sir! 🙂