unsloth/studio/backend
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
..
assets feat(studio): expose provider_type selector in model provider dialog (#4277) 2026-06-15 15:07:06 +02:00
auth Studio: add --secure Cloudflare-only mode and revamp API usage examples (#6300) 2026-06-15 04:18:15 -07:00
core Studio: reserve the duplicated MTP target KV context for MLA models (GLM-5.2 OOM) (#6447) 2026-06-18 10:24:33 -07:00
hub Studio: free chat model VRAM at training start only when the GPU is tight (#6243) 2026-06-18 09:04:01 -07:00
loggers Studio: trim serving-log noise and surface llama-server engine stats (#6377) 2026-06-17 05:37:57 -07:00
models 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
plugins Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
requirements studio: select torchao version from the installed torch (#6400) 2026-06-17 04:09:40 -07:00
routes 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
state Studio: Add inline confirmation (Allow/Always allow/Deny) for tool calls (#5869) 2026-06-12 10:55:26 +02:00
storage feat: implement thread forking functionality with associated database… (#5810) 2026-06-15 14:57:39 +01:00
tests Studio: reserve the duplicated MTP target KV context for MLA models (GLM-5.2 OOM) (#6447) 2026-06-18 10:24:33 -07:00
utils Studio: show an actionable message when the GGUF runtime is missing (#6327) 2026-06-18 06:00:00 -07:00
__init__.py Final cleanup 2026-03-12 18:28:04 +00:00
_platform_compat.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
cloudflare_tunnel.py Reap Studio child processes when the parent dies abnormally (#6425) 2026-06-18 05:51:22 -07:00
colab.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
main.py Studio: pin CUDA_DEVICE_ORDER=PCI_BUS_ID and list GPUs at startup (#6353) 2026-06-17 22:40:04 -07:00
run.py Keep server-side tools enabled under --secure (#6403) 2026-06-18 05:52:40 -07:00
startup_banner.py Studio: enable stdio MCP servers on a loopback bind (#6295) 2026-06-15 03:02:32 +01:00