From 796f8497e79843aec1afdd3f5190d5ec15177201 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 20 Jul 2026 06:47:08 -0700 Subject: [PATCH] Studio (Windows): keep prompt caching on full GPU offload (#7260) * Studio (Windows): keep prompt caching on full GPU offload (#5692 follow-up) The #5692 full-offload tuning also added --no-cache-prompt, which disables in-VRAM prompt-prefix reuse. That is unrelated to the host-RAM KV checkpoints #5692 fixed (--cache-ram 0 / --ctx-checkpoints 0): a fully offloaded model keeps its KV cache in VRAM, so reusing a common prefix does not copy to system RAM and does not cause the PCI-E overhead. --no-cache-prompt only forces every request to re-prefill the whole prompt, which is small for short chats but severe for large stable system prompts reused across calls (coding agents, long multi-turn chats). Remove --no-cache-prompt; keep the checkpoint disables and the thread/OMP tuning. _prompt_cache_disabled stays False (its default), so slot save/restore is intact. Verified on a fully offloaded gemma GGUF: an identical repeated prompt reprefills 1 token instead of 2220. * Guard against re-adding --no-cache-prompt to any llama-server command Add a backend-wide test that AST-scans studio/backend and fails if --no-cache-prompt is appended/extended/+= into a command. This locks in the #7260 fix across every code path, not just load_model. Detecting the flag or honouring a user-supplied one stays allowed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 10 +--- .../tests/test_llama_cpp_mtp_detection.py | 55 ++++++++++++++++++- 2 files changed, 57 insertions(+), 8 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index dee507fae2..2c7433f7a4 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -7563,8 +7563,9 @@ class LlamaCppBackend: else: self._api_key = None - # Windows + full offload: disable KV checkpoints (WDDM/PCI-E - # overhead). CPU/partial offload keeps prompt caching. #5692. + # Windows + full offload: drop the host-RAM KV checkpoints that cause + # WDDM/PCI-E overhead, but keep prompt caching (in-VRAM prefix reuse) so + # a repeated prompt is not re-prefilled on every request. #5692. if sys.platform == "win32" and full_offload_tuning_active: unsupported_cache_flags: list[str] = [] if server_caps.get("supports_cache_ram"): @@ -7575,11 +7576,6 @@ class LlamaCppBackend: cmd.extend(["--ctx-checkpoints", "0"]) else: unsupported_cache_flags.append("--ctx-checkpoints") - if server_caps.get("supports_no_cache_prompt"): - cmd.append("--no-cache-prompt") - self._prompt_cache_disabled = True - else: - unsupported_cache_flags.append("--no-cache-prompt") if unsupported_cache_flags: logger.info( "Skipping unsupported Windows cache flags for llama-server: %s", diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 68b706ebf9..1d15647967 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -9,6 +9,7 @@ the _already_in_target_state mirror that prevents needless reloads. from __future__ import annotations +import ast import inspect import os import struct @@ -345,10 +346,62 @@ def test_windows_full_offload_flags_use_current_llama_server_args(): stale_checkpoint_flag = "--checkpoint-" + "every-n-tokens" assert '"--cache-ram"' in src assert '"--ctx-checkpoints"' in src - assert '"--no-cache-prompt"' in src + # Prompt caching stays on (in-VRAM prefix reuse); #5692 only needed the host-RAM + # checkpoints (--cache-ram / --ctx-checkpoints) disabled, not prompt reuse. + assert '"--no-cache-prompt"' not in src assert stale_checkpoint_flag not in src +# Backend-wide guard: Unsloth must never inject --no-cache-prompt into a llama-server +# command. It disables in-VRAM prompt-prefix reuse, re-prefilling every repeated prompt +# (#5692 only needed --cache-ram / --ctx-checkpoints off; #7260 dropped the stray flag). +# Detecting it (_is_real) or honouring a user-supplied one (_prompt_cache_off) is fine. +_NO_CACHE_PROMPT_FLAG = "--no-cache-prompt" +_LIST_MUTATORS = frozenset({"append", "extend", "insert"}) + + +def _has_flag_literal(node: ast.AST) -> bool: + return any( + isinstance(n, ast.Constant) and n.value == _NO_CACHE_PROMPT_FLAG for n in ast.walk(node) + ) + + +def _no_cache_prompt_injections(source: str, filename: str) -> list[tuple[str, int]]: + """(file, lineno) for each spot adding --no-cache-prompt to a list.""" + hits: list[tuple[str, int]] = [] + for node in ast.walk(ast.parse(source, filename = filename)): + # cmd.append/extend/insert(... flag ...) or cmd += [... flag ...] + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in _LIST_MUTATORS + and any(_has_flag_literal(a) for a in node.args) + ) or ( + isinstance(node, ast.AugAssign) + and isinstance(node.op, ast.Add) + and _has_flag_literal(node.value) + ): + hits.append((filename, node.lineno)) + return hits + + +def test_unsloth_never_injects_no_cache_prompt_into_any_command(): + root = Path(_BACKEND_DIR) + files = [p for p in root.rglob("*.py") if "tests" not in p.relative_to(root).parts] + violations: list[tuple[str, int]] = [] + for path in files: + try: + violations += _no_cache_prompt_injections(path.read_text(encoding = "utf-8"), str(path)) + except (OSError, UnicodeDecodeError, SyntaxError): + continue + assert files, "no backend source files were scanned" + assert violations == [], ( + "Unsloth must never add --no-cache-prompt to a llama-server command " + "(it disables prompt-prefix reuse); detecting or honouring a user-supplied " + f"one is fine. Offending sites: {violations}" + ) + + def test_load_model_sets_threads_once(): src = inspect.getsource(LlamaCppBackend.load_model) assert src.count('cmd.extend(["--threads", str(') == 1