unsloth/studio/backend/models
Daniel Han 27d4aced59
studio: add --spec-draft-n-max toggle for MTP speculative decoding (#5582)
* studio: add --spec-draft-n-max toggle for MTP speculative decoding

Surface llama-server's --spec-draft-n-max as a first-class
LoadRequest field so users can tune the MTP draft tree size from
the chat settings panel. Default behaviour is unchanged: when the
caller omits spec_draft_n_max, the existing platform defaults still
apply (6 on GPU, 3 on CPU/Mac).

Why this matters: on context-constrained loads the draft KV cache
competes with the target model's KV cache for VRAM. Lowering
spec_draft_n_max reduces that pressure, lets a larger user context
fit, and recovers throughput; raising it pays off when draft
acceptance is high enough to amortise the extra cache.

Backend
- LoadRequest gains an optional spec_draft_n_max: int (1..16).
- LlamaCppBackend.load_model accepts and persists the override on
  self._spec_draft_n_max, used in place of the hardcoded 6/3 in the
  MTP emit branch.
- LoadResponse and InferenceStatusResponse echo the active value
  (None when the platform default is in effect) so the UI can
  hydrate the input on refresh.
- _already_in_target_state and _request_matches_loaded_settings
  compare spec_draft_n_max alongside speculative_type so a value
  change triggers a reload rather than no-op'ing.
- strip_shadowing_flags now strips inherited --spec-* extras when
  either speculative_type or spec_draft_n_max is in fields_set, so
  an inherited --spec-draft-n-max cannot last-wins-override a fresh
  request's first-class field.

Frontend
- LoadModelRequest, LoadModelResponse, InferenceStatusResponse
  TypeScript shapes get spec_draft_n_max.
- chat-runtime-store gains specDraftNMax / loadedSpecDraftNMax and
  a setter, hydrated from /v1/status and /v1/load.
- chat-settings-sheet renders a "Draft Tokens" numeric input
  directly under the Speculative Decoding switch when that switch
  is on. Toggling the switch off clears the override; the Reset
  button restores the loaded value.

Tests
- Four new regression tests cover _already_in_target_state with
  matching / mismatching / non-MTP / unset spec_draft_n_max.
- Existing test_llama_server_args.py and test_llama_cpp_mtp_detection.py
  green: 141 passed locally.

* studio: add --spec-draft-p-min and --spec-draft-p-split to spec strip set

llama.cpp server documents --spec-draft-p-min (default 0.75, min draft
acceptance probability) and --spec-draft-p-split (default 0.10). Both
are first-class spec-decoding knobs that should travel with the rest
of the --spec-* family when an Apply re-sets speculative_type, so an
inherited override doesn't leak across a fresh load.

* studio/tests: skip MTP capability-probe tests on Windows

The four probe_server_capabilities tests use a bash stub written to
tmp_path/llama-server, which Windows' subprocess can't execute
directly (no shebang resolution, .bat / .cmd would be needed). Mark
them skipif sys.platform == 'win32' so the rest of the MTP plumbing
suite stays green on Windows CI. Unix coverage is unchanged.

* studio: lower MTP GPU default --spec-draft-n-max from 6 to 2

Bench on B200 / Qwen3.6-27B-MTP-GGUF UD-Q4_K_XL across five prompt
types (essay, code, story, math, science) with greedy temp=0:

  prompt    OFF    n=1    n=2    n=3    n=6
  essay    79.1   93.4   93.8   84.7   64.6
  code     79.1  104.4  116.6  113.5  103.0
  story    79.1   99.2  105.7  101.8   88.9
  math     79.1  100.8  110.8  111.8   98.2
  science  79.1  100.1  110.8  110.8  102.9

The previous hardcoded GPU default of 6 was 17% SLOWER than spec-off
on the essay prompt (64.6 vs 79.1 t/s) and 11-50% slower than n=2 on
the rest. n=2 wins on 4/5 prompts with a 1.18x-1.47x speedup vs OFF;
n=3 wins on the math prompt by a hair. n=6 collapses once acceptance
rate drops past n=3 -- wasted draft decode dominates the per-step
budget.

Matches the dataset README ("n_max=2 is the sweet spot for 36 of 42
quants"). Keeps CPU/Mac default at 3, which empirically tracks the
narrower ngram+MTP chained budget on those platforms.

Users who want the old behaviour can pass spec_draft_n_max in
LoadRequest (the toggle this PR also adds) or --spec-draft-n-max via
llama_extra_args.

* studio: skip MTP auto-promote on sub-2B models, backfill chat usage

Two MTP-visibility fixes uncovered while bisecting llama.cpp post-#22673
on Qwen3.6-27B-MTP-GGUF UD-Q4_K_XL on B200.

Size gate. Direct llama-server bench (no Studio measurement loop) at
n_predict=192 across 9 prompts shows MTP regresses vs spec-off on
sub-2B dense models because draft cost exceeds savings:

  Qwen3.5-0.8B Q4_K_XL   GPU: 452.0 OFF -> 283.4 t/s n=2  (0.63x)
                         CPU: 84.5  OFF -> 64.9  t/s n=3  (0.77x)
  Qwen3.5-4B  Q4_K_XL    GPU: 241.0 OFF -> 258.2 t/s n=2  (1.07x)
  Qwen3.5-9B  Q4_K_XL    GPU: 201.6 OFF -> 228.9 t/s n=2  (1.14x)
  Qwen3.5-27B Q4_K_XL    GPU:  78.8 OFF -> 113.6 t/s n=2  (1.44x)
  Qwen3.6-27B Q4_K_XL    GPU:  78.8 OFF -> 113.6 t/s n=2  (1.44x)
  Qwen3.6-35B-A3B Q4     GPU: 192.3 OFF -> 223.2 t/s n=2  (1.16x)

The 2B inflection is sharp. Skip auto-promote to draft-mtp when the
identifier reports <2.0B params; users can still force via --spec-type
or the Speculative Decoding toggle. Mirror the gate in the
reload-skip check so a sub-2B reload-with-default does not bounce a
spec-off backend.

Chat-completions usage. llama-server's final SSE chunk emits both an
OpenAI-style usage block and a custom timings block. timings.predicted_n
is always populated, but usage.completion_tokens is zero on some
server builds. The Studio chat UI computes generation t/s from
meta.usage.completion_tokens / totalStreamTime, so a zero
completion_tokens makes the UI fall back to wall-clock time
(including SSE / proxy / template overhead) which dilutes MTP gains and
makes ON look the same as OFF.

Add _backfill_usage_from_timings: if usage.completion_tokens is missing
or zero AND timings has predicted_n/prompt_n, synthesize a complete
usage dict. Apply at the streaming metadata yield in
generate_chat_completion and at the three accumulator/yield sites in
generate_chat_completion_with_tools so per-iteration counts are not
silently lost across tool calls.

Tests cover both the gate (sub-2B skips, 2B+ promotes) and the
backfill (zero usage filled, real usage preserved, empty timings
passthrough).

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

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

* studio: probe + emit legacy ngram-mod flags for pre-rename llama-server

llama.cpp upstream renamed the ngram-mod tuning knobs:

  --draft-max         -> --spec-ngram-mod-n-max  (and --spec-draft-n-max)
  --draft-min         -> --spec-ngram-mod-n-min  (and --spec-draft-n-min)
  --spec-ngram-size-n -> --spec-ngram-mod-n-match

The new names are real flags on post-rename builds and stub removal
entries on the same builds (with description "argument has been
removed"). Pre-rename builds only carry the legacy names as real
flags. Studio was emitting the new names unconditionally, so a user
running a pre-rename llama-server (e.g. an older prebuilt or a
hand-installed binary) would see "unknown argument" errors when the
ngram-mod path engages, or silent drop of the ngram knobs.

Extend `probe_server_capabilities` to parse the help text into
per-flag description blocks and tell real flags apart from removal
stubs by the "argument has been removed" marker. Add three new probe
fields: `ngram_mod_flavor` ("new" / "legacy" / None),
`supports_ngram_mod`, and `spec_draft_n_max_flag` (the actual n_max
flag the binary accepts). Cached by (path, mtime) the same way as
`mtp_token`.

Add `_build_ngram_mod_flags(caps, ...)` that picks the right flag
set, returning [] when neither is usable so callers can drop ngram
chaining entirely on minimal binaries.

Wire both call sites to use the probe-driven flag set:
- CPU/Mac MTP comma-chain (--spec-type ngram-mod,draft-mtp) emits
  legacy or new knobs as appropriate. If neither set is available,
  degrade to MTP-only (warn but still engage spec).
- Standalone --spec-type ngram-mod branch uses the same helper.

Tests cover post-rename detection, legacy detection, removal-stub
discrimination, minimal-binary case, and all three branches of
`_build_ngram_mod_flags` plus custom n_match/n_min/n_max values.

Verified against three real binaries (Studio bundled 726704a, my
build of 45b455e HEAD, and the MTP merge baseline 2555826) all
correctly reporting ngram_mod_flavor=new.

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

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

* studio: sub-3B MTP falls back to ngram-mod, not off

Earlier sub-2B gate disabled speculative decoding entirely for tiny
dense MTP models because the MTP draft head's per-token cost exceeds
the acceptance savings at that scale. The "fully off" fallback was
conservative -- ngram-mod has near-zero idle cost on diverse content
and consistently outperforms both off and draft-mtp at sub-3B.

Clean-methodology bench (each of 9 distinct prompts run once after
two unrelated warmup prompts so the ngram-mod hash pool is
realistically populated but never holds the exact deterministic
output we're about to measure):

  Q4_K_XL on B200:
    0.8B  OFF=451  draft-mtp n=2=263 (0.58x)  ngram-only=498 (1.10x)
    2B    OFF=377  draft-mtp n=2=308 (0.82x)  ngram-only=369 (1.00x)
    4B    OFF=240  draft-mtp n=2=260 (1.08x)  -- 4B+ wins with MTP

  Q4_K_XL on x86 48 cores:
    0.8B  OFF= 80  chained n=2= 69 (0.86x)  ngram-only= 95 (1.19x)
    2B    OFF= 62  chained n=2= 51 (0.83x)  ngram-only= 63 (1.01x)
    4B    OFF= 31  chained n=2= 41 (1.33x)

Change:
- Raise the MTP-skip threshold from 2.0B to 3.0B (2B falls below it).
- When skipping the MTP head, fall back to --spec-type ngram-mod via
  the probe-driven _build_ngram_mod_flags helper. Works on both
  post-rename and pre-rename llama-server builds.
- If the binary advertises neither ngram-mod flavor, fall back to
  spec-off (older binaries that don't support ngram-mod at all).
- Mirror the same fallback in _already_in_target_state so a sub-3B
  reload-with-default does not bounce a ngram-mod backend.

Tests updated: monkeypatch probe_server_capabilities so the gate
behavior is deterministic regardless of which llama-server happens
to be on the host. +1 new test for the "binary has no ngram-mod
support" branch; renamed prior 2B/0.8B tests to reflect new semantics.

This generalizes the size gate to be probe-driven instead of a hard
"disable spec" branch.

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

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

* studio: 5-mode Speculative Decoding dropdown (Auto / MTP / Ngram / MTP+Ngram / Off)

Replace the Chat Settings Speculative Decoding on/off Switch with a 5-option
Select. Auto preserves today's platform-aware resolver (MTP on MTP GGUFs,
ngram-mod fallback for sub-3B, --spec-default for non-MTP). The other 3 modes
force the user's choice on BOTH GPU and CPU: MTP emits draft-mtp only (no
ngram chain on CPU), Ngram emits ngram-mod only, MTP+Ngram emits the
ngram-mod,draft-mtp chain on both platforms. Off is the existing fully-off
state, kept so the Switch's "disable" capability isn't lost.

Backend
- New module-level _canonicalize_spec_mode(value) maps any accepted input
  (canonical, legacy "default" / "draft-mtp" / "ngram-mod" / "ngram-simple",
  or comma-chained "ngram-mod,draft-mtp") onto one of auto / mtp / ngram /
  mtp+ngram / off / ngram-simple / None. Lets external callers and old
  persisted UI state round-trip without breaking.
- LlamaCppBackend grows a _requested_spec_mode field + requested_spec_mode
  property storing the canonical UI mode the user requested. Status
  responses round-trip this instead of the resolved internal flag, so the
  dropdown restores the picked value after reload / refresh (Auto on a 27B
  MTP GGUF resolves to draft-mtp internally but the dropdown stays on
  "Auto").
- The resolver block in load_model is extracted into a unit-testable
  _build_speculative_flags method. Forced MTP / MTP+Ngram on a sub-3B or
  non-MTP GGUF logs a warning and engages anyway (user override > the
  Auto-path sub-3B fallback).
- _already_in_target_state and routes/inference._request_matches_loaded_settings
  now compare canonical-requested mode, dropping the old auto-promotion
  mirror. spec_draft_n_max still gates on the resolved spec so Auto + a
  changed n_max still bounces a reload.

Frontend
- chat-settings-sheet.tsx: Switch swapped for Select modeled on the KV
  Cache Dtype Select. Items: Auto / MTP / Ngram / MTP+Ngram / Off. Draft
  Tokens input only visible when speculativeType is "mtp" or "mtp+ngram".
- chat-runtime-store.ts: initial value flips from "default" to "auto".
- use-chat-model-runtime.ts normalizeSpeculativeType mirrors the backend
  canonicaliser so persisted "default" / "draft-mtp" / "ngram-mod" / chain
  values hydrate to the right dropdown option.
- types/api.ts: docs the canonical wire vocabulary.

Tests
- 53 new assertions in test_llama_cpp_mtp_detection.py: full
  _canonicalize_spec_mode table, a 23-row resolver matrix across
  (requested mode) x (GPU/CPU) x (model size class), plus n_max override,
  user-extra-args precedence, requested-mode round-trip, and graceful
  degrade on an outdated llama-server without an MTP token.
- 165 existing backend tests still green. 218 total in the MTP /
  server-args / reload-inheritance suite.

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

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

* studio: reset Speculative Decoding to Auto on model switch

When the user switches from model A to a different model B, clear the
runtime store's speculativeType + specDraftNMax (and their loaded*
shadows). The new load request then carries null, the backend
canonicalises that to "auto", and its platform-aware resolver runs
fresh for the new model.

Without this, a non-MTP model loaded with "Off" carried the Off choice
into a subsequent MTP load, suppressing MTP auto-promotion (and the
sub-3B ngram-mod fallback) until the user manually opened settings and
flipped the dropdown back to Auto. The clean-sweep deep probe caught
it as anomaly A-1.

The reset only fires when currentCheckpoint != modelId, so a
same-model reapply or forceReload still honours the user's current
spec choice. End-to-end probe on Qwen3.5-4B-GGUF (non-MTP, Off) ->
Qwen3.5-0.8B-MTP confirms: dropdown shows Auto, /api/inference/status
returns speculative_type=auto, studio.log shows the Auto sub-3B
fallback emitted --spec-type ngram-mod.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-05-19 06:17:04 -07:00
..
.gitkeep fix: restore models directory files deleted during restructure 2026-02-02 19:36:30 +00:00
__init__.py Studio: Dark theme refactor, right sidebar redesign, and chat UI polish (#5150) 2026-05-07 14:33:31 +04:00
auth.py studio: security and hardening pass (auth rate-limit, sandbox, path containment, schema validation, headers) (#5375) 2026-05-13 06:12:18 -07:00
data_recipe.py feat(studio): multi-file unstructured seed upload with better backend extraction (#4468) 2026-03-20 13:22:42 -07:00
datasets.py Final cleanup 2026-03-12 18:28:04 +00:00
export.py studio: security and hardening pass (auth rate-limit, sandbox, path containment, schema validation, headers) (#5375) 2026-05-13 06:12:18 -07:00
inference.py studio: add --spec-draft-n-max toggle for MTP speculative decoding (#5582) 2026-05-19 06:17:04 -07:00
models.py Studio: add folder browser modal for Custom Folders (#5035) 2026-04-15 08:04:33 -07:00
providers.py Studio: make API key optional for local providers (llama.cpp/vLLM/Ollama) (#5457) 2026-05-15 23:33:22 +04:00
responses.py Final cleanup 2026-03-12 18:28:04 +00:00
training.py studio: drop unused max_grad_value schema + route plumbing (#5424) 2026-05-14 05:43:58 -07:00
users.py fix: remove old comments (#4292) 2026-03-14 16:50:13 +04:00