ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: enable MTP for sub-3B Gemma separate-drafter GGUFs (#6191)

Browse source 
* Studio: enable MTP for sub-3B Gemma separate-drafter GGUFs

The sub-3B auto-drop to ngram-mod was tuned for an embedded draft head
(Qwen), whose per-token cost regresses below 3B. Gemma ships the head as a
separate root mtp-*.gguf drafter, a tiny standalone model that is cheap
enough to win below 3B: B200 Q4_K_XL bench, draft-mtp n=2 vs spec-off,
gemma-4-E2B (2B) = 1.21x (accept ~0.65) while ngram-mod is 1.00x.

Exempt a separate drafter from the sub-3B gate everywhere the threshold is
applied: the resolver (_mtp_too_small), the auto-fit VRAM reserve, the
drafter auto-download decision, and the reload-skip mirror via a
has_separate_drafter flag on _auto_mode_drops_mtp. Embedded sub-3B heads
(Qwen) still drop to ngram-mod. A drafter the binary cannot build (older
prebuilt, or a CUDA kernel limit) still aborts the spawn and the load
retries once without speculative decoding.

Adds the full Qwen3.5 + Gemma-4 (regular and QAT) auto/off/forced resolver
matrix, plus explicit sub-3B exemption tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Always compare the separate drafter in the reload-skip mirror

The sub-3B wrapper around the drafter compare could skip it when the drafter
was deleted out from under a running sub-3B server (detected None, stored set),
leaving a stale launch. The resolved-path compare is cheap and already handles
every case, so drop the _auto_mode_drops_mtp guard (and its now-unused imports)
and always compare when the mode can use a drafter and the user does not own
--spec-type. Addresses review feedback on #6191.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-11 05:27:46 -07:00 • committed by GitHub
commit c3604d01f7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 230 additions and 36 deletions
Download patch file Download diff file Expand all files Collapse all files

60
studio/backend/core/inference/llama_cpp.py
View file

@ -514,9 +514,18 @@ _CTX_FIT_VRAM_FRACTION = 0.90
_MTP_VRAM_RESERVE_FRAC = 0.05
def _auto_mode_drops_mtp(req_mode: Optional[str], size_b: Optional[float]) -> bool:
"""Auto mode drops MTP below _MTP_MIN_SIZE_B (draft-mtp regresses there);
forced mtp / mtp+ngram engage regardless of size."""
def _auto_mode_drops_mtp(
req_mode: Optional[str],
size_b: Optional[float],
*,
has_separate_drafter: bool = False,
) -> bool:
"""Auto mode drops MTP below _MTP_MIN_SIZE_B for an embedded draft head
(its per-token cost regresses there); a separate drafter (Gemma) is a tiny
standalone model that still speeds up below 3B, so it never drops. Forced
mtp / mtp+ngram engage regardless of size."""
if has_separate_drafter:
return False
return req_mode == "auto" and size_b is not None and size_b < _MTP_MIN_SIZE_B
@ -2910,19 +2919,14 @@ class LlamaCppBackend:
# Auto-download the separate MTP drafter (e.g. Gemma) when
# the requested spec mode can use it. Repos with the head
# baked into the main GGUF (Qwen) have no mtp- sibling and
# this no-ops. Skipped when the user disabled MTP, drives
# --spec-type manually via extra_args, or in auto mode on a
# sub-3B model (e.g. Gemma E2B) where the resolver drops
# MTP anyway -- no point fetching a drafter it never uses.
# Forced mtp / mtp+ngram still download (user override).
# this no-ops, so the size gate stays out of it: a separate
# drafter speeds up even sub-3B (Gemma E2B), and the resolver
# below decides the final emission. Skipped only when the
# user disabled MTP or drives --spec-type manually.
_spec_canon = _canonicalize_spec_mode(speculative_type) or "auto"
_auto_drops_mtp = _auto_mode_drops_mtp(
_spec_canon, _extract_model_size_b(model_identifier)
)
if (
not mtp_draft_path
and _spec_canon in ("auto", "mtp", "mtp+ngram")
and not _auto_drops_mtp
and not _extra_args_set_spec_type(extra_args)
):
mtp_draft_path = self._download_mtp(
@ -3002,8 +3006,12 @@ class LlamaCppBackend:
_mtp_canonical = _canonicalize_spec_mode(speculative_type)
_mtp_effective = _mtp_canonical or "auto"
_mtp_size_for_fit = _extract_model_size_b(model_identifier)
# Sub-3B drops MTP only for an embedded head; a separate
# drafter (Gemma) engages and needs its VRAM reserved.
_mtp_sub_3b_for_fit = (
_mtp_size_for_fit is not None and _mtp_size_for_fit < _MTP_MIN_SIZE_B
_mtp_size_for_fit is not None
and _mtp_size_for_fit < _MTP_MIN_SIZE_B
and not bool(mtp_draft_path)
)
_mtp_will_engage = bool(
not _extra_args_set_spec_type(extra_args)
@ -3771,18 +3779,26 @@ class LlamaCppBackend:
https://github.com/ggml-org/llama.cpp/pull/18471
MTP guide: unsloth.ai/docs/models/qwen3.6#mtp-guide
Sub-3B dense MTP regresses vs spec-off: the draft head's per-token
cost exceeds the acceptance savings at this scale. Q4_K_XL clean
bench (each prompt once after an unrelated warmup) on B200 + x86 CPU:
Sub-3B dense MTP regresses vs spec-off when the head is baked into the
main GGUF (Qwen): the draft head's per-token cost exceeds the
acceptance savings at this scale. Q4_K_XL clean bench (each prompt once
after an unrelated warmup) on B200 + x86 CPU:
0.8B GPU: draft-mtp n=2 = 0.58x vs OFF; ngram-only = 1.10x
2B GPU: draft-mtp n=2 = 0.82x vs OFF; OFF or ngram = 1.00x
0.8B CPU: chained n=2 = 0.86x vs OFF; ngram-only = 1.19x
2B CPU: chained n=2 = 0.83x vs OFF; ngram-only = 1.01x
4B+ GPU/CPU: spec on is a net win (1.08x-1.46x).
A separate drafter (Gemma's root mtp-*.gguf) is a different, cheaper
mechanism that wins even below 3B, so it is exempt from the sub-3B drop
(``mtp_draft_path`` set -> not too small). B200 Q4_K_XL bench, draft-mtp
n=2 vs OFF: gemma-4-E2B (2B) = 1.21x, accept ~0.65 (vs ngram = 1.00x);
gemma-4-E4B (4B) and 12B engage as usual.
Auto falls back to ngram-mod (zero-VRAM, near-zero idle cost on
diverse content); forced MTP on a model with no head/drafter defaults
back (mtp -> spec-default, mtp+ngram -> ngram-mod) since llama-server
aborts otherwise; sub-3B real-MTP engages with a warning.
diverse content) for an embedded sub-3B head; forced MTP on a model
with no head/drafter defaults back (mtp -> spec-default, mtp+ngram ->
ngram-mod) since llama-server aborts otherwise; a drafter the binary
cannot build (older prebuilt, or a CUDA kernel limit) aborts the spawn
and the load retries once without speculative decoding.
"""
flags: List[str] = []
# Reset; emit branches re-set on the resolved emission.
@ -3801,7 +3817,11 @@ class LlamaCppBackend:
)
user_owns_spec_type = _extra_args_set_spec_type(extra_args)
_mtp_size_b = _extract_model_size_b(model_identifier)
_mtp_too_small = _mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B
# The sub-3B regression is an embedded-head cost; a separate drafter
# (Gemma) is a cheap standalone model that wins below 3B, so exempt it.
_mtp_too_small = (
_mtp_size_b is not None and _mtp_size_b < _MTP_MIN_SIZE_B and not bool(mtp_draft_path)
)
if user_owns_spec_type:
# User --spec-type wins outright; suppress auto-emit to avoid a

26
studio/backend/routes/inference.py
View file

@ -417,7 +417,6 @@ try:
LlamaCppBackend,
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_T_MAX_PREDICT_MS,
_auto_mode_drops_mtp,
_canonicalize_spec_mode,
_extra_args_set_spec_type,
_hf_offline_if_dns_dead,
@ -431,7 +430,6 @@ try:
from utils.inference import load_inference_config
from utils.models.model_config import (
detect_mtp_file,
extract_model_size_b,
load_model_defaults,
)
from utils.native_path_leases import (
@ -450,7 +448,6 @@ except ImportError:
LlamaCppBackend,
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_T_MAX_PREDICT_MS,
_auto_mode_drops_mtp,
_canonicalize_spec_mode,
_extra_args_set_spec_type,
_hf_offline_if_dns_dead,
@ -464,7 +461,6 @@ except ImportError:
from utils.inference import load_inference_config
from utils.models.model_config import (
detect_mtp_file,
extract_model_size_b,
load_model_defaults,
)
from utils.native_path_leases import (
@ -1025,14 +1021,15 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC
else:
if list(request.llama_extra_args) != backend_extra:
return False
# A drafter that appeared next to the loaded weights since the last load
# changes the launch command (--model-draft) when the mode can use it;
# without this, a duplicate /load is deduped and MTP can't engage short
# of an unload. Runs last: it stats the filesystem (two dir scans against
# the resolved weight path -- covers local dirs and HF cache snapshots
# alike), so every pure-memory comparison above short-circuits first.
# Skipped when auto drops MTP anyway (sub-3B) or the user owns
# --spec-type, where a drafter changes nothing. Resolve both sides: the
# A separate drafter (Gemma's root mtp-*.gguf) appearing or disappearing
# next to the loaded weights changes the launch command (--model-draft),
# so a duplicate /load must reload rather than dedupe. Always compare the
# detected vs stored drafter when the mode can use one and the user does
# not own --spec-type: the resolved-path compare is cheap and handles all
# four cases (both None -> match; one None -> reload; equal -> match;
# different -> reload), including a drafter deleted out from under a
# running server. Runs last: it stats the filesystem, so every pure-memory
# comparison above short-circuits first. Resolve both sides since the
# stored launch path may be a snapshot symlink while detect_mtp_file
# returns the resolved blob.
if req_mode in ("auto", "mtp", "mtp+ngram") and llama_backend.gguf_path:
@ -1041,10 +1038,7 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC
if request.llama_extra_args is not None
else llama_backend.extra_args
)
size_b = extract_model_size_b(llama_backend.model_identifier or "")
if not _auto_mode_drops_mtp(req_mode, size_b) and not _extra_args_set_spec_type(
effective_extras
):
if not _extra_args_set_spec_type(effective_extras):
detected = detect_mtp_file(llama_backend.gguf_path)
stored = llama_backend.mtp_draft_path
try:

180
studio/backend/tests/test_llama_cpp_mtp_detection.py
View file

@ -1219,3 +1219,183 @@ def test_forced_mtp_ngram_on_non_mtp_model_keeps_ngram(monkeypatch):
assert parsed.get("--spec-type") == "ngram-mod"
assert backend.speculative_type == "ngram-mod"
assert backend.requested_spec_mode == "mtp+ngram"
# ── Full named-repo resolver matrix (the shipping Studio families) ─────
#
# Locks auto / off / forced-mtp routing for every Qwen3.5 (MTP + plain) and
# gemma-4 (regular + QAT) GGUF repo, including the giant MoEs that stay
# resolver-only (122B-A10B / 397B-A17B). Expectations are derived from the
# same signals load_model uses -- _extract_model_size_b (active>effective>
# total, so E2B->2, A3B->3, A10B->10, A17B->17), _is_mtp_model_name, and the
# separate-drafter flag -- so each row mirrors what the loader emits on a
# B200 (GPU default, n=2). gemma carries no -MTP marker; its MTP comes from
# the root mtp-*.gguf drafter, modelled here by passing mtp_draft_path.
#
# auto_spec: "draft-mtp" = head/drafter engaged (>=3B MTP, or any size with a
# separate drafter); "ngram-mod" = embedded sub-3B drop (zero-VRAM); None =
# non-MTP -> llama-server --spec-default.
_GEMMA_DRAFTER = "/snap/mtp-gemma-4-it.gguf" # stand-in separate drafter
_REAL_REPO_MATRIX = [
# repo, drafter, auto_spec, auto_ngram_knobs
("unsloth/Qwen3.5-0.8B-MTP-GGUF", None, "ngram-mod", True),
("unsloth/Qwen3.5-2B-MTP-GGUF", None, "ngram-mod", True),
("unsloth/Qwen3.5-4B-MTP-GGUF", None, "draft-mtp", False),
("unsloth/Qwen3.5-9B-MTP-GGUF", None, "draft-mtp", False),
("unsloth/Qwen3.5-27B-MTP-GGUF", None, "draft-mtp", False),
("unsloth/Qwen3.5-35B-A3B-MTP-GGUF", None, "draft-mtp", False),
("unsloth/Qwen3.5-122B-A10B-MTP-GGUF", None, "draft-mtp", False),
("unsloth/Qwen3.5-397B-A17B-MTP-GGUF", None, "draft-mtp", False),
("unsloth/Qwen3.5-0.8B-GGUF", None, None, False),
("unsloth/Qwen3.5-2B-GGUF", None, None, False),
("unsloth/Qwen3.5-4B-GGUF", None, None, False),
("unsloth/Qwen3.5-9B-GGUF", None, None, False),
# E2B is 2B but ships a separate drafter -> exempt from the sub-3B drop.
("unsloth/gemma-4-E2B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-E4B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-12b-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-26B-A4B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-31B-it-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-E2B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-E4B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-12b-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-26B-A4B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
("unsloth/gemma-4-31B-it-qat-GGUF", _GEMMA_DRAFTER, "draft-mtp", False),
]
def _resolve_real(monkeypatch, repo, drafter, mode):
backend = _resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = mode,
spec_draft_n_max = None,
extra_args = None,
model_identifier = repo,
model_path = None,
gpus = True, # B200 default
binary = "/fake/llama-server",
mtp_draft_path = drafter,
)
return backend, flags, _flags_dict(flags)
@pytest.mark.parametrize(
"repo, drafter, auto_spec, auto_ngram_knobs",
_REAL_REPO_MATRIX,
ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX],
)
def test_real_repo_auto_routing(monkeypatch, repo, drafter, auto_spec, auto_ngram_knobs):
# Auto is the default mode the dropdown ships with.
backend, flags, parsed = _resolve_real(monkeypatch, repo, drafter, "auto")
if auto_spec is None:
# Non-MTP: no draft-mtp, hand off to llama-server's own default.
assert "--spec-type" not in parsed
assert "--spec-default" in flags
assert backend.speculative_type == "default"
elif auto_spec == "draft-mtp":
assert parsed.get("--spec-type") == "draft-mtp"
assert parsed.get("--spec-draft-n-max") == "2"
assert backend.speculative_type == "draft-mtp"
# gemma ships a separate drafter; Qwen bakes the head into the GGUF.
assert (
(parsed.get("--model-draft") == drafter) if drafter else ("--model-draft" not in parsed)
)
else: # ngram-mod (sub-3B MTP drop)
assert parsed.get("--spec-type") == "ngram-mod"
assert "--model-draft" not in parsed # draft head dropped
assert backend.speculative_type == "ngram-mod"
if auto_ngram_knobs:
assert "--spec-ngram-mod-n-match" in parsed
assert backend.requested_spec_mode == "auto"
@pytest.mark.parametrize(
"repo, drafter",
[(r[0], r[1]) for r in _REAL_REPO_MATRIX],
ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX],
)
def test_real_repo_off_emits_nothing(monkeypatch, repo, drafter):
# Off must suppress speculative decoding for every family.
backend, flags, _ = _resolve_real(monkeypatch, repo, drafter, "off")
assert flags == []
assert backend.speculative_type is None
assert backend.requested_spec_mode == "off"
@pytest.mark.parametrize(
"repo, drafter",
[(r[0], r[1]) for r in _REAL_REPO_MATRIX],
ids = [r[0].split("/")[-1] for r in _REAL_REPO_MATRIX],
)
def test_real_repo_forced_mtp_never_aborts(monkeypatch, repo, drafter):
# Forcing MTP on the dropdown: real MTP models (name marker or separate
# drafter) engage draft-mtp even below 3B; non-MTP models default back to
# --spec-default instead of emitting a draft-mtp llama-server will abort on.
backend, flags, parsed = _resolve_real(monkeypatch, repo, drafter, "mtp")
is_real_mtp = _is_mtp_model_name(repo) or bool(drafter)
if is_real_mtp:
assert parsed.get("--spec-type") == "draft-mtp"
assert backend.speculative_type == "draft-mtp"
assert (
(parsed.get("--model-draft") == drafter) if drafter else ("--model-draft" not in parsed)
)
else:
assert "--spec-type" not in parsed
assert "--spec-default" in flags
assert backend.speculative_type == "default"
assert backend.requested_spec_mode == "mtp"
# ── Sub-3B separate-drafter exemption (Gemma) ─────────────────────────
#
# The sub-3B MTP drop is an embedded-head cost (Qwen). A separate drafter
# (Gemma's root mtp-*.gguf) is a cheap standalone model that wins below 3B
# (B200 Q4_K_XL: gemma-4-E2B draft-mtp n=2 = 1.21x vs OFF), so it is exempt.
def test_sub3b_gemma_separate_drafter_engages_mtp(monkeypatch):
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-E2B-it-GGUF", # 2B
model_path = None,
gpus = True,
binary = "/fake/llama-server",
mtp_draft_path = "/snap/mtp-gemma-4-E2B-it.gguf", # separate drafter
)
parsed = _flags_dict(flags)
assert parsed.get("--spec-type") == "draft-mtp"
assert parsed.get("--model-draft") == "/snap/mtp-gemma-4-E2B-it.gguf"
assert "--spec-ngram-mod-n-match" not in parsed
assert backend.speculative_type == "draft-mtp"
def test_sub3b_qwen_embedded_head_still_drops_to_ngram(monkeypatch):
backend = _resolver_backend(monkeypatch)
flags = backend._build_speculative_flags(
speculative_type = "auto",
spec_draft_n_max = None,
extra_args = None,
model_identifier = "unsloth/Qwen3.5-2B-MTP-GGUF", # 2B, embedded head
model_path = None,
gpus = True,
binary = "/fake/llama-server",
mtp_draft_path = None, # no separate drafter
)
parsed = _flags_dict(flags)
assert parsed.get("--spec-type") == "ngram-mod"
assert "--model-draft" not in parsed
assert backend.speculative_type == "ngram-mod"
def test_auto_mode_drops_mtp_exempts_separate_drafter():
from core.inference.llama_cpp import _auto_mode_drops_mtp
assert _auto_mode_drops_mtp("auto", 2.0) is True
assert _auto_mode_drops_mtp("auto", 2.0, has_separate_drafter = True) is False
assert _auto_mode_drops_mtp("auto", 4.0) is False
assert _auto_mode_drops_mtp("mtp", 2.0) is False # forced engages regardless