From cc3f0430e6e62a1cb0b720de11d66db538046bfb Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Thu, 23 Jul 2026 20:01:06 -0700 Subject: [PATCH] Studio: detect local MTP subdirectory drafters --- .../tests/test_mtp_drafter_companion.py | 47 +++++++++++++++ studio/backend/utils/models/model_config.py | 58 ++++++++++++++++++- .../src/features/chat/chat-settings-sheet.tsx | 6 +- 3 files changed, 107 insertions(+), 4 deletions(-) diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 02230632b6..999e1bd0ca 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -211,6 +211,53 @@ def test_detect_mtp_file_search_root(tmp_path): assert found is not None and found.endswith("mtp-gemma-4-12b-it.gguf") +def test_detect_mtp_file_falls_back_to_new_scheme_subdir(tmp_path): + weight = tmp_path / "gemma-4-E4B-it-qat-Q4_0.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + (sub / "mtp-gemma-4-E4B-it-BF16.gguf").write_bytes(b"x") + q4 = sub / "mtp-gemma-4-E4B-it-Q4_0.gguf" + q4.write_bytes(b"x") + + found = detect_mtp_file(str(weight)) + assert found == str(q4.resolve()) + + +def test_detect_mtp_file_falls_back_to_old_scheme_subdir(tmp_path): + weight = tmp_path / "gemma-4-12b-it-Q4_K_M.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + drafter = sub / "gemma-4-12b-it-Q8_0-MTP.gguf" + drafter.write_bytes(b"x") + + found = detect_mtp_file(str(weight)) + assert found == str(drafter.resolve()) + + +def test_detect_mtp_file_root_still_wins_over_subdir(tmp_path): + weight = tmp_path / "gemma-4-E4B-it-qat-Q4_0.gguf" + weight.write_bytes(b"x") + root = tmp_path / "mtp-gemma-4-E4B-it.gguf" + root.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + (sub / "mtp-gemma-4-E4B-it-Q4_0.gguf").write_bytes(b"x") + + assert detect_mtp_file(str(weight)) == str(root.resolve()) + + +def test_detect_mtp_file_subdir_skips_foreign_drafter(tmp_path): + weight = tmp_path / "gemma-4-E4B-it-qat-Q4_0.gguf" + weight.write_bytes(b"x") + sub = tmp_path / "MTP" + sub.mkdir() + (sub / "mtp-gemma-4-12b-it-Q4_0.gguf").write_bytes(b"x") + + assert detect_mtp_file(str(weight)) is None + + # ── Reload dedup includes the drafter ──────────────────────────────── diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 4897f05ce4..6fa1b85c08 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1455,8 +1455,36 @@ def detect_mtp_file(path: str, search_root: Optional[str] = None) -> Optional[st unsloth names the drafter ``mtp-.gguf`` where ```` prefixes the weight filename across all Gemma 4 repos (e.g. ``mtp-gemma-4-12B-it.gguf`` next to ``gemma-4-12B-it-qat-Q4_0.gguf``). - An unmatched drafter is skipped (fail-safe: no MTP). + If the root drafter is absent, also accept its precision copy under the + repository's ``MTP/`` directory. An unmatched drafter is skipped. """ + + def _pairing_stem(name: str) -> str: + stem = Path(name).stem.lower() + if stem.startswith("mtp-"): + stem = stem[len("mtp-") :] + if stem.endswith("-mtp"): + stem = stem[: -len("-mtp")] + return re.sub(r"-(?:q[0-9]+_[0-9]+|bf16|f16)$", "", stem) + + def _matches_weight(candidate: Path) -> bool: + if weight_name is None: + return True + stem = _pairing_stem(candidate.name) + return bool(stem) and weight_name.startswith(stem) + + def _precision_rank(candidate: Path) -> tuple[int, str]: + name = candidate.name.lower() + if "-q4_0" in name: + rank = 0 + elif "-q8_0" in name: + rank = 1 + elif "-bf16" in name or "-f16" in name: + rank = 2 + else: + rank = 3 + return rank, name + p = Path(path) weight_name = p.name.lower() if p.suffix.lower() == ".gguf" else None start_dir = p.parent if p.is_file() else p @@ -1472,14 +1500,38 @@ def detect_mtp_file(path: str, search_root: Optional[str] = None) -> Optional[st name = f.name.lower() if not (name.startswith("mtp-") and name.endswith(".gguf")): continue - stem = name[len("mtp-") : -len(".gguf")] - if not stem or (weight_name is not None and not weight_name.startswith(stem)): + if not _matches_weight(f): continue try: if f.is_file(): return str(f.resolve()) except OSError: continue + + subdir_candidates: list[Path] = [] + for d in dirs: + mtp_dir = d / "MTP" + try: + entries = sorted(mtp_dir.iterdir()) + except OSError: + continue + for f in entries: + rel = f"MTP/{f.name}" + if not _is_mtp_drafter(rel) or not _matches_weight(f): + continue + try: + if f.is_file(): + subdir_candidates.append(f) + except OSError: + continue + + for candidate in sorted(subdir_candidates, key = _precision_rank): + try: + resolved = candidate.resolve() + except OSError: + continue + logger.info(f"Detected MTP subdirectory drafter: {resolved}") + return str(resolved) return None diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index c3a59e9860..e52a6e15cc 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -378,6 +378,8 @@ export function ChatSettingsPanel({ const isMobile = useIsMobile(); const isLoadedGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null; const currentCheckpoint = params.checkpoint; + const isDirectLocalGguf = + currentCheckpoint?.toLowerCase().endsWith(".gguf") ?? false; const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); // Direct-file / custom-folder GGUFs load without a variant label but still // report a GGUF context, so detect them via the context and the checkpoint @@ -812,7 +814,9 @@ export function ChatSettingsPanel({ : specFallbackReason === "runtime_error" ? "MTP could not start for this model on the installed llama.cpp build, so it is running without speculative decoding." : 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." + ? isDirectLocalGguf + ? "This local model supports MTP, but no matching drafter file was found. Place its mtp-*.gguf beside the model or in its MTP folder, then reload the model." + : "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."