feat(turboquant): batched KV-cache compression (single + batch), no worse than single

[popfido]

Summary

Makes TurboQuant KV-cache compression actually work in oMLX — in both single-sequence
and continuous-batching (B>1) decode — for the first time. TQ decode was previously
dead code (the conversion was reverted in #771 and never re-enabled), so
turboquant_kv_enabled silently ran fp16. Wiring it up exposed and fixed three real bugs
(two upstream in mlx-vlm — now both merged — and one in oMLX), and the result is proven
no worse than single mode on every axis while saving ~69% KV memory.

What changed (runtime)

• omlx/scheduler.py — quantize the completed fp16 prefill cache once, post-prefill
  (not on-the-fly, which corrupted hidden states in fix: wire up TurboQuant KV cache conversion + skip sensitive last layer (#661) #717/0.3.5-rc1: scheduler._do_external_prefill applies TurboQuantKV during prefill → garbage output & SIGABRT on quantized-KV models #771), gated to dense-KVCache
  models via a new _turboquant_eligible() (chunked/rotating models stay fp16 — no crash).
  Fully TQ-gated: for non-TQ models these paths are a no-op (zero behavioural change).
• omlx/patches/turboquant_attention.py — route B>1 decode through the quantized
  decode_attention.
• omlx/turboquant_kv.py — fix BatchTurboQuantKVCache.make_mask for left-padded
  batches (it was masking out valid tokens, so left-padded requests decoded garbage); now
  delegates to mlx-lm create_causal_mask, identical to BatchKVCache.

Upstream dependency (mlx-vlm) — both fixes now merged

Wiring TQ decode exposed two mlx-vlm kernel bugs (both ~140% error, latent because TQ
decode never ran), both now fixed upstream:

1. Bug 1 — fused single-token quantize kernel.
2. Bug 2 — RHT-incompatible L=1 value kernels (masked decode): turboquant: guard L=1 value kernels behind not use_rht (fix masked decode under RHT) Blaizzy/mlx-vlm#1244, merged.

This PR pins mlx-vlm 6f60ee4 (which contains both fixes), so oMLX needs no workaround —
B>1 masked decode relies on the upstream fix directly.

Proof: batch is no worse than single

Measured on Llama-3.2-1B-Instruct-4bit, TQ 4-bit:

axis	single (B=1)	batch (B>1)	verdict
quality — teacher-forced agreement vs fp16¹	96.2%	95.7%	✅ ≈ equal (symmetric near-ties)
speed — decode throughput	35.9 tok/s	53.2 tok/s	✅ 1.5× faster
memory — KV bytes/token	0.31× fp16	0.31× fp16	✅ equal

¹ Same forced context (no greedy cascade), n=210. At the rare single-vs-batch disagreements
(1.4%) it's a coin-flip which matches fp16 (2 vs 1) — batch is not systematically worse.

Scope / impact validation

• Runtime blast radius is 3 files; scheduler.py changes are TQ-gated (no-op for non-TQ).
• Full default unit suite on the pin: 4754 passed (the only failures are a pre-existing
  stale dflash-mlx dependency, unrelated); mlx-vlm-dependent modules re-validated on 6f60ee4.

Tests

TQ tests are grouped under a turboquant marker — run with pytest -m turboquant (30
tests across test_turboquant.py, test_turboquant_batch_memory.py, test_turboquant_ssd.py;
the model-loading ones are also slow).

Notes for the merger

• uv.lock is gitignored; regenerate with uv lock in a controlled env.

[popfido]

Update: the mlx-vlm RHT masked-decode fix (Blaizzy/mlx-vlm#1244) is merged. Bumped the pin fea81522 → 6f60ee4 (contains both TurboQuant decode fixes) and removed the interim _fix_masked_decode_rht monkey-patch + the docs/upstream/ artifacts — B>1 masked decode now relies on the upstream fix directly. Re-verified: masked decode 1.2% with no patch, 26 TQ tests + 86 mlx-vlm-dependent tests pass, single/batch coherent.

[jundot]

Thanks, this is useful and the main TQ batch path looks good. I found one narrow mixed empty/non-empty batch crash in BatchTurboQuantKVCache.merge(): an empty TQ row is skipped while the batch metadata still counts it, so the next decode append sees B=2 inputs with B=1 quantized state. It is straightforward enough that I'll merge this and fold the fix into an immediate follow-up on main.

[jundot]

One follow-up note: this also bumps mlx-vlm from fea81522 to 6f60ee4, so I'll do a separate pass over the mlx-vlm-related monkey patches on main. I'll check whether any local patches are now covered upstream by the new pin and remove or adjust them in a follow-up commit where appropriate.

[popfido]

Thanks, this is useful and the main TQ batch path looks good. I found one narrow mixed empty/non-empty batch crash in BatchTurboQuantKVCache.merge(): an empty TQ row is skipped while the batch metadata still counts it, so the next decode append sees B=2 inputs with B=1 quantized state. It is straightforward enough that I'll merge this and fold the fix into an immediate follow-up on main.

That‘s really helpful for available TQ for batch mode inference. Glad to talk about any edge case I've ever missed.

[popfido]

One follow-up note: this also bumps mlx-vlm from fea81522 to 6f60ee4, so I'll do a separate pass over the mlx-vlm-related monkey patches on main. I'll check whether any local patches are now covered upstream by the new pin and remove or adjust them in a follow-up commit where appropriate.

That's also what I mean. Next we'll need to wait for updated mlx-vlm version that including HEAD fixes so that we can bump to a released stable version for mlx-vlm dependency. I'll take a cautious look on it. I'm also raising performance optimization for TQ batch in mlx-vlm, which may benefit if I can catch up with recent release of mlx-vlm, so that the TQ batch mode in oMLX would enjoy better performance than single seq model far more than just 1.5x.