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

@@ -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.

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(

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)

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

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.

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."
+                                : "")}
                       </p>
                       {mtpUpdatable && llamaUpdateStatus?.update_available && (
                         <Button