From 9e83399f9e04cfdb3c3ca1190afdae9459f91e24 Mon Sep 17 00:00:00 2001 From: UmranPros <152087084+umran666@users.noreply.github.com> Date: Sat, 20 Jun 2026 17:13:37 +0530 Subject: [PATCH] Studio: fix Gemma 4 separate-drafter MTP detection and fallback (#6459) Recognise the Gemma 4 separate-drafter MTP family, auto-download the drafter with retry, fall back to n-gram with a clear reason when it cannot be resolved, and retry the download on reload. Gemma 3n (ships no drafter) and embedded-MTP models (Qwen) are unaffected. Fixes #6406 --- studio/backend/core/inference/llama_cpp.py | 114 ++++++++++++++-- studio/backend/models/inference.py | 13 +- studio/backend/routes/inference.py | 14 +- .../tests/test_llama_cpp_mtp_detection.py | 125 ++++++++++++++++++ studio/backend/tests/test_mtp_vram_budget.py | 14 ++ .../src/features/chat/chat-settings-sheet.tsx | 10 +- 6 files changed, 270 insertions(+), 20 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 16b312d8f6..0d32a7cc38 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -657,11 +657,32 @@ def detect_reasoning_flags( return flags +# Gemma 4 ships MTP as a separate drafter (no "-mtp" in the name). Gemma 3n +# ships no drafter, so it is excluded -- it takes the normal non-MTP path. +_GEMMA_MTP_FAMILY_RE = re.compile(r"gemma[-_]?4[-_]", re.IGNORECASE) + + +def _is_gemma_mtp_family(name: Optional[str]) -> bool: + """Match Gemma 4 by name.""" + return bool(name) and bool(_GEMMA_MTP_FAMILY_RE.search(name)) + + +def _is_gemma_mtp_name(model_identifier: Optional[str], gguf_path: Optional[str] = None) -> bool: + """Match Gemma 4 by id or GGUF filename.""" + return _is_gemma_mtp_family(model_identifier) or _is_gemma_mtp_family( + Path(gguf_path).name if gguf_path else None + ) + + def _is_mtp_model_name(model_identifier: Optional[str], gguf_path: Optional[str] = None) -> bool: """Name-based MTP detector. Fallback for the metadata signal.""" for cand in (model_identifier, Path(gguf_path).name if gguf_path else None): if cand and "-mtp" in cand.lower(): return True + # Recognise Gemma 4 too, so a failed drafter download surfaces a + # fallback reason instead of silently defaulting. + if cand and _is_gemma_mtp_family(cand): + return True return False @@ -1374,6 +1395,11 @@ class LlamaCppBackend: def gguf_path(self) -> Optional[str]: return self._gguf_path + @property + def hf_repo(self) -> Optional[str]: + """HF repo of the loaded model, or None for local/native file loads.""" + return self._hf_repo + @property def mtp_draft_path(self) -> Optional[str]: return self._mtp_draft_path @@ -3813,11 +3839,31 @@ class LlamaCppBackend: return None target: Optional[str] = None - try: - from huggingface_hub import list_repo_files - target = pick(list_repo_files(hf_repo, token = hf_token)) - except Exception as e: - logger.debug(f"Could not list repo files for {label}: {e}") + from huggingface_hub import list_repo_files + + # Retry a transient listing blip; permanent repo/auth errors and offline + # mode are not retried (offline raises at once -> fall through to cache). + for attempt in range(3): + if self._cancel_event.is_set(): + return None + try: + target = pick(list_repo_files(hf_repo, token = hf_token)) + break + except Exception as e: + if type(e).__name__ in ( + "RepositoryNotFoundError", + "GatedRepoError", + "RevisionNotFoundError", + "EntryNotFoundError", + "OfflineModeIsEnabled", + ): + logger.debug(f"Could not list repo files for {label}: {e}") + break + logger.debug( + f"Could not list repo files for {label} " f"(attempt {attempt + 1}/3): {e}" + ) + if attempt < 2: + self._cancel_event.wait(2**attempt) if target is None: try: @@ -4831,6 +4877,12 @@ class LlamaCppBackend: self._nextn_predict_layers or _is_mtp_model_name(model_identifier, model_path) or bool(mtp_draft_path) + ) and not ( + # Drafterless Gemma falls back to ngram-mod; reserve no + # drafter VRAM for it (mirrors the launch resolver). + _is_gemma_mtp_name(model_identifier, model_path) + and not mtp_draft_path + and not self._nextn_predict_layers ) _mtp_binary_ok = True _mtp_probe_raised = False @@ -6202,6 +6254,13 @@ 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) ) + # Drafterless Gemma (name-only MTP, no embedded head): emitting MTP + # would abort llama-server, so every mode below falls back instead. + _mtp_drafter_missing = ( + _is_gemma_mtp_name(model_identifier, model_path) + and not mtp_draft_path + and not self._nextn_predict_layers + ) # 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 @@ -6307,6 +6366,20 @@ class LlamaCppBackend: logger.info("Spec decoding: ngram-mod") return True + def _fallback_drafter_not_found() -> None: + """Drafterless Gemma: use ngram-mod (or spec-default) and record why.""" + logger.warning( + "Model %s is MTP-capable but no drafter or head was found; " + "falling back. Check network or run `unsloth studio update`.", + model_identifier, + ) + if self.probe_server_capabilities(binary).get("supports_ngram_mod"): + _emit_ngram_mod() + else: + flags.append("--spec-default") + self._speculative_type = "default" + self._spec_fallback_reason = "drafter_not_found" + if effective_mode == "off": return flags # nothing to emit if effective_mode == "ngram-simple": @@ -6327,6 +6400,10 @@ class LlamaCppBackend: flags.append("--spec-default") self._speculative_type = "default" return flags + if _mtp_drafter_missing: + # Drafterless: draft-mtp would abort llama-server, so fall back. + _fallback_drafter_not_found() + return flags if _mtp_too_small: logger.warning( f"Forcing MTP on a {_mtp_size_b:.1f}B model; " @@ -6345,6 +6422,10 @@ class LlamaCppBackend: ) _emit_ngram_mod() return flags + if _mtp_drafter_missing: + # No head/drafter: keep ngram-mod, drop the draft-mtp chain. + _fallback_drafter_not_found() + return flags if _mtp_too_small: logger.warning( f"Forcing MTP+Ngram on a {_mtp_size_b:.1f}B model; " @@ -6380,13 +6461,18 @@ class LlamaCppBackend: ) # 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) + if _mtp_drafter_missing: + # Name-only MTP, drafter did not resolve (download failed/absent). + _fallback_drafter_not_found() + else: + # GPU: MTP-only. CPU/Mac: chain ngram-mod + MTP. + _emit_mtp(chain_ngram = not gpus) elif is_mtp_model and _mtp_too_small: # Sub-3B fallback: drop the MTP draft head, keep ngram-mod when # the binary supports it. - _small_caps = self.probe_server_capabilities(binary) - if _small_caps.get("supports_ngram_mod"): + if _mtp_drafter_missing: + _fallback_drafter_not_found() + elif self.probe_server_capabilities(binary).get("supports_ngram_mod"): logger.info( f"MTP GGUF detected but model size {_mtp_size_b:.1f}B " "is below the 3B speedup threshold; using ngram-mod " @@ -6476,6 +6562,16 @@ class LlamaCppBackend: if req_mode != backend_mode: return False + # Prior HF load fell back with drafter_not_found; a same-settings reload + # must retry the download in load_model, not dedupe to the stale fallback + # (HF loads resolve the drafter there, so gguf_path is None here). + if ( + self._spec_fallback_reason == "drafter_not_found" + and gguf_path is None + and req_mode in ("auto", "mtp", "mtp+ngram") + ): + return False + # spec_draft_n_max only matters when an MTP variant is engaged. Compare # on the resolved spec so an Auto request promoted to draft-mtp still # bounces a reload when n_max changes. diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 0d33cfa976..b8432f588c 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -412,12 +412,13 @@ 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. '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." + "current build could not run it; 'drafter_not_found' -> the model's " + "separate MTP drafter could not be resolved; '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( diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4b068941b8..eb597caa0a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1829,6 +1829,16 @@ def _request_matches_loaded_settings( backend_mode = llama_backend.requested_spec_mode or "auto" if req_mode != backend_mode: return False + # Prior HF load fell back with drafter_not_found: a same-settings reload must + # retry the download, not dedupe to the stale fallback. HF only (hf_repo set); + # local/native loads have no download to retry (handled by the path compare). + if ( + llama_backend.hf_repo + and llama_backend.spec_fallback_reason == "drafter_not_found" + and req_mode in ("auto", "mtp", "mtp+ngram") + and not _extra_args_set_spec_type(effective_extra) + ): + return False # spec_draft_n_max only matters with an MTP variant; 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: @@ -2206,7 +2216,9 @@ async def load_model( and llama_backend.model_identifier.lower() == model_identifier.lower() # Match runtime settings so Apply isn't dropped (#5401). and _request_matches_loaded_settings( - request, llama_backend, effective_chat_template_override + request, + llama_backend, + effective_chat_template_override, ) # Skip if a prior audio probe failed -- let load_model retry. and getattr(llama_backend, "_audio_probed", True) diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index a81a49acd5..3f9d2a8f50 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -1653,6 +1653,8 @@ _REAL_REPO_MATRIX = [ def _resolve_real(monkeypatch, repo, drafter, mode): backend = _resolver_backend(monkeypatch) + if "qwen" in repo.lower() and "-mtp" in repo.lower(): + backend._nextn_predict_layers = 1 flags = backend._build_speculative_flags( speculative_type = mode, spec_draft_n_max = None, @@ -1843,3 +1845,126 @@ def test_spec_fallback_reason_reset_on_off(monkeypatch): binary = "/fake/llama-server", ) assert backend.spec_fallback_reason is None + + +def test_is_gemma_mtp_family(): + from core.inference.llama_cpp import _is_gemma_mtp_family + + assert _is_gemma_mtp_family("unsloth/gemma-4-E4B-it-GGUF") is True + assert _is_gemma_mtp_family("unsloth/gemma-4-12b-it-GGUF") is True + # gemma-3n ships no separate drafter, so it is not a drafter family. + assert _is_gemma_mtp_family("unsloth/gemma-3n-E2B-it-GGUF") is False + assert _is_gemma_mtp_family("unsloth/Qwen3.5-35B-A3B-MTP-GGUF") is False + assert _is_gemma_mtp_family("unsloth/llama-3-8b") is False + + +def test_gemma_3n_without_drafter_is_not_mtp(monkeypatch): + # gemma-3n ships no drafter; it must take the normal non-MTP path, not + # drafter_not_found (which would make every reload retry a missing drafter). + backend = _resolver_backend(monkeypatch) + backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = "unsloth/gemma-3n-E4B-it-GGUF", + model_path = None, + gpus = True, + binary = "/fake/llama-server", + mtp_draft_path = None, + ) + assert backend.spec_fallback_reason is None + + +def test_spec_fallback_reason_drafter_not_found(monkeypatch): + # Drafterless Gemma should fall back to ngram-mod + drafter_not_found. + 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-E4B-it-GGUF", + model_path = None, + gpus = True, + binary = "/fake/llama-server", + mtp_draft_path = None, # Drafter download failed + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "ngram-mod" + assert backend.speculative_type == "ngram-mod" + assert backend.spec_fallback_reason == "drafter_not_found" + + +def test_is_gemma_mtp_name_none_safe(): + # model_identifier=None (local load) must not raise; recognise via filename. + from core.inference.llama_cpp import _is_gemma_mtp_family, _is_gemma_mtp_name + + assert _is_gemma_mtp_family(None) is False + assert _is_gemma_mtp_name(None, "/models/gemma-4-E4B-it-Q4_K_M.gguf") is True + assert _is_gemma_mtp_name("unsloth/Qwen3.5-4B-MTP-GGUF", None) is False + + +@pytest.mark.parametrize("mode", ["mtp", "mtp+ngram"]) +def test_forced_mtp_gemma_without_drafter_falls_back(monkeypatch, mode): + # Forced MTP on a drafterless Gemma must fall back, not emit draft-mtp. + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = mode, + spec_draft_n_max = None, + extra_args = None, + model_identifier = "unsloth/gemma-4-E4B-it-GGUF", + model_path = None, + gpus = True, + binary = "/fake/llama-server", + mtp_draft_path = None, + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "ngram-mod" + assert "--model-draft" not in parsed + assert backend.spec_fallback_reason == "drafter_not_found" + + +def test_local_gemma_gguf_without_identifier_falls_back(monkeypatch): + # Local Gemma GGUF (family only in filename) must not crash; falls back. + backend = _resolver_backend(monkeypatch) + flags = backend._build_speculative_flags( + speculative_type = "auto", + spec_draft_n_max = None, + extra_args = None, + model_identifier = None, + model_path = "/models/gemma-4-E4B-it-Q4_K_M.gguf", + gpus = True, + binary = "/fake/llama-server", + mtp_draft_path = None, + ) + parsed = _flags_dict(flags) + assert parsed.get("--spec-type") == "ngram-mod" + assert backend.spec_fallback_reason == "drafter_not_found" + + +def _drafter_not_found_kwargs(): + return dict( + model_identifier = "unsloth/gemma-4-E4B-it-GGUF", + 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, + gguf_path = None, # HF load: drafter resolves inside load_model + ) + + +def test_already_in_target_state_retries_after_hf_drafter_not_found(): + # Recoverable drafter_not_found must not dedupe; reload re-attempts download. + backend = _mtp_backend( + _model_identifier = "unsloth/gemma-4-E4B-it-GGUF", + _speculative_type = "ngram-mod", + _spec_fallback_reason = "drafter_not_found", + _mtp_draft_path = None, + _gguf_path = None, + ) + assert backend._already_in_target_state(**_drafter_not_found_kwargs()) is False + # Sanity: with no fallback reason the same request still dedupes (matches). + ok = _mtp_backend(_model_identifier = "unsloth/gemma-4-E4B-it-GGUF", _gguf_path = None) + assert ok._already_in_target_state(**_drafter_not_found_kwargs()) is True diff --git a/studio/backend/tests/test_mtp_vram_budget.py b/studio/backend/tests/test_mtp_vram_budget.py index e7c6b02232..02f0c37b4d 100644 --- a/studio/backend/tests/test_mtp_vram_budget.py +++ b/studio/backend/tests/test_mtp_vram_budget.py @@ -829,6 +829,20 @@ class TestExtraArgsMtpDetection: "request.tensor_parallel,llama_backend.tensor_parallel)" in body ) + def test_route_matcher_retries_after_drafter_not_found(self): + # drafter_not_found must not report "already loaded" or the reload never + # retries the download (#6459). Read source: importing routes pulls deps. + routes_src = ( + Path(__file__).resolve().parent.parent / "routes" / "inference.py" + ).read_text() + start = routes_src.index("def _request_matches_loaded_settings") + end = routes_src.index("\ndef ", start + 1) + body = "".join(routes_src[start:end].split()) + assert 'llama_backend.spec_fallback_reason=="drafter_not_found"' in body + assert "not_extra_args_set_spec_type(effective_extra)" in body + # HF-only (hf_repo): local/native loads have no download to retry. + assert "llama_backend.hf_repo" in body + def test_extra_args_main_cache_type_heavier_axis(self): # Asymmetric --cache-type-k/-v must budget the heavier axis (extras win # per axis at launch), not the last-wins single type that under-reserves. diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 935f7e1bc6..2ef4c1573e 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -1009,10 +1009,12 @@ export function ChatSettingsPanel({ ? "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 - ? " Update llama.cpp to enable it." - : "")} + : specFallbackReason === "drafter_not_found" + ? "This model supports MTP, but its drafter file could not be downloaded, so MTP is off and it falls back to n-gram speculative decoding where the llama.cpp build supports it. Check your network connection or Hugging Face access, then reload the model to retry the drafter." + : "MTP is not available in the installed llama.cpp build, so this model is running without it." + + (llamaUpdateStatus?.update_available + ? " Update llama.cpp to enable it." + : "")}

{mtpUpdatable && llamaUpdateStatus?.update_available && (