ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

studio: tighten MTP reload guards and asymmetric spec flags for #5582

Browse source 
Five issues surfaced after #5582 merged. All addressed with matching
pytest coverage (15 new tests, 147 total green).

Bug A -- route guard compared against the requested UI mode rather than
the backend's resolved spec mode. A user request setting
``spec_draft_n_max=2`` against a backend that was auto-promoted from
``auto`` -> ``draft-mtp`` saw ``requested_spec_mode == "auto"`` (not in
``("mtp", "mtp+ngram")``) and skipped the comparison, returning
``already_loaded`` with the stale value still active. Now mirrors the
backend-side guard's check against ``speculative_type == "draft-mtp"``.

Bug B -- both reload guards short-circuited the n_max comparison when
the request value was ``None``, treating it as a wildcard. A backend
loaded with an explicit override of 8 could never be cleared back to
the platform default without swapping the model. Both guards now treat
the ``(None vs explicit)`` flip as a difference: clear-to-default and
set-from-default both bounce a reload, while ``(None == None)`` and
``(N == N)`` continue to match.

Bug C -- chained MTP+ngram on a legacy llama-server (pre arg-rename)
emitted ``--draft-max`` twice: once for MTP's draft length (e.g. 2),
once for ngram-mod's size-N max (e.g. 64). llama-server's last-wins
parsing clobbered the MTP value with 64, defeating the
``--spec-draft-n-max`` slider. ``_build_ngram_mod_flags`` now takes a
``chain_with_mtp`` kwarg that suppresses ``--draft-max`` on the legacy
flavor when MTP will emit it; the post-rename flavor uses distinct
``--spec-ngram-mod-*`` names that cannot collide.

Bug D -- a forced ``speculative_type="ngram"`` request emitted
``--spec-type ngram-mod`` even on binaries that did not advertise
ngram-mod support, causing llama-server to refuse to start. The auto
path already checked ``supports_ngram_mod`` before emitting; the
forced path now mirrors that check and loads without spec (with a
warning that matches the MTP-token-missing path).

Bug E -- ``speculative_type="none"`` is llama.cpp's own explicit-disable
spelling, and external API callers commonly use ``"disable"`` /
``"disabled"``. None of these were in the canonical spec mode set or
the legacy alias map, so they fell through to ``"auto"`` and silently
re-enabled MTP -- the opposite of the user's intent. Added all three
to ``_LEGACY_SPEC_MODE_MAP`` as aliases for ``"off"``.

Tests
-----
- test_canonicalize_spec_mode_none_aliases_map_to_off (6 cases via
  parametrize): "none"/"None"/"NONE"/"  none  "/"disable"/"Disabled"
  all canonicalise to "off".
- test_build_ngram_mod_flags_legacy_chained_omits_draft_max +
  test_build_ngram_mod_flags_legacy_standalone_keeps_draft_max +
  test_build_ngram_mod_flags_new_flavor_always_emits_distinct_names:
  the chain_with_mtp kwarg suppresses only the legacy flavor's
  --draft-max, never the new-flavor knobs.
- test_build_speculative_flags_chained_mtp_ngram_legacy_no_duplicate_draft_max:
  end-to-end check that the assembled spec block has exactly one
  --draft-max carrying the MTP draft length.
- test_build_speculative_flags_forced_ngram_without_support_skips_spec
  + test_build_speculative_flags_forced_ngram_with_support_emits_spec:
  forced ngram refuses on a binary lacking ngram-mod support; still
  emits cleanly on a supporting binary.
- test_already_in_target_state_{clear,set}_explicit_n_max_*_forces_reload:
  backend-side guard covers both clear-to-default and set-from-default.
- test_route_guard_auto_promoted_mtp_{bounces,matches,clear_*}: route
  guard now compares against resolved spec mode and handles the None
  flip symmetrically.
- test_route_guard_ignores_n_max_when_resolved_spec_is_not_mtp:
  non-MTP resolved spec (e.g. ngram-mod) still ignores n_max.

147/147 spec/reload test suites green.
This commit is contained in:
Daniel Han 2026-05-22 13:10:08 +00:00 • committed by Daniel Han
commit 1b621cd9e7
3 changed files with 399 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

76
studio/backend/core/inference/llama_cpp.py
View file

