From b5aef63c03a1927ebec6d0e662b6c95c1a41d9ae Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 06:46:00 -0700 Subject: [PATCH] Studio: resolve the repo-root MTP drafter after the MTP/ GGUF rename (#7031) * Studio: resolve the repo-root MTP drafter after the MTP/ GGUF rename The Gemma 4 QAT GGUF repos renamed the higher-precision MTP/ subdir copies from gemma-4-...--MTP.gguf to mtp-gemma-4-...-.gguf, so their basenames now start with the same mtp- prefix as the small repo-root drafter (mtp-gemma-4-E4B-it.gguf). The drafter selectors filtered candidates by a mtp- basename prefix and took the first in sort order. With the new names the MTP/ copies also match, and because MTP/ (uppercase) sorts before the lowercase root file, selection flipped to the large BF16 copy under MTP/ instead of the root drafter both functions document they should pick. Restrict both selectors, and the companion byte estimate, to root-level mtp-*.gguf so the MTP/ copies stay explicit-selection only: - core/inference/llama_cpp.py _pick_mtp (loader auto-download) - hub/utils/gguf_plan.py preferred_mtp_sibling (Hub variant plans) - routes/inference.py _remote_gguf_companion_bytes (VRAM headroom) Also reuse a drafter already in the local cache before downloading, so a device that already holds a copy on disk does not re-fetch it. Old-scheme names keep working (they have no root-level mtp- sibling to mis-select). Adds regression tests for the new naming, both selection paths, and the on-disk reuse. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: gate MTP drafter cache reuse to offline mode Reuse the cached drafter only when HF is offline. Online, route back through _download_companion_gguf/hf_hub_download so the current revision is checked (etag) and a changed drafter is refetched, matching the offline-only cross-snapshot reuse already used for the main GGUF. This avoids pairing freshly downloaded weights with a stale cached draft. Make the reuse tests offline and add an online-skips-reuse test. * Studio: prefer a root MTP drafter across all cached snapshots Offline reuse scanned snapshots one at a time and returned the first snapshot that held any drafter, only preferring root within it. A newer partial snapshot with just the MTP/ copy could shadow the small root drafter in an older snapshot. Collect drafters across all snapshots and prefer any repo-root file before an MTP/ copy. * Studio: keep newest-first snapshot order when reusing cached drafters Collecting root candidates and sorting by absolute snapshot path could pick a drafter from an older snapshot. _iter_hf_cache_snapshots yields newest first and the main GGUF is resolved in that order, so preserve it (root still preferred over MTP/ copies) to avoid pairing a fresh main weight with a stale drafter revision. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 39 +++- studio/backend/hub/utils/gguf_plan.py | 13 +- studio/backend/routes/inference.py | 8 +- .../tests/test_mtp_drafter_companion.py | 183 +++++++++++++++++- 4 files changed, 233 insertions(+), 10 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3ba9eff857..7d782fbbdd 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -4673,6 +4673,29 @@ class LlamaCppBackend: cancel_event = cancel_event, ) + def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]: + """A drafter already in this repo's local HF cache, reused offline when a + fresh copy can't be fetched. Prefers a repo-root ``mtp-*.gguf`` across all + cached snapshots; else an existing ``MTP/`` copy (any precision -- the + target verifies every drafted token). None if none is cached.""" + try: + from utils.models.model_config import _iter_hf_cache_snapshots + + roots: list[Path] = [] + subdirs: list[Path] = [] + for snap in _iter_hf_cache_snapshots(hf_repo): # newest first + for f in sorted(_gguf_snapshot_files(snap)): + if _is_companion_gguf_path(f) and "mmproj" not in f.lower(): + (roots if "/" not in f else subdirs).append(snap / f) + # Keep snapshot order (newest first), root before any MTP/ copy, so a + # newer main GGUF pairs with the newest cached drafter, not a stale one. + for cand in roots + subdirs: + if cand.is_file(): + return str(cand) + except Exception as e: + logger.debug("Cached MTP drafter lookup failed for %s: %s", hf_repo, e) + return None + def _download_mtp( self, *, @@ -4689,11 +4712,25 @@ class LlamaCppBackend: are intentionally skipped. Returns the local path, or None. """ + # Offline, reuse any drafter already on disk (a fresh copy can't be + # fetched). Online, _download_companion_gguf/hf_hub_download reuse the + # current cached file and refetch a changed one, so skip the probe here + # rather than pair new weights with a stale draft. + if _hf_env_offline(): + cached = self._cached_repo_mtp_drafter(hf_repo) + if cached: + logger.info(f"Reusing cached MTP drafter (offline): {cached}") + return cached + def _pick_mtp(candidates: list[str]) -> Optional[str]: + # Root-level only: MTP/ subdir copies now share the mtp- prefix but + # are explicit-selection, not auto-fetch (they'd sort ahead of root). mtp_files = sorted( f for f in candidates - if f.lower().endswith(".gguf") and Path(f).name.lower().startswith("mtp-") + if f.lower().endswith(".gguf") + and "/" not in f + and Path(f).name.lower().startswith("mtp-") ) return mtp_files[0] if mtp_files else None diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py index 2abdb0fb79..18daa4f84e 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -102,16 +102,17 @@ def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]: def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]: """The separate MTP drafter to fetch with every variant: the repo-root ``mtp-*.gguf`` copy unsloth ships for llama.cpp ``-hf`` auto-discovery - (Gemma 4). Same pick as the loader's drafter resolution (``mtp-`` basename - prefix, first in sort order) so download and load resolve the same file; - the higher-precision ``MTP/`` subdir copies are for explicit selection and - are not auto-fetched. None for repos with the head baked into the main - GGUF (Qwen).""" + (Gemma 4). Same pick as the loader's drafter resolution (root-level + ``mtp-`` prefix, first in sort order) so download and load resolve the same + file; the higher-precision ``MTP/`` subdir copies are for explicit + selection and are not auto-fetched. None for repos with the head baked into + the main GGUF (Qwen).""" + # Root-level only: the MTP/ subdir copies now share the mtp- prefix too. candidates = sorted( ( s for s in siblings - if (name := _gguf_rfilename(s)) and name.lower().rsplit("/", 1)[-1].startswith("mtp-") + if (name := _gguf_rfilename(s)) and "/" not in name and name.lower().startswith("mtp-") ), key = lambda s: getattr(s, "rfilename"), ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4bb9ce655e..1e39c6d129 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -3075,10 +3075,14 @@ def _remote_gguf_companion_bytes( info = model_info(repo, token = hf_token, files_metadata = True) total = 0 for sibling in info.siblings or []: - base = Path(sibling.rfilename or "").name.lower() + name = sibling.rfilename or "" + base = Path(name).name.lower() if not base.endswith(".gguf"): continue - if base.startswith("mtp-") or (include_mmproj and "mmproj" in base): + # Root-level mtp- only: -hf auto-fetches the repo-root drafter, not + # the MTP/ subdir copies (which now share the mtp- prefix too). + is_root_mtp = "/" not in name and base.startswith("mtp-") + if is_root_mtp or (include_mmproj and "mmproj" in base): total += getattr(sibling, "size", 0) or 0 return total except Exception as e: diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index d5e8d13652..86e528ae67 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -23,7 +23,11 @@ if _BACKEND_DIR not in sys.path: from hub.utils.download_manifest import ExpectedFile from hub.utils.gguf import is_mtp_drafter_path -from hub.utils.gguf_plan import build_gguf_variant_plans, plan_from_expected_files +from hub.utils.gguf_plan import ( + build_gguf_variant_plans, + plan_from_expected_files, + preferred_mtp_sibling, +) from utils.models.model_config import ( _is_mtp_drafter, detect_gguf_model, @@ -37,6 +41,8 @@ from utils.models.model_config import ( DRAFTER_CASES = [ ("mtp-gemma-4-12b-it.gguf", True), ("MTP/gemma-4-12b-it-Q8_0-MTP.gguf", True), + # New-scheme MTP/ copies carry the mtp- basename prefix too. + ("MTP/mtp-gemma-4-E4B-it-BF16.gguf", True), ("foo/MTP/bar.gguf", True), ("gemma-4-12b-it-Q8_0.gguf", False), # Baked-in Qwen MTP repos: the head is inside the main GGUF, the file @@ -274,3 +280,178 @@ def test_detect_gguf_model_rejects_mtp_subdir_copy(tmp_path): assert detect_gguf_model(str(copy)) is None # Selecting the MTP dir itself must not surface the copies as models. assert detect_gguf_model(str(sub)) is None + + +# ── Root drafter wins over new-scheme MTP/ copies ──────────────────── +# The MTP/ copies were renamed to share the mtp- basename prefix (e.g. +# MTP/mtp-gemma-4-E4B-it-BF16.gguf). Auto-fetch/load must still resolve the +# small repo-root drafter, not a sort-first MTP/ copy (uppercase precedes +# lowercase, so the subdir path would otherwise win). + +NEW_SCHEME_SIBLINGS = [ + _sib("gemma-4-12b-it-Q4_K_M.gguf", 4_000, "main-q4"), + _sib("gemma-4-12b-it-Q8_0.gguf", 8_000, "main-q8"), + _sib("mtp-gemma-4-12b-it.gguf", 100, "drafter"), + _sib("MTP/mtp-gemma-4-12b-it-Q8_0.gguf", 100, "mtp-sub-q8"), + _sib("MTP/mtp-gemma-4-12b-it-BF16.gguf", 200, "mtp-sub-bf16"), + _sib("mmproj-F16.gguf", 500, "mmproj"), +] + + +def test_preferred_mtp_sibling_prefers_root_over_new_scheme_copies(): + picked = preferred_mtp_sibling(NEW_SCHEME_SIBLINGS) + assert picked is not None and picked.rfilename == "mtp-gemma-4-12b-it.gguf" + + +def test_variant_plans_new_scheme_uses_root_drafter(): + plans = build_gguf_variant_plans(NEW_SCHEME_SIBLINGS) + assert set(plans) == {"q4_k_m", "q8_0"} + for plan in plans.values(): + assert "mtp-gemma-4-12b-it.gguf" in plan.target_filenames + assert not any("MTP/" in name for name in plan.target_filenames) + assert "drafter" in plan.companion_hashes + # Download size = main + mmproj + root drafter (not the 200-byte BF16 copy). + assert plans["q4_k_m"].download_size_bytes == 4_600 + + +def test_download_mtp_prefers_root_over_new_scheme_copies(monkeypatch): + # _pick_mtp is nested; capture it via the companion-download seam. + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) # online: skip reuse probe + captured = {} + + def _fake_companion( + *, + hf_repo, + hf_token, + pick, + label, + cancel_event = None, + ): + captured["pick"] = pick + return None + + b = LlamaCppBackend() + b._download_companion_gguf = _fake_companion + b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + + repo_files = [ + "MTP/mtp-gemma-4-E4B-it-BF16.gguf", + "MTP/mtp-gemma-4-E4B-it-Q4_0.gguf", + "MTP/mtp-gemma-4-E4B-it-Q8_0.gguf", + "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf", + "mmproj-F16.gguf", + "mtp-gemma-4-E4B-it.gguf", + ] + assert captured["pick"](repo_files) == "mtp-gemma-4-E4B-it.gguf" + + +# ── Reuse an on-disk drafter offline; fetch fresh online ───────────── + + +def _seed_snapshot(tmp_path, names): + snap = tmp_path / "snap" + for rel in names: + f = snap / rel + f.parent.mkdir(parents = True, exist_ok = True) + f.write_bytes(b"x") + return snap + + +def test_download_mtp_reuses_cached_root_drafter_offline(tmp_path, monkeypatch): + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snap = _seed_snapshot( + tmp_path, + [ + "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf", + "mtp-gemma-4-E4B-it.gguf", + "MTP/mtp-gemma-4-E4B-it-BF16.gguf", + "mmproj-F16.gguf", + ], + ) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf" + + +def test_download_mtp_reuses_cached_subdir_copy_when_no_root_offline(tmp_path, monkeypatch): + # Pre-fix build may have fetched only the MTP/ copy; reuse it offline. + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snap = _seed_snapshot( + tmp_path, + [ + "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf", + "MTP/mtp-gemma-4-E4B-it-BF16.gguf", + ], + ) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it-BF16.gguf" + + +def test_download_mtp_prefers_root_across_snapshots_offline(tmp_path, monkeypatch): + # A newer partial snapshot holds only the MTP/ copy; an older one has the + # root. Must still return the small root, not the large subdir copy. + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snap_partial = _seed_snapshot(tmp_path / "new", ["MTP/mtp-gemma-4-E4B-it-BF16.gguf"]) + snap_full = _seed_snapshot(tmp_path / "old", ["mtp-gemma-4-E4B-it.gguf"]) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap_partial, snap_full]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf" + + +def test_download_mtp_reuse_follows_snapshot_order_offline(tmp_path, monkeypatch): + # Two snapshots both hold a root drafter; newest-first order must win so a + # fresh main GGUF is not paired with a stale drafter revision. + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + newest = _seed_snapshot(tmp_path / "newest", ["mtp-gemma-4-E4B-it.gguf"]) + oldest = _seed_snapshot(tmp_path / "oldest", ["mtp-gemma-4-E4B-it.gguf"]) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [newest, oldest]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).parent.parent.name == "newest" + + +def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch): + # Online, do not reuse a cached copy: go to the download path so a changed + # drafter is refetched (hf_hub_download checks the current revision). + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + snap = _seed_snapshot(tmp_path, ["mtp-gemma-4-E4B-it.gguf"]) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) + + reached = {} + + def _fake_companion( + *, + hf_repo, + hf_token, + pick, + label, + cancel_event = None, + ): + reached["hit"] = True + return None + + b = LlamaCppBackend() + b._download_companion_gguf = _fake_companion + assert b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") is None + assert reached.get("hit") is True