ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/backend/tests/test_llama_cpp_context_fit.py
Matt Van Horn 15d70a1d7b
 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>
2026-05-28 11:34:35 +04:00

579 lines
20 KiB
Python
Raw Permalink Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for the GGUF load-time context auto-fit decision.
Guards two regressions in ``LlamaCppBackend.load_model``:
1. **Auto mode on weights-exceed-VRAM** (``n_ctx == 0``): when the model
weights alone exceed 90% of every GPU subset's free memory, the
auto-pick loop used to exit without matching, leaving
``effective_ctx`` at the model's native context (e.g. 196608 for
MiniMax-M2.7). The intended default per Studio's UI spec is 4096 so
the slider lands on a usable value; the user can still drag higher
and trigger ``--fit on`` with a warning.
2. **Explicit ctx silently shrunk when KV overflows**: with fittable
weights but a requested ctx whose KV cache pushes total memory over
90% of VRAM, the old code binary-searched a smaller ctx and emitted
``-c <capped> -ngl -1`` without informing the caller. The UI had
already surfaced its "might be slower" warning and expects the user's
explicit ctx to be honored with ``--fit on`` flexing ``-ngl`` instead.
Tests avoid GPU probing, subprocess spawning, and GGUF I/O by driving the
post-metadata decision block directly against a stubbed instance.
Requires no GPU, network, or external libraries beyond pytest.
Cross-platform: Linux, macOS, Windows, WSL.
"""
from __future__ import annotations
import sys
import types as _types
from pathlib import Path
import pytest
# ---------------------------------------------------------------------------
# Stub heavy / unavailable external dependencies before importing the
# module under test. Same pattern as test_kv_cache_estimation.py.
# ---------------------------------------------------------------------------
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# loggers
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
# structlog
_structlog_stub = _types.ModuleType("structlog")
sys.modules.setdefault("structlog", _structlog_stub)
# httpx
_httpx_stub = _types.ModuleType("httpx")
for _exc_name in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
):
setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
class _FakeTimeout:
def __init__(self, *a, **kw):
pass
_httpx_stub.Timeout = _FakeTimeout
_httpx_stub.Client = type(
"Client",
(),
{
"__init__": lambda self, **kw: None,
"__enter__": lambda self: self,
"__exit__": lambda self, *a: None,
},
)
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
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
GIB = 1024**3
FALLBACK_CTX = 4096
def _make_backend(
native_ctx = 131072,
n_layers = 80,
n_kv_heads = 8,
n_heads = 64,
kv_key_length = 128,
kv_value_length = 128,
):
"""Create a LlamaCppBackend instance with GGUF metadata fields set and
the helpers used by the decision block stubbed out."""
inst = LlamaCppBackend.__new__(LlamaCppBackend)
inst._context_length = native_ctx
inst._n_layers = n_layers
inst._n_kv_heads = n_kv_heads
inst._n_heads = n_heads
inst._embedding_length = 8192
inst._kv_key_length = kv_key_length
inst._kv_value_length = kv_value_length
inst._kv_lora_rank = None
inst._sliding_window = None
inst._sliding_window_pattern = None
inst._ssm_inner_size = None
inst._full_attention_interval = None
inst._key_length_mla = None
inst._n_kv_heads_by_layer = None
inst._kv_key_length_swa = None
inst._kv_value_length_swa = None
return inst
def _drive(
n_ctx,
model_gib,
gpus,
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.
Mirrors the decision block at llama_cpp.py:1137-1296 so we can assert
the command that would be built, without subprocesses or GPU probes.
"""
inst = _make_backend(native_ctx = native_ctx)
model_size = int(model_gib * GIB)
cache_type_kv = None
def fake_estimate(n_ctx_, _type = None, **_kwargs):
return 0 if n_ctx_ <= 0 else n_ctx_ * kv_per_token_bytes
inst._estimate_kv_cache_bytes = fake_estimate
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 = requested_ctx if requested_ctx > 0 else (context_length or 0)
max_available_ctx = context_length or effective_ctx
if requested_ctx > 0:
effective_ctx = requested_ctx
elif context_length is not None:
effective_ctx = context_length
else:
effective_ctx = 0
original_ctx = effective_ctx
max_available_ctx = context_length or effective_ctx
gpu_indices, use_fit = None, True
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
if native_ctx_for_cap > 0:
ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True)
best_cap = 0
for n_gpus in range(1, len(ranked_for_cap) + 1):
subset = ranked_for_cap[:n_gpus]
pool_mib = sum(free for _, free in subset)
capped = inst._fit_context_to_vram(
native_ctx_for_cap,
pool_mib,
model_size,
cache_type_kv,
)
kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv)
total_mib = (model_size + kv) / (1024 * 1024)
if total_mib <= pool_mib * 0.90:
best_cap = max(best_cap, capped)
if best_cap > 0:
max_available_ctx = best_cap
if explicit_ctx:
requested_total = model_size + inst._estimate_kv_cache_bytes(
effective_ctx, cache_type_kv
)
gpu_indices, use_fit = inst._select_gpus(requested_total, gpus)
else:
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
matched = False
pin_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION
for n_gpus in range(1, len(ranked) + 1):
subset = ranked[:n_gpus]
pool_mib = sum(free for _, free in subset)
capped = inst._fit_context_to_vram(
effective_ctx,
pool_mib,
model_size,
cache_type_kv,
)
kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv)
total_mib = (model_size + kv) / (1024 * 1024)
if total_mib <= pool_mib * pin_fraction:
effective_ctx = capped
gpu_indices = sorted(idx for idx, _ in subset)
use_fit = False
matched = True
break
if not matched:
effective_ctx = min(FALLBACK_CTX, effective_ctx)
# Mirror llama_cpp.py: re-check fit at FALLBACK_CTX.
if effective_ctx > 0:
for n_gpus in range(1, len(ranked) + 1):
subset = ranked[:n_gpus]
pool_mib = sum(free for _, free in subset)
kv = inst._estimate_kv_cache_bytes(effective_ctx, cache_type_kv)
total_mib = (model_size + kv) / (1024 * 1024)
if total_mib <= pool_mib * pin_fraction:
gpu_indices = sorted(idx for idx, _ in subset)
use_fit = False
break
elif gpus:
gpu_indices, use_fit = inst._select_gpus(model_size, gpus)
if use_fit and not explicit_ctx:
effective_ctx = (
min(FALLBACK_CTX, effective_ctx) if effective_ctx > 0 else FALLBACK_CTX
)
return {
"c_arg": effective_ctx if effective_ctx > 0 else 0,
"use_fit": use_fit,
"gpu_indices": gpu_indices,
"max_available_ctx": max_available_ctx,
"original_ctx": original_ctx,
"ctx_override": ctx_override,
}
# ---------------------------------------------------------------------------
# Auto mode, model weights exceed VRAM (Bug A guard)
# ---------------------------------------------------------------------------
class TestAutoModeWeightsExceedVRAM:
"""``n_ctx == 0`` on a model whose weights don't fit anywhere."""
def test_minimax_like_single_gpu(self):
plan = _drive(
n_ctx = 0,
model_gib = 131,
gpus = [(0, 97_000)],
native_ctx = 196608,
)
assert plan["c_arg"] == FALLBACK_CTX
assert plan["use_fit"] is True
assert plan["gpu_indices"] is None
# UI slider ceiling stays at native: user can still drag higher
# and get the "might be slower" path.
assert plan["max_available_ctx"] == 196608
def test_multi_gpu_all_subsets_fail(self):
plan = _drive(
n_ctx = 0,
model_gib = 400,
gpus = [(0, 80_000), (1, 80_000), (2, 80_000), (3, 80_000)],
native_ctx = 131072,
)
assert plan["c_arg"] == FALLBACK_CTX
assert plan["use_fit"] is True
assert plan["gpu_indices"] is None
def test_no_kv_metadata_auto(self):
"""File-size-only fallback path also defaults to 4096."""
plan = _drive(
n_ctx = 0,
model_gib = 131,
gpus = [(0, 97_000)],
native_ctx = 196608,
can_estimate_kv = False,
)
assert plan["c_arg"] == FALLBACK_CTX
assert plan["use_fit"] is True
# ---------------------------------------------------------------------------
# Explicit ctx, KV overflows fittable weights (Bug B guard)
# ---------------------------------------------------------------------------
class TestExplicitCtxRespectsUser:
"""``n_ctx > 0`` must never be silently shrunk."""
def test_fittable_weights_oversized_kv(self):
# 8 GB weights + 131k ctx KV on 24 GB VRAM.
# Budget = 21.6 GB, KV at 131k >> 13.6 GB remaining, so
# _select_gpus flips use_fit=True.
plan = _drive(
n_ctx = 131072,
model_gib = 8,
gpus = [(0, 24_000)],
native_ctx = 131072,
)
assert plan["c_arg"] == 131072
assert plan["use_fit"] is True
assert plan["gpu_indices"] is None
def test_explicit_that_fits_uses_ngl(self):
plan = _drive(
n_ctx = 8192,
model_gib = 8,
gpus = [(0, 24_000)],
native_ctx = 131072,
)
assert plan["c_arg"] == 8192
assert plan["use_fit"] is False
assert plan["gpu_indices"] == [0]
def test_explicit_on_weights_exceed_vram(self):
# User drags the slider to 32k on a too-big model: honored.
plan = _drive(
n_ctx = 32768,
model_gib = 131,
gpus = [(0, 97_000)],
native_ctx = 196608,
)
assert plan["c_arg"] == 32768
assert plan["use_fit"] is True
def test_explicit_at_fallback_on_too_big(self):
plan = _drive(
n_ctx = FALLBACK_CTX,
model_gib = 131,
gpus = [(0, 97_000)],
native_ctx = 196608,
)
assert plan["c_arg"] == FALLBACK_CTX
assert plan["use_fit"] is True
def test_explicit_below_floor_honored(self):
# 2048 is below --fit-ctx default; still honored since user set it.
plan = _drive(
n_ctx = 2048,
model_gib = 8,
gpus = [(0, 24_000)],
)
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
# ---------------------------------------------------------------------------
class TestFittableAutoPickRegressions:
def test_small_model_one_gpu(self):
plan = _drive(
n_ctx = 0,
model_gib = 8,
gpus = [(0, 24_000)],
native_ctx = 131072,
kv_per_token_bytes = 8192,
)
assert plan["use_fit"] is False
assert plan["gpu_indices"] == [0]
assert plan["c_arg"] > FALLBACK_CTX
def test_medium_model_needs_multi_gpu(self):
plan = _drive(
n_ctx = 0,
model_gib = 60,
gpus = [(0, 40_000), (1, 40_000)],
native_ctx = 131072,
kv_per_token_bytes = 8192,
)
assert plan["use_fit"] is False
assert plan["gpu_indices"] == [0, 1]
def test_no_kv_metadata_fittable_auto(self):
plan = _drive(
n_ctx = 0,
model_gib = 8,
gpus = [(0, 24_000)],
native_ctx = 131072,
can_estimate_kv = False,
)
assert plan["use_fit"] is False
assert plan["gpu_indices"] == [0]
# ---------------------------------------------------------------------------
# #5106 regression: 91-95% utilization must still pin GPU.
# ---------------------------------------------------------------------------
class TestTightFitPinsToGPU:
"""Models that fit at 91-95% of free VRAM must use the GPU."""
def test_rtx_4090_qwen_24gb_class(self):
# noahterbest's #5106 log: 20.8 GB model on 22805 MiB free
# GPU, ctx=4096 -> ~94% utilization, ~1.4 GiB headroom.
plan = _drive(
n_ctx = 0,
model_gib = 20.8,
gpus = [(0, 22_805)],
native_ctx = 131072,
kv_per_token_bytes = 25_000,
)
assert plan["use_fit"] is False
assert plan["gpu_indices"] == [0]
def test_explicit_ctx_at_94_pct_pins_to_gpu(self):
# Explicit-ctx branch must agree with auto-ctx on headroom.
plan = _drive(
n_ctx = 4096,
model_gib = 20.8,
gpus = [(0, 22_805)],
native_ctx = 131072,
kv_per_token_bytes = 25_000,
)
assert plan["use_fit"] is False
assert plan["gpu_indices"] == [0]
def test_genuine_overflow_still_uses_fit(self):
# Beyond 95% must still defer to --fit on.
plan = _drive(
n_ctx = 4096,
model_gib = 23,
gpus = [(0, 22_000)],
native_ctx = 131072,
kv_per_token_bytes = 25_000,
)
assert plan["use_fit"] is True
assert plan["gpu_indices"] is None
# ---------------------------------------------------------------------------
# Platform-agnostic input shape
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("platform_tag", ["linux", "windows", "mac", "rocm"])
def test_identical_decision_across_platforms(platform_tag):
"""The decision function takes ``[(gpu_idx, free_mib), ...]`` regardless
of how upstream (nvidia-smi / nvidia-smi.exe / Metal / rocm-smi) produced
it. Identical inputs must yield identical plans."""
plan_a = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)])
plan_b = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)])
assert plan_a == plan_b, platform_tag
# ---------------------------------------------------------------------------
# _classify_gpu_offload: detect silent CPU fallback (#5106).
# ---------------------------------------------------------------------------
class TestClassifyGpuOffload:
def _backend(self, stdout_lines):
inst = LlamaCppBackend.__new__(LlamaCppBackend)
inst._stdout_lines = list(stdout_lines)
return inst
def test_cuda_buffer_present_returns_true(self):
inst = self._backend(
[
"load_tensors: offloaded 33/33 layers to GPU",
"load_tensors: CUDA0 model buffer size = 21000.0 MiB",
"load_tensors: CPU_Mapped model buffer size = 0.6 MiB",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
def test_cpu_only_buffer_returns_false(self):
# llama-server printed buffer lines but only CPU buffers --
# this is the silent CPU fallback symptom we want to catch.
inst = self._backend(
[
"load_tensors: CPU_Mapped model buffer size = 21000.0 MiB",
"load_tensors: CPU model buffer size = 0.6 MiB",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
def test_no_buffer_lines_returns_none(self):
# If we can't see buffer-allocation lines at all, don't guess.
inst = self._backend(
[
"INFO [main] starting server",
"load_tensors: file format = GGUF V3",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is None
def test_no_gpus_detected_returns_none(self):
# CPU-only systems are valid; suppress the warning entirely.
inst = self._backend(
[
"load_tensors: CPU_Mapped model buffer size = 21000.0 MiB",
]
)
assert inst._classify_gpu_offload(False, []) is None
def test_user_did_not_intend_gpu_returns_none(self):
# Studio called start_llama_server without expecting GPU use;
# don't warn.
inst = self._backend(
[
"load_tensors: CPU_Mapped model buffer size = 21000.0 MiB",
]
)
assert inst._classify_gpu_offload(False, [(0, 22805)]) is None
def test_rocm_buffer_marker_returns_true(self):
inst = self._backend(
[
"load_tensors: ROCm0 model buffer size = 21000.0 MiB",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
def test_metal_buffer_marker_returns_true(self):
inst = self._backend(
[
"load_tensors: Metal model buffer size = 8000.0 MiB",
]
)
assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
Reference in a new issue View git blame Copy permalink