Studio: warn when llama.cpp prebuilt is too old for MTP (#5528)
* Studio: warn when llama.cpp prebuilt is too old for MTP Layered on #5527. Adds a one-shot llama-server --help capability probe so users get a clear signal when their prebuilt is missing MTP support, plus a graceful fallback if they load an MTP GGUF against an outdated binary. What's surfaced: 1. Startup log + stderr line in main.py:lifespan() if MTP isn't advertised: WARNING: llama.cpp prebuilt is missing MTP support (--spec-type mtp / draft-mtp). Run `unsloth studio update` to refresh it. MTP GGUFs will load without speculative decoding. 2. Load-time graceful fallback in load_model's spec block: skip the auto-emit and log a clear warning instead of letting llama-server fail with an unknown-flag error. 3. /api/inference/status now returns llama_cpp_supports_mtp: bool so the frontend can show a banner / popup. Probe internals: - Class-level cache keyed on (binary_path, mtime). One subprocess call the first time, instant thereafter. Touching the binary (e.g. via `unsloth studio update`) invalidates the cache automatically because the mtime changes, so the new build is picked up without restarting the server. - Recognises both upstream naming forms: the original draft-mtp from llama.cpp PR #22673 and the renamed mtp variant in later commits. - Spec block uses whichever token the binary accepts so we emit the right value regardless of which release the user has. Tests: - 6 new cases in test_llama_cpp_mtp_detection.py covering each probe variant (draft-mtp, renamed mtp, pre-MTP build, missing binary, mtime-based cache invalidation). - Existing 38 MTP detection cases still pass; broader 188-test regression suite (server args, reload inheritance, gguf metadata, load progress, context fit, model validation) still green. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
5a6e94d422
commit
fc04809bfe
5 changed files with 220 additions and 27 deletions
|
|
@ -915,6 +915,61 @@ class LlamaCppBackend:
|
|||
|
||||
return None
|
||||
|
||||
# ── llama-server capability probe ─────────────────────────────
|
||||
|
||||
# Cached on (path, mtime); `unsloth studio update` bumps mtime.
|
||||
_capability_cache: dict[tuple[str, int], dict[str, object]] = {}
|
||||
|
||||
@classmethod
|
||||
def probe_server_capabilities(
|
||||
cls, binary: Optional[str] = None
|
||||
) -> dict[str, object]:
|
||||
"""Parse `llama-server --help` for feature flags. Returns
|
||||
{found, mtp_token, supports_mtp}. mtp_token is "draft-mtp"
|
||||
(older) or "mtp" (renamed upstream), or None."""
|
||||
bin_path = binary or cls._find_llama_server_binary()
|
||||
if not bin_path or not Path(bin_path).is_file():
|
||||
return {"found": False, "mtp_token": None, "supports_mtp": False}
|
||||
try:
|
||||
mtime = int(Path(bin_path).stat().st_mtime)
|
||||
except OSError:
|
||||
mtime = 0
|
||||
cache_key = (bin_path, mtime)
|
||||
cached = cls._capability_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
mtp_token: Optional[str] = None
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[bin_path, "--help"],
|
||||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 10,
|
||||
check = False,
|
||||
)
|
||||
help_text = (result.stdout or "") + "\n" + (result.stderr or "")
|
||||
spec_line = ""
|
||||
for line in help_text.splitlines():
|
||||
if "--spec-type" in line:
|
||||
spec_line = line
|
||||
break
|
||||
# PR #22673 used draft-mtp; later renamed to mtp.
|
||||
if "draft-mtp" in spec_line:
|
||||
mtp_token = "draft-mtp"
|
||||
elif re.search(r"[|,\[]mtp[|,\]]", spec_line):
|
||||
mtp_token = "mtp"
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
logger.debug(f"llama-server --help probe failed: {exc}")
|
||||
|
||||
info = {
|
||||
"found": True,
|
||||
"mtp_token": mtp_token,
|
||||
"supports_mtp": mtp_token is not None,
|
||||
}
|
||||
cls._capability_cache[cache_key] = info
|
||||
return info
|
||||
|
||||
# ── GPU allocation ────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -2569,36 +2624,50 @@ class LlamaCppBackend:
|
|||
cmd.append("--spec-default")
|
||||
self._speculative_type = "default"
|
||||
elif normalized_spec == "draft-mtp":
|
||||
if gpus:
|
||||
cmd.extend(
|
||||
[
|
||||
"--spec-type",
|
||||
"draft-mtp",
|
||||
"--spec-draft-n-max",
|
||||
"6",
|
||||
]
|
||||
# Probe binary; fail gracefully on outdated prebuilts.
|
||||
# Use whichever token the binary advertises
|
||||
# (older: draft-mtp; renamed upstream: mtp).
|
||||
caps = self.probe_server_capabilities(binary)
|
||||
mtp_token = caps.get("mtp_token") if caps else None
|
||||
if not mtp_token:
|
||||
logger.warning(
|
||||
"MTP GGUF detected but llama-server lacks "
|
||||
"--spec-type mtp/draft-mtp; run "
|
||||
"`unsloth studio update`. Loading without "
|
||||
"speculative decoding."
|
||||
)
|
||||
self._speculative_type = None
|
||||
else:
|
||||
cmd.extend(
|
||||
[
|
||||
"--spec-type",
|
||||
"draft-mtp",
|
||||
"--spec-draft-n-max",
|
||||
"3",
|
||||
"--spec-type",
|
||||
"ngram-mod",
|
||||
"--spec-ngram-mod-n-match",
|
||||
"24",
|
||||
"--spec-ngram-mod-n-min",
|
||||
"48",
|
||||
"--spec-ngram-mod-n-max",
|
||||
"6",
|
||||
]
|
||||
if gpus:
|
||||
cmd.extend(
|
||||
[
|
||||
"--spec-type",
|
||||
mtp_token,
|
||||
"--spec-draft-n-max",
|
||||
"6",
|
||||
]
|
||||
)
|
||||
else:
|
||||
cmd.extend(
|
||||
[
|
||||
"--spec-type",
|
||||
mtp_token,
|
||||
"--spec-draft-n-max",
|
||||
"3",
|
||||
"--spec-type",
|
||||
"ngram-mod",
|
||||
"--spec-ngram-mod-n-match",
|
||||
"24",
|
||||
"--spec-ngram-mod-n-min",
|
||||
"48",
|
||||
"--spec-ngram-mod-n-max",
|
||||
"6",
|
||||
]
|
||||
)
|
||||
self._speculative_type = "draft-mtp"
|
||||
logger.info(
|
||||
f"Spec decoding: {mtp_token} ({'GPU' if gpus else 'CPU/Mac'})"
|
||||
)
|
||||
self._speculative_type = "draft-mtp"
|
||||
logger.info(
|
||||
f"Spec decoding: draft-mtp ({'GPU' if gpus else 'CPU/Mac'})"
|
||||
)
|
||||
elif normalized_spec in _valid_spec_types:
|
||||
cmd.extend(["--spec-type", normalized_spec])
|
||||
if normalized_spec == "ngram-mod":
|
||||
|
|
|
|||
|
|
@ -198,6 +198,29 @@ async def lifespan(app: FastAPI):
|
|||
# Detect hardware first — sets DEVICE global used everywhere
|
||||
detect_hardware()
|
||||
|
||||
# llama.cpp capability probe; warns if the prebuilt lacks MTP support.
|
||||
try:
|
||||
from core.inference.llama_cpp import LlamaCppBackend
|
||||
|
||||
_caps = LlamaCppBackend.probe_server_capabilities()
|
||||
app.state.llama_cpp_capabilities = _caps
|
||||
if _caps.get("found") and not _caps.get("supports_mtp"):
|
||||
import structlog as _structlog
|
||||
|
||||
_msg = (
|
||||
"llama.cpp prebuilt lacks MTP support "
|
||||
"(--spec-type mtp/draft-mtp). Run `unsloth studio update`. "
|
||||
"MTP GGUFs will load without speculative decoding."
|
||||
)
|
||||
_structlog.get_logger(__name__).warning(_msg)
|
||||
print(f"WARNING: {_msg}", flush = True)
|
||||
except Exception as _probe_exc:
|
||||
import structlog as _structlog
|
||||
|
||||
_structlog.get_logger(__name__).debug(
|
||||
"llama.cpp capability probe failed: %s", _probe_exc
|
||||
)
|
||||
|
||||
from storage.studio_db import cleanup_orphaned_runs
|
||||
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -342,6 +342,13 @@ class InferenceStatusResponse(BaseModel):
|
|||
None,
|
||||
description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled",
|
||||
)
|
||||
llama_cpp_supports_mtp: bool = Field(
|
||||
True,
|
||||
description = (
|
||||
"Whether llama.cpp supports MTP (--spec-type mtp/draft-mtp). "
|
||||
"False -> recommend `unsloth studio update`."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# =====================================================================
|
||||
|
|
|
|||
|
|
@ -1282,6 +1282,13 @@ async def get_status(
|
|||
try:
|
||||
llama_backend = get_llama_cpp_backend()
|
||||
|
||||
# MTP capability probe (cached). Drives the UI update banner.
|
||||
try:
|
||||
_caps = type(llama_backend).probe_server_capabilities()
|
||||
_supports_mtp = bool(_caps.get("supports_mtp", False))
|
||||
except Exception:
|
||||
_supports_mtp = True # fail open
|
||||
|
||||
# If a GGUF model is loaded via llama-server, report that
|
||||
if llama_backend.is_loaded:
|
||||
_model_id = llama_backend.model_identifier
|
||||
|
|
@ -1324,6 +1331,7 @@ async def get_status(
|
|||
cache_type_kv = llama_backend.cache_type_kv,
|
||||
chat_template_override = llama_backend.chat_template_override,
|
||||
speculative_type = llama_backend.speculative_type,
|
||||
llama_cpp_supports_mtp = _supports_mtp,
|
||||
)
|
||||
|
||||
# Otherwise, report Unsloth backend status
|
||||
|
|
@ -1384,6 +1392,7 @@ async def get_status(
|
|||
supports_preserve_thinking = False,
|
||||
supports_tools = False,
|
||||
chat_template = chat_template,
|
||||
llama_cpp_supports_mtp = _supports_mtp,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -409,3 +409,88 @@ def test_unload_resets_nextn_predict_layers():
|
|||
backend._nextn_predict_layers = 1
|
||||
backend.unload_model()
|
||||
assert backend._nextn_predict_layers is None
|
||||
|
||||
|
||||
# llama-server capability probe.
|
||||
|
||||
|
||||
def _make_fake_llama_server(path: Path, help_text: str) -> Path:
|
||||
"""Bash stub that prints `help_text` on --help."""
|
||||
path.write_text("#!/usr/bin/env bash\n" f"cat <<'EOF'\n{help_text}\nEOF\n")
|
||||
path.chmod(0o755)
|
||||
return path
|
||||
|
||||
|
||||
def _clear_caps_cache():
|
||||
LlamaCppBackend._capability_cache.clear()
|
||||
|
||||
|
||||
def test_probe_server_capabilities_detects_draft_mtp(tmp_path):
|
||||
# Original naming from llama.cpp #22673.
|
||||
fake = _make_fake_llama_server(
|
||||
tmp_path / "llama-server",
|
||||
"--spec-type none,draft-simple,draft-eagle3,draft-mtp,"
|
||||
"ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache",
|
||||
)
|
||||
_clear_caps_cache()
|
||||
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
|
||||
assert caps["found"] is True
|
||||
assert caps["mtp_token"] == "draft-mtp"
|
||||
assert caps["supports_mtp"] is True
|
||||
|
||||
|
||||
def test_probe_server_capabilities_detects_renamed_mtp(tmp_path):
|
||||
# Renamed upstream: draft-mtp -> mtp.
|
||||
fake = _make_fake_llama_server(
|
||||
tmp_path / "llama-server",
|
||||
"--spec-type [none|mtp|ngram-cache|ngram-simple|ngram-map-k|"
|
||||
"ngram-map-k4v|ngram-mod]",
|
||||
)
|
||||
_clear_caps_cache()
|
||||
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
|
||||
assert caps["mtp_token"] == "mtp"
|
||||
assert caps["supports_mtp"] is True
|
||||
|
||||
|
||||
def test_probe_server_capabilities_reports_outdated_binary(tmp_path):
|
||||
# Pre-MTP llama.cpp: only ngram variants.
|
||||
fake = _make_fake_llama_server(
|
||||
tmp_path / "llama-server",
|
||||
"--spec-type none,ngram-simple,ngram-mod",
|
||||
)
|
||||
_clear_caps_cache()
|
||||
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
|
||||
assert caps["found"] is True
|
||||
assert caps["mtp_token"] is None
|
||||
assert caps["supports_mtp"] is False
|
||||
|
||||
|
||||
def test_probe_server_capabilities_handles_missing_binary():
|
||||
_clear_caps_cache()
|
||||
caps = LlamaCppBackend.probe_server_capabilities("/no/such/llama-server")
|
||||
assert caps["found"] is False
|
||||
assert caps["supports_mtp"] is False
|
||||
|
||||
|
||||
def test_probe_server_capabilities_caches_by_mtime(tmp_path):
|
||||
# Same (path, mtime) -> cache hit. Bumped mtime -> re-probe.
|
||||
fake = _make_fake_llama_server(
|
||||
tmp_path / "llama-server",
|
||||
"--spec-type none,ngram-mod",
|
||||
)
|
||||
_clear_caps_cache()
|
||||
caps1 = LlamaCppBackend.probe_server_capabilities(str(fake))
|
||||
assert caps1["supports_mtp"] is False
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
_make_fake_llama_server(
|
||||
fake,
|
||||
"--spec-type none,draft-mtp,ngram-mod",
|
||||
)
|
||||
new_mtime = int(time.time()) + 2
|
||||
os.utime(fake, (new_mtime, new_mtime))
|
||||
caps2 = LlamaCppBackend.probe_server_capabilities(str(fake))
|
||||
assert caps2["mtp_token"] == "draft-mtp"
|
||||
assert caps2["supports_mtp"] is True
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue