Studio: Auto disables MTP for MLA models (GLM-5.2 et al.); UNSLOTH_MLA_MTP_ENABLED to re-enable (#6468)

* Studio: Auto disables MTP for MLA models (GLM-5.2 et al.); UNSLOTH_MLA_MTP_ENABLED to re-enable

Studio's Auto speculative mode promotes any embedded-MTP model >=3B to
--spec-type draft-mtp. For MLA models (GLM-5.2/DeepSeek/Kimi) that is a
regression: llama.cpp's MLA/DSA MTP path keeps a duplicated full target-KV
context and recomputes the sparse-attention indexer every draft step, so it
runs ~2x slower than no speculation (GLM-5.2 UD-IQ1_S bench: 27 vs 45 tok/s,
flat across draft depth 1..6 and 96-100% acceptance, on both prose and code).
vLLM/SGLang get a speedup from the same model, so this is a llama.cpp
implementation gap, not a model property.

Auto now drops embedded MTP for MLA models and falls back to ngram-mod (or
spec-off when the binary lacks ngram-mod), mirroring the existing sub-3B
fallback. The metadata separator is kv_lora_rank: it is present on MLA models
and absent on non-MLA embedded-MTP models (Qwen3.x-MTP), whose MTP module is
structurally identical but fast, so a "full layer" heuristic cannot tell them
apart. Qwen MTP, separate drafters (Gemma, --model-draft), and non-MTP models
are unchanged.

Explicit overrides still engage the slower MTP route: choosing MTP / MTP+Ngram
in Settings, or passing --spec-type in extra args. UNSLOTH_MLA_MTP_ENABLED=1
re-enables Auto promotion for MLA once the upstream path is optimized.

A new spec_fallback_reason value "mla_mtp_disabled" surfaces this as an
Auto-mode policy downgrade (not a binary/update problem), with a settings
banner that points users at the MTP override. It is deliberately kept out of
the "Update llama.cpp" affordance since updating does not help.

Tests: resolver-matrix rows for MLA->ngram-mod / MLA-no-ngram->off /
non-MLA-Qwen->draft-mtp / MLA-separate-drafter->draft-mtp /
non-MTP-MLA->default / forced mtp|mtp+ngram on MLA->draft-mtp / env flag;
kv_lora_rank metadata fixtures; and reload-skip coverage (Auto ngram-mod is
idempotent, forced mtp bounces a reload).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-19 05:40:16 -07:00 committed by GitHub
commit 76a2b9edf1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 342 additions and 5 deletions

View file

@ -851,6 +851,23 @@ def _auto_mode_drops_mtp(
return req_mode == "auto" and size_b is not None and size_b < _MTP_MIN_SIZE_B
def _mla_mtp_auto_enabled() -> bool:
"""Whether Auto may pick embedded MTP for an MLA model (GLM-5.2/DeepSeek/Kimi).
Off by default: llama.cpp's MLA/DSA MTP path keeps a duplicated full target-KV
context and recomputes the sparse-attention indexer every draft step, so it runs
~2x slower than no speculation (GLM-5.2 bench: 27 vs 45 tok/s, flat across draft
depth and 96-100% acceptance) -- the opposite of the vLLM/SGLang speedup on the
same model. Set UNSLOTH_MLA_MTP_ENABLED=1 to let Auto promote MLA MTP again once
that path is optimized upstream. Forced mtp / mtp+ngram ignore this gate."""
return os.environ.get("UNSLOTH_MLA_MTP_ENABLED", "0").strip().lower() in (
"1",
"true",
"yes",
"on",
)
def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool:
"""User passed --spec-type / --spec-default? llama-server takes one
--spec-type (comma-separated to chain), so suppress auto-emit."""
@ -6169,6 +6186,17 @@ class LlamaCppBackend:
_mtp_too_small = (
_mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B and not bool(mtp_draft_path)
)
# Embedded MTP head on an MLA model (GLM-5.2/DeepSeek/Kimi, detected by
# kv_lora_rank): llama.cpp's MLA/DSA MTP path is ~2x slower than no spec,
# so Auto drops it (override via the Settings dropdown / forced mtp, or
# UNSLOTH_MLA_MTP_ENABLED=1). Separate drafters (Gemma, mtp_draft_path) and
# non-MLA embedded heads (Qwen, no kv_lora_rank) are unaffected.
_auto_mla_embedded_mtp = (
bool(self._nextn_predict_layers)
and self._kv_lora_rank is not None
and not bool(mtp_draft_path)
and not _mla_mtp_auto_enabled()
)
if user_owns_spec_type:
# User --spec-type wins outright; suppress auto-emit to avoid a
@ -6312,7 +6340,30 @@ class LlamaCppBackend:
# effective_mode == "auto": the promotion path. llama.cpp #22673:
# MTP is compatible with mmproj, so there's no vision gate.
if is_mtp_model and not _mtp_too_small:
if _auto_mla_embedded_mtp:
# MLA embedded-MTP (GLM-5.2 et al.): the MTP path regresses vs spec-off
# on llama.cpp today, so Auto drops it and falls back to ngram-mod (or
# spec-off if unsupported), mirroring the sub-3B branch. Forced mtp /
# mtp+ngram (handled above) still engage; UNSLOTH_MLA_MTP_ENABLED=1
# re-enables this promotion once upstream optimizes the path.
self._spec_fallback_reason = "mla_mtp_disabled"
_mla_caps = self.probe_server_capabilities(binary)
if _mla_caps.get("supports_ngram_mod"):
logger.info(
"Auto: MLA embedded-MTP model detected; llama.cpp's MLA/DSA "
"MTP path is slower than no speculation, so using ngram-mod "
"instead. Override via the Studio Speculative Decoding "
"dropdown or UNSLOTH_MLA_MTP_ENABLED=1."
)
_emit_ngram_mod()
else:
logger.info(
"Auto: MLA embedded-MTP model detected; disabling speculative "
"decoding (this llama-server does not advertise ngram-mod). "
"Override via the dropdown or UNSLOTH_MLA_MTP_ENABLED=1."
)
# spec-off: emit nothing, mirroring the sub-3B no-ngram path.
elif is_mtp_model and not _mtp_too_small:
# GPU: MTP-only. CPU/Mac: chain ngram-mod + MTP.
_emit_mtp(chain_ngram = not gpus)
elif is_mtp_model and _mtp_too_small:

View file

@ -412,8 +412,12 @@ class InferenceStatusResponse(BaseModel):
"(auto on an MTP model, or forced mtp / mtp+ngram). "
"'binary_no_mtp' / 'binary_outdated' -> a newer prebuilt would "
"re-enable it (show the update affordance); 'runtime_error' -> the "
"current build could not run it. None when MTP engaged or was not "
"requested."
"current build could not run it. 'mla_mtp_disabled' -> an Auto-mode "
"policy downgrade: the model is MLA (GLM-5.2 et al.) whose llama.cpp "
"MTP path runs slower than no speculation, so Auto used ngram-mod or "
"spec-off instead -- updating won't help; choose MTP in Settings (or "
"set UNSLOTH_MLA_MTP_ENABLED=1) to force it. None when MTP engaged or "
"was not requested."
),
)
llama_cpp_prebuilt_stale: bool = Field(

View file

@ -62,6 +62,7 @@ from core.inference.llama_cpp import (
_extra_args_set_any_flag,
_extra_args_set_spec_type,
_is_mtp_model_name,
_mla_mtp_auto_enabled,
)
@ -1329,6 +1330,282 @@ def test_forced_mtp_ngram_on_non_mtp_model_keeps_ngram(monkeypatch):
assert backend.requested_spec_mode == "mtp+ngram"
# ── Auto drops embedded MTP for MLA models (GLM-5.2 et al.) ───────────
#
# llama.cpp's MLA/DSA MTP path runs ~2x slower than no speculation (GLM-5.2
# bench), so Auto downgrades it to ngram-mod (or spec-off). The clean
# metadata separator from non-MLA MTP (Qwen, kept on draft-mtp) is
# self._kv_lora_rank. Forced mtp / mtp+ngram and separate drafters (Gemma)
# stay on draft-mtp; UNSLOTH_MLA_MTP_ENABLED=1 re-enables Auto promotion.
# GLM-5.2's repo name has no "MTP" marker, so its MTP signal is metadata-only
# (nextn_predict_layers) -- exactly the embedded-MLA case we gate.
_GLM_MLA_MODEL = "unsloth/GLM-5.2-GGUF"
def _mla_resolver_backend(
monkeypatch,
*,
ngram_supported = True,
kv_lora_rank = 512,
nextn = 1,
):
"""Resolver backend posing as an embedded-MTP MLA model (kv_lora_rank set)."""
backend = _resolver_backend(monkeypatch, ngram_supported = ngram_supported)
backend._nextn_predict_layers = nextn
backend._kv_lora_rank = kv_lora_rank
return backend
@pytest.mark.parametrize("gpus", [True, False])
def test_auto_mla_embedded_mtp_falls_back_to_ngram(monkeypatch, gpus):
# Auto + MLA embedded MTP + ngram supported -> ngram-mod on BOTH platforms
# (the CPU chain ngram-mod,draft-mtp is dropped: no draft-mtp for MLA).
backend = _mla_resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = "auto",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _GLM_MLA_MODEL,
model_path = None,
gpus = gpus,
binary = "/fake/llama-server",
)
parsed = _flags_dict(flags)
assert parsed.get("--spec-type") == "ngram-mod"
assert "--spec-draft-n-max" not in parsed
assert "--spec-ngram-mod-n-match" in parsed
assert backend.speculative_type == "ngram-mod"
assert backend.requested_spec_mode == "auto"
assert backend.spec_fallback_reason == "mla_mtp_disabled"
assert backend.spec_draft_n_max is None
def test_auto_mla_embedded_mtp_no_ngram_disables_spec(monkeypatch):
# Auto + MLA embedded MTP + no ngram-mod support -> emit nothing (spec-off),
# mirroring the sub-3B no-ngram path. Still flagged as a policy downgrade.
backend = _mla_resolver_backend(monkeypatch, ngram_supported = False)
flags = backend._build_speculative_flags(
speculative_type = "auto",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _GLM_MLA_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
assert "--spec-type" not in flags
assert backend.speculative_type is None
assert backend.requested_spec_mode == "auto"
assert backend.spec_fallback_reason == "mla_mtp_disabled"
def test_auto_non_mla_embedded_mtp_keeps_draft_mtp(monkeypatch):
# Auto + embedded MTP + NON-MLA (kv_lora_rank None, e.g. Qwen) -> unchanged:
# still draft-mtp at the platform default. No policy downgrade.
backend = _mla_resolver_backend(monkeypatch, kv_lora_rank = None)
flags = backend._build_speculative_flags(
speculative_type = "auto",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
parsed = _flags_dict(flags)
assert parsed.get("--spec-type") == "draft-mtp"
assert parsed.get("--spec-draft-n-max") == "2"
assert backend.speculative_type == "draft-mtp"
assert backend.spec_fallback_reason is None
def test_auto_mla_separate_drafter_keeps_mtp(monkeypatch):
# Auto + MLA + a separate drafter (mtp_draft_path) -> the drafter exemption
# wins over the MLA gate: still draft-mtp (Gemma-style external drafter is
# not the slow embedded MLA/DSA path).
backend = _mla_resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = "auto",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _GLM_MLA_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
mtp_draft_path = "/fake/mtp-draft.gguf",
)
parsed = _flags_dict(flags)
assert parsed.get("--spec-type") == "draft-mtp"
assert backend.speculative_type == "draft-mtp"
assert backend.spec_fallback_reason is None
def test_auto_non_mtp_mla_model_unaffected(monkeypatch):
# Auto + MLA but NO embedded MTP head (kv_lora_rank set, nextn None, e.g.
# GLM-4.7-Flash) -> non-MTP default; no accidental ngram drop.
backend = _mla_resolver_backend(monkeypatch, nextn = None)
flags = backend._build_speculative_flags(
speculative_type = "auto",
spec_draft_n_max = None,
extra_args = None,
model_identifier = "unsloth/GLM-4.7-Flash-GGUF",
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
assert "--spec-default" in flags
assert "ngram-mod" not in flags
assert backend.speculative_type == "default"
assert backend.spec_fallback_reason is None
@pytest.mark.parametrize(
"mode, expect_spec_type, expect_n_max",
[
("mtp", "draft-mtp", "2"),
("mtp+ngram", "ngram-mod,draft-mtp", "2"),
],
)
def test_forced_mtp_on_mla_still_engages(monkeypatch, mode, expect_spec_type, expect_n_max):
# Explicit override engages the deliberately-slower MTP route on MLA models,
# regardless of the Auto gate. No policy downgrade reason.
backend = _mla_resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = mode,
spec_draft_n_max = None,
extra_args = None,
model_identifier = _GLM_MLA_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
parsed = _flags_dict(flags)
assert parsed.get("--spec-type") == expect_spec_type
assert parsed.get("--spec-draft-n-max") == expect_n_max
assert backend.speculative_type == "draft-mtp"
assert backend.requested_spec_mode == mode
assert backend.spec_fallback_reason is None
def test_env_flag_reenables_auto_mla_mtp(monkeypatch):
# UNSLOTH_MLA_MTP_ENABLED=1 -> Auto promotes MLA embedded MTP to draft-mtp
# again (the forward hook for when llama.cpp optimizes the path).
monkeypatch.setenv("UNSLOTH_MLA_MTP_ENABLED", "1")
backend = _mla_resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = "auto",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _GLM_MLA_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
parsed = _flags_dict(flags)
assert parsed.get("--spec-type") == "draft-mtp"
assert backend.speculative_type == "draft-mtp"
assert backend.spec_fallback_reason is None
@pytest.mark.parametrize("value", ["1", "true", "yes", "on", "TRUE", "On"])
def test_mla_mtp_auto_enabled_truthy_values(monkeypatch, value):
monkeypatch.setenv("UNSLOTH_MLA_MTP_ENABLED", value)
assert _mla_mtp_auto_enabled() is True
@pytest.mark.parametrize("value", ["0", "false", "no", "off", "", " ", "bogus"])
def test_mla_mtp_auto_disabled_default_and_falsy(monkeypatch, value):
monkeypatch.setenv("UNSLOTH_MLA_MTP_ENABLED", value)
assert _mla_mtp_auto_enabled() is False
def test_mla_mtp_auto_disabled_when_unset(monkeypatch):
monkeypatch.delenv("UNSLOTH_MLA_MTP_ENABLED", raising = False)
assert _mla_mtp_auto_enabled() is False
def test_read_gguf_metadata_captures_kv_lora_rank(tmp_path):
# GLM-5.2-style header: MLA (kv_lora_rank) + embedded MTP (nextn) populate
# both fields, so the Auto gate sees an MLA embedded-MTP model.
gguf = _write_minimal_gguf(
tmp_path / "model.gguf",
arch = "glm-dsa",
nextn = 1,
extra_uint32 = {
"glm-dsa.block_count": 4,
"glm-dsa.attention.kv_lora_rank": 512,
},
)
backend = LlamaCppBackend()
backend._read_gguf_metadata(str(gguf))
assert backend._nextn_predict_layers == 1
assert backend._kv_lora_rank == 512
def test_read_gguf_metadata_qwen_mtp_has_no_kv_lora_rank(tmp_path):
# Qwen MTP header: embedded MTP but non-MLA, so kv_lora_rank stays None and
# Auto keeps it on draft-mtp.
gguf = _write_minimal_gguf(
tmp_path / "model.gguf",
arch = "qwen35moe",
nextn = 1,
extra_uint32 = {"qwen35moe.block_count": 4},
)
backend = LlamaCppBackend()
backend._read_gguf_metadata(str(gguf))
assert backend._nextn_predict_layers == 1
assert backend._kv_lora_rank is None
def test_reload_skip_auto_mla_ngram_is_idempotent():
# A GLM model resolved to ngram-mod under Auto must not churn: a duplicate
# Auto /load at the same settings is already-satisfied.
backend = _mtp_backend(
_model_identifier = _GLM_MLA_MODEL,
_speculative_type = "ngram-mod",
_requested_spec_mode = "auto",
)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = _GLM_MLA_MODEL,
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = "auto",
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is True
)
def test_reload_forced_mtp_bounces_auto_mla():
# Overriding Auto (ngram-mod) with a forced mtp request must reload (to the
# slower draft-mtp route), not dedup against the running ngram-mod server.
backend = _mtp_backend(
_model_identifier = _GLM_MLA_MODEL,
_speculative_type = "ngram-mod",
_requested_spec_mode = "auto",
)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = _GLM_MLA_MODEL,
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = "mtp",
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is False
)
# ── Full named-repo resolver matrix (the shipping Studio families) ─────
#
# Locks auto / off / forced-mtp routing for every Qwen3.5 (MTP + plain) and

