Summary

Multi-head Latent Attention (MLA) models crash the engine loop when TurboQuant is enabled:

omlx.scheduler   - INFO  - TurboQuant: converted 46/47 cache layers to 4.0-bit, skipped last KVCache layer
omlx.engine_core - ERROR - Engine loop error: 'TurboQuantMSEState' object has no attribute 'swapaxes'

This PR makes Scheduler._turboquant_eligible() recognise MLA architectures and keep them on an fp16 KV cache — no crash, no TurboQuant — with a one-time info log. Fixes #1613.

Root cause

MLA attention (DeepSeek-V2/V3, GLM-4-MoE / GLM-4.7-Flash, MiniCPM3, …) does not delegate everything to scaled_dot_product_attention(cache=...). It reads the fetched cache tensors directly:

kv_latent, k_pe = cache.update_and_fetch(kv_latent, k_pe)
pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)   # <-- crash

(mlx_lm/models/glm4_moe_lite.py, deepseek_v3.py, …)

TurboQuant replaces each KV layer's state with a per-head quantized TurboQuantMSEState NamedTuple (norms, indices) that has no array methods, so .swapaxes (and the absorbed-MLA embed_q/unembed_out matmuls) raise AttributeError. MLA also stores the latent (kv_lora_rank, e.g. 512) and rope key (qk_rope_head_dim, e.g. 64) with mismatched head dims and a single head, which the per-head codec was never designed for.

This is consistent with the rest of the ecosystem — TurboQuant / KV-cache vector quant is standard-attention (MHA/GQA) only in mlx-vlm, vLLM, and SGLang; MLA always uses a separate KV path (FP8 grouped latent + bf16 rope) or none at all (MLX). So excluding MLA matches how every implementation treats it.

The fix

A memoized _model_uses_mla() detector with three signals, called at the top of _turboquant_eligible():

1. Config — kv_lora_rank (int), including nested text_config/llm_config and VLM-adapter _language_model/language_model delegation.
2. Architecture — module walk for kv_a_proj_with_mqa / kv_a_layernorm / int kv_lora_rank.

_turboquant_eligible() returns False for MLA, gating both conversion call sites (empty + post-prefill convert).

Scope / safety — does not affect other functionality

• The change is a single early-return in the existing eligibility gate. For non-MLA models the detector returns False immediately and the original cache-type logic runs byte-for-byte identically — TurboQuant behaviour for dense/GQA models is unchanged.
• Detection is memoized (model never changes per scheduler), so there is no per-request cost; the module walk only runs once and only if the config signal misses.
• No public API, settings, or cache-format changes. MLA models simply keep the fp16 KV cache they already used before TurboQuant was toggled on.

Coverage

Verified against all MLA families in the installed mlx-lm: deepseek_v2, deepseek_v3, deepseek_v32, glm4_moe_lite, glm_moe_dsa, kimi_linear, kimi_vl, longcat_flash, longcat_flash_ngram, minicpm3, youtu_llm — each caught by the config signal (top-level or nested) and/or the architecture walk. No false positives: only MLA models define kv_lora_rank.

Test plan

New tests/test_scheduler.py::TestTurboQuantMLAGuard (6 cases):

• MLA detected by config, by architecture, by nested text_config, and via VLM-adapter _language_model
• standard MHA/GQA model stays eligible
• detection is memoized (single module walk)

pytest tests/test_scheduler.py::TestTurboQuantMLAGuard tests/test_turboquant.py tests/test_scheduler.py
# 134 passed
pytest            # full default suite: 4892 passed, 37 skipped

(The only failures in the full run are 4 pre-existing test_dflash_engine.py cases failing on a dflash_mlx.runtime.config import in this environment — unrelated to this change.)