@ -487,6 +487,8 @@ def _build_ngram_mod_flags(
n_match: int = 24,
n_min: int = 48,
n_max: int = 64,
*,
chain_with_mtp: bool = False,
) -> list[str]:
"""Emit the right ngram-mod knob flags for the running llama-server.
@ -496,9 +498,19 @@ def _build_ngram_mod_flags(
``probe_server_capabilities``; ``ngram_mod_flavor`` tells us which
set is real (vs a removal-stub entry). Returns ``[]`` when neither
set is available so the caller can drop ngram-mod entirely.
``chain_with_mtp`` is set by the chained MTP+ngram emit path on
legacy builds, where ``--draft-max`` is shared between MTP (number
of draft tokens per pass; small, e.g. 2) and ngram-mod (max ngram
size N; large, e.g. 64). Emitting both produces a duplicate flag
and last-wins clobbers the MTP draft-length the user picked. When
set, skip the ngram-mod ``--draft-max`` so MTP's emission wins.
"""
flavor = caps.get("ngram_mod_flavor") if caps else None
if flavor == "new":
# New-flavor knobs are distinct flag names (``--spec-ngram-mod-*``),
# so they never collide with MTP's ``--spec-draft-n-max``; emit all
# three regardless of chain_with_mtp.
return [
"--spec-ngram-mod-n-match",
str(n_match),
@ -511,14 +523,18 @@ def _build_ngram_mod_flags(
# Legacy llama.cpp before the spec arg rename: same knobs lived
# under --spec-ngram-size-n (lookup length) and the generic
# --draft-min / --draft-max (ngram size N range).
return [
out = [
"--spec-ngram-size-n",
str(n_match),
"--draft-min",
str(n_min),
"--draft-max",
str(n_max),
]
if not chain_with_mtp:
# Only safe to set --draft-max here when MTP is NOT in the
# same emission; otherwise this duplicates MTP's --draft-max
# and last-wins clobbers the MTP draft length.
out.extend(["--draft-max", str(n_max)])
return out
return []
@ -531,6 +547,15 @@ _LEGACY_SPEC_MODE_MAP = {
"default": "auto",
"draft-mtp": "mtp",
"ngram-mod": "ngram",
# llama.cpp's own ``--spec-type none`` spelling, plus the common
# English ``disable`` / ``disabled`` aliases external API callers
# use, all mean "do not engage speculative decoding". Without an
# explicit mapping these strings fall through the comma-parser
# below to ``auto``, which silently enables MTP -- the opposite of
# the user's intent. Map them to canonical ``off``.
"none": "off",
"disable": "off",
"disabled": "off",
}
@ -3504,7 +3529,10 @@ class LlamaCppBackend:
draft_n_max = _resolved_draft_n_max()
n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max"
if chain_ngram:
ngram_knobs = _build_ngram_mod_flags(caps)
# chain_with_mtp suppresses ngram-mod's --draft-max on
# legacy builds so it does not collide with MTP's draft
# length set just above.
ngram_knobs = _build_ngram_mod_flags(caps, chain_with_mtp = True)
if ngram_knobs:
spec_value = f"ngram-mod,{mtp_token}"
else:
@ -3539,9 +3567,27 @@ class LlamaCppBackend:
def _emit_ngram_mod() -> bool:
"""Append --spec-type ngram-mod + flag-set knobs."""
ngram_caps = self.probe_server_capabilities(binary)
# Skip emission when the binary advertises no ngram-mod
# support at all (no knobs probed under either flavor).
# llama-server rejects ``--spec-type ngram-mod`` on those
# builds, so emitting it would refuse to start the server
# instead of silently disabling spec. Mirror the auto-path
# fallback at the elif _mtp_too_small branch below: log and
# return False, leaving the caller to load without spec.
if not (ngram_caps and ngram_caps.get("supports_ngram_mod")):
logger.warning(
"Requested ngram-mod speculative decoding but "
"llama-server does not advertise ngram-mod support; "
"run `unsloth studio update`. Loading without "
"speculative decoding."
)
return False
ngram_knobs = _build_ngram_mod_flags(ngram_caps)
flags.extend(["--spec-type", "ngram-mod"])
if not ngram_knobs:
# supports_ngram_mod is True but knobs returned empty
# (shouldn't happen given the flavor check, but kept
# for defensive logging).
logger.warning(
"llama-server lacks ngram-mod tuning "
"flags; loading without --spec-ngram-mod-* knobs"
@ -3688,12 +3734,22 @@ class LlamaCppBackend:
# engaged. Compare on the resolved spec rather than the requested
# mode so an Auto request that auto-promoted to draft-mtp under
# the hood still bounces a reload when the user changes n_max.
if (
self._speculative_type == "draft-mtp"
and spec_draft_n_max is not None
and int(spec_draft_n_max) != (self._spec_draft_n_max or 0)
):
return False
#
# ``_spec_draft_n_max`` is None when the backend was loaded with
# the platform default (no explicit override); ``None`` on either
# side means "platform default". Reload when:
# * both sides explicit and differ, OR
# * request is None but backend has an explicit value (user is
# clearing the override back to default), OR
# * request is explicit but backend is on default (user is
# setting a fresh override).
if self._speculative_type == "draft-mtp":
req_n = spec_draft_n_max
backend_n = self._spec_draft_n_max
if (req_n is None) != (backend_n is None):
return False
if req_n is not None and int(req_n) != int(backend_n):
return False
if (self._chat_template_override or None) != (chat_template_override or None):
return False

19
studio/backend/routes/inference.py
View file

@ -504,10 +504,21 @@ def _request_matches_loaded_settings(
backend_mode = llama_backend.requested_spec_mode or "auto"
if req_mode != backend_mode:
return False
# spec_draft_n_max only matters when an MTP variant is engaged; None
# means "platform default" and matches whatever the backend chose.
if backend_mode in ("mtp", "mtp+ngram") and request.spec_draft_n_max is not None:
if int(request.spec_draft_n_max) != (llama_backend.spec_draft_n_max or 0):
# spec_draft_n_max only matters when an MTP variant is ACTUALLY
# engaged. Mirror the backend-side guard's check against the
# RESOLVED spec mode (``llama_backend.speculative_type``) rather
# than the requested UI mode -- an Auto request that auto-promoted
# to ``draft-mtp`` has ``requested_spec_mode = "auto"`` but still
# honours user n_max changes. ``None`` on either side means
# "platform default"; reload when the explicit/default state flips
# (clear-to-default or set-from-default) or when both explicit
# values differ.
if llama_backend.speculative_type == "draft-mtp":
req_n = request.spec_draft_n_max
backend_n = llama_backend.spec_draft_n_max
if (req_n is None) != (backend_n is None):
return False
if req_n is not None and int(req_n) != int(backend_n):
return False
if (request.chat_template_override or None) != (
llama_backend.chat_template_override or None

318
studio/backend/tests/test_llama_cpp_mtp_detection.py
View file

@ -34,6 +34,9 @@ for _exc in (
"ReadError",
"RemoteProtocolError",
"CloseError",
"HTTPError",
"HTTPStatusError",
"RequestError",
):
setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None})
@ -46,6 +49,19 @@ _httpx_stub.Client = type(
"__exit__": lambda s, *a: None,
},
)
# AsyncClient instantiated at module load time in
# core/inference/external_provider.py; tests that import routes.inference
# need this stub. The methods do not have to work -- nothing in the
# route reload guard exercises them.
_httpx_stub.AsyncClient = type(
"AC",
(),
{
"__init__": lambda s, **kw: None,
"__aenter__": lambda s: s,
"__aexit__": lambda s, *a: None,
},
)
sys.modules.setdefault("httpx", _httpx_stub)
import pytest
@ -1198,3 +1214,305 @@ def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch):
# _requested_spec_mode still reflects the user's choice.
assert backend.requested_spec_mode == "mtp"
assert backend.speculative_type is None
# ---------------------------------------------------------------------------
# Followup regression tests for #5582.
#
# Bug A: route guard at routes/inference.py:509 compared `request.spec_draft_
# n_max` against the requested UI mode (which is "auto" for an
# auto-promoted draft-mtp load), so a slider change while running
# under Auto-promoted MTP returned `already_loaded` and kept the
# stale value.
# Bug B: both reload guards short-circuited the n_max comparison when the
# request cleared back to None ("platform default"), so an explicit
# override could never be cleared without a model swap.
# Bug C: legacy chained MTP+ngram emitted --draft-max twice (MTP's draft
# length, then ngram's size-N max); last-wins clobbered the MTP
# value.
# Bug D: forced ngram standalone emitted --spec-type ngram-mod even on
# binaries that did not advertise ngram-mod support, causing
# llama-server to refuse to start.
# Bug E: speculative_type="none" (the spelling llama.cpp itself uses for
# the disable case, as well as "disable" / "disabled" from
# external API callers) fell through the canonicaliser to "auto",
# silently enabling MTP when the user said disable.
# ---------------------------------------------------------------------------
# ---- Bug E: "none"/"disable" canonicalise to "off" ----
import pytest as _pytest
@_pytest.mark.parametrize(
"value",
["none", "None", "NONE", " none ", "disable", "Disabled", "DISABLED"],
)
def test_canonicalize_spec_mode_none_aliases_map_to_off(value):
"""Without the followup these all fell through to "auto" and silently
re-enabled MTP. They must canonicalise to "off"."""
assert _canonicalize_spec_mode(value) == "off"
# ---- Bug C: legacy chained MTP+ngram does not duplicate --draft-max ----
def test_build_ngram_mod_flags_legacy_chained_omits_draft_max():
"""When chaining ngram-mod alongside MTP on a legacy llama-server,
--draft-max collides with MTP's --draft-max (same flag, different
semantic in old builds). The chain_with_mtp flag must suppress the
ngram-mod --draft-max so MTP's draft length wins."""
caps = {"ngram_mod_flavor": "legacy"}
chained = _build_ngram_mod_flags(caps, chain_with_mtp = True)
assert "--draft-max" not in chained, (
f"chain_with_mtp=True must drop --draft-max on legacy; got {chained}"
)
# The other two knobs must still be present so ngram-mod actually
# tunes the chain.
assert "--spec-ngram-size-n" in chained
assert "--draft-min" in chained
def test_build_ngram_mod_flags_legacy_standalone_keeps_draft_max():
"""When NOT chaining, --draft-max is the ngram-mod size-N max and must
still be emitted on legacy."""
caps = {"ngram_mod_flavor": "legacy"}
standalone = _build_ngram_mod_flags(caps, chain_with_mtp = False)
assert "--draft-max" in standalone
def test_build_ngram_mod_flags_new_flavor_always_emits_distinct_names():
"""The post-rename flavor uses --spec-ngram-mod-* names that never
collide with MTP's --spec-draft-n-max, so chain_with_mtp does not
matter -- the full knob set is always emitted."""
caps = {"ngram_mod_flavor": "new"}
chained = _build_ngram_mod_flags(caps, chain_with_mtp = True)
standalone = _build_ngram_mod_flags(caps, chain_with_mtp = False)
assert chained == standalone
assert "--spec-ngram-mod-n-max" in chained
def test_build_speculative_flags_chained_mtp_ngram_legacy_no_duplicate_draft_max(
monkeypatch,
):
"""End-to-end: forced mtp+ngram on a legacy llama-server must emit
--draft-max exactly once (the MTP draft length), not twice with a
last-wins overwrite from ngram-mod."""
fake = {
"found": True,
"mtp_token": "mtp",
"supports_mtp": True,
"ngram_mod_flavor": "legacy",
"supports_ngram_mod": True,
"spec_draft_n_max_flag": "--draft-max", # legacy n_max flag
}
monkeypatch.setattr(
LlamaCppBackend,
"probe_server_capabilities",
classmethod(lambda cls, binary = None: fake),
)
backend = LlamaCppBackend()
backend._nextn_predict_layers = 1 # is_mtp_model
flags = backend._build_speculative_flags(
speculative_type = "mtp+ngram",
spec_draft_n_max = 2,
extra_args = None,
model_identifier = _MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
occurrences = [i for i, t in enumerate(flags) if t == "--draft-max"]
assert len(occurrences) == 1, (
f"--draft-max must appear exactly once on legacy chained MTP+ngram; "
f"got {len(occurrences)}: {flags}"
)
# And the value must be MTP's choice (2), not the ngram size-N max (64).
idx = occurrences[0]
assert flags[idx + 1] == "2", (
f"the single --draft-max must carry the MTP draft length, not "
f"ngram's size-N max; got {flags[idx + 1]!r} in {flags}"
)
# ---- Bug D: forced ngram refuses on binaries with no ngram-mod support ----
def test_build_speculative_flags_forced_ngram_without_support_skips_spec(monkeypatch):
"""A forced ``speculative_type="ngram"`` request on a binary that
does not advertise ngram-mod support must NOT emit --spec-type
ngram-mod (llama-server would refuse to start); load without spec
instead, mirroring the auto-path sub-3B fallback."""
backend = _resolver_backend(monkeypatch, ngram_supported = False)
backend._nextn_predict_layers = None
flags = backend._build_speculative_flags(
speculative_type = "ngram",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _NON_MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
assert "--spec-type" not in flags, (
f"forced ngram must not emit --spec-type when binary lacks "
f"ngram-mod support; got {flags}"
)
assert backend.speculative_type is None
# User's UI choice is preserved on the requested-mode round-trip.
assert backend.requested_spec_mode == "ngram"
def test_build_speculative_flags_forced_ngram_with_support_emits_spec(monkeypatch):
"""Sanity check the positive case: forced ngram on a supporting
binary still emits --spec-type ngram-mod plus the knob set."""
backend = _resolver_backend(monkeypatch, ngram_supported = True)
backend._nextn_predict_layers = None
flags = backend._build_speculative_flags(
speculative_type = "ngram",
spec_draft_n_max = None,
extra_args = None,
model_identifier = _NON_MTP_MODEL,
model_path = None,
gpus = True,
binary = "/fake/llama-server",
)
parsed = _flags_dict(flags)
assert parsed.get("--spec-type") == "ngram-mod"
assert backend.speculative_type == "ngram-mod"
# ---- Bug B: clear-to-None forces reload on the backend guard ----
def test_already_in_target_state_clear_explicit_n_max_to_none_forces_reload():
"""Backend loaded with explicit ``spec_draft_n_max=8``; new request
clears the value to None (platform default). Without the followup
the guard short-circuited on ``spec_draft_n_max is not None`` and
returned True, leaving the old 8 in effect. Must now return False."""
backend = _mtp_backend(_spec_draft_n_max = 8)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
spec_draft_n_max = None,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is False
)
def test_already_in_target_state_set_n_max_from_default_forces_reload():
"""Mirror: backend loaded on default (None); new request adds an
explicit 8. Must reload."""
backend = _mtp_backend(_spec_draft_n_max = None)
assert (
backend._already_in_target_state(
gguf_path = None,
model_identifier = "unsloth/Qwen3.6-27B-MTP-GGUF",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = None,
spec_draft_n_max = 8,
chat_template_override = None,
extra_args = None,
is_vision = False,
)
is False
)
# ---- Bug A + Bug B: route-level guard checks resolved mode + handles None ----
class _FakeLoadRequest:
"""Minimal LoadRequest stand-in for the route guard test."""
def __init__(self, **kw):
self.max_seq_length = kw.get("max_seq_length", 8192)
self.cache_type_kv = kw.get("cache_type_kv", None)
self.speculative_type = kw.get("speculative_type", "auto")
self.spec_draft_n_max = kw.get("spec_draft_n_max", None)
self.chat_template_override = kw.get("chat_template_override", None)
self.llama_extra_args = kw.get("llama_extra_args", None)
def _auto_promoted_mtp_backend(**overrides):
"""Backend that's running under draft-mtp because Auto auto-promoted
it (requested_spec_mode == "auto", speculative_type == "draft-mtp").
"""
backend = _mtp_backend(
_requested_spec_mode = "auto",
_speculative_type = "draft-mtp",
_spec_draft_n_max = 8,
**overrides,
)
backend._extra_args = None
return backend
def test_route_guard_auto_promoted_mtp_bounces_on_n_max_change():
"""Without the followup, route guard compared n_max only when
``backend_mode in ("mtp", "mtp+ngram")`` -- but an Auto-promoted
backend has ``requested_spec_mode == "auto"``, so a slider change
from 8 -> 2 returned ``already_loaded`` and kept the stale value.
Now must compare against the RESOLVED speculative_type."""
from routes.inference import _request_matches_loaded_settings
backend = _auto_promoted_mtp_backend()
request = _FakeLoadRequest(
speculative_type = "auto",
spec_draft_n_max = 2, # was 8
)
assert _request_matches_loaded_settings(request, backend) is False
def test_route_guard_auto_promoted_mtp_matches_when_n_max_unchanged():
"""Same setup; same n_max value (8) on both sides; must still match."""
from routes.inference import _request_matches_loaded_settings
backend = _auto_promoted_mtp_backend()
request = _FakeLoadRequest(
speculative_type = "auto",
spec_draft_n_max = 8,
)
assert _request_matches_loaded_settings(request, backend) is True
def test_route_guard_clear_explicit_n_max_to_none_forces_reload():
"""Backend has explicit 8 (auto-promoted to draft-mtp); request
clears to None. Route guard must return False so the reload path
runs and re-resolves the default value."""
from routes.inference import _request_matches_loaded_settings
backend = _auto_promoted_mtp_backend()
request = _FakeLoadRequest(
speculative_type = "auto",
spec_draft_n_max = None,
)
assert _request_matches_loaded_settings(request, backend) is False
def test_route_guard_ignores_n_max_when_resolved_spec_is_not_mtp():
"""Backend resolved to ngram-mod (not draft-mtp); n_max is MTP-only
and must not force a reload."""
from routes.inference import _request_matches_loaded_settings
backend = _mtp_backend(
_requested_spec_mode = "ngram",
_speculative_type = "ngram-mod",
_spec_draft_n_max = None,
)
backend._extra_args = None
request = _FakeLoadRequest(
speculative_type = "ngram",
spec_draft_n_max = 8, # would force reload if we checked, but shouldn't
)
assert _request_matches_loaded_settings(request, backend) is True