From 15d70a1d7bd2557fa42672a9e5e2672f71a9c18c Mon Sep 17 00:00:00 2001 From: Matt Van Horn Date: Thu, 28 May 2026 00:34:35 -0700 Subject: [PATCH 1/2] fix: honor --ctx-size and other forwarded args from `unsloth studio run` in Studio's context-fit logic (#5815) * fix: honor --ctx-size and other forwarded args from `unsloth studio run` in Studio's context-fit logic * refactor: extract resolve_requested_ctx as single source of truth The test helper was reimplementing the two-line 'ctx_override = parse_ctx_override(...); requested_ctx = ctx_override if ctx_override is not None else n_ctx' pattern locally, so the test asserted against its own reimplementation rather than production logic. Extract the conditional into resolve_requested_ctx and have both the production caller and the test use it. * fix(studio): honor pass-through cache type flags in KV VRAM estimate Studio's KV cache VRAM estimate computed from the first-class cache_type_kv even when the user passed -ctk/--cache-type-k/-ctv/ --cache-type-v via extras. Those flags reached llama-server fine (last-wins on the CLI) but the pre-launch estimate kept using the default f16 bytes-per-element, so GPU placement decisions could be off when the user lowered cache precision via pass-through. Adds parse_cache_override + resolve_cache_type_kv in llama_server_args.py (mirroring parse_ctx_override / resolve_requested_ctx), wires both into load_model alongside the existing ctx resolution, and adds focused unit tests for the parser + resolver. Follow-up to @rolandtannous review on #5815. --------- Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 34 +++++- .../core/inference/llama_server_args.py | 113 ++++++++++++++++++ .../tests/test_llama_cpp_context_fit.py | 58 ++++++++- .../backend/tests/test_llama_server_args.py | 83 +++++++++++++ 4 files changed, 278 insertions(+), 10 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 0325620b2d..2d95112d6d 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -28,6 +28,12 @@ from urllib.parse import urlparse import httpx +from core.inference.llama_server_args import ( + parse_cache_override, + parse_ctx_override, + resolve_cache_type_kv, + resolve_requested_ctx, +) from core.tool_healing import ( _TC_END_TAG_RE, _TC_FUNC_CLOSE_RE, @@ -2724,7 +2730,23 @@ class LlamaCppBackend: # Select GPU(s) based on model size + estimated KV cache. # Seed safe defaults before GPU probing so the except path # still has valid state to publish. - effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0) + ctx_override = parse_ctx_override(extra_args) + requested_ctx = resolve_requested_ctx(extra_args, n_ctx) + cache_override = parse_cache_override(extra_args) + cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv) + if ctx_override is not None and ctx_override > 0: + logger.info( + f"User --ctx-size {ctx_override} honored; " + "skipping auto-reduce" + ) + if cache_override is not None: + logger.info( + f"User --cache-type-k/-v {cache_override} " + "honored for KV estimate" + ) + effective_ctx = ( + requested_ctx if requested_ctx > 0 else (self._context_length or 0) + ) max_available_ctx = self._context_length or effective_ctx gpus: list[tuple[int, int]] = [] try: @@ -2734,8 +2756,8 @@ class LlamaCppBackend: # Resolve effective context: 0 means let llama-server use the # model's native length. Only expand to a known native length # if metadata is available; otherwise preserve 0 as a sentinel. - if n_ctx > 0: - effective_ctx = n_ctx + if requested_ctx > 0: + effective_ctx = requested_ctx elif self._context_length is not None: effective_ctx = self._context_length else: @@ -2788,7 +2810,7 @@ class LlamaCppBackend: # since multi-GPU is slower and the user didn't ask for a # specific context length. gpu_indices, use_fit = None, True - explicit_ctx = n_ctx > 0 + explicit_ctx = requested_ctx > 0 if gpus and self._can_estimate_kv() and effective_ctx > 0: # Compute the largest hardware-aware cap from the model's @@ -2845,7 +2867,7 @@ class LlamaCppBackend: gpu_indices, use_fit = self._select_gpus( requested_total, gpus ) - # No silent shrink: effective_ctx stays == n_ctx. + # No silent shrink: effective_ctx stays == requested_ctx. else: # Auto context: prefer fewer GPUs, cap context # to fit. Same headroom threshold as @@ -2934,7 +2956,7 @@ class LlamaCppBackend: except Exception as e: logger.warning(f"GPU selection failed ({e}), using --fit on") gpu_indices, use_fit = None, True - effective_ctx = n_ctx # fall back to original + effective_ctx = requested_ctx # fall back to original launch_mmproj_path = self._resolve_launch_mmproj_path( model_path = model_path, diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 4f528a689f..b299e1ee9e 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -110,6 +110,8 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]: f"and cannot be passed as an extra arg" ) out.append(token) + parse_ctx_override(out) + parse_cache_override(out) return out @@ -166,6 +168,117 @@ _BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset( ) +def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]: + """Return the last user-supplied ``-c`` / ``--ctx-size`` value. + + Mirrors llama.cpp's last-wins flag parsing for the one pass-through + numeric knob Studio's load-time fit logic needs to see. + """ + if not args: + return None + + tokens = [str(a) for a in args] + override: Optional[int] = None + i, n = 0, len(tokens) + while i < n: + tok = tokens[i] + flag = _flag_name(tok) + if flag is None or flag not in _CONTEXT_FLAGS: + i += 1 + continue + + if "=" in tok: + raw_value = tok.split("=", 1)[1] + i += 1 + else: + if i + 1 >= n or _flag_name(tokens[i + 1]) is not None: + raise ValueError( + f"llama-server flag '{flag}' requires an integer value" + ) + raw_value = tokens[i + 1] + i += 2 + + try: + value = int(str(raw_value).strip()) + except ValueError as exc: + raise ValueError( + f"llama-server flag '{flag}' requires an integer value" + ) from exc + if value < 0: + raise ValueError( + f"llama-server flag '{flag}' requires a non-negative integer value" + ) + override = value + + return override + + +def resolve_requested_ctx( + args: Optional[Iterable[str]], + fallback_n_ctx: int, +) -> int: + """Return the context size load_model should treat as requested. + + Single source of truth for the two-line ``ctx_override = parse_ctx_override(...); + requested_ctx = ctx_override if ctx_override is not None else n_ctx`` pattern + used by ``load_model`` so tests don't have to reimplement the conditional + locally and then assert against their own reimplementation. + """ + override = parse_ctx_override(args) + return override if override is not None else fallback_n_ctx + + +def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: + """Return the last-wins cache type if extras pass cache flags. + + Mirrors parse_ctx_override but for cache type. Recognises both -ctk + (key) and -ctv (value). When both flags appear, returns the last-wins + value, treating key and value cache flags as the same setting because + Studio's KV estimate has a single cache_type_kv knob. + """ + if not args: + return None + + tokens = [str(a) for a in args] + override: Optional[str] = None + i, n = 0, len(tokens) + while i < n: + tok = tokens[i] + flag = _flag_name(tok) + if flag is None or flag not in _CACHE_FLAGS: + i += 1 + continue + + if "=" in tok: + raw_value = tok.split("=", 1)[1] + i += 1 + else: + if i + 1 >= n or _flag_name(tokens[i + 1]) is not None: + raise ValueError(f"llama-server flag '{flag}' requires a value") + raw_value = tokens[i + 1] + i += 2 + + value = str(raw_value).strip() + if not value: + raise ValueError(f"llama-server flag '{flag}' requires a non-empty value") + override = value + + return override + + +def resolve_cache_type_kv( + args: Optional[Iterable[str]], + fallback_cache_type_kv: Optional[str], +) -> Optional[str]: + """Return the cache type load_model should treat as requested. + + Single source of truth for the cache override conditional used by + ``load_model``. + """ + override = parse_cache_override(args) + return override if override is not None else fallback_cache_type_kv + + def strip_shadowing_flags( args: Iterable[str], *, diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py index 1ea76edd15..6fe5372147 100644 --- a/studio/backend/tests/test_llama_cpp_context_fit.py +++ b/studio/backend/tests/test_llama_cpp_context_fit.py @@ -84,6 +84,7 @@ _httpx_stub.Client = type( sys.modules.setdefault("httpx", _httpx_stub) from core.inference.llama_cpp import LlamaCppBackend +from core.inference.llama_server_args import parse_ctx_override, resolve_requested_ctx # --------------------------------------------------------------------------- @@ -131,6 +132,7 @@ def _drive( native_ctx = 131072, kv_per_token_bytes = 325_000, can_estimate_kv = True, + extra_args = None, ): """Drive the post-metadata portion of load_model with stubbed inputs. @@ -148,11 +150,16 @@ def _drive( inst._can_estimate_kv = lambda: can_estimate_kv context_length = inst._context_length + # Use the production helper instead of reimplementing the conditional + # locally; reimplementing makes the test pass for the test's own logic + # rather than production's, and silent drift won't be caught. + ctx_override = parse_ctx_override(extra_args) + requested_ctx = resolve_requested_ctx(extra_args, n_ctx) - effective_ctx = n_ctx if n_ctx > 0 else (context_length or 0) + effective_ctx = requested_ctx if requested_ctx > 0 else (context_length or 0) max_available_ctx = context_length or effective_ctx - if n_ctx > 0: - effective_ctx = n_ctx + if requested_ctx > 0: + effective_ctx = requested_ctx elif context_length is not None: effective_ctx = context_length else: @@ -161,7 +168,7 @@ def _drive( max_available_ctx = context_length or effective_ctx gpu_indices, use_fit = None, True - explicit_ctx = n_ctx > 0 + explicit_ctx = requested_ctx > 0 if gpus and inst._can_estimate_kv() and effective_ctx > 0: native_ctx_for_cap = context_length or effective_ctx @@ -236,6 +243,7 @@ def _drive( "gpu_indices": gpu_indices, "max_available_ctx": max_available_ctx, "original_ctx": original_ctx, + "ctx_override": ctx_override, } @@ -349,6 +357,48 @@ class TestExplicitCtxRespectsUser: assert plan["c_arg"] == 2048 +# --------------------------------------------------------------------------- +# Pass-through --ctx-size participates in context fit (#5676). +# --------------------------------------------------------------------------- + + +class TestExtraArgsCtxOverride: + def test_ctx_size_extra_honored_over_auto(self): + plan = _drive( + n_ctx = 0, + model_gib = 131, + gpus = [(0, 97_000)], + native_ctx = 196608, + extra_args = ["--ctx-size", "128000"], + ) + assert plan["ctx_override"] == 128000 + assert plan["original_ctx"] == 128000 + assert plan["c_arg"] == 128000 + assert plan["use_fit"] is True + + def test_ctx_size_short_alias_honored_over_auto(self): + plan = _drive( + n_ctx = 0, + model_gib = 131, + gpus = [(0, 97_000)], + native_ctx = 196608, + extra_args = ["-c", "128000"], + ) + assert plan["c_arg"] == 128000 + assert plan["use_fit"] is True + + def test_ctx_size_extra_wins_over_first_class_field(self): + plan = _drive( + n_ctx = 4096, + model_gib = 8, + gpus = [(0, 24_000)], + native_ctx = 131072, + extra_args = ["--ctx-size", "128000"], + ) + assert plan["original_ctx"] == 128000 + assert plan["c_arg"] == 128000 + + # --------------------------------------------------------------------------- # Non-regression: fittable + auto still auto-picks largest fitting ctx # --------------------------------------------------------------------------- diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index 02a272ba3e..2f7431497d 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -30,6 +30,9 @@ _spec = importlib.util.spec_from_file_location("_lsa_test_only", _LSA_PATH) _lsa = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(_lsa) is_managed_flag = _lsa.is_managed_flag +parse_cache_override = _lsa.parse_cache_override +parse_ctx_override = _lsa.parse_ctx_override +resolve_cache_type_kv = _lsa.resolve_cache_type_kv strip_shadowing_flags = _lsa.strip_shadowing_flags validate_extra_args = _lsa.validate_extra_args @@ -410,6 +413,86 @@ def test_is_managed_flag_false_for_mtp_pass_through(): assert is_managed_flag("--spec-ngram-mod-n-max") is False +# ── parse_ctx_override ─────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "args,expected", + [ + (None, None), + ([], None), + (["--top-k", "20"], None), + (["--ctx-size", "128000"], 128000), + (["--ctx-size=128000"], 128000), + (["-c", "128000"], 128000), + (["-c=128000"], 128000), + (["-c", "4096", "--ctx-size", "128000"], 128000), + ], +) +def test_parse_ctx_override(args, expected): + assert parse_ctx_override(args) == expected + + +@pytest.mark.parametrize( + "args", + [ + ["--ctx-size"], + ["--ctx-size", "--top-k"], + ["--ctx-size", "abc"], + ["--ctx-size=abc"], + ["-c", "-1"], + ], +) +def test_parse_ctx_override_rejects_malformed_values(args): + with pytest.raises(ValueError, match = "ctx-size|'-c'"): + parse_ctx_override(args) + + +def test_validate_extra_args_rejects_malformed_ctx_override(): + with pytest.raises(ValueError, match = "ctx-size"): + validate_extra_args(["--ctx-size", "abc"]) + + +# ── parse_cache_override ───────────────────────────────────────────── + + +@pytest.mark.parametrize( + "args,expected", + [ + (None, None), + ([], None), + (["--top-k", "20"], None), + (["--cache-type-k", "q8_0"], "q8_0"), + (["-ctk", "q4_0"], "q4_0"), + (["-ctv", "q4_0"], "q4_0"), + (["--cache-type-k=q4_0"], "q4_0"), + (["-ctk", "f16", "-ctk", "q8_0"], "q8_0"), + ], +) +def test_parse_cache_override(args, expected): + assert parse_cache_override(args) == expected + + +@pytest.mark.parametrize( + "args", + [ + ["-ctk"], + ["-ctk", "-c", "4096"], + ], +) +def test_parse_cache_override_rejects_malformed_values(args): + with pytest.raises(ValueError, match = "cache-type|'-ctk'"): + parse_cache_override(args) + + +def test_resolve_cache_type_kv_uses_override_when_present(): + assert resolve_cache_type_kv(["--cache-type-k", "q8_0"], "f16") == "q8_0" + + +def test_resolve_cache_type_kv_uses_fallback_without_override(): + assert resolve_cache_type_kv(["--top-k", "20"], "f16") == "f16" + + def test_strip_shadowing_flags_boolean_does_not_consume_next_token(): # `--spec-default` is boolean; drop just the flag, keep the next token. out = strip_shadowing_flags(["--spec-default", "ngram-mod"], strip_spec = True) From 185ff00c62853c376d1a7dbeac98d0e60517b2fc Mon Sep 17 00:00:00 2001 From: alkinun Date: Thu, 28 May 2026 12:28:52 +0300 Subject: [PATCH 2/2] Fix non-streaming GGUF chat completion usage (#5781) * Fix GGUF non-stream chat completion usage * Handle nullable GGUF completion usage --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> --- studio/backend/routes/inference.py | 14 +++- .../tests/test_gguf_completion_usage.py | 78 +++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) create mode 100644 studio/backend/tests/test_gguf_completion_usage.py diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 259337616c..7d1c7b2488 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -3091,9 +3091,12 @@ async def openai_chat_completions( else: try: full_text = "" + completion_usage = None for token in gguf_generate(): if isinstance(token, dict): - continue # skip metadata dict in non-streaming path + if token.get("type") == "metadata": + completion_usage = token.get("usage") + continue full_text = token response = ChatCompletion( @@ -3106,6 +3109,15 @@ async def openai_chat_completions( finish_reason = "stop", ) ], + usage = CompletionUsage( + prompt_tokens = (completion_usage or {}).get("prompt_tokens") + or 0, + completion_tokens = (completion_usage or {}).get( + "completion_tokens" + ) + or 0, + total_tokens = (completion_usage or {}).get("total_tokens") or 0, + ), ) return JSONResponse(content = response.model_dump()) diff --git a/studio/backend/tests/test_gguf_completion_usage.py b/studio/backend/tests/test_gguf_completion_usage.py new file mode 100644 index 0000000000..b8cfaee7c7 --- /dev/null +++ b/studio/backend/tests/test_gguf_completion_usage.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Regression tests for GGUF non-streaming chat completion usage.""" + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from auth.authentication import get_current_subject +import routes.inference as inference_route + + +class _GgufBackend: + is_loaded = True + model_identifier = "test/model.gguf" + _is_audio = False + is_vision = False + supports_tools = False + + def __init__(self, usage): + self.usage = usage + + def generate_chat_completion(self, **kwargs): + yield "answer" + yield { + "type": "metadata", + "usage": self.usage, + "timings": {"prompt_n": 23, "predicted_n": 1283}, + } + + +def _request_completion(monkeypatch, usage): + monkeypatch.setattr( + inference_route, "get_llama_cpp_backend", lambda: _GgufBackend(usage) + ) + monkeypatch.setattr( + inference_route, "_effective_enable_tools", lambda payload: False + ) + + app = FastAPI() + app.include_router(inference_route.router) + app.dependency_overrides[get_current_subject] = lambda: "test-user" + + return TestClient(app).post( + "/chat/completions", + json = { + "messages": [{"role": "user", "content": "Why is the sky blue?"}], + "stream": False, + }, + ) + + +def test_non_streaming_gguf_completion_includes_generated_usage(monkeypatch): + response = _request_completion( + monkeypatch, + {"prompt_tokens": 23, "completion_tokens": 1283, "total_tokens": 1306}, + ) + + assert response.status_code == 200 + assert response.json()["usage"] == { + "prompt_tokens": 23, + "completion_tokens": 1283, + "total_tokens": 1306, + } + + +def test_non_streaming_gguf_completion_defaults_nullable_usage_to_zero(monkeypatch): + response = _request_completion( + monkeypatch, + {"prompt_tokens": None, "completion_tokens": 1283, "total_tokens": None}, + ) + + assert response.status_code == 200 + assert response.json()["usage"] == { + "prompt_tokens": 0, + "completion_tokens": 1283, + "total_tokens": 0, + }