unsloth/studio
Daniel Han 1390e721cb
Studio: reserve the duplicated MTP target KV context for MLA models (GLM-5.2 OOM) (#6447)
* Studio: reserve the duplicated MTP target KV context for MLA models

GLM-5.2 UD-IQ1_S advertised its native 1,048,576-token context, loaded, then
crashed cublasCreate on the first generation with "CUDA error: the resource
allocation failed" on a 2x B200 box. The model loaded fine; the decode OOMed.

Cause: when MTP speculative decoding is engaged, llama.cpp keeps a second full
copy of the target model's KV context for draft verification (ctx_tgt=yes in the
spec log), at f16. On an MLA model that copy is ~the main KV again -- for GLM-5.2
at 1M ctx llama.cpp sized it at ~97.5 GiB -- but the auto-fit reserve only
counted the tiny embedded draft head (~2 GiB), 46x too low. So weights (~202 GiB)
+ main KV (~83 GiB) + a 2 GiB reserve looked like it fit in 2x182 GiB, when the
real footprint with the ~97 GiB MTP copy is ~382 GiB and overruns the cards.
Disabling speculative decoding removed the copy and the same context ran fine.

_estimate_mtp_overhead_bytes now adds the duplicated target context (the main KV
re-estimated at f16) for MLA models, so auto-fit backs the context off (or selects
more GPUs) instead of advertising one that OOMs. It is gated strictly on MLA
(kv_lora_rank present), which is exactly the family that keeps the extra copy
(GLM-5.x, DeepSeek, Kimi-K2); non-MLA MTP (Qwen, Gemma) is byte-for-byte
unchanged. The reserve stays deterministic from GGUF dims, matching #6312.

test_mtp_mla_target_ctx.py covers it: the MLA reserve includes the f16 target
copy and dominates the draft head, the copy is f16 regardless of the main cache
type and scales with context, non-MLA embedded heads keep overhead == draft KV,
and _fit_context_to_vram on the GLM-5.2 / 2x B200 budget now returns a context
below the requested 1M where the old draft-only reserve kept the full 1M.

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

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

* Studio: stub structlog/loggers in MTP-MLA test so it is import-order-independent

The new test imports core.inference.llama_cpp, which pulls in orchestrator ->
structlog. In the lightweight test env structlog is absent, so when this file is
collected before test_mtp_vram_budget.py (it sorts first) or run directly,
collection aborted with ModuleNotFoundError. Install the same loggers/structlog
(+ conditional httpx) stubs the sibling MTP tests use before the import, matching
the established per-file convention.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-18 10:24:33 -07:00
..
backend Studio: reserve the duplicated MTP target KV context for MLA models (GLM-5.2 OOM) (#6447) 2026-06-18 10:24:33 -07:00
frontend Studio: Bypass Permissions menu fix, decimal GB sizes, and GLM-5.2 high/max/disabled thinking (#6444) 2026-06-18 10:10:01 -07:00
src-tauri feat(studio): Hub + Download Manager (#5916) 2026-06-09 04:11:24 -07:00
__init__.py Final cleanup 2026-03-12 18:28:04 +00:00
install_llama_prebuilt.py Studio: reach the published source asset when a mix build's commit 404s (#6314) 2026-06-18 05:59:42 -07:00
install_python_stack.py studio: select torchao version from the installed torch (#6400) 2026-06-17 04:09:40 -07:00
LICENSE.AGPL-3.0 Add AGPL-3.0 license to studio folder 2026-03-09 19:36:25 +00:00
package-lock.json ci: advisory lockfile supply-chain audit (no install-script changes) (#5604) 2026-05-19 05:56:56 -07:00
package.json ci: advisory lockfile supply-chain audit (no install-script changes) (#5604) 2026-05-19 05:56:56 -07:00
setup.bat Final cleanup 2026-03-12 18:28:04 +00:00
setup.ps1 Make the Studio installers (sh + ps1) resilient to transient uv download failures (#6281) 2026-06-16 03:48:01 -07:00
setup.sh Make the Studio installers (sh + ps1) resilient to transient uv download failures (#6281) 2026-06-16 03:48:01 -07:00
Unsloth_Studio_Colab.ipynb Fix/studio colab proxy and iframe - Unsloth Studio not loading in Colab (iframe "refused to connect" and wrong URL) (#5844) 2026-05-28 23:54:48 -07:00