diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 85ca32eaf5..7d67097d09 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -514,9 +514,18 @@ _CTX_FIT_VRAM_FRACTION = 0.90 _MTP_VRAM_RESERVE_FRAC = 0.05 -def _auto_mode_drops_mtp(req_mode: Optional[str], size_b: Optional[float]) -> bool: - """Auto mode drops MTP below _MTP_MIN_SIZE_B (draft-mtp regresses there); - forced mtp / mtp+ngram engage regardless of size.""" +def _auto_mode_drops_mtp( + req_mode: Optional[str], + size_b: Optional[float], + *, + has_separate_drafter: bool = False, +) -> bool: + """Auto mode drops MTP below _MTP_MIN_SIZE_B for an embedded draft head + (its per-token cost regresses there); a separate drafter (Gemma) is a tiny + standalone model that still speeds up below 3B, so it never drops. Forced + mtp / mtp+ngram engage regardless of size.""" + if has_separate_drafter: + return False return req_mode == "auto" and size_b is not None and size_b < _MTP_MIN_SIZE_B @@ -2910,19 +2919,14 @@ class LlamaCppBackend: # Auto-download the separate MTP drafter (e.g. Gemma) when # the requested spec mode can use it. Repos with the head # baked into the main GGUF (Qwen) have no mtp- sibling and - # this no-ops. Skipped when the user disabled MTP, drives - # --spec-type manually via extra_args, or in auto mode on a - # sub-3B model (e.g. Gemma E2B) where the resolver drops - # MTP anyway -- no point fetching a drafter it never uses. - # Forced mtp / mtp+ngram still download (user override). + # this no-ops, so the size gate stays out of it: a separate + # drafter speeds up even sub-3B (Gemma E2B), and the resolver + # below decides the final emission. Skipped only when the + # user disabled MTP or drives --spec-type manually. _spec_canon = _canonicalize_spec_mode(speculative_type) or "auto" - _auto_drops_mtp = _auto_mode_drops_mtp( - _spec_canon, _extract_model_size_b(model_identifier) - ) if ( not mtp_draft_path and _spec_canon in ("auto", "mtp", "mtp+ngram") - and not _auto_drops_mtp and not _extra_args_set_spec_type(extra_args) ): mtp_draft_path = self._download_mtp( @@ -3002,8 +3006,12 @@ class LlamaCppBackend: _mtp_canonical = _canonicalize_spec_mode(speculative_type) _mtp_effective = _mtp_canonical or "auto" _mtp_size_for_fit = _extract_model_size_b(model_identifier) + # Sub-3B drops MTP only for an embedded head; a separate + # drafter (Gemma) engages and needs its VRAM reserved. _mtp_sub_3b_for_fit = ( - _mtp_size_for_fit is not None and _mtp_size_for_fit < _MTP_MIN_SIZE_B + _mtp_size_for_fit is not None + and _mtp_size_for_fit < _MTP_MIN_SIZE_B + and not bool(mtp_draft_path) ) _mtp_will_engage = bool( not _extra_args_set_spec_type(extra_args) @@ -3771,18 +3779,26 @@ class LlamaCppBackend: https://github.com/ggml-org/llama.cpp/pull/18471 MTP guide: unsloth.ai/docs/models/qwen3.6#mtp-guide - Sub-3B dense MTP regresses vs spec-off: the draft head's per-token - cost exceeds the acceptance savings at this scale. Q4_K_XL clean - bench (each prompt once after an unrelated warmup) on B200 + x86 CPU: + Sub-3B dense MTP regresses vs spec-off when the head is baked into the + main GGUF (Qwen): the draft head's per-token cost exceeds the + acceptance savings at this scale. Q4_K_XL clean bench (each prompt once + after an unrelated warmup) on B200 + x86 CPU: 0.8B GPU: draft-mtp n=2 = 0.58x vs OFF; ngram-only = 1.10x 2B GPU: draft-mtp n=2 = 0.82x vs OFF; OFF or ngram = 1.00x 0.8B CPU: chained n=2 = 0.86x vs OFF; ngram-only = 1.19x 2B CPU: chained n=2 = 0.83x vs OFF; ngram-only = 1.01x 4B+ GPU/CPU: spec on is a net win (1.08x-1.46x). + A separate drafter (Gemma's root mtp-*.gguf) is a different, cheaper + mechanism that wins even below 3B, so it is exempt from the sub-3B drop + (``mtp_draft_path`` set -> not too small). B200 Q4_K_XL bench, draft-mtp + n=2 vs OFF: gemma-4-E2B (2B) = 1.21x, accept ~0.65 (vs ngram = 1.00x); + gemma-4-E4B (4B) and 12B engage as usual. Auto falls back to ngram-mod (zero-VRAM, near-zero idle cost on - diverse content); forced MTP on a model with no head/drafter defaults - back (mtp -> spec-default, mtp+ngram -> ngram-mod) since llama-server - aborts otherwise; sub-3B real-MTP engages with a warning. + diverse content) for an embedded sub-3B head; forced MTP on a model + with no head/drafter defaults back (mtp -> spec-default, mtp+ngram -> + ngram-mod) since llama-server aborts otherwise; a drafter the binary + cannot build (older prebuilt, or a CUDA kernel limit) aborts the spawn + and the load retries once without speculative decoding. """ flags: List[str] = [] # Reset; emit branches re-set on the resolved emission. @@ -3801,7 +3817,11 @@ class LlamaCppBackend: ) user_owns_spec_type = _extra_args_set_spec_type(extra_args) _mtp_size_b = _extract_model_size_b(model_identifier) - _mtp_too_small = _mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B + # The sub-3B regression is an embedded-head cost; a separate drafter + # (Gemma) is a cheap standalone model that wins below 3B, so exempt it. + _mtp_too_small = ( + _mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B and not bool(mtp_draft_path) + ) if user_owns_spec_type: # User --spec-type wins outright; suppress auto-emit to avoid a diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4cc76e7c2a..ef67f8bb59 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -417,7 +417,6 @@ try: LlamaCppBackend, _DEFAULT_MAX_TOKENS_FLOOR, _DEFAULT_T_MAX_PREDICT_MS, - _auto_mode_drops_mtp, _canonicalize_spec_mode, _extra_args_set_spec_type, _hf_offline_if_dns_dead, @@ -431,7 +430,6 @@ try: from utils.inference import load_inference_config from utils.models.model_config import ( detect_mtp_file, - extract_model_size_b, load_model_defaults, ) from utils.native_path_leases import ( @@ -450,7 +448,6 @@ except ImportError: LlamaCppBackend, _DEFAULT_MAX_TOKENS_FLOOR, _DEFAULT_T_MAX_PREDICT_MS, - _auto_mode_drops_mtp, _canonicalize_spec_mode, _extra_args_set_spec_type, _hf_offline_if_dns_dead, @@ -464,7 +461,6 @@ except ImportError: from utils.inference import load_inference_config from utils.models.model_config import ( detect_mtp_file, - extract_model_size_b, load_model_defaults, ) from utils.native_path_leases import ( @@ -1025,14 +1021,15 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC else: if list(request.llama_extra_args) != backend_extra: return False - # A drafter that appeared next to the loaded weights since the last load - # changes the launch command (--model-draft) when the mode can use it; - # without this, a duplicate /load is deduped and MTP can't engage short - # of an unload. Runs last: it stats the filesystem (two dir scans against - # the resolved weight path -- covers local dirs and HF cache snapshots - # alike), so every pure-memory comparison above short-circuits first. - # Skipped when auto drops MTP anyway (sub-3B) or the user owns - # --spec-type, where a drafter changes nothing. Resolve both sides: the + # A separate drafter (Gemma's root mtp-*.gguf) appearing or disappearing + # next to the loaded weights changes the launch command (--model-draft), + # so a duplicate /load must reload rather than dedupe. Always compare the + # detected vs stored drafter when the mode can use one and the user does + # not own --spec-type: the resolved-path compare is cheap and handles all + # four cases (both None -> match; one None -> reload; equal -> match; + # different -> reload), including a drafter deleted out from under a + # running server. Runs last: it stats the filesystem, so every pure-memory + # comparison above short-circuits first. Resolve both sides since the # stored launch path may be a snapshot symlink while detect_mtp_file # returns the resolved blob. if req_mode in ("auto", "mtp", "mtp+ngram") and llama_backend.gguf_path: @@ -1041,10 +1038,7 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC if request.llama_extra_args is not None else llama_backend.extra_args ) - size_b = extract_model_size_b(llama_backend.model_identifier or "") - if not _auto_mode_drops_mtp(req_mode, size_b) and not _extra_args_set_spec_type( - effective_extras - ): + if not _extra_args_set_spec_type(effective_extras): detected = detect_mtp_file(llama_backend.gguf_path) stored = llama_backend.mtp_draft_path try: diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 25ce71a774..a7f30d2951 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -1219,3 +1219,183 @@ def test_forced_mtp_ngram_on_non_mtp_model_keeps_ngram(monkeypatch): assert parsed.get("--spec-type") == "ngram-mod" assert backend.speculative_type == "ngram-mod" assert backend.requested_spec_mode == "mtp+ngram" + + +# ── Full named-repo resolver matrix (the shipping Studio families) ───── +# +# Locks auto / off / forced-mtp routing for every Qwen3.5 (MTP + plain) and +# gemma-4 (regular + QAT) GGUF repo, including the giant MoEs that stay +# resolver-only (122B-A10B / 397B-A17B). Expectations are derived from the +# same signals load_model uses -- _extract_model_size_b (active>effective> +# total, so E2B->2, A3B->3, A10B->10, A17B->17), _is_mtp_model_name, and the +# separate-drafter flag -- so each row mirrors what the loader emits on a +# B200 (GPU default, n=2). gemma carries no -MTP marker; its MTP comes from +# the root mtp-*.gguf drafter, modelled here by passing mtp_draft_path. +# +# auto_spec: "draft-mtp" = head/drafter engaged (>=3B MTP, or any size with a +# separate drafter); "ngram-mod" = embedded sub-3B drop (zero-VRAM); None = +# non-MTP -> llama-server --spec-default. + +_GEMMA_DRAFTER = "/snap/mtp-gemma-4-it.gguf" # stand-in separate drafter + +_REAL_REPO_MATRIX = [ + # repo, drafter, auto_spec, auto_ngram_knobs + ("unsloth/Qwen3.5-0.8B-MTP-GGUF", None, "ngram-mod", True), + ("unsloth/Qwen3.5-2B-MTP-GGUF", None, "ngram-mod", True), + ("unsloth/Qwen3.5-4B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-9B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-27B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-35B-A3B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-122B-A10B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-397B-A17B-MTP-GGUF", None, "draft-mtp", False), + ("unsloth/Qwen3.5-0.8B-GGUF", None, None, False), + ("unsloth/Qwen3.5-2B-GGUF", None, None, False), + ("unsloth/Qwen3.5-4B-GGUF", None, None, False), + ("unsloth/Qwen3.5-9B-GGUF", None, None, False), + # E2B is 2B but ships a separate drafter -> exempt from the sub-3B drop. + ("unsloth/gemma-4-E2B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-E4B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-12b-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-26B-A4B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-31B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-E2B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-E4B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-12b-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-26B-A4B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), + ("unsloth/gemma-4-31B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False), +] + + +def _resolve_real(monkeypatch, repo, drafter, mode): + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = mode, + spec_draft_n_max = None, + extra_args = None, + model_identifier = repo, + model_path = None, + gpus = True, # B200 default + binary = "/fake/llama-server", + mtp_draft_path = drafter, + ) + return backend, flags, _flags_dict(flags) + + +@pytest.mark.parametrize( + "repo, drafter, auto_spec, auto_ngram_knobs", + _REAL_REPO_MATRIX, + ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX], +) +def test_real_repo_auto_routing(monkeypatch, repo, drafter, auto_spec, auto_ngram_knobs): + # Auto is the default mode the dropdown ships with. + backend, flags, parsed = _resolve_real(monkeypatch, repo, drafter, "auto") + if auto_spec is None: + # Non-MTP: no draft-mtp, hand off to llama-server's own default. + assert "--spec-type" not in parsed + assert "--spec-default" in flags + assert backend.speculative_type == "default" + elif auto_spec == "draft-mtp": + assert parsed.get("--spec-type") == "draft-mtp" + assert parsed.get("--spec-draft-n-max") == "2" + assert backend.speculative_type == "draft-mtp" + # gemma ships a separate drafter; Qwen bakes the head into the GGUF. + assert ( + (parsed.get("--model-draft") == drafter) if drafter else ("--model-draft" not in parsed) + ) + else: # ngram-mod (sub-3B MTP drop) + assert parsed.get("--spec-type") == "ngram-mod" + assert "--model-draft" not in parsed # draft head dropped + assert backend.speculative_type == "ngram-mod" + if auto_ngram_knobs: + assert "--spec-ngram-mod-n-match" in parsed + assert backend.requested_spec_mode == "auto" + + +@pytest.mark.parametrize( + "repo, drafter", + [(r[0], r[1]) for r in _REAL_REPO_MATRIX], + ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX], +) +def test_real_repo_off_emits_nothing(monkeypatch, repo, drafter): + # Off must suppress speculative decoding for every family. + backend, flags, _ = _resolve_real(monkeypatch, repo, drafter, "off") + assert flags == [] + assert backend.speculative_type is None + assert backend.requested_spec_mode == "off" + + +@pytest.mark.parametrize( + "repo, drafter", + [(r[0], r[1]) for r in _REAL_REPO_MATRIX], + ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX], +) +def test_real_repo_forced_mtp_never_aborts(monkeypatch, repo, drafter): + # Forcing MTP on the dropdown: real MTP models (name marker or separate + # drafter) engage draft-mtp even below 3B; non-MTP models default back to + # --spec-default instead of emitting a draft-mtp llama-server will abort on. + backend, flags, parsed = _resolve_real(monkeypatch, repo, drafter, "mtp") + is_real_mtp = _is_mtp_model_name(repo) or bool(drafter) + if is_real_mtp: + assert parsed.get("--spec-type") == "draft-mtp" + assert backend.speculative_type == "draft-mtp" + assert ( + (parsed.get("--model-draft") == drafter) if drafter else ("--model-draft" not in parsed) + ) + else: + assert "--spec-type" not in parsed + assert "--spec-default" in flags + assert backend.speculative_type == "default" + assert backend.requested_spec_mode == "mtp" + + +# ── Sub-3B separate-drafter exemption (Gemma) ───────────────────────── +# +# The sub-3B MTP drop is an embedded-head cost (Qwen). A separate drafter +# (Gemma's root mtp-*.gguf) is a cheap standalone model that wins below 3B +# (B200 Q4_K_XL: gemma-4-E2B draft-mtp n=2 = 1.21x vs OFF), so it is exempt. + + +def test_sub3b_gemma_separate_drafter_engages_mtp(monkeypatch): + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = "unsloth/gemma-4-E2B-it-GGUF", # 2B + model_path = None, + gpus = True, + binary = "/fake/llama-server", + mtp_draft_path = "/snap/mtp-gemma-4-E2B-it.gguf", # separate drafter + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "draft-mtp" + assert parsed.get("--model-draft") == "/snap/mtp-gemma-4-E2B-it.gguf" + assert "--spec-ngram-mod-n-match" not in parsed + assert backend.speculative_type == "draft-mtp" + + +def test_sub3b_qwen_embedded_head_still_drops_to_ngram(monkeypatch): + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = "unsloth/Qwen3.5-2B-MTP-GGUF", # 2B, embedded head + model_path = None, + gpus = True, + binary = "/fake/llama-server", + mtp_draft_path = None, # no separate drafter + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "ngram-mod" + assert "--model-draft" not in parsed + assert backend.speculative_type == "ngram-mod" + + +def test_auto_mode_drops_mtp_exempts_separate_drafter(): + from core.inference.llama_cpp import _auto_mode_drops_mtp + + assert _auto_mode_drops_mtp("auto", 2.0) is True + assert _auto_mode_drops_mtp("auto", 2.0, has_separate_drafter = True) is False + assert _auto_mode_drops_mtp("auto", 4.0) is False + assert _auto_mode_drops_mtp("mtp", 2.0) is False # forced engages regardless