View file

@ -1005,7 +1005,9 @@ export function ChatSettingsPanel({
speculativeType === "mtp+ngram") && (
<div className="rounded-lg bg-amber-500/[0.08] px-3 py-2 text-[12px] leading-[1.4] text-nav-fg/80">
<p>
{specFallbackReason === "runtime_error"
{specFallbackReason === "mla_mtp_disabled"
? "MTP is disabled by default for this model architecture because it currently runs slower than standard decoding. Select MTP above to force it."
: specFallbackReason === "runtime_error"
? "MTP could not start for this model on the installed llama.cpp build, so it is running without speculative decoding."
: "MTP is not available in the installed llama.cpp build, so this model is running without it." +
(llamaUpdateStatus?.update_available

View file

@ -196,7 +196,10 @@ export interface InferenceStatusResponse {
/**
* Why MTP was disabled on the loaded model despite being requested.
* "binary_no_mtp" / "binary_outdated" -> updating llama.cpp would re-enable
* it; "runtime_error" -> the current build could not run it. Null otherwise.
* it; "runtime_error" -> the current build could not run it;
* "mla_mtp_disabled" -> an Auto-mode policy downgrade for MLA models
* (GLM-5.2 et al.) whose llama.cpp MTP path is slower than no speculation
* (updating won't help; choose MTP in Settings to force it). Null otherwise.
*/
spec_fallback_reason?: string | null;
}