From e6ab7f69522b7a6f31ef5ec09e37aefbc17ce862 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 27 Jun 2026 09:04:06 +0000 Subject: [PATCH 1/2] Studio: size separate-drafter MTP draft KV at full context (swa_full) For a separate drafter (Gemma mtp-*.gguf) the draft caches against the full target window, so its sliding-window layers are not window-capped at runtime and the draft KV grows with n_ctx. _estimate_kv_cache_bytes defaulted to the SWA window cap, which sized the drafter at a near-flat ~16 MiB and under- reserved on long contexts. Pass swa_full = True so the draft KV is sized as a safe upper bound. Measured draft overhead on a 12b drafter is 580 MiB at 32768, 868 MiB at 131072 and 1252 MiB at 262144, all under-reserved before. Embedded MTP heads (Qwen) are unaffected. All fit, KV, MTP and context tests pass. --- studio/backend/core/inference/llama_cpp.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 8ec11fa79f..90e8aeba75 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2907,7 +2907,11 @@ class LlamaCppBackend: # The drafter is served under the same --parallel slot count as the # main model, so price its KV per slot too: a sliding-window drafter # (Gemma) grows KV with slots and would otherwise be under-reserved. - kv = db._estimate_kv_cache_bytes(n_ctx, heavier, n_parallel = n_parallel) + # swa_full: the draft caches against the full target window, so its SWA + # layers are not window-capped and grow with n_ctx (safe upper bound). + kv = db._estimate_kv_cache_bytes( + n_ctx, heavier, n_parallel = n_parallel, swa_full = True + ) return kv or None nextn = self._nextn_predict_layers or 0 n_kv = self._n_kv_heads or self._n_heads From 6b19a8984aa83c592dbb96c6a9ef5426027bf44e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 27 Jun 2026 09:04:54 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/llama_cpp.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 90e8aeba75..06cf4363e9 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2909,9 +2909,7 @@ class LlamaCppBackend: # (Gemma) grows KV with slots and would otherwise be under-reserved. # swa_full: the draft caches against the full target window, so its SWA # layers are not window-capped and grow with n_ctx (safe upper bound). - kv = db._estimate_kv_cache_bytes( - n_ctx, heavier, n_parallel = n_parallel, swa_full = True - ) + kv = db._estimate_kv_cache_bytes(n_ctx, heavier, n_parallel = n_parallel, swa_full = True) return kv or None nextn = self._nextn_predict_layers or 0 n_kv = self._n_kv_heads or self._n_heads