diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 6c97ef8cb2..35933e6685 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -433,6 +433,8 @@ class LlamaCppBackend: self._hf_variant: Optional[str] = None self._is_vision: bool = False self._healthy = False + # Set by _classify_gpu_offload after _wait_for_health. + self._gpu_offload_active: Optional[bool] = None self._context_length: Optional[int] = None self._effective_context_length: Optional[int] = None self._max_context_length: Optional[int] = None @@ -956,6 +958,13 @@ class LlamaCppBackend: logger.debug(f"torch GPU probe failed: {e}") return [] + # Free-VRAM fraction at which Studio pins the GPU directly instead + # of deferring to ``--fit on``. 5% headroom covers CUDA context + + # compute buffers; 0.90 was too conservative and dropped 91-94% + # fits to CPU offload (#5106). The fork's --fit on still catches + # the truly-too-large case. + _GPU_PIN_VRAM_FRACTION = 0.95 + @staticmethod def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]: """Return DLL dirs from pip-installed CUDA wheels under @@ -1024,11 +1033,11 @@ class LlamaCppBackend: """Pick GPU(s) for a model based on estimated VRAM and free memory. ``model_size_bytes`` should include both model weights and estimated - KV cache. The 90% threshold provides headroom for compute buffers, - CUDA context, and other runtime overhead. + KV cache. The ``_GPU_PIN_VRAM_FRACTION`` threshold provides headroom + for compute buffers, CUDA context, and other runtime overhead. Returns (gpu_indices, use_fit): - - ([1], False) model fits on 1 GPU at 90% of free + - ([1], False) model fits on 1 GPU at the headroom threshold - ([1, 2], False) model needs 2 GPUs - (None, True) model too large, let --fit handle it """ @@ -1036,12 +1045,13 @@ class LlamaCppBackend: return None, True model_size_mib = model_size_bytes / (1024 * 1024) + usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION # Sort GPUs by free memory descending ranked = sorted(gpus, key = lambda g: g[1], reverse = True) - # Try fitting on 1 GPU (90% of free memory threshold) - if ranked[0][1] * 0.90 >= model_size_mib: + # Try fitting on 1 GPU at the usable-VRAM threshold. + if ranked[0][1] * usable_fraction >= model_size_mib: return [ranked[0][0]], False # Try fitting on N GPUs (accumulate free memory from most-free) @@ -1049,7 +1059,7 @@ class LlamaCppBackend: selected = [] for idx, free_mib in ranked: selected.append(idx) - cumulative += free_mib * 0.90 + cumulative += free_mib * usable_fraction if cumulative >= model_size_mib: return sorted(selected), False @@ -1282,10 +1292,11 @@ class LlamaCppBackend: ) -> int: """Return the largest context length that fits in GPU VRAM. - Uses 90% of available VRAM as the budget (matching _select_gpus - threshold -- 10% reserved for compute buffers, CUDA context, - scratch space, flash-attn workspace, etc.). - If the model weights alone don't fit, returns min_ctx unchanged. + Uses 90% of available VRAM as the ctx-fit budget. Tighter than + ``_GPU_PIN_VRAM_FRACTION`` on purpose: over-promising context + OOMs at runtime, while pinning conservatively just defers to + --fit on. If the weights alone don't fit, returns + ``requested_ctx`` unchanged. ``kv_on_gpu`` mirrors ``--kv-offload`` (default on). When False the KV cache lives in CPU RAM and doesn't compete with weights @@ -2031,6 +2042,7 @@ class LlamaCppBackend: # still has valid state to publish. effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0) max_available_ctx = self._context_length or effective_ctx + gpus: list[tuple[int, int]] = [] try: model_size = self._get_gguf_size_bytes(model_path) gpus = self._get_gpu_free_memory() @@ -2114,8 +2126,11 @@ class LlamaCppBackend: gpu_indices, use_fit = self._select_gpus(requested_total, gpus) # No silent shrink: effective_ctx stays == n_ctx. else: - # Auto context: prefer fewer GPUs, cap context to fit. + # Auto context: prefer fewer GPUs, cap context + # to fit. Same headroom threshold as + # _select_gpus (#5106). ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + pin_fraction = self._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) @@ -2130,18 +2145,31 @@ class LlamaCppBackend: capped, cache_type_kv, n_parallel = n_parallel ) total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * 0.90: + if total_mib <= pool_mib * pin_fraction: effective_ctx = capped gpu_indices = sorted(idx for idx, _ in subset) use_fit = False break else: - # No subset can host the weights (weights alone - # exceed 90% of every pool). Per spec, default - # the UI-visible context to 4096 and let - # --fit on flex -ngl so llama-server offloads - # layers to CPU RAM. + # Native ctx doesn't fit. Drop to 4096 and + # re-check before deferring to --fit on: + # a model that overflows at 131k may pin + # comfortably with a 4096 KV cache (#5106). effective_ctx = min(4096, effective_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 = self._estimate_kv_cache_bytes( + effective_ctx, + cache_type_kv, + n_parallel = n_parallel, + ) + 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: # Can't estimate KV -- fall back to file-size-only check. @@ -2570,12 +2598,54 @@ class LlamaCppBackend: self._healthy = True + # Catch silent CPU fallback when GPU was intended (#5106). + self._gpu_offload_active = self._classify_gpu_offload( + gpu_indices is not None or use_fit, gpus or [] + ) + if self._gpu_offload_active is False: + logger.warning( + "llama-server appears to have loaded the model entirely " + "on CPU even though Studio detected at least one GPU. " + "This usually means the prebuilt binary's GPU backend " + "failed to load -- on Windows, cudart64_X.dll / " + "cublas64_X.dll could not be resolved. Reinstall the " + "Studio llama.cpp prebuilt or install a matching CUDA " + "toolkit (issue unslothai/unsloth#5106).", + ) + logger.info( f"llama-server ready on port {self._port} " f"for model '{model_identifier}'" ) return True + def _classify_gpu_offload( + self, + expected_gpu: bool, + detected_gpus: list[tuple[int, int]], + ) -> Optional[bool]: + """True if a GPU model buffer was allocated, False if only CPU + buffers landed despite GPU intent, None when there's no signal + (no GPU detected, no buffer-size lines, etc.).""" + if not detected_gpus or not expected_gpu: + return None + # llama-server logs one ``... model buffer size = N MiB`` line + # per backend buffer; CUDA0 / ROCm0 / Metal / Vulkan0 / + # OpenCL0 / SYCL0 are GPU, CPU / CPU_Mapped are not. + gpu_markers = ("CUDA", "ROCm", "Metal", "Vulkan", "OpenCL", "SYCL") + saw_buffer_line = False + saw_gpu_buffer = False + for line in self._stdout_lines: + if "model buffer size" not in line: + continue + saw_buffer_line = True + if any(marker in line for marker in gpu_markers): + saw_gpu_buffer = True + break + if not saw_buffer_line: + return None + return saw_gpu_buffer + def unload_model(self) -> bool: """Terminate the llama-server subprocess and cancel any in-flight download.""" self._cancel_event.set() diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py index caa6397901..1ea76edd15 100644 --- a/studio/backend/tests/test_llama_cpp_context_fit.py +++ b/studio/backend/tests/test_llama_cpp_context_fit.py @@ -192,6 +192,7 @@ def _drive( 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) @@ -203,7 +204,7 @@ def _drive( ) kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv) total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * 0.90: + if total_mib <= pool_mib * pin_fraction: effective_ctx = capped gpu_indices = sorted(idx for idx, _ in subset) use_fit = False @@ -211,6 +212,17 @@ def _drive( 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: @@ -378,6 +390,52 @@ class TestFittableAutoPickRegressions: 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 # --------------------------------------------------------------------------- @@ -391,3 +449,81 @@ def test_identical_decision_across_platforms(platform_tag): 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