unsloth/studio/backend/core/inference
Daniel Han bef6da59aa
studio: reserve VRAM headroom for the MTP draft cache in auto-fit (#5585)
* studio: reserve VRAM headroom for the MTP draft cache in auto-fit

When MTP is going to engage on this load, _fit_context_to_vram now
budgets 0.85 of available VRAM instead of 0.90, leaving room for
llama.cpp's secondary MTP draft KV cache + compute graph buffers.

Motivation: a user report on RTX 5090 (32 GB) showed Qwen3.6-27B-MTP-GGUF
UD-Q4_K_XL at native auto-context running roughly half the speed of
the same model with a slightly smaller context. The most parsimonious
explanation is a VRAM cliff: at native context the target's KV
already eats the 90% budget, then llama-server allocates the draft
cache + draft graph on top and spills into a slower partial-offload
path. Reducing the budget by 5% on MTP loads avoids the spill without
penalising non-MTP loads. On hardware with abundant VRAM (B200, etc.)
the fit is unchanged because the requested context already fits in
the tighter budget too.

MTP detection mirrors the auto-promotion logic in load_model: the
GGUF advertises nextn_predict_layers, or the model identifier /
local path matches the -MTP marker, and the user has not explicitly
opted out via speculative_type="off" or --spec-type extra args.

Tests: two new cases in test_kv_cache_estimation.py verify that
mtp_engaged=True yields a context less-than-or-equal-to the
non-MTP path on a tight budget, and that kv_on_gpu=False still
short-circuits regardless of mtp_engaged.

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

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

* studio: gate _mtp_will_engage on canonical-mode resolver

After PR #5582 introduced the 5-mode Speculative Decoding dropdown plus
_canonicalize_spec_mode, the auto-fit MTP-engaged predicate becomes:
  * forced mtp / mtp+ngram -> always engage MTP (extra VRAM needed)
  * auto + MTP GGUF (>= 3B) -> engages MTP via auto-promotion
  * auto + MTP GGUF (sub-3B) -> falls back to ngram-mod (no extra VRAM)
  * ngram / ngram-simple / off -> never engage MTP
  * user --spec-type in extra_args -> resolver suppressed; no headroom

The old gate triggered on "anything but off", so it over-reserved the
0.85 budget when the user explicitly picked Ngram (no MTP) or when
Auto fell back to ngram-mod on a sub-3B MTP model. The 5% headroom
cost was minor but unnecessary.

Mirrors the same logic already encoded in _build_speculative_flags so
the auto-fit budget and the actual emission agree on whether MTP is
running.

All 361 backend tests pass.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-05-19 06:19:02 -07:00
..
__init__.py Final cleanup 2026-03-12 18:28:04 +00:00
_html_to_md.py fix: studio web search SSL failures and empty page content (#4754) 2026-04-01 06:12:02 -07:00
anthropic_compat.py Studio: support images on /v1/messages (Anthropic-compat) (#5128) 2026-04-22 03:25:07 +04:00
audio_codecs.py Add native GGUF intake to Studio (#5246) 2026-05-04 11:46:18 +02:00
defaults.py studio: engage draft-mtp on vision MTP GGUFs (drop incorrect vision gate) (#5560) 2026-05-18 08:42:55 -07:00
external_provider.py fix(studio): handle expired OpenAI shell-tool containers without surfacing error in chat (#5547) 2026-05-18 05:47:57 -07:00
inference.py Pin bitsandbytes to continuous-release_main on ROCm (4-bit decode fix) (#4954) 2026-04-10 06:25:39 -07:00
key_exchange.py studio: API external provider support for chat (OpenAI, Mistral, Gemini, Cohere, Anthropic, OpenRouter, DeepSeek, custom providers) (#4706) 2026-05-14 16:13:59 +04:00
llama_cpp.py studio: reserve VRAM headroom for the MTP draft cache in auto-fit (#5585) 2026-05-19 06:19:02 -07:00
llama_server_args.py studio: add --spec-draft-n-max toggle for MTP speculative decoding (#5582) 2026-05-19 06:17:04 -07:00
mlx_inference.py MLX training support for Studio on Apple Silicon (#5340) 2026-05-14 05:24:20 -07:00
orchestrator.py Add native GGUF intake to Studio (#5246) 2026-05-04 11:46:18 +02:00
providers.py Polish/cloud to providers (#5450) 2026-05-15 19:29:21 +04:00
tools.py studio: tighten sandbox blocklist precision (bash, hf upload, NOFILE) (#5487) 2026-05-18 00:01:17 -07:00
worker.py studio: load cached GGUF models when fully offline (#5505) 2026-05-17 21:25:39 -07:00