# 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 (``n_ctx == 0``) when weights exceed every GPU subset's free memory: auto-pick should fall back to 4096 (a usable slider value) rather than leaving native ctx. User can still drag higher onto ``--fit on``. 2. Explicit ctx must never be silently shrunk: when KV overflows fittable weights, honor the explicit ctx with ``--fit on`` flexing ``-ngl``. Drives the post-metadata decision block against a stubbed instance: no GPU, network, subprocess, or GGUF I/O. Cross-platform. """ from __future__ import annotations import sys import types as _types from pathlib import Path import pytest # --------------------------------------------------------------------------- # Stub heavy/unavailable deps before importing the module under test. # --------------------------------------------------------------------------- _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) _loggers_stub = _types.ModuleType("loggers") _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) _structlog_stub = _types.ModuleType("structlog") sys.modules.setdefault("structlog", _structlog_stub) _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 ( _APPLE_UNIFIED_MEMORY_FRACTION, _CTX_FIT_VRAM_FRACTION, LlamaCppBackend, classify_gpu_offload_lines, ) 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, ): """LlamaCppBackend with GGUF metadata set and decision helpers stubbed.""" 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, apple_budget_mib = 0, flat_mtp_reserve = 0.0, ): """Drive the post-metadata portion of load_model with stubbed inputs. Mirrors llama_cpp.py:1137-1296 to assert the built command, 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, not a reimplementation, to avoid testing our own logic. 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 * _CTX_FIT_VRAM_FRACTION: 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 elif apple_budget_mib > 0 and effective_ctx > 0: # Mirrors the Apple unified-memory branch in load_model: flat MTP reserve # off the budget up front (no-op at 0), sparse-KV floors to FALLBACK_CTX, # only auto context shrinks. native_ctx_for_cap = context_length or effective_ctx apple_fit_budget_mib = int(apple_budget_mib * max(0.0, 1.0 - flat_mtp_reserve)) if inst._can_estimate_kv(): cap = inst._fit_context_to_vram( native_ctx_for_cap, apple_fit_budget_mib, model_size, cache_type_kv, budget_frac = 1.0, ) cap_footprint_mib = (model_size + inst._estimate_kv_cache_bytes(cap, cache_type_kv)) / ( 1024 * 1024 ) max_available_ctx = ( cap if cap_footprint_mib <= apple_fit_budget_mib else min(FALLBACK_CTX, native_ctx_for_cap) ) else: max_available_ctx = min(FALLBACK_CTX, native_ctx_for_cap) if not explicit_ctx: effective_ctx = max_available_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 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; 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): """Decision takes ``[(gpu_idx, free_mib), ...]`` regardless of source; 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): # Buffer lines printed but only CPU buffers -- 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; 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 def test_offloaded_zero_count_returns_false(self): # Authoritative count overrides any GPU-looking buffer line. inst = self._backend( [ "load_tensors: offloaded 0/33 layers to GPU", "load_tensors: CUDA0 model buffer size = 21000.0 MiB", ] ) assert inst._classify_gpu_offload(True, [(0, 22805)]) is False def test_offloaded_draft_then_main_returns_true(self): # A small draft model (0/2) does not mask the main model (33/33). inst = self._backend( [ "load_tensors: offloaded 0/2 layers to GPU", "load_tensors: offloaded 33/33 layers to GPU", ] ) assert inst._classify_gpu_offload(True, [(0, 22805)]) is True def test_main_on_cpu_with_draft_on_gpu_returns_false(self): # MTP: the small drafter fits on GPU (1/1) but the main model is on CPU # (0/33). Decide on the largest model, so the warning still fires. inst = self._backend( [ "load_tensors: offloaded 0/33 layers to GPU", "load_tensors: offloaded 1/1 layers to GPU", ] ) assert inst._classify_gpu_offload(True, [(0, 22805)]) is False def test_main_on_gpu_with_draft_on_cpu_returns_true(self): # Reverse: main model on GPU (33/33), drafter on CPU (0/1) -> no warning. inst = self._backend( [ "load_tensors: offloaded 33/33 layers to GPU", "load_tensors: offloaded 0/1 layers to GPU", ] ) assert inst._classify_gpu_offload(True, [(0, 22805)]) is True def test_cuda_host_buffer_excluded_returns_false(self): # CUDA_Host is CPU-pinned memory, not a model offload. inst = self._backend( [ "load_tensors: CUDA_Host model buffer size = 500.0 MiB", "load_tensors: CPU model buffer size = 21000.0 MiB", ] ) assert inst._classify_gpu_offload(True, [(0, 22805)]) is False def test_device_info_gpu_row_alone_is_inconclusive(self): # device_info lists available devices, not where the model loaded, so a # GPU row alone is not proof of offload. inst = self._backend( [ "print_info: device_info:", " - CUDA0 : 24564 MiB free", ] ) assert inst._classify_gpu_offload(True, [(0, 22805)]) is None def test_cpu_buffers_with_gpu_device_row_returns_false(self): # Definite CPU-only buffers must win over a GPU device-inventory row. inst = self._backend( [ "load_tensors: CPU model buffer size = 21000.0 MiB", "print_info: device_info:", " - CUDA0 : 24564 MiB free", ] ) assert inst._classify_gpu_offload(True, [(0, 22805)]) is False def test_device_info_cpu_only_returns_false(self): inst = self._backend( [ "print_info: device_info:", " - CPU : 64000 MiB free", ] ) assert inst._classify_gpu_offload(True, [(0, 22805)]) is False def test_system_info_cuda_before_device_info_does_not_count(self): # A compiled-in backend named in system_info is not proof of offload; # only the device_info table (here CPU only) decides. inst = self._backend( [ "system_info: CUDA : ARCHS = 890 | n_threads = 8", "print_info: device_info:", " - CPU : 64000 MiB free", ] ) assert inst._classify_gpu_offload(True, [(0, 22805)]) is False @pytest.mark.parametrize( "marker", ["CUDA0", "ROCm0", "HIP0", "Metal", "Vulkan0", "OpenCL0", "SYCL0", "MUSA0", "CANN0"], ) def test_all_gpu_buffer_markers_return_true(self, marker): assert ( classify_gpu_offload_lines([f"load_tensors: {marker} model buffer size = 8000.0 MiB"]) is True ) def test_module_level_no_signal_returns_none(self): assert classify_gpu_offload_lines(["INFO starting server"]) is None def test_select_gpus_ranks_by_usable_not_raw_free(): # 80 GB card (30 GB free -> 25.9 GB usable) vs 32 GB card (29 GB free -> 27.4 # GB usable). A 27 GB model fits the 32 GB card alone; raw-free ranking would # try the 80 GB card first and split across both. Usable ranking picks [1]. gpus = [(0, 30000), (1, 29000)] totals = {0: 81920, 1: 32607} model = int(27000 * 1024 * 1024) idxs, use_fit = LlamaCppBackend._select_gpus(model, gpus, total_by_idx = totals) assert idxs == [1] and use_fit is False def test_select_gpus_reserves_per_device_overhead(): # Two 16 GB cards, ~15181 MiB usable each at 0.95 -> 30362 MiB pooled. A 30000 # MiB model fits the pool with no per-device overhead, but a layer split also # pays ~1 GiB/extra-GPU; that pushes the 2-GPU need to 31024 MiB > pool, so a # pin would OOM -> must fall back to --fit. Single-GPU fits add no overhead # (Finding F1, the explicit/file-size multi-GPU pin gap). gpus = [(0, 16000), (1, 16000)] totals = {0: 16384, 1: 16384} gib = 1024 * 1024 * 1024 model = int(30000 * 1024 * 1024) idxs, use_fit = LlamaCppBackend._select_gpus(model, gpus, total_by_idx = totals) assert idxs == [0, 1] and use_fit is False # fits 2 GPUs without overhead idxs2, use_fit2 = LlamaCppBackend._select_gpus( model, gpus, total_by_idx = totals, per_device_overhead_bytes = gib ) assert idxs2 is None and use_fit2 is True # overhead tips it past the pool # A single-GPU fit is unchanged by the overhead (k=1 adds nothing). small = int(15000 * 1024 * 1024) a, _ = LlamaCppBackend._select_gpus(small, gpus, total_by_idx = totals) b, _ = LlamaCppBackend._select_gpus( small, gpus, total_by_idx = totals, per_device_overhead_bytes = gib ) assert a == [0] and b == [0] # --------------------------------------------------------------------------- # Apple Silicon unified-memory context cap (#5118, #6529): no discrete GPU on # Metal, so the auto context defaulted to native and over-committed unified # memory. The fix budgets and caps the auto context (explicit stays verbatim). # --------------------------------------------------------------------------- def _force_apple(monkeypatch): import platform as _platform monkeypatch.setattr(_platform, "system", lambda: "Darwin") monkeypatch.setattr(_platform, "machine", lambda: "arm64") def _install_fake_mlx(monkeypatch, working_set_bytes): """Minimal mlx.core stub exposing metal.is_available() and device_info().""" mlx = _types.ModuleType("mlx") mlx_core = _types.ModuleType("mlx.core") mlx_core.metal = _types.SimpleNamespace(is_available = lambda: True) mlx_core.device_info = lambda: {"max_recommended_working_set_size": working_set_bytes} mlx.core = mlx_core monkeypatch.setitem(sys.modules, "mlx", mlx) monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) class TestAppleUnifiedMemoryBudget: def test_zero_off_apple_silicon(self, monkeypatch): import platform as _platform monkeypatch.setattr(_platform, "system", lambda: "Linux") monkeypatch.setattr(_platform, "machine", lambda: "x86_64") assert LlamaCppBackend._apple_metal_memory_budget_bytes() == 0 def test_uses_metal_working_set(self, monkeypatch): _force_apple(monkeypatch) ws = 27 * GIB # ~recommended working set on a 36 GB Mac _install_fake_mlx(monkeypatch, ws) assert LlamaCppBackend._apple_metal_memory_budget_bytes() == int( ws * _APPLE_UNIFIED_MEMORY_FRACTION ) def test_falls_back_to_total_ram_without_mlx(self, monkeypatch): _force_apple(monkeypatch) monkeypatch.setitem(sys.modules, "mlx", None) # import mlx.core -> ImportError fake_psutil = _types.ModuleType("psutil") fake_psutil.virtual_memory = lambda: _types.SimpleNamespace(total = 36 * GIB) monkeypatch.setitem(sys.modules, "psutil", fake_psutil) assert LlamaCppBackend._apple_metal_memory_budget_bytes() == int( 36 * GIB * _APPLE_UNIFIED_MEMORY_FRACTION ) def test_zero_when_no_budget_resolvable(self, monkeypatch): _force_apple(monkeypatch) monkeypatch.setitem(sys.modules, "mlx", None) monkeypatch.setitem(sys.modules, "psutil", None) assert LlamaCppBackend._apple_metal_memory_budget_bytes() == 0 class TestAppleContextCap: """The real ``_fit_context_to_vram`` against the reporter's M3 Pro case.""" def test_caps_native_context_into_unified_budget(self): # ~15.7 GB weights at native 262144 (~16 GB KV) -> ~32 GB on a 36 GB M3 # Pro (~23 GB budget); the fit must reduce the context to fit. inst = _make_backend(native_ctx = 262144) inst._can_estimate_kv = lambda: True inst._estimate_kv_cache_bytes = ( lambda n, *a, **k: 0 if n <= 0 else int(n * 64_000) # ~16 GB @ 262144 ) model_size_fit = int(15.7 * GIB) budget_mib = int(27 * GIB * _APPLE_UNIFIED_MEMORY_FRACTION) // (1024 * 1024) # The native footprint over-commits the budget -- this is the bug. native_footprint_mib = (model_size_fit + inst._estimate_kv_cache_bytes(262144)) // ( 1024 * 1024 ) assert native_footprint_mib > budget_mib capped = inst._fit_context_to_vram( 262144, budget_mib, model_size_fit, None, budget_frac = 1.0 ) assert capped < 262144 capped_footprint_mib = (model_size_fit + inst._estimate_kv_cache_bytes(capped)) // ( 1024 * 1024 ) assert capped_footprint_mib <= budget_mib class TestAppleBranchEndToEnd: """Drive the Apple elif glue (cap / floor / explicit) via _drive, no GPU.""" def test_auto_context_capped_below_native(self): plan = _drive( n_ctx = 0, model_gib = 15.7, gpus = [], native_ctx = 262144, kv_per_token_bytes = 64_000, apple_budget_mib = 23_000, # ~22 GB: weights fit, native KV doesn't ) assert 0 < plan["c_arg"] < 262144 assert plan["use_fit"] is True # --fit on still ships as a backstop assert plan["gpu_indices"] is None # no CUDA device pinning on Metal assert plan["max_available_ctx"] == plan["c_arg"] def test_floors_to_fallback_when_weights_exceed_budget(self): # Weights alone exceed budget: ctx can't help, so floor to 4096. plan = _drive( n_ctx = 0, model_gib = 100, gpus = [], native_ctx = 262144, apple_budget_mib = 20_000, ) assert plan["c_arg"] == FALLBACK_CTX assert plan["use_fit"] is True assert plan["gpu_indices"] is None def test_explicit_context_honored_verbatim(self): # Explicit context is never shrunk, but the UI ceiling still tightens. plan = _drive( n_ctx = 200_000, model_gib = 15.7, gpus = [], native_ctx = 262144, kv_per_token_bytes = 64_000, apple_budget_mib = 23_000, ) assert plan["c_arg"] == 200_000 # launch context honored verbatim assert plan["use_fit"] is True # Ceiling reflects the budget so the over-budget warning still fires. assert plan["max_available_ctx"] < 262144 class TestAppleMtpFlatReserve: """Apple cap reserves the flat MTP fraction up front (like _pin_fraction) so an unsized MTP draft (Qwen3.6-MTP, #6529) can't over-commit.""" def test_flat_reserve_keeps_draft_within_budget(self): # No reserve -> cap fills the budget, leaving nothing for the ~5% draft. kw = dict( n_ctx = 0, model_gib = 15.7, gpus = [], native_ctx = 262144, kv_per_token_bytes = 64_000, apple_budget_mib = 23_000, ) no_reserve = _drive(**kw, flat_mtp_reserve = 0.0) with_reserve = _drive(**kw, flat_mtp_reserve = 0.05) def footprint_mib(ctx): return (15.7 * GIB + ctx * 64_000) / (1024 * 1024) # No reserve: main footprint + 5% draft exceeds the budget. assert footprint_mib(no_reserve["c_arg"]) + 0.05 * 23_000 > 23_000 # With reserve: the cap is smaller and the full footprint fits. assert with_reserve["c_arg"] < no_reserve["c_arg"] assert footprint_mib(with_reserve["c_arg"]) + 0.05 * 23_000 <= 23_000 def test_no_reserve_is_a_noop_when_mtp_absent(self): # flat_mtp_reserve == 0 (the common, non-MTP case) must not change the cap. kw = dict( n_ctx = 0, model_gib = 15.7, gpus = [], native_ctx = 262144, kv_per_token_bytes = 64_000, apple_budget_mib = 23_000, ) assert _drive(**kw, flat_mtp_reserve = 0.0) == _drive(**kw) class TestAppleNoKvMetadataFloor: """Sparse KV metadata floors the auto context to FALLBACK_CTX (like the discrete file-size-only fallback) instead of launching at native.""" def test_sparse_kv_floors_auto_context(self): plan = _drive( n_ctx = 0, model_gib = 15.7, gpus = [], native_ctx = 262144, can_estimate_kv = False, apple_budget_mib = 23_000, ) assert plan["c_arg"] == FALLBACK_CTX # not native 262144 assert plan["use_fit"] is True assert plan["gpu_indices"] is None def test_sparse_kv_still_honors_explicit_context(self): plan = _drive( n_ctx = 100_000, model_gib = 15.7, gpus = [], native_ctx = 262144, can_estimate_kv = False, apple_budget_mib = 23_000, ) assert plan["c_arg"] == 100_000 # explicit honored even without KV sizing