unsloth/studio/backend/core
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
..
data_recipe ci: advisory lockfile supply-chain audit (no install-script changes) (#5604) 2026-05-19 05:56:56 -07:00
export feat(studio): MLX training tab on Apple Silicon (LoRA / full FT, VLM, export) (#5265) 2026-05-05 23:54:58 -07:00
inference studio: reserve VRAM headroom for the MTP draft cache in auto-fit (#5585) 2026-05-19 06:19:02 -07:00
training studio: install flash-linear-attention and tilelang for Qwen3.5 family (#5434) 2026-05-18 03:49:06 -07:00
__init__.py [Studio] Show non exported models in chat UI (#4892) 2026-04-14 15:03:58 +04:00
tool_healing.py studio: extract tool-call XML parser into a reusable helper module (#5583) 2026-05-19 05:06:17 -07:00