diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index bf8a3c04df..bc84c5f9f0 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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,19 @@ 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 [ - "--spec-ngram-size-n", - str(n_match), - "--draft-min", - str(n_min), - "--draft-max", - str(n_max), - ] + out = ["--spec-ngram-size-n", str(n_match)] + if not chain_with_mtp: + # Gate --draft-min AND --draft-max together. Both flags are + # generic on legacy binaries, so emitting either in the same + # invocation as MTP would race with MTP's own --draft-min / + # --draft-max (typically 2/3 from spec_draft_n_max). Earlier + # we suppressed only --draft-max, which produced an inverted + # legacy range (--draft-min 48 --draft-max 2/3) and disabled + # ngram-mod entirely. Keeping the pair together preserves a + # valid range when ngram is standalone and avoids the clobber + # when it is chained with MTP. + out.extend(["--draft-min", str(n_min), "--draft-max", str(n_max)]) + return out return [] @@ -531,6 +548,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 +3530,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 +3568,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 +3735,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 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 02270ab405..1d397cdc94 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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 diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 4a8276adc0..8bb890c0b1 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -26,27 +26,54 @@ _structlog_stub = _types.ModuleType("structlog") _structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") sys.modules.setdefault("structlog", _structlog_stub) -_httpx_stub = _types.ModuleType("httpx") -for _exc in ( - "ConnectError", - "TimeoutException", - "ReadTimeout", - "ReadError", - "RemoteProtocolError", - "CloseError", -): - setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) -_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) -_httpx_stub.Client = type( - "C", - (), - { - "__init__": lambda s, **kw: None, - "__enter__": lambda s: s, - "__exit__": lambda s, *a: None, - }, -) -sys.modules.setdefault("httpx", _httpx_stub) +# Prefer real httpx when it is installed: route-guard tests import +# ``routes.inference`` which transitively pulls ``huggingface_hub`` +# (uses ``from httpx import Response``) and a stub that lacks +# ``Response`` would break that import chain. Fall back to the stub +# only when httpx is genuinely missing (keeps the legacy local-dev +# fast path for the unit-level tests that never touch the route module). +try: + import httpx as _real_httpx # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "HTTPStatusError", + "RequestError", + ): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) + _httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) + _httpx_stub.Response = type("R", (), {"__init__": lambda s, *a, **k: None}) + _httpx_stub.Request = type("Req", (), {"__init__": lambda s, *a, **k: None}) + _httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__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["httpx"] = _httpx_stub import pytest @@ -1198,3 +1225,314 @@ 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_min_and_max(): + """Legacy chained ngram+MTP must drop BOTH --draft-min and --draft-max. + Both are generic flags on legacy binaries, so MTP's own --draft-min / + --draft-max (e.g. 2/3 from spec_draft_n_max) would race with ngram's + larger size-N range. Suppressing only --draft-max produced an inverted + legacy range (--draft-min 48 --draft-max 2/3) that disabled ngram-mod + on affected builds; gate them together. --spec-ngram-size-n is a + distinct flag and must remain so ngram-mod still tunes the chain.""" + 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}" + assert ( + "--draft-min" not in chained + ), f"chain_with_mtp=True must drop --draft-min on legacy; got {chained}" + assert "--spec-ngram-size-n" in chained + + +def test_build_ngram_mod_flags_legacy_standalone_keeps_draft_min_and_max(): + """Standalone ngram on legacy must emit --draft-min AND --draft-max as + a valid range (size-N min/max) so ngram-mod actually has a window.""" + caps = {"ngram_mod_flavor": "legacy"} + standalone = _build_ngram_mod_flags(caps, chain_with_mtp = False) + assert "--draft-min" in standalone + assert "--draft-max" in standalone + # The range must not be inverted: min <= max. + i_min = standalone.index("--draft-min") + i_max = standalone.index("--draft-max") + assert int(standalone[i_min + 1]) <= int(standalone[i_max + 1]) + + +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