From 5f1f30ec82d097d92d1093d1d557427dcbf079e6 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Sun, 19 Jul 2026 09:46:22 -0300 Subject: [PATCH] Studio: GPU memory configuration for GGUF models (#6414) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Studio: GPU memory dropdown — llama.cpp --fit on and manual gpu-layers/cpu-moe * Studio: simplify GPU memory changes (reuse ParamSlider, GPU_LAYERS_ALL, loadedGpuMemoryFields helper) * Studio: GPU picker — choose which GPUs a GGUF model loads on (gpu_ids) * Studio: simplify GPU picker (share /api/system fetch, validate gpu_ids) * Studio: GPU picker review fixes (gate relative indices, no cross-model leak, validate, types) * Studio: group GPU controls under a collapsible GPU section * Studio: GPU feature review fixes (fix fit-ctx test, behavior-test the floor, comment accuracy) * Studio: make GPU a top-level settings section (not nested under Model) * Studio: flatten GPU controls into the Model section, group by GPU/context/generation * Studio: move GPU Memory to the bottom of Model with its dependent controls beneath it * Studio: move GPU Memory below Tensor Parallelism and GPUs below GPU Memory * Studio: tighten GPU Memory and GPU Layers tooltip copy * Studio: fix fit-mode context slider track-click, restore GPU Memory tooltip, shorten fit dropdown label * Studio: GPU Memory tooltip one mode per line, briefer * Studio: note HIP_VISIBLE_DEVICES (ROCm) in the GPUs picker tooltip * Studio: narrow the GPU Memory dropdown to fit the shortened label * Studio: use 'llama.cpp --fit' in the GPU Memory tooltip for consistency * Studio: allow Tensor Parallelism in Manual GPU mode * Studio: graduated MoE-on-CPU offload (--n-cpu-moe) replacing the all-or-nothing toggle * Studio: size the MoE-offload slider for staged (deferred-load) models * Studio: share one GGUF header walk for the context-length and MoE-count readers * Studio: size the GPU Layers slider for staged models (one staged-header read) * Studio: move Tensor Parallelism below the GPUs picker * Studio: GPU split (--tensor-split) per-GPU model share in Manual mode * Studio: tolerate whitespace in GPU split input, move it below GPU Layers * Studio: rename the GPU split control to "Split ratio" * Studio: Split ratio sends explicit even input; fix blank=free-VRAM (not even) copy * Studio: tighten llama.cpp --fit VRAM margin with --fit-target 512 * Studio: GPU memory review fixes (rollback re-baseline, single-GPU TP gate, accurate copy) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: move Split ratio below MoE Layers on CPU * Studio: address PR review (fix GPU-info hydration race, share fit context-length across load paths) * Studio: address codex review (manual single-GPU TP guard, GPU-aware spec defaults in fit/manual, GGUF-only context/preference) * Studio: address codex review round 2 (gpu_present seed, single-GPU tensor-split guard, staged manual-knob reset, strip inherited offload flags) * Studio: address codex review round 3 (strip inherited --n-cpu-moe, CPU-fallback warning in Manual mode) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: address codex review round 4 (preserve pinned fit context across a later Apply) * Studio: address codex review round 5 (honor GPU picker for diffusion GGUFs, clear fit pin on cross-model switch) * Studio: preserve the pending GPU Memory mode when staging a model * Studio: pin diffusion GPU device order and reset GPU-memory state for diffusion loads * Studio: address codex review round 6 (fit-Auto rollback context, preserve manual non-tensor split modes, persist GPU mode on load not select) * Studio: persist the applied GPU Memory mode, not the requested one (skip diffusion loads) * Studio: replace Manual-mode split-ratio field with per-GPU layer sliders * Studio: clarify per-GPU layer split hint for tensor-parallel mode * Studio: address codex review round 7 (allow GGUF gpu_ids past the legacy guard, replay GPU-memory fields on respawn) * Studio: address codex review round 8 (size the validate preflight like the load in fit mode, across both load paths) * Studio: skip the training-OOM guard for llama.cpp --fit GGUF loads (they spill to RAM) * Studio: drop the now-redundant compare-path validate sizing (the --fit guard skip makes it moot) * Studio: address codex review round 9 (keep the training guard for fit loads, forward gpu_ids to validate, strip inherited manual tensor-split) * Studio: address codex review round 10 (gate GPU-memory adoption on is_gguf, record manual knobs only in Manual mode) * Studio: handle diffusion GGUFs symmetrically in the GPU Memory controls (preserve the standing mode preference, hide the inapplicable mode/TP controls) * Studio: remember the GPU Memory settings per model * Studio: consolidate --fit mode and Manual mode into a single Manual mode * Studio: preserve the per-GPU layer split across GPU Layers changes * Studio: trim overly long GPU Memory comments * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address GPU memory config review comments * trim redundant GPU memory tests * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reconcile manual-mode TP drops with the #6659 drop-site invariants * Preserve quantized KV in manual --fit, charge GGUF companions in full, reconcile GPU pick on load * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Clear stale GPU baseline on non-GGUF loads so it can't read as dirty * Fix no-context-shift test for the conditional -c flag * Credit manual GPU-layer offload for cached HF GGUFs * Reset per-model load knobs on GGUF quant switch * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Strip inherited tensor-split when manual ratio is cleared * Match auto-load validation to safetensors placement * Reset editable manual knobs after Auto GGUF loads * Record a single device for diffusion GPU picks * Reset per-model GPU knobs before applying saved settings * Address review comments * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Guard manual tensor splits and keep remembered context on auto-load * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Snapshot compare knobs, seed splits from free VRAM, flag zero-offload loads * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Exempt CPU-only loads from the guard floor and harden compare and reseed paths * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reach full offload from the layers slider and charge extras drafters in the guard * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Warm the GPU device cache before pick reconciles and disable staged GPU controls * Align the training guard with inherited extras, spec mode, and compare targets * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Hide GPUs from companion-less zero-offload loads * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Size diffusion picks per device, own manual offload flags, reject XPU picks * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop tensor flags at zero layers and exempt CPU-pinned drafters * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Allowlist the zero-layer tensor parallel drop site * Keep validate and load guards on the same extras and refresh stale baselines * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop mismatched manual tensor splits before launch * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate XPU picks on the real backend field and harden split and hydration paths * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Weight full GPUs as zero, clamp split shares, and refine the zero-layer mask gate * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Carry fit context across mode changes and align drafter and picker gates * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Catch variant switches, uncached diffusion repos, and text-only mmproj skips * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Check companions on the first device and size native and remote zero-layer loads * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Replace the training guard's precise VRAM modeling with a conservative bound * Baseline context pins on non-GGUF hydration and reprobe list-seeded staged GGUFs * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Size manual splits by their largest share and preserve resolved context from Default * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Default-deny unsized required companions and price KV at the effective cache dtype * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reserve MTP draft KV and MLA target-copy in the training guard * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Size tensor-parallel loads per device and show GPU controls for native GGUFs * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reserve MTP overhead for uncached remote GGUFs and the mmproj runtime factor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop the training-coexistence VRAM estimation this PR added * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate remembered load settings to GGUF picks * Lock the remaining load-time controls during a staged load * Clear the stale native-path token on compare loads * Drop a stale guard reference from the zero-offload masking comment * Seed GPU baselines from the rollback response and drop never-emitted offload flags * Match validate's training guard to load and keep the native reload token * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim verbose GPU-memory comments * Thread the variants header walk off the event loop, honor device pins on zero-offload, and hold staged GPU edits * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Honor manual placement and classify pinned zero-offload loads * Close diffusion admission and status hydration gaps * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Check the actual diffusion GPU during training * Align staged baselines and manual reload dedupe * Fix GGUF placement and rollback state * Harden manual GGUF placement boundaries * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove unused resolve_tensor_parallel import in llama_cpp.py The name is used only in llama_server_args.py, routes/inference.py, and tests, not in llama_cpp.py; the unused hoisted import trips the import-hoist verifier in the source-lint CI job. * Fix diffusion GPU dedup and training guard for non-numeric device tokens The diffusion runner drives only its single lowest device and the backend records that one device (self._gpu_ids = [sorted(gpu_ids)[0]]), but the reload dedupe compared it against the full requested list, so a multi-GPU pick that resolves to the same device forced a needless reload. Normalize the request the same way for a loaded diffusion model in both _already_in_target_state and the route _request_matches_loaded_settings. The chat-during-training coexistence guard called int() on the single-device token and hard-rejected when it could not parse. A non-numeric token (a CUDA UUID / MIG handle) now sizes against the whole visible pool like the GGUF guard instead of falsely blocking the load, and an empty token (a CPU-only runner such as a CPU diffusion GGUF) is allowed outright since it uses no GPU VRAM. * Tighten comments added by the GPU memory config changes * Harden GGUF placement from independent review: VRAM sizing, diffusion TP reset, tensor_split validation - Training coexistence guard: a single-device runner pinned through an unresolvable UUID/MIG token was sized against the aggregate visible-VRAM pool, so a load could pass on capacity it cannot use and then OOM active training. Size against the worst-case visible device (min free) instead, keeping the guard's documented default-deny contract. The empty-token (CPU-only runner) allow path is unchanged. - Diffusion startup: _start_diffusion_server now resets self._tensor_parallel to False alongside the other placement resets. A prior tensor-parallel chat load (process killed but not fully unload-reset) otherwise left /status misreporting tensor parallelism and made an identical diffusion re-Apply reload against the stale state. - tensor_split: reject negative / non-finite / all-zero splits up front. They were dropped at launch but still compared raw in the reload dedupe, so an identical Apply reloaded indefinitely. - Tests: the shared httpx stub was incomplete and, installed via setdefault before real httpx loaded, broke a combined pytest run (collection errors on httpx.Response). Import the real installed httpx instead. * [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> Co-authored-by: danielhanchen Co-authored-by: danielhanchen --- studio/backend/core/inference/llama_cpp.py | 715 +++++++++++++- .../core/inference/llama_server_args.py | 27 +- studio/backend/main.py | 14 + studio/backend/models/inference.py | 144 ++- studio/backend/routes/inference.py | 518 ++++++++--- studio/backend/routes/models.py | 6 +- studio/backend/routes/training_vram.py | 42 +- .../tests/test_chat_load_during_training.py | 333 ++++++- studio/backend/tests/test_gguf_metadata.py | 73 ++ studio/backend/tests/test_gpu_memory_mode.py | 879 ++++++++++++++++++ studio/backend/tests/test_gpu_selection.py | 18 +- .../tests/test_llama_cpp_no_context_shift.py | 12 +- .../tests/test_llama_cpp_props_readback.py | 31 +- .../backend/tests/test_llama_server_args.py | 45 + studio/backend/tests/test_tensor_parallel.py | 5 +- .../tests/test_tp_vision_regression.py | 17 +- studio/backend/utils/models/gguf_metadata.py | 109 ++- .../remembered-load-settings.ts | 24 +- .../src/features/chat/api/chat-adapter.ts | 128 ++- .../src/features/chat/api/chat-api.ts | 38 +- .../frontend/src/features/chat/chat-page.tsx | 13 +- .../src/features/chat/chat-settings-sheet.tsx | 450 ++++++++- .../chat/hooks/use-chat-model-runtime.ts | 217 ++++- .../hooks/use-staged-model-preparation.ts | 46 +- .../lib/apply-inference-status-to-store.ts | 117 ++- .../features/chat/presets/preset-policy.ts | 31 + .../src/features/chat/shared-composer.tsx | 109 ++- .../chat/stores/chat-runtime-store.ts | 369 +++++++- .../frontend/src/features/chat/types/api.ts | 38 + studio/frontend/src/hooks/use-gpu-info.ts | 186 ++-- studio/frontend/src/hooks/use-system.ts | 3 + 31 files changed, 4356 insertions(+), 401 deletions(-) create mode 100644 studio/backend/tests/test_gpu_memory_mode.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index c9ab7eb83b..d7c7eed518 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -11,6 +11,7 @@ import atexit import contextlib import functools import json +import math import os import re import struct @@ -23,11 +24,22 @@ import sys import threading import time from pathlib import Path -from typing import Callable, Collection, Generator, Iterable, List, Mapping, Optional, Union +from typing import ( + Callable, + Collection, + Generator, + Iterable, + List, + Literal, + Mapping, + Optional, + Union, +) import httpx from core.inference.llama_server_args import ( + _LAYER_OFFLOAD_FLAGS, _effective_tensor_parallel, _tensor_parallel_matches_loaded, extra_args_disable_mmproj, @@ -234,8 +246,7 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]": return out -# Plan-without-action re-prompt state (intent signal, caps, message) now lives -# in tool_call_parser, imported above under its old aliases. +# Plan-without-action re-prompt state now lives in tool_call_parser (imported above). # Default max_tokens to the effective context when known. The floor is high # enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. @@ -1431,7 +1442,10 @@ def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool: return _extra_args_set_any_flag(extra_args, {"--spec-type", "--spec-default"}) -_GPU_OFFLOAD_OVERRIDE_FLAGS = frozenset({"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"}) +# Layer-offload override detection. Single-sourced from llama_server_args, which +# also strips these (plus the MoE flags) from inherited extras; sharing the layer +# set keeps detection and stripping from drifting. +_GPU_OFFLOAD_OVERRIDE_FLAGS = _LAYER_OFFLOAD_FLAGS _THREAD_OVERRIDE_FLAGS = frozenset({"-t", "--threads"}) @@ -1895,6 +1909,17 @@ class LlamaCppBackend: self._cache_type_kv: Optional[str] = None # Whether --split-mode tensor was applied on the active load. self._tensor_parallel: bool = False + # GPU memory strategy applied on the active load ("auto"/"manual"). + self._gpu_memory_mode: str = "auto" + # Manual-mode load options (echoed back so the UI round-trips them). + self._gpu_layers: int = -1 + # MoE expert layers to keep on CPU (--n-cpu-moe); 0 = none. + self._n_cpu_moe: int = 0 + # Relative model share per GPU (--tensor-split), in GPU order; None = + # default (llama.cpp splits by free VRAM). + self._tensor_split: Optional[List[float]] = None + # User-picked physical GPU indices (None = automatic selection). + self._gpu_ids: Optional[List[int]] = None # Layer load kept multi-GPU only to honor a downgraded tensor request, so a # later explicit tensor-off reloads instead of deduping to it (#6659). self._layer_preserves_tensor_intent: bool = False @@ -1909,6 +1934,11 @@ class LlamaCppBackend: self._spec_draft_n_max: Optional[int] = None # KV-cache estimation fields (populated by _read_gguf_metadata) self._n_layers: Optional[int] = None + # MoE metadata (populated by _read_gguf_metadata): expert count (>0 = + # MoE) and leading dense-layer count (offsets --n-cpu-moe, which counts + # from layer 0). See the n_moe_layers property. + self._n_experts: Optional[int] = None + self._leading_dense_block_count: Optional[int] = None self._n_kv_heads: Optional[int] = None self._n_kv_heads_by_layer: Optional[list[int]] = None self._n_heads: Optional[int] = None @@ -2329,6 +2359,79 @@ class LlamaCppBackend: """Whether --split-mode tensor is active on the loaded server.""" return self._tensor_parallel + @property + def gpu_memory_mode(self) -> str: + """Active GPU memory strategy: 'auto' or 'manual' (gpu_layers < 0 = Auto/--fit, >= 0 = pinned).""" + return self._gpu_memory_mode + + @property + def gpu_layers(self) -> int: + """Requested --gpu-layers for manual mode (-1 when not manual).""" + return self._gpu_layers + + @property + def n_cpu_moe(self) -> int: + """MoE expert layers manual mode kept on CPU (--n-cpu-moe); 0 = none.""" + return self._n_cpu_moe + + @property + def tensor_split(self) -> Optional[List[float]]: + """Manual-mode relative model share per GPU (--tensor-split); None = + default (split by free VRAM).""" + return self._tensor_split + + @property + def gpu_ids(self) -> Optional[List[int]]: + """User-picked physical GPU indices, or None for automatic selection.""" + return self._gpu_ids + + @property + def n_layers(self) -> Optional[int]: + """Model layer count (GGUF block_count), or None if unknown.""" + return self._n_layers + + @property + def n_moe_layers(self) -> int: + """Number of MoE expert layers (the --n-cpu-moe ceiling), 0 if not MoE. + + block_count minus the leading dense layers (which carry no experts): + --n-cpu-moe counts from layer 0, so those dense layers are no-ops. + """ + if not self._n_experts or not self._n_layers: + return 0 + return max(0, self._n_layers - (self._leading_dense_block_count or 0)) + + @staticmethod + def _resolve_cpu_moe_flag( + n_cpu_moe: int, n_moe_layers: int, leading_dense: int + ) -> Optional[int]: + """The --n-cpu-moe value (absolute first-N layers), or None to omit it. + + Clamps the requested count to the model's MoE layers, then offsets past + the leading dense layers (--n-cpu-moe counts from layer 0). Returns None + for nothing-to-offload (0 requested) or a non-MoE model. + """ + if n_cpu_moe <= 0 or n_moe_layers <= 0: + return None + return leading_dense + min(n_cpu_moe, n_moe_layers) + + @staticmethod + def _sanitize_tensor_split(tensor_split: Optional[List[float]]) -> List[float]: + """Per-GPU shares with negative and non-finite entries clamped to 0. + + A direct caller's negative entry would launch a placement different + from the ratio the UI showed, and inf would pass a plain ``> 0`` total + gate and emit ``--tensor-split inf,...``. Returns [] for input that + can't be read as floats (the length gate at the call site then drops + the split). + """ + try: + return [ + x if math.isfinite(x) and x > 0.0 else 0.0 for x in (float(v) for v in tensor_split) + ] + except (TypeError, ValueError, OverflowError): + return [] + @property def layer_preserves_tensor_intent(self) -> bool: """True when a downgraded tensor request kept this layer load multi-GPU.""" @@ -2530,6 +2633,7 @@ class LlamaCppBackend: "spec_draft_n_max_flag": None, "supports_kv_unified": False, "supports_fit_ctx": False, + "supports_fit_target": False, "supports_cache_ram": False, "supports_ctx_checkpoints": False, "supports_no_cache_prompt": False, @@ -2549,6 +2653,7 @@ class LlamaCppBackend: spec_draft_n_max_flag: Optional[str] = None supports_kv_unified = False supports_fit_ctx = False + supports_fit_target = False supports_cache_ram = False supports_ctx_checkpoints = False supports_no_cache_prompt = False @@ -2646,6 +2751,7 @@ class LlamaCppBackend: supports_kv_unified = _is_real("--kv-unified") supports_fit_ctx = _is_real("--fit-ctx") + supports_fit_target = _is_real("--fit-target") supports_cache_ram = _is_real("--cache-ram") supports_ctx_checkpoints = _is_real("--ctx-checkpoints") supports_no_cache_prompt = _is_real("--no-cache-prompt") @@ -2662,6 +2768,7 @@ class LlamaCppBackend: "spec_draft_n_max_flag": spec_draft_n_max_flag, "supports_kv_unified": supports_kv_unified, "supports_fit_ctx": supports_fit_ctx, + "supports_fit_target": supports_fit_target, "supports_cache_ram": supports_cache_ram, "supports_ctx_checkpoints": supports_ctx_checkpoints, "supports_no_cache_prompt": supports_no_cache_prompt, @@ -2746,6 +2853,57 @@ class LlamaCppBackend: except ValueError: return None + @staticmethod + def _emit_child_gpu_visibility(env: dict, pinned: str) -> None: + """Write the child's GPU visibility mask (CUDA, plus the HIP mirror on + ROCm, where narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child + seeing the full set). Do NOT also set ROCR_VISIBLE_DEVICES: ROCR and HIP + mask at different layers, so the same indices apply twice -- ROCR reduces + and re-indexes from 0, then a non-zero HIP pin points out of range, HIP + enumerates 0 devices, and llama.cpp falls back to CPU. The HIP mask alone + narrows correctly; clear any inherited ROCR mask so it can't double up.""" + env["CUDA_VISIBLE_DEVICES"] = pinned + try: + import torch as _torch + if getattr(_torch.version, "hip", None) is not None: + env["HIP_VISIBLE_DEVICES"] = pinned + env.pop("ROCR_VISIBLE_DEVICES", None) + except Exception as e: + logger.debug("Failed to set ROCm visibility env vars for child: %s", e) + + @staticmethod + def _pin_visible_gpu_order_for_split(env: dict) -> None: + """Pin the child's GPU enumeration to the picker's order for a manual + ``--tensor-split`` across the whole visible set. CUDA's default + FASTEST_FIRST enumeration applies the shares to the wrong cards on + heterogeneous hosts (#5025), and CUDA_DEVICE_ORDER only fixes the + numbering base: an inherited numeric visibility mask ALSO defines + enumeration order, so a reordered parent mask (CUDA_VISIBLE_DEVICES=3,1) + would still hand the shares to the wrong cards. The UI built the split + positionally over get_backend_visible_gpu_info's device list (ascending + physical via nvidia-smi, inherited mask order on the torch fallback), so + re-emit the same set in that report order -- not an assumed ascending + sort. The visible set itself never changes. No mask, an empty mask, or a + UUID/MIG mask (which resolves to None) is left alone -- the multi-GPU + controls are hidden for the latter.""" + env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + inherited = LlamaCppBackend._resolve_visible_physical_ids() + if not inherited: + return + order = None + try: + from utils.hardware import get_backend_visible_gpu_info + info = get_backend_visible_gpu_info() + if info.get("available") and info.get("index_kind") == "physical": + reported = [d["index"] for d in info.get("devices", [])] + if sorted(reported) == sorted(inherited): + order = reported + except Exception as e: + logger.debug("Could not read reported GPU order for split pin: %s", e) + if order is None: + order = sorted(inherited) + LlamaCppBackend._emit_child_gpu_visibility(env, ",".join(str(i) for i in order)) + @staticmethod def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool: """True only for AMD unified-memory APUs (gfx1150/gfx1151), where @@ -3262,6 +3420,20 @@ class LlamaCppBackend: # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + # Main-model placement settings that Manual mode owns. They must not leak + # from Studio's parent environment into llama-server and silently override + # the command assembled from the current request. Draft-model placement is + # intentionally separate and remains available to speculative decoding. + _MANUAL_PLACEMENT_ENV_VARS = ( + "LLAMA_ARG_CPU_MOE", + "LLAMA_ARG_N_CPU_MOE", + "LLAMA_ARG_N_GPU_LAYERS", + "LLAMA_ARG_TENSOR_SPLIT", + "LLAMA_ARG_FIT", + "LLAMA_ARG_FIT_TARGET", + "LLAMA_ARG_FIT_CTX", + ) + # (binary, mtime, model) that aborted on --split-mode tensor this process (#6415 # geometry limit, e.g. MQA n_head_kv=1). Model-keyed so one model's abort doesn't # skip tensor for others; tensor is tried by default, recorded only on a real abort. @@ -3426,6 +3598,12 @@ class LlamaCppBackend: return env + @classmethod + def _clear_manual_placement_env(cls, env: dict[str, str]) -> None: + """Remove inherited main-model placement owned by Manual mode.""" + for name in cls._MANUAL_PLACEMENT_ENV_VARS: + env.pop(name, None) + @staticmethod def _select_gpus( model_size_bytes: int, @@ -4246,6 +4424,8 @@ class LlamaCppBackend: self._supports_preserve_thinking = False self._supports_tools = False self._n_layers = None + self._n_experts = None + self._leading_dense_block_count = None self._n_kv_heads = None self._n_kv_heads_by_layer = None self._n_heads = None @@ -4335,6 +4515,8 @@ class LlamaCppBackend: arch_keys = { f"{arch}.context_length": "context_length", f"{arch}.block_count": "n_layers", + f"{arch}.expert_count": "n_experts", + f"{arch}.leading_dense_block_count": "leading_dense_block_count", f"{arch}.attention.head_count_kv": "n_kv_heads", f"{arch}.attention.head_count": "n_heads", f"{arch}.embedding_length": "embedding_length", @@ -4523,6 +4705,28 @@ class LlamaCppBackend: return None + @staticmethod + def _diffusion_gpu_arg(gpu_ids: Optional[List[int]], *, cpu_only: bool = False) -> str: + """Device token passed to the diffusion visual-server child. + + The visual engine replaces its child's CUDA visibility mask with this + token, so an unpinned load must carry forward the first token from the + parent's mask rather than turning a parent-relative ordinal into a new + physical selection. + """ + if gpu_ids: + return str(sorted(gpu_ids)[0]) + if cpu_only: + return "" + if "DG_GPU" in os.environ: + return os.environ["DG_GPU"] + parent_mask = os.environ.get("CUDA_VISIBLE_DEVICES") + if parent_mask: + first = next((token.strip() for token in parent_mask.split(",") if token.strip()), "") + if first and first != "-1": + return first + return "0" + def _start_diffusion_server( self, *, @@ -4533,6 +4737,7 @@ class LlamaCppBackend: model_identifier: str, n_ctx: int, extra_args: Optional[List[str]], + gpu_ids: Optional[List[int]] = None, ) -> bool: """Launch the OpenAI-compat diffusion shim (which drives the on-device visual decoder) and wait for health. Presents the same /v1 + /health @@ -4558,7 +4763,11 @@ class LlamaCppBackend: # CUDA_VISIBLE_DEVICES="" to force CPU serving. Keep the visual-server child # CPU-masked (empty --gpu) so the shim does not re-expose GPU 0 via its default. cpu_only = self._effective_gpu_count() == 0 - gpu = "" if cpu_only else os.environ.get("DG_GPU", "0") + # Honor the GPU picker first: the diffusion runner takes a single device, + # so use the lowest selected GPU (matches the sorted set recorded below, so + # the device used == the echoed gpu_ids[0]). With no pick, fall back to the + # CPU-only mask, else DG_GPU / 0. + gpu = self._diffusion_gpu_arg(gpu_ids, cpu_only = cpu_only) cmd = list(shim_cmd) + [ "--gguf", @@ -4586,6 +4795,11 @@ class LlamaCppBackend: env.setdefault("UNSLOTH_ALLOW_CPU", "1") env["DG_VISUAL_BIN"] = visual_bin env["DG_GPU"] = gpu + if gpu_ids: + # The visual server remasks via CUDA_VISIBLE_DEVICES=; pin PCI + # order (as the llama-server path does) so the picked physical id maps + # to the GPU the picker showed, not CUDA's default fastest-first order. + env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH. # (The zoo-package shim is an installed module and needs no PYTHONPATH change.) if extra_pythonpath: @@ -4631,6 +4845,23 @@ class LlamaCppBackend: self._model_identifier = model_identifier self._cache_type_kv = None self._gpu_offload_active = True + # Diffusion doesn't use the llama.cpp GPU-memory knobs; reset them to + # defaults (the picked device is still recorded below) so /load, /status + # and reload dedup don't report a previous GGUF's manual settings. + self._gpu_memory_mode = "auto" + self._gpu_layers = -1 + self._n_cpu_moe = 0 + self._tensor_split = None + # Diffusion is never tensor-parallel; clear any state left by a prior TP + # chat load (load_model phase 1 only kills the process, it doesn't run + # the unload reset) so /status doesn't misreport TP and an identical + # re-Apply doesn't reload against stale tensor-parallel state. + self._tensor_parallel = False + # Record only the single device the runner actually uses (the lowest + # selected GPU, chosen above) -- not the whole pick. The diffusion runner + # is single-device, so echoing a multi-GPU list would misreport placement + # in /status and let a re-Apply dedup against GPUs the runner never used. + self._gpu_ids = [sorted(gpu_ids)[0]] if gpu_ids else None if hf_variant: self._hf_variant = hf_variant elif gguf_path: @@ -5721,6 +5952,11 @@ class LlamaCppBackend: speculative_type: Optional[str] = None, spec_draft_n_max: Optional[int] = None, tensor_parallel: bool = False, + gpu_memory_mode: Literal["auto", "manual"] = "auto", + gpu_layers: int = -1, + n_cpu_moe: int = 0, + tensor_split: Optional[List[float]] = None, + gpu_ids: Optional[List[int]] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # caller compat, unused n_parallel: int = 1, @@ -5753,6 +5989,14 @@ class LlamaCppBackend: "speculative_type": speculative_type, "spec_draft_n_max": spec_draft_n_max, "tensor_parallel": tensor_parallel, + # GPU-memory placement: replayed on respawn so a server SIGKILL'd by + # GPU/RAM pressure reloads onto the same devices with the same + # offload, not the auto defaults. + "gpu_memory_mode": gpu_memory_mode, + "gpu_layers": gpu_layers, + "n_cpu_moe": n_cpu_moe, + "tensor_split": list(tensor_split) if tensor_split is not None else None, + "gpu_ids": list(gpu_ids) if gpu_ids is not None else None, "n_threads": n_threads, "n_gpu_layers": n_gpu_layers, "n_parallel": n_parallel, @@ -5779,6 +6023,11 @@ class LlamaCppBackend: speculative_type = speculative_type, spec_draft_n_max = spec_draft_n_max, tensor_parallel = tensor_parallel, + gpu_memory_mode = gpu_memory_mode, + gpu_layers = gpu_layers, + n_cpu_moe = n_cpu_moe, + tensor_split = tensor_split, + gpu_ids = gpu_ids, chat_template_override = chat_template_override, extra_args = extra_args, is_vision = is_vision, @@ -5899,6 +6148,7 @@ class LlamaCppBackend: model_identifier = model_identifier, n_ctx = n_ctx, extra_args = extra_args, + gpu_ids = gpu_ids, ) if not binary: @@ -5960,6 +6210,59 @@ class LlamaCppBackend: # use the same helper so a healthy env-driven tensor server matches. split_mode_override = parse_split_mode_override(extra_args) tensor_parallel = _effective_tensor_parallel(extra_args, tensor_parallel) + # gpu_layers=0 leaves nothing to split, yet --split-mode tensor or + # a per-GPU ratio still launches tensor mode -- and under the + # CPU-only mask below (no visible devices) that aborts the server + # instead of loading on CPU. Drop both here (nothing to split). + if gpu_memory_mode == "manual" and gpu_layers == 0: + if tensor_parallel or tensor_split: + logger.info( + "Manual gpu_layers=0: dropping tensor split/parallel " + "flags (nothing to split on the GPU)" + ) + tensor_parallel = False + tensor_split = None + # Record the requested strategy for /status and the load + # response. 'manual' has no fallback, so the request value is the + # value actually applied. + self._gpu_memory_mode = gpu_memory_mode + # The layer/MoE/split knobs apply only with an explicit offload + # (manual + gpu_layers >= 0); else record defaults so /status and + # /load don't report knobs the server never applied. + if gpu_memory_mode == "manual" and gpu_layers >= 0: + self._gpu_layers = gpu_layers + self._n_cpu_moe = n_cpu_moe + self._tensor_split = tensor_split + else: + self._gpu_layers = -1 + self._n_cpu_moe = 0 + self._tensor_split = None + self._gpu_ids = sorted(gpu_ids) if gpu_ids else None + # Manual offload skips the TP planner but still emits --split-mode + # tensor at launch; drop it when fewer than 2 GPUs are in use -- + # tensor split is a no-op there and aborts on some architectures. + # Done before the cache-drop below so a quantized KV survives. + if ( + tensor_parallel + and gpu_memory_mode == "manual" + and gpu_layers >= 0 + and self._effective_gpu_count(sorted(gpu_ids) if gpu_ids else None) < 2 + ): + logger.info( + "Tensor parallelism requested in manual mode but fewer " + "than 2 GPUs are in use; ignoring (needs >= 2)." + ) + tensor_parallel = False + # Drop TP for manual + Auto layers before the cache-drop below (like + # the <2-GPU guard above), so a requested quantized KV survives into + # the --fit load rather than being stripped for a tensor attempt. + if tensor_parallel and gpu_memory_mode == "manual" and gpu_layers < 0: + logger.info( + "Manual mode with Auto layers hands memory management to " + "llama.cpp --fit, which is incompatible with tensor " + "parallelism; ignoring the tensor split." + ) + tensor_parallel = False # Tensor mode aborts on a quantized KV cache, so drop it for the # tensor attempt (and strip any inherited/explicit --cache-type # that would re-impose it when appended last). Layer split does @@ -6040,6 +6343,12 @@ class LlamaCppBackend: "Vision-capable GGUF loaded without a usable mmproj; " "image input will be disabled for this session" ) + # Seed before the try: the except (GPU-selection failure -> + # --fit on) falls through to the launch which reads this, and the + # probe that assigns it may throw first. Captured before manual + # empty `gpus` so the speculative defaults stay GPU-aware and the + # CPU-fallback check still knows GPUs were present. + _detected_gpus: list[tuple[int, int]] = [] model_size = None # set in the fit try; used by the APU RAM guard # Layer-fallback min GPUs; raised below on a tensor downgrade. Bound # before the try so the --fit-on except path still has it (no UnboundLocal). @@ -6057,6 +6366,18 @@ class LlamaCppBackend: _gpu_mem = self._get_gpu_memory(binary) gpus = [(idx, free) for idx, free, _t in _gpu_mem] total_by_idx = {idx: total for idx, _f, total in _gpu_mem} + # GPU picker: restrict every mode to the chosen devices, so + # auto selection only considers them and manual mask to + # them (the env block below pins CUDA/HIP_VISIBLE_DEVICES). + if gpu_ids: + _picked = set(gpu_ids) + gpus = [g for g in gpus if g[0] in _picked] + + # GPUs the model will run on -- captured before manual + # empty `gpus` to bypass the planner. bool() drives the + # GPU-aware speculative defaults; the list feeds the + # CPU-fallback check. + _detected_gpus = list(gpus) def _gpu_usable(g, frac = _CTX_FIT_VRAM_FRACTION): # Per-GPU usable budget for ranking: free - (1-frac)*total. @@ -6088,6 +6409,44 @@ class LlamaCppBackend: # GPU/VRAM-fit logic below may shrink it on limited HW. max_available_ctx = self._context_length or effective_ctx + # Manual + Auto layers (the Manual default): hand memory + # management to llama.cpp's --fit. Emptying the probed GPU set + # no-ops the selection/TP planning below, leaving gpu_indices + # None (an explicit gpu_ids pick still pins below) and use_fit + # True. An explicit context is honored (--fit optimizes around + # it); 0 lets --fit size it. + if gpu_memory_mode == "manual" and gpu_layers < 0: + # Tensor parallelism was already dropped above (before the + # cache-drop), so a quantized KV survives into this --fit load. + gpus = [] + effective_ctx = requested_ctx if requested_ctx > 0 else 0 + original_ctx = effective_ctx + # --fit aborts under --split-mode tensor; a raw extras + # --split-mode/--tensor-split (appended last) would + # otherwise reach llama-server. Strip it like the TP + # downgrade does. + extra_args = strip_split_mode_only(extra_args) + elif gpu_memory_mode == "manual": + # Manual offload (--gpu-layers + --fit off): no automatic + # device masking (a gpu_ids pick still pins below) or + # context cap -- the user owns both. tensor_parallel is + # honored but skips the memory-based planner (gpus = []); + # the toggle just emits --split-mode tensor (split by free + # VRAM, or by the Split ratio if set). + gpus = [] + effective_ctx = ( + requested_ctx if requested_ctx > 0 else (self._context_length or 0) + ) + original_ctx = effective_ctx + # Strip the user --split-mode when the toggle owns the split + # (TP engaged -> Studio emits --split-mode tensor) or when the + # user asked for tensor (which aborts on a single GPU even if + # the manual <2-GPU guard downgraded TP). Otherwise keep their + # non-tensor mode (row/none/layer) -- the toggle can't express + # those. + if tensor_parallel or split_mode_override == "tensor": + extra_args = strip_split_mode_only(extra_args) + # Will MTP engage? If so, auto-fit reserves draft-model VRAM. # Mirrors _build_speculative_flags: forced mtp/mtp+ngram always # engage; auto only on an MTP model >= 3B; ngram/off never. A @@ -6175,7 +6534,10 @@ class LlamaCppBackend: _extra_n_max = _extra_args_spec_draft_n_max(extra_args) _mtp_eff_n_max = _extra_n_max if _extra_n_max is not None else spec_draft_n_max if _mtp_eff_n_max is None: - _mtp_eff_n_max = 2 if gpus else 3 + # _detected_gpus (not gpus) so manual -- which empty + # gpus to bypass the planner -- keep the GPU draft depth the + # launch flags also use, instead of the CPU default. + _mtp_eff_n_max = 2 if _detected_gpus else 3 # Separate-drafter weights live on GPU (an embedded head is # already in model_size). Size the drafter the launch loads, by # precedence: extras --model-draft (last-wins), else Unsloth's @@ -6313,7 +6675,8 @@ class LlamaCppBackend: # honor it, cap only if it fits no combination. Auto (native): # prefer fewer GPUs with reduced context (multi-GPU is slower). gpu_indices, use_fit = None, True - # Per-GPU weight proportions for tensor mode (None = even). + # Per-GPU weight proportions for tensor mode (None lets + # llama.cpp split by free VRAM). tp_tensor_split: Optional[list[int]] = None explicit_ctx = requested_ctx > 0 # Flat MTP reserve fraction: used only as the fallback when the @@ -6388,7 +6751,12 @@ class LlamaCppBackend: # GPUs below that reserve from the set up front (gpu_indices # becomes the CUDA_VISIBLE_DEVICES mask, fully excluding them). tp_gpus = gpus - if tensor_parallel: + # Manual mode owns the layer count and context, so it skips + # the memory-based planner; its toggle still emits + # --split-mode tensor below (split by free VRAM, or by the + # Split ratio if set). auto plans here. + plan_tp = tensor_parallel and gpu_memory_mode != "manual" + if plan_tp: # Deterministic per-device compute buffer (replicated on # every device in tensor mode); flat fallback when dims # are unavailable. _plan_tensor_parallel uses the same. @@ -6407,7 +6775,7 @@ class LlamaCppBackend: # free yet have no budget left. tp_gpus = [g for g in gpus if _gpu_usable(g) >= reserve_mib] - if tensor_parallel and len(tp_gpus) < 2: + if plan_tp and len(tp_gpus) < 2: # Tensor parallelism needs >= 2 usable GPUs. On a single # GPU --split-mode tensor is a no-op; with 0 GPUs (CPU-only # or probe failed) it must not reach llama-server; and a @@ -6823,6 +7191,12 @@ class LlamaCppBackend: tp_tensor_split = None effective_ctx = requested_ctx # fall back to original + # GPU picker: when no narrower subset was chosen (manual, or + # a failed/file-size selection), pin the whole picked set so the + # model can't spill onto an unpicked GPU. + if gpu_ids and gpu_indices is None: + gpu_indices = sorted(gpu_ids) + # Unified-memory APUs load weights into system RAM (under WSL the VM # cap, not the ROCm-reported VRAM, is the real ceiling); refuse an # oversize load the OS would otherwise kill mid-flight. Base model @@ -6859,8 +7233,6 @@ class LlamaCppBackend: model_path, "--port", str(self._port), - "-c", - str(effective_ctx) if effective_ctx > 0 else "0", "--parallel", str(n_parallel), "--flash-attn", @@ -6868,6 +7240,17 @@ class LlamaCppBackend: # Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length". "--no-context-shift", ] + # A positive context is always passed (in auto-fit, --fit then + # optimizes the gpu-layer offload around it). When auto-fit has + # no explicit context, omit -c so --fit sizes it to fit VRAM: + # "-c 0" would instead pin the FULL native context (llama.cpp's + # -c handler sets fit_params_min_ctx = UINT32_MAX on value 0, + # disabling --fit's reduction). See gpu_memory_mode. + auto_fit = gpu_memory_mode == "manual" and gpu_layers < 0 + if effective_ctx > 0: + cmd.extend(["-c", str(effective_ctx)]) + elif not auto_fit: + cmd.extend(["-c", "0"]) # Report a clean public model id (matching GET /v1/models) rather # than the raw -m path in llama-server's own /v1/models and the @@ -6879,7 +7262,63 @@ class LlamaCppBackend: cmd.extend(["--alias", _alias]) fully_gpu_offloaded = False - if use_fit: + # Set when a positional --tensor-split is emitted, so the env block + # can pin CUDA to PCI order even without a GPU subset (see below). + manual_tensor_split_emitted = False + if gpu_memory_mode == "manual" and gpu_layers >= 0: + # Pin the user's layer count and disable auto-fit. --fit off + # also means _ctx_integrity_flags must not add --fit-ctx. + use_fit = False + cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"]) + # Keep the first n_cpu_moe MoE layers' experts on CPU. + moe_flag = self._resolve_cpu_moe_flag( + n_cpu_moe, + self.n_moe_layers, + self._leading_dense_block_count or 0, + ) + if moe_flag is not None: + cmd.extend(["--n-cpu-moe", str(moe_flag)]) + elif n_cpu_moe: + # Requested on a dense model: nothing was emitted, so + # don't report a count llama-server never received. + self._n_cpu_moe = 0 + # Distribute the model across GPUs by the user's per-GPU shares + # (--tensor-split). Works with layer split and tensor + # parallelism; --fit off means no fit/tensor abort. Only emit + # when >1 GPU is in use AND the list length matches that count: + # the field is hidden (not cleared) when the picker narrows to + # one, and a direct caller can send a stale ratio for a different + # GPU set. Studio drops any mismatch to the free-VRAM default + # (llama.cpp would silently zero-pad a short list, or abort past + # its 16-device cap). + _split_gpus = self._effective_gpu_count(gpu_indices) + if tensor_split and _split_gpus > 1: + # An all-zero/non-positive sanitized split assigns nothing + # anywhere, so fall through to the free-VRAM default in + # that case. + _sanitized_split = self._sanitize_tensor_split(tensor_split) + _split_total = sum(_sanitized_split) + if len(_sanitized_split) == _split_gpus and _split_total > 0: + cmd.extend( + ["--tensor-split", ",".join(f"{x:g}" for x in _sanitized_split)] + ) + self._tensor_split = _sanitized_split + manual_tensor_split_emitted = True + else: + logger.warning( + "Dropping manual --tensor-split (%d entries for " + "%d GPUs, sanitized total %s); llama.cpp's " + "free-VRAM split applies instead", + len(tensor_split), + _split_gpus, + _split_total, + ) + self._tensor_split = None + elif tensor_split: + # Single effective GPU: the split is never emitted, so + # don't report it as active via /status and /load. + self._tensor_split = None + elif use_fit: cmd.extend(["--fit", "on"]) elif gpu_indices is not None: # Fits on selected GPU(s) -- force all layers on GPU. --fit off is @@ -6897,6 +7336,7 @@ class LlamaCppBackend: self._ctx_integrity_flags( n_parallel, use_fit, + auto_fit, requested_ctx, effective_ctx, server_caps, @@ -6960,9 +7400,11 @@ class LlamaCppBackend: self._cache_type_kv = None # Tensor parallelism: split the model across GPUs by tensor - # rather than by layer. Multi-GPU only -- a no-op on a single - # GPU. Default (layer split) is left implicit by omitting the - # flag. See llama.cpp --split-mode. + # rather than by layer. The UI only offers it on multi-GPU; a + # direct single-GPU caller is redundant (supported archs no-op, + # unsupported ones abort and the /load path retries layer split). + # Default (layer split) is left implicit by omitting the flag. + # See llama.cpp --split-mode. if tensor_parallel: cmd.extend(["--split-mode", "tensor"]) if tp_tensor_split and len(tp_tensor_split) > 1: @@ -6994,7 +7436,7 @@ class LlamaCppBackend: extra_args = extra_args, model_identifier = model_identifier, model_path = model_path, - gpus = bool(gpus), + gpus = bool(_detected_gpus), binary = binary, mtp_draft_path = launch_mtp_draft_path, ) @@ -7112,6 +7554,8 @@ class LlamaCppBackend: # Library paths so llama-server finds its shared libs and CUDA DLLs. env = self._llama_server_env_for_binary(binary) + if gpu_memory_mode == "manual": + self._clear_manual_placement_env(env) # Omitting --threads relies on llama.cpp's physical-core default, so # drop an inherited LLAMA_ARG_THREADS that would otherwise feed the # arg handler and silently force hardware_concurrency(). #5692 @@ -7170,28 +7614,39 @@ class LlamaCppBackend: # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device # (above), not here. - if gpu_indices is not None and not is_vulkan_backend: - pinned = ",".join(str(i) for i in gpu_indices) - env["CUDA_VISIBLE_DEVICES"] = pinned - try: - import torch as _torch - if getattr(_torch.version, "hip", None) is not None: - env["HIP_VISIBLE_DEVICES"] = pinned - # Do NOT also set ROCR_VISIBLE_DEVICES to the same - # value. ROCR_VISIBLE_DEVICES filters at the HSA/ROCr - # layer and HIP_VISIBLE_DEVICES at the HIP layer, so - # setting both with the same physical indices applies - # the mask twice: ROCR reduces the visible set and - # re-indexes it from 0, then HIP indexes into the - # already-reduced set. A single non-zero pin (e.g. - # "1") then points out of range at the HIP layer, HIP - # enumerates 0 devices, and llama.cpp falls back to - # CPU ("ggml_cuda_init: no ROCm-capable device is - # detected"). The HIP mask alone narrows correctly; - # clear any inherited ROCR mask so it can't double up. - env.pop("ROCR_VISIBLE_DEVICES", None) - except Exception as e: - logger.debug("Failed to set ROCm visibility env vars for child: %s", e) + # A deliberate zero-offload load with no GPU companions runs + # entirely on CPU, yet a visible CUDA device still costs the child + # ~0.5 GB (context + compute scratch) that the CPU-only + # classification below reports as free. Hide the GPUs so the load + # is exactly what it claims: zero VRAM (verified: GPU stays at idle + # baseline and generation runs). Companion loads keep the normal + # masking, and a user device pin (in extras or an inherited + # LLAMA_ARG_DEVICE) keeps control of its own devices -- the child + # aborts on a pin it can't see. The draft-device forms count too: + # llama-server parses them even with no drafter loaded. + _cpu_only_zero_offload = ( + gpu_memory_mode == "manual" + and gpu_layers == 0 + and not is_vulkan_backend + and not self._zero_offload_keeps_gpu_visible(cmd, env) + ) + if _cpu_only_zero_offload: + self._emit_child_gpu_visibility(env, "-1") + elif gpu_indices is not None and not is_vulkan_backend: + # When the user picked GPUs by index, align CUDA's ordering + # with the PCI-bus order the picker enumerated (nvidia-smi), + # so "GPU 1" in the UI is GPU 1 to llama.cpp -- not CUDA's + # default FASTEST_FIRST order (#5025). + if gpu_ids: + env["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + self._emit_child_gpu_visibility(env, ",".join(str(i) for i in gpu_indices)) + elif manual_tensor_split_emitted and not is_vulkan_backend: + # A manual per-GPU ratio across ALL GPUs (no explicit pick, so + # no CUDA_VISIBLE_DEVICES mask above): the UI built the + # --tensor-split list in ascending physical/PCI index order, + # so pin the child's enumeration to that order too. The whole + # visible set stays in use; only its ordering is fixed. + self._pin_visible_gpu_order_for_split(env) # Captured before any text-only fallback strips it from cmd. launched_with_mmproj = "--mmproj" in cmd @@ -7350,7 +7805,6 @@ class LlamaCppBackend: self._effective_context_length = ( effective_ctx if effective_ctx > 0 else self._context_length ) - self._reconcile_effective_ctx_with_server() self._max_context_length = ( max_available_ctx if max_available_ctx > 0 else self._effective_context_length ) @@ -7535,6 +7989,10 @@ class LlamaCppBackend: "session; run 'unsloth studio update' to enable vision." ) cmd = self._strip_mmproj_args(_last_spawn_cmd) + # This retry bypasses _spawn_and_wait, so refresh the + # launched-argv snapshot itself -- the zero-offload + # classification below must not see the stripped --mmproj. + _last_spawn_cmd = list(cmd) self._is_vision = False self._mmproj_has_audio = False self._start_llama_process(cmd, env) @@ -7566,6 +8024,13 @@ class LlamaCppBackend: self._healthy = True self._commit_effective_parallel_slots(n_parallel) + # Server is up: adopt the real per-request context it allocated + # -- the length --fit chose, or a --parallel slot split -- so the + # reported context_length matches reality. (Querying /props + # before the spawn above always failed; the seeded value was the + # requested/native length.) + self._reconcile_effective_ctx_with_server() + # Commit caller intent only after _healthy=True so a failed start # can't poison the next inheritance check. None keeps prior, [] # clears, list sets. Source records hf_variant for the route's @@ -7580,11 +8045,24 @@ class LlamaCppBackend: self._mtp_runtime_fallback_active = _mtp_active_for_launched_server self._start_mtp_crash_watchdog() - # 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: + # Catch silent CPU fallback when GPU was intended (#5106). Manual + # offload (no picker) leaves gpu_indices None and use_fit False, so + # include its GPU-layer intent; use the preserved probe since + # auto-layers/manual empty `gpus`. A deliberate zero-offload load + # classifies by its launched argv instead: the main model is + # CPU-only by construction and must read False (not None), or + # training needlessly unloads a server holding no VRAM. + _deliberate_cpu_only = gpu_memory_mode == "manual" and gpu_layers == 0 + if _deliberate_cpu_only: + self._gpu_offload_active = self._zero_offload_gpu_flag( + _last_spawn_cmd, _detected_gpus, env + ) + else: + self._gpu_offload_active = self._classify_gpu_offload( + gpu_indices is not None or use_fit or gpu_memory_mode == "manual", + _detected_gpus, + ) + if self._gpu_offload_active is False and not _deliberate_cpu_only: logger.warning( "llama-server appears to have loaded the model entirely " "on CPU even though Unsloth detected at least one GPU. " @@ -7947,6 +8425,11 @@ class LlamaCppBackend: gguf_path: Optional[str] = None, spec_draft_n_max: Optional[int] = None, tensor_parallel: bool = False, + gpu_memory_mode: Literal["auto", "manual"] = "auto", + gpu_layers: int = -1, + n_cpu_moe: int = 0, + tensor_split: Optional[List[float]] = None, + gpu_ids: Optional[List[int]] = None, mtp_draft_path: Optional[str] = None, preserve_multi_gpu_on_layer: bool = False, ) -> bool: @@ -8003,6 +8486,38 @@ class LlamaCppBackend: ): return False + # The diffusion runner is mode-agnostic (always "auto", ignores the + # layer/MoE/split knobs), so a standing manual preference in the + # request must not force a needless reload -- only the GPU pick matters. + if not self._is_diffusion: + # A GPU-memory-mode flip (Unsloth / manual) must always reload. + if self._gpu_memory_mode != gpu_memory_mode: + return False + # Manual: a layer-count change always reloads (covers Auto(-1) <-> a + # pinned count); MoE/split only matter with an explicit offload. + if gpu_memory_mode == "manual" and ( + self._gpu_layers != gpu_layers + or ( + gpu_layers >= 0 + and ( + self._n_cpu_moe != n_cpu_moe + or (self._tensor_split or None) != (tensor_split or None) + ) + ) + ): + return False + # A changed GPU pick must reload (compare order-insensitively; None/[] + # both mean automatic). The diffusion runner collapses a multi-GPU pick + # to its single lowest device, so self._gpu_ids holds just that device; + # normalize the request the same way, or a multi-GPU pick that resolves + # to the same device needlessly reloads. + if self._is_diffusion: + requested_gpu_pick = [sorted(gpu_ids)[0]] if gpu_ids else None + else: + requested_gpu_pick = sorted(gpu_ids) if gpu_ids else None + if (self._gpu_ids or None) != requested_gpu_pick: + return False + # Compare on the canonical requested mode. With --spec-type in # extra_args the backend stores None; mirror that here. if _extra_args_set_spec_type(extra_args): @@ -8071,6 +8586,78 @@ class LlamaCppBackend: return None return classify_gpu_offload_lines(self._stdout_lines) + @staticmethod + def _cmd_has_gpu_companion(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + """True when the argv/env carries a GPU companion: any --mmproj form, or + a drafter (Studio's --model-draft, the extras aliases, or the + LLAMA_ARG_SPEC_DRAFT_* env) -- these offload to the GPU regardless of + the main ``--gpu-layers``. A drafter explicitly forced to CPU + (--spec-draft-ngl 0 / --spec-draft-device cpu) doesn't count.""" + if any(str(a).startswith("--mmproj") for a in cmd): + return True + if _extra_args_mtp_draft_path(cmd, env) is None: + return False + return not _extra_args_draft_offloaded_to_cpu(cmd, env) + + @staticmethod + def _zero_offload_keeps_gpu_visible(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + """Whether a zero-layer launch still has a reason to use visible GPUs. + + Keep this shared by child masking and post-launch residency bookkeeping: + a device pin, surviving tensor mode, mmproj, or GPU drafter prevents the + launch from being a confirmed zero-VRAM server. + """ + return ( + LlamaCppBackend._cmd_has_gpu_device_pin(cmd, env) + or _effective_tensor_parallel(cmd, False, env) + or LlamaCppBackend._cmd_has_gpu_companion(cmd, env) + ) + + @staticmethod + def _cmd_has_gpu_device_pin(cmd: list, env: Optional[Mapping[str, str]] = None) -> bool: + """True when the effective main or draft ``--device`` pin names a GPU.""" + main_flags = {"--device", "-dev"} + draft_flags = {"--spec-draft-device", "-devd", "--device-draft"} + last_main: Optional[str] = None + last_draft: Optional[str] = None + args = [str(arg) for arg in cmd] + for index, raw in enumerate(args): + flag, equals, inline = raw.partition("=") + if flag not in main_flags and flag not in draft_flags: + continue + value = inline if equals else (args[index + 1] if index + 1 < len(args) else "") + if flag in main_flags: + last_main = value + else: + last_draft = value + if last_main is None: + last_main = (env or {}).get("LLAMA_ARG_DEVICE") + + def _names_gpu(value: Optional[str]) -> bool: + if value is None: + return False + devices = [item.strip().lower() for item in value.split(",") if item.strip()] + return not devices or any(item not in ("cpu", "none") for item in devices) + + return _names_gpu(last_main) or _names_gpu(last_draft) + + @staticmethod + def _zero_offload_gpu_flag( + spawn_cmd: list, + detected_gpus: list, + env: Optional[Mapping[str, str]] = None, + ) -> Optional[bool]: + """GPU-residency flag for a deliberate manual zero-offload load. The + main model is CPU-only by construction, but device pins, tensor mode, + mmproj, and GPU drafters can still make the server hold VRAM. The counted + offload classifier cannot see those allocations. This uses the same + predicate as the launch-time zero-VRAM mask; None means no GPU signal.""" + if not detected_gpus: + return None + if LlamaCppBackend._is_vulkan_backend(): + return True + return LlamaCppBackend._zero_offload_keeps_gpu_visible(spawn_cmd, env) + def load_cancelled(self) -> bool: """True if a load was cancelled (e.g. via unload/_cancel_event) and not yet consumed by the next load_model. Lets the tensor->layer fallback @@ -8114,11 +8701,18 @@ class LlamaCppBackend: self._supports_tools = False self._cache_type_kv = None self._tensor_parallel = False + self._gpu_memory_mode = "auto" + self._gpu_layers = -1 + self._n_cpu_moe = 0 + self._tensor_split = None + self._gpu_ids = None self._layer_preserves_tensor_intent = False self._speculative_type = None self._requested_spec_mode = None self._spec_draft_n_max = None self._n_layers = None + self._n_experts = None + self._leading_dense_block_count = None self._n_kv_heads = None self._n_kv_heads_by_layer = None self._n_heads = None @@ -8181,6 +8775,10 @@ class LlamaCppBackend: # Clear healthy so a /load during the replacement's warm-up can't # short-circuit against the previous server's health (#5401). self._healthy = False + # Reset to unknown so the training guard treats the next (still + # loading) server as VRAM-resident rather than reading the killed + # server's stale zero-offload flag until the health probe reclassifies. + self._gpu_offload_active = None # Drives _wait_for_vram_settle in the next load_model; set in finally # so both in-process and frontend Apply paths record the kill. self._last_kill_monotonic = time.monotonic() @@ -8785,7 +9383,12 @@ class LlamaCppBackend: @staticmethod def _ctx_integrity_flags( - n_parallel: int, use_fit: bool, requested_ctx: int, effective_ctx: int, caps: dict + n_parallel: int, + use_fit: bool, + auto_fit: bool, + requested_ctx: int, + effective_ctx: int, + caps: dict, ) -> list[str]: """Flags that keep the per-request window equal to the advertised ctx. @@ -8793,14 +9396,28 @@ class LlamaCppBackend: ``--kv-unified`` default, silently splitting ``-c`` into per-slot windows of ``-c / N``; restore the shared pool so one request can use the full context. With ``--fit on``, ``--fit-ctx`` floors the fit step - at an explicitly requested ctx (default floor is 4096) so it offloads - or fails instead of silently shrinking the window. + at an explicitly requested ctx so it offloads or fails instead of + silently shrinking the window. The 8192 auto-floor and the tighter + ``--fit-target`` margin apply only under Manual + Auto (``auto_fit``), + which omits ``-c``: on the legacy auto path ``-c 0`` already pins the + native window and ``--fit-ctx 8192`` would override it down to 8192. """ flags: list[str] = [] if n_parallel > 1 and caps.get("supports_kv_unified"): flags.append("--kv-unified") - if use_fit and requested_ctx > 0 and effective_ctx > 0 and caps.get("supports_fit_ctx"): - flags.extend(["--fit-ctx", str(effective_ctx)]) + if use_fit and caps.get("supports_fit_ctx"): + if requested_ctx > 0 and effective_ctx > 0: + # Floor the fit step at the explicitly requested ctx. + flags.extend(["--fit-ctx", str(effective_ctx)]) + elif auto_fit: + # Manual + Auto omits -c, so floor at 8192 so --fit doesn't + # shrink the window below a usable size. + flags.extend(["--fit-ctx", "8192"]) + if use_fit and auto_fit and caps.get("supports_fit_target"): + # llama.cpp's --fit leaves 1 GiB free per device by default; + # tighten that to 512 MiB so it packs more of the model onto + # the GPU before spilling to system RAM. + flags.extend(["--fit-target", "512"]) return flags def _query_server_n_ctx(self) -> Optional[int]: diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 70d0dc774d..e72e10e071 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -186,12 +186,25 @@ _SPLIT_MODE_FLAGS: frozenset[str] = frozenset({"-sm", "--split-mode"}) _TENSOR_SPLIT_FLAGS: frozenset[str] = frozenset({"-ts", "--tensor-split"}) _SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS +# GPU-offload flags. Stripped only when the GPU Memory mode owns offload +# (manual emits --fit / --gpu-layers / --n-cpu-moe); in auto, a user's +# inherited -ngl is respected (the offload_overridden path), so this group is +# opt-in, not default. Layer flags are shared with llama_cpp's override +# detection; the MoE flags are strip-only (manual's --n-cpu-moe slider owns them). +_LAYER_OFFLOAD_FLAGS: frozenset[str] = frozenset( + {"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"} +) +_MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset({"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"}) +_OFFLOAD_SHADOWING_FLAGS: frozenset[str] = _LAYER_OFFLOAD_FLAGS | _MOE_OFFLOAD_FLAGS + _SHADOWING_FLAGS: frozenset[str] = ( _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS ) # Shadowing flags that take no value -- strip the flag only, not the next token. -_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"}) +_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset( + {"--spec-default", "--jinja", "--no-jinja", "-cmoe", "--cpu-moe"} +) def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]: @@ -424,6 +437,8 @@ def strip_shadowing_flags( strip_spec: bool = True, strip_template: bool = True, strip_split_mode: bool = True, + strip_tensor_split: bool = False, + strip_offload: bool = False, ) -> list[str]: """Strip flags that shadow first-class Unsloth settings. @@ -432,6 +447,12 @@ def strip_shadowing_flags( (same for cache / spec / template / split-mode). Each ``strip_*`` toggle controls one group; the route only strips groups whose first-class field the caller actually supplied. + + ``strip_split_mode`` removes both ``--split-mode`` and the coupled + ``--tensor-split`` (the Tensor Parallelism toggle owns the whole split). + ``strip_tensor_split`` removes ``--tensor-split`` *alone*, so manual mode can + replace an inherited per-GPU ratio while leaving the user's ``--split-mode`` + row/none/layer choice intact. """ shadowing: set[str] = set() if strip_context: @@ -444,6 +465,10 @@ def strip_shadowing_flags( shadowing |= _TEMPLATE_FLAGS if strip_split_mode: shadowing |= _SPLIT_SHADOWING_FLAGS + if strip_tensor_split: + shadowing |= _TENSOR_SPLIT_FLAGS + if strip_offload: + shadowing |= _OFFLOAD_SHADOWING_FLAGS tokens = [str(a) for a in (args or [])] out: list[str] = [] diff --git a/studio/backend/main.py b/studio/backend/main.py index 4797764ce7..81d4c16e52 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -1156,9 +1156,23 @@ def _get_cached_system_gpu_info(logger) -> dict[str, Any]: enriched_dev["vram_utilization_pct"] = util.get("vram_utilization_pct") enriched_devices.append(enriched_dev) + # Whether GGUF loads accept an explicit gpu_ids pick: /load and + # /validate 400 picks on XPU hosts (no visibility mask speaks torch-xpu + # ordinals) and on Vulkan-only builds (--device pins ggml's own + # ordinals), so the picker must not offer them. + try: + from core.inference.llama_cpp import LlamaCppBackend + from utils.hardware import DeviceType, get_device + gpu_ids_supported = ( + get_device() != DeviceType.XPU and not LlamaCppBackend._is_vulkan_backend() + ) + except Exception as e: + logger.debug(f"Could not resolve gpu_ids support: {e}") + gpu_ids_supported = True gpu_info = { "available": visibility_info.get("available", False), "devices": enriched_devices, + "gguf_gpu_ids_supported": gpu_ids_supported, } _system_gpu_cache = (time.monotonic(), gpu_info) return gpu_info diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index f3ae0f70df..d51d35189b 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -64,7 +64,7 @@ class LoadRequest(BaseModel): ) gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.", + description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. For GGUF models the picked devices are pinned via CUDA/HIP_VISIBLE_DEVICES.", ) speculative_type: Optional[str] = Field( None, @@ -100,6 +100,66 @@ class LoadRequest(BaseModel): "No effect on a single GPU. Ignored for non-GGUF models." ), ) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = ( + "GPU memory strategy for GGUF models. 'auto' (default): Unsloth " + "selects GPUs and caps context to fit VRAM. 'manual': you own the " + "offload. Leave gpu_layers at -1 (Auto) to hand memory management to " + "llama.cpp's --fit (no device masking, no context auto-reduce, no " + "gpu-layer/tensor-split planning); set gpu_layers >= 0 to pin layers " + "and n_cpu_moe yourself (--fit off), with tensor_parallel still " + "applying (split by free VRAM unless tensor_split is set, no planner). " + "Ignored for non-GGUF." + ), + ) + gpu_layers: int = Field( + -1, + ge = -1, + description = ( + "Manual mode only: number of layers to offload to the GPU " + "(--gpu-layers, with --fit off). A value >= the model's layer count " + "offloads all of them. -1 = Auto: hand layer + context sizing to " + "llama.cpp's --fit. Ignored unless gpu_memory_mode is 'manual'." + ), + ) + n_cpu_moe: int = Field( + 0, + ge = 0, + description = ( + "Manual mode only: keep the first N MoE expert layers on the CPU " + "(--n-cpu-moe) to save VRAM on MoE models. 0 = none, N = number of " + "MoE layers offloaded (the backend offsets past any leading dense " + "layers). Ignored unless gpu_memory_mode is 'manual' with gpu_layers >= 0." + ), + ) + tensor_split: Optional[List[float]] = Field( + None, + description = ( + "Manual mode only: relative share of the model per GPU (--tensor-split), " + "in the order of the GPUs in use, e.g. [2, 1] for 2:1. Omit it to let " + "llama.cpp use its default, which splits by free VRAM. Any list given is " + "passed through as-is, so send [1, 1] to force an even split. Ignored " + "unless gpu_memory_mode is 'manual' with gpu_layers >= 0." + ), + ) + + @field_validator("tensor_split") + @classmethod + def _reject_degenerate_tensor_split(cls, value: Optional[List[float]]) -> Optional[List[float]]: + # A negative / non-finite / all-zero split is silently dropped at launch + # (stored as None) yet still compared raw in the reload dedupe, so an + # identical Apply reloads forever. Reject it up front; [] = no split. + if not value: + return value + import math + + if any((not math.isfinite(v)) or v < 0 for v in value): + raise ValueError("tensor_split entries must be finite and non-negative") + if sum(value) <= 0: + raise ValueError("tensor_split must have a positive total") + return value + llama_extra_args: Optional[List[str]] = Field( None, description = ( @@ -133,6 +193,14 @@ class ValidateModelRequest(BaseModel): max_seq_length: int = Field(0, ge = 0, le = 1048576) load_in_4bit: bool = Field(True) gpu_ids: Optional[List[int]] = Field(None) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = ( + "GGUF GPU-memory strategy intended for the follow-up load. Manual " + "placement bypasses the training coexistence estimate: Auto layers " + "delegate fitting to llama.cpp, while explicit layers are user-owned." + ), + ) include_context_length: bool = Field( False, description = "Also read the native context length from the local GGUF header. " @@ -188,6 +256,16 @@ class ValidateModelResponse(BaseModel): description = "Native training context length, read from the GGUF header when the file " "is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.", ) + layer_count: Optional[int] = Field( + None, + description = "Total layer count (GGUF block_count), the manual gpu-layers ceiling, read " + "from the header alongside context_length; None when not read.", + ) + moe_layer_count: Optional[int] = Field( + None, + description = "MoE expert-layer count (the manual --n-cpu-moe ceiling), read from the GGUF " + "header alongside context_length; 0 for dense models, None when not read.", + ) # Additive fields; the consuming consent dialog ships in a follow-up frontend PR. requires_transformers_upgrade: bool = Field( False, @@ -333,6 +411,34 @@ class LoadResponse(BaseModel): False, description = "Whether tensor-parallel split (--split-mode tensor) is active.", ) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = "Active GPU memory strategy ('auto' or 'manual').", + ) + gpu_layers: int = Field( + -1, + description = "Manual mode: requested --gpu-layers value (-1 = Auto/--fit, or when not manual).", + ) + n_cpu_moe: int = Field( + 0, + description = "Manual mode: MoE expert layers pinned to CPU (--n-cpu-moe); 0 = none.", + ) + tensor_split: Optional[List[float]] = Field( + None, + description = "Manual mode: relative model share per GPU (--tensor-split); None = default (split by free VRAM).", + ) + n_layers: Optional[int] = Field( + None, + description = "Model's layer count (GGUF block_count), for the manual gpu-layers ceiling.", + ) + n_moe_layers: int = Field( + 0, + description = "Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not an MoE model.", + ) + gpu_ids: Optional[List[int]] = Field( + None, + description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + ) class UnloadResponse(BaseModel): @@ -461,6 +567,42 @@ class InferenceStatusResponse(BaseModel): False, description = "Whether tensor-parallel split (--split-mode tensor) is active.", ) + gpu_memory_mode: Literal["auto", "manual"] = Field( + "auto", + description = "Active GPU memory strategy ('auto' or 'manual').", + ) + gpu_layers: int = Field( + -1, + description = "Manual mode: requested --gpu-layers value (-1 = Auto/--fit, or when not manual).", + ) + n_cpu_moe: int = Field( + 0, + description = "Manual mode: MoE expert layers pinned to CPU (--n-cpu-moe); 0 = none.", + ) + tensor_split: Optional[List[float]] = Field( + None, + description = "Manual mode: relative model share per GPU (--tensor-split); None = default (split by free VRAM).", + ) + requested_context_length: Optional[int] = Field( + None, + description = ( + "The n_ctx the active GGUF load was invoked with (0 = Auto). Lets the " + "UI re-seed a Manual + Auto-layers context pin on hydration, where " + "context_length only exposes the resolved value. None for non-GGUF." + ), + ) + n_layers: Optional[int] = Field( + None, + description = "Model's layer count (GGUF block_count), for the manual gpu-layers ceiling.", + ) + n_moe_layers: int = Field( + 0, + description = "Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not an MoE model.", + ) + gpu_ids: Optional[List[int]] = Field( + None, + description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + ) llama_cpp_supports_mtp: bool = Field( True, description = ( diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3d527bf317..136e4f7645 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -13,7 +13,7 @@ from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.responses import StreamingResponse, JSONResponse, Response from starlette.requests import ClientDisconnect -from typing import Any, Callable, List, Optional, Union +from typing import Any, Callable, List, Literal, Optional, Union import json import httpx from loggers import get_logger @@ -3115,13 +3115,16 @@ def _normalise_settings_str(value: Optional[str]) -> Optional[str]: def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[str]]) -> bool: - """Whether an inherited --split-mode should be stripped on reload. + """Whether an inherited --split-mode (and its coupled --tensor-split) should + be stripped on reload. The binary Tensor Parallelism toggle can't carry --split-mode's row/none/ layer modes, so only strip when the toggle overrides it: tensor being turned on, or the inherited mode is tensor (toggle turning it off). Non-tensor modes - survive. Shared by the inheritance strip and the already-loaded stale check - so they agree on what reload would do. + survive. A manual per-GPU ratio is handled by _should_strip_tensor_split, + which strips only --tensor-split so the inherited mode is kept. Shared by the + inheritance strip and the already-loaded stale check so they agree on what + reload would do. """ fields_set = getattr(request, "model_fields_set", set()) return "tensor_parallel" in fields_set and ( @@ -3129,6 +3132,25 @@ def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[ ) +def _should_strip_tensor_split(request: LoadRequest) -> bool: + """Whether an inherited --tensor-split alone should be stripped on reload. + + Manual explicit offload (gpu_layers >= 0) owns the per-GPU split: with a ratio + it emits its own --tensor-split (an inherited one, appended last, would + override it), and with the ratio cleared it wants llama.cpp's default + free-VRAM split. Either way an inherited --tensor-split must go, else the + cleared case silently keeps the stale ratio while status reports None. + Unlike _should_strip_split_mode this leaves --split-mode untouched, so a + user's row/none/layer mode survives a Studio split-ratio edit. When the + Tensor Parallelism toggle IS overriding the mode, _should_strip_split_mode + (called alongside this at every site) strips --split-mode anyway. + """ + return ( + getattr(request, "gpu_memory_mode", "auto") == "manual" + and getattr(request, "gpu_layers", -1) >= 0 + ) + + def _carry_preserved_tensor_intent( *, preserved: bool, same_model: bool, explicit_drop: bool ) -> bool: @@ -3187,12 +3209,44 @@ def _request_matches_loaded_settings( else strip_shadowing_flags( backend_extra, strip_split_mode = _should_strip_split_mode(request, backend_extra), + strip_tensor_split = _should_strip_tensor_split(request), + strip_offload = request.gpu_memory_mode == "manual", ) ) if not _tensor_parallel_matches_loaded( effective_extra, request.tensor_parallel, llama_backend.tensor_parallel ): return False + # The diffusion runner is mode-agnostic (it always reports "auto" and ignores + # the layer/MoE/split knobs), so a standing manual preference in the request + # must not force a needless reload -- only the GPU pick matters. + if not llama_backend.is_diffusion: + if request.gpu_memory_mode != llama_backend.gpu_memory_mode: + return False + # Manual: a layer-count change always reloads; MoE/split only matter with + # an explicit offload (gpu_layers >= 0), so a leftover value under Auto + # must not force one. Mirrors LlamaCppBackend._already_in_target_state. + if request.gpu_memory_mode == "manual" and ( + request.gpu_layers != llama_backend.gpu_layers + or ( + request.gpu_layers >= 0 + and ( + request.n_cpu_moe != llama_backend.n_cpu_moe + or (request.tensor_split or None) != (llama_backend.tensor_split or None) + ) + ) + ): + return False + # A changed GPU pick must reload. The diffusion runner collapses a multi-GPU + # request to its single lowest device (it drives one device only), so the + # backend records just that device; compare the request the same way, or a + # multi-GPU pick that resolves to the same device needlessly reloads. + if llama_backend.is_diffusion: + _req_gpu_ids = [sorted(request.gpu_ids)[0]] if request.gpu_ids else None + else: + _req_gpu_ids = sorted(request.gpu_ids) if request.gpu_ids else None + if _req_gpu_ids != llama_backend.gpu_ids: + return False # Preserved tensor->layer fallback (both report tensor=off, so the check above # matches): if the user now explicitly drops tensor intent, reload so placement # re-selects instead of keeping the all-GPU mask (#6659). The effective check @@ -3235,14 +3289,17 @@ def _request_matches_loaded_settings( # contain any shadow flag, so the reload path strips them rather than # leaving a stale override in effect. (backend_extra computed above.) if request.llama_extra_args is None: - # Mirror the reload's conditional split-mode strip, so a preserved - # non-tensor mode (row/none/layer) isn't seen as stale and doesn't - # trigger a needless reload of a healthy server. + # Mirror the reload's conditional strips, so a preserved non-tensor mode + # (row/none/layer) isn't seen as stale and doesn't trigger a needless + # reload of a healthy server, while an inherited offload/ratio flag that + # the reload *would* strip is correctly seen as stale. if ( backend_extra and strip_shadowing_flags( backend_extra, strip_split_mode = _should_strip_split_mode(request, backend_extra), + strip_tensor_split = _should_strip_tensor_split(request), + strip_offload = request.gpu_memory_mode == "manual", ) != backend_extra ): @@ -3861,6 +3918,46 @@ def _estimate_gguf_required_gb( return None +def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]: + """Classify a GGUF as diffusion, normal, or unknown before it is loaded. + + ``None`` is important here: a remote GGUF whose header is not cached can + still be routed to the single-GPU diffusion runner after download. Treating + that case as normal would let Manual mode skip the training guard even + though the runner ignores Manual's llama-server placement controls. + """ + identity = " ".join( + str(getattr(config, attr, "") or "") for attr in ("identifier", "gguf_hf_repo", "gguf_file") + ).lower() + if "diffusion" in identity: + return True + + try: + main = getattr(config, "gguf_file", None) + if not (main and Path(main).is_file()): + repo = getattr(config, "gguf_hf_repo", None) + variant = getattr(config, "gguf_variant", None) + if repo and variant: + from hub.utils.gguf import resolve_local_gguf_path + main = resolve_local_gguf_path(repo, variant) + if not main or not Path(main).is_file(): + return None + + probe = LlamaCppBackend() + probe._read_gguf_metadata(str(main)) + if probe.is_diffusion: + return True + # A successfully decoded architecture proves that this is a normal + # llama-server GGUF. No architecture means the lightweight probe could + # not establish the routing decision, so preserve the unknown state. + if getattr(probe, "_architecture", None): + return False + return None + except Exception as e: + logger.debug("Could not identify diffusion GGUF for training guard: %s", e) + return None + + def _guard_chat_load_against_training( config: ModelConfig, *, @@ -3871,11 +3968,19 @@ def _guard_chat_load_against_training( requested_gpu_ids: Optional[List[int]], llama_extra_args: Optional[list[str]] = None, n_parallel: int = 1, + gpu_memory_mode: Literal["auto", "manual"] = "auto", ) -> None: - """Refuse loading a local chat model that would OOM an active training run. + """Protect active training from automatically placed chat-model loads. + No-op when training is inactive or unknown. `load_in_4bit` must be the - effective quantization (see _effective_load_in_4bit). Raises HTTP 409 when the - model would not fit alongside training.""" + effective quantization (see _effective_load_in_4bit). Manual chat-GGUF + placement is an explicit override: Auto layers delegate fitting to + llama.cpp's ``--fit`` and pinned layers are owned by the user, so neither is + estimated here. Diffusion is still guarded because its mode-agnostic runner + ignores those controls and uses one GPU. An unclassified GGUF is guarded as + potentially diffusion until its local header proves otherwise. Other loads + raise HTTP 409 when they would not fit beside training. + """ from core.training import get_training_backend from routes.training_vram import can_load_chat_during_training @@ -3887,6 +3992,19 @@ def _guard_chat_load_against_training( return is_gguf = bool(getattr(config, "is_gguf", False)) + diffusion_kind = _classify_diffusion_gguf(config) if is_gguf else False + if is_gguf and gpu_memory_mode == "manual" and diffusion_kind is False: + return + + diffusion_gpu = None + if is_gguf and diffusion_kind is not False: + # Use the same token selection as the runner: an explicit pick wins, + # followed by DG_GPU, the first parent-visible token, then GPU 0. + diffusion_gpu = LlamaCppBackend._diffusion_gpu_arg( + requested_gpu_ids, + cpu_only = LlamaCppBackend._effective_gpu_count() == 0, + ) + required_override_gb = ( _estimate_gguf_required_gb( config, @@ -3907,6 +4025,7 @@ def _guard_chat_load_against_training( requested_gpu_ids = requested_gpu_ids, is_gguf = is_gguf, required_override_gb = required_override_gb, + single_device_gpu = diffusion_gpu, ) if ok: return @@ -3934,6 +4053,98 @@ def _guard_chat_load_against_training( raise HTTPException(status_code = 409, detail = detail) +def _resolve_inherited_extra_args( + request, + config: ModelConfig, + model_identifier: str, + extra_llama_args: Optional[list[str]], + effective_chat_template_override: Optional[str] = None, +) -> Optional[list[str]]: + """Effective pass-through extras for a GGUF request that omitted the field: + the previous same-model load's extras, shadow-stripped, so a settings-Apply + reload (which does not round-trip the extras field) keeps them (#5401).""" + if getattr(request, "llama_extra_args", None) is not None: + return extra_llama_args + if not getattr(config, "is_gguf", False): + return extra_llama_args + llama_backend = get_llama_cpp_backend() + if not llama_backend.extra_args: + return extra_llama_args + # Inherit the previous load's extras (the chat-settings Apply path doesn't + # round-trip them; an explicit [] still clears). Gated on (model_identifier, + # hf_variant) to refuse cross-model pickup, and shadowing flags are + # stripped so an inherited override can't win the last-wins CLI + # parse against a freshly-supplied first-class field. + source = llama_backend.extra_args_source + # Compare against the resolved variant, not the request field: callers + # commonly omit gguf_variant for local ``.gguf`` paths and HF auto-pick + # flows. ``config.gguf_variant`` is the variant load_model was actually + # invoked with, so both sides of the comparison key off the same string. + resolved_variant = (config.gguf_variant or "").lower() + request_variant = (request.gguf_variant or "").lower() + stored_variant = (source[1] or "").lower() if source else "" + same_model = bool(source and source[0] and source[0].lower() == model_identifier.lower()) + if request.gguf_variant: + variant_mismatch = request_variant != stored_variant + else: + variant_mismatch = bool(stored_variant and resolved_variant != stored_variant) + same_source = same_model and not variant_mismatch + if not same_source: + logger.info( + "Not inheriting llama_extra_args: stored args came from %s, loading %s", + source, + (model_identifier, resolved_variant), + ) + # Cross-model: clear explicitly so the backend doesn't + # inherit via "no opinion" semantics. + extra_llama_args = [] + else: + # Strip only the groups whose first-class field was set by the caller, so + # an inherited --chat-template-file survives an Apply that omits + # chat_template_override. A bundled family template (e.g. gemma-4) counts as + # a first-class template even when the request omits chat_template_override, + # so strip the inherited --chat-template-file then too -- else the stale arg + # (appended last) shadows the bundled template while Studio reports its caps. + fields_set = getattr(request, "model_fields_set", set()) + stripped = strip_shadowing_flags( + llama_backend.extra_args, + strip_context = "max_seq_length" in fields_set, + strip_cache = "cache_type_kv" in fields_set, + strip_spec = ("speculative_type" in fields_set or "spec_draft_n_max" in fields_set), + strip_template = ( + "chat_template_override" in fields_set + or effective_chat_template_override is not None + ), + strip_split_mode = _should_strip_split_mode(request, llama_backend.extra_args), + # manual + per-GPU ratio emits its own --tensor-split; drop + # an inherited one (appended last would override it) while + # keeping the user's --split-mode row/none/layer choice. + strip_tensor_split = _should_strip_tensor_split(request), + # manual emits its own --fit/--gpu-layers, so an inherited offload flag + # must not last-wins-override it. auto leaves a user's inherited -ngl + # alone. getattr: a validate request reuses this resolver, no offload fields. + strip_offload = getattr(request, "gpu_memory_mode", "auto") == "manual", + ) + try: + extra_llama_args = validate_extra_args(stripped) + except ValueError: + # Shouldn't happen on already-validated args; degrade to + # no-extras rather than 400 if managed flags changed. + logger.warning( + "Stored llama_extra_args failed revalidation; loading without them: %s", + stripped, + ) + extra_llama_args = [] + else: + if extra_llama_args: + logger.info( + "Inheriting llama_extra_args from previous " + "load (same model, shadow-stripped): %s", + extra_llama_args, + ) + return extra_llama_args + + def _model_json_response(model, status_code: int = 200) -> Response: """Serialize a pydantic response once via pydantic-core. @@ -4040,6 +4251,35 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre None if request.llama_extra_args is None else extra_llama_args ) + # Manual mode owns the offload flags: strip them from EXPLICIT extras + # too (the inherited path already does), or a last-wins --gpu-layers / + # --fit in extras re-enables GPU offload on a load status reports as + # CPU-only. Manual + per-GPU ratio owns --tensor-split the same way. + if request.gpu_memory_mode == "manual" and extra_llama_args: + _stripped_explicit = strip_shadowing_flags( + extra_llama_args, + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + strip_tensor_split = _should_strip_tensor_split(request), + strip_offload = True, + ) + if _stripped_explicit != extra_llama_args: + logger.info( + "Manual GPU memory owns the offload flags; stripping them " + "from explicit llama_extra_args: %s -> %s", + extra_llama_args, + _stripped_explicit, + ) + extra_llama_args = _stripped_explicit + + # Keep every downstream consumer on the normalized explicit list. In + # particular, the already-loaded comparator must not compare the raw + # request's managed offload flags against the stripped launch state. + request = request.model_copy(update = {"llama_extra_args": extra_llama_args}) + model_identifier, model_log_label, native_grant_backed = ( _resolve_model_identifier_for_request(request, operation = "load-model") ) @@ -4121,6 +4361,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, tensor_parallel = llama_backend.tensor_parallel, + gpu_memory_mode = llama_backend.gpu_memory_mode, + gpu_layers = llama_backend.gpu_layers, + n_cpu_moe = llama_backend.n_cpu_moe, + tensor_split = llama_backend.tensor_split, + n_layers = llama_backend.n_layers, + n_moe_layers = llama_backend.n_moe_layers, + gpu_ids = llama_backend.gpu_ids, ) else: if ( @@ -4187,12 +4434,41 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre # Normalize gpu_ids: empty list means auto-selection, same as None effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # Reject GGUF + gpu_ids first so the guard can't mask it with a VRAM 409. + # GGUF supports gpu_ids: validate the pick up front (before the training + # guard) so a bad pick is a clean 400, not masked by a VRAM 409. Rejects + # negative / out-of-range / duplicate ids and UUID/MIG parents. XPU hosts + # are rejected outright: the picker's indices are torch-xpu ordinals neither + # applicator speaks (CUDA/HIP masks don't apply, the Vulkan --device pin + # uses ggml's own Vulkan ordinals), so a pick could land on the wrong device. if config.is_gguf and effective_gpu_ids is not None: - raise HTTPException( - status_code = 400, - detail = "gpu_ids is not supported for GGUF models yet.", - ) + from utils.hardware import DeviceType, get_device + from utils.hardware.hardware import resolve_requested_gpu_ids + + if get_device() == DeviceType.XPU: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported on Intel XPU. " + "Omit gpu_ids to use all devices." + ), + ) + # Same reasoning for a Vulkan-only build: --device pins ggml's own + # Vulkan ordinals, so a physical pick can land on the wrong card on + # masked or non-contiguous hosts. + if LlamaCppBackend._is_vulkan_backend(): + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported with a Vulkan " + "llama.cpp build: physical GPU ids have no defined " + "mapping to Vulkan device ordinals. Omit gpu_ids to use " + "all devices." + ), + ) + try: + resolve_requested_gpu_ids(effective_gpu_ids) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc if not config.is_gguf and _mlx_distributed_launch_detected(): raise HTTPException( status_code = 400, @@ -4222,8 +4498,20 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre "architectures)" ) - # Refuse a load that would OOM active training, before the unload step below - # frees the resident model. Off-loop: guard does sync nvidia-smi / HF work. + # Inherit the previous same-model load's pass-through extras when this + # request omits the field (a settings-Apply reload doesn't round-trip + # them); shadow-stripped so an inherited flag can't override a + # first-class field the caller did set (#5401). + extra_llama_args = _resolve_inherited_extra_args( + request, + config, + model_identifier, + extra_llama_args, + effective_chat_template_override, + ) + + # Apply the training coexistence policy before the unload step below + # frees the resident model. Off-loop: the default-mode guard does sync work. await asyncio.to_thread( _guard_chat_load_against_training, config, @@ -4234,6 +4522,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre requested_gpu_ids = effective_gpu_ids, llama_extra_args = extra_llama_args, n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1), + gpu_memory_mode = request.gpu_memory_mode, ) # ── GGUF path: load via llama-server ────────────────────── @@ -4245,84 +4534,6 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre from core.inference.llama_cpp import gguf_load_in_flight gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo)) - # Inherit llama_extra_args from the previous load when the request - # omits the field (the chat-settings Apply path doesn't round-trip - # them; explicit [] still clears). Gated on (model_identifier, - # hf_variant) to refuse cross-model pickup, and shadowing flags are - # stripped so an inherited override can't win the last-wins CLI - # parse against a freshly-supplied first-class field. - if request.llama_extra_args is None and llama_backend.extra_args: - source = llama_backend.extra_args_source - # Compare against the resolved variant, not the request - # field: callers commonly omit gguf_variant for local - # ``.gguf`` paths and HF auto-pick flows. ``config.gguf_ - # variant`` is the variant load_model was actually - # invoked with (see the HF / local branches below), so - # both sides of the comparison key off the same string. - resolved_variant = (config.gguf_variant or "").lower() - request_variant = (request.gguf_variant or "").lower() - stored_variant = (source[1] or "").lower() if source else "" - same_model = bool( - source and source[0] and source[0].lower() == model_identifier.lower() - ) - if request.gguf_variant: - variant_mismatch = request_variant != stored_variant - else: - variant_mismatch = bool(stored_variant and resolved_variant != stored_variant) - same_source = same_model and not variant_mismatch - if not same_source: - logger.info( - "Not inheriting llama_extra_args: stored args came from %s, loading %s", - source, - (model_identifier, resolved_variant), - ) - # Cross-model: clear explicitly so the backend doesn't - # inherit via "no opinion" semantics. - extra_llama_args = [] - else: - # Strip only the groups whose first-class field was set by - # the caller, so an inherited --chat-template-file survives - # an Apply that omits chat_template_override. A bundled family - # template (e.g. the gemma-4 override) is an effective - # first-class template setting even when the raw request - # omits chat_template_override, so strip the inherited - # --chat-template-file in that case too -- otherwise the stale - # extra arg (appended last) shadows the bundled template while - # Unsloth reports the bundled template's capabilities. - fields_set = getattr(request, "model_fields_set", set()) - stripped = strip_shadowing_flags( - llama_backend.extra_args, - strip_context = "max_seq_length" in fields_set, - strip_cache = "cache_type_kv" in fields_set, - strip_spec = ( - "speculative_type" in fields_set or "spec_draft_n_max" in fields_set - ), - strip_template = ( - "chat_template_override" in fields_set - or effective_chat_template_override is not None - ), - strip_split_mode = _should_strip_split_mode( - request, llama_backend.extra_args - ), - ) - try: - extra_llama_args = validate_extra_args(stripped) - except ValueError: - # Shouldn't happen on already-validated args; degrade to - # no-extras rather than 400 if managed flags changed. - logger.warning( - "Stored llama_extra_args failed revalidation; loading without them: %s", - stripped, - ) - extra_llama_args = [] - else: - if extra_llama_args: - logger.info( - "Inheriting llama_extra_args from previous " - "load (same model, shadow-stripped): %s", - extra_llama_args, - ) - # Block cache writes that would race the download manager. This runs # after pass-through argument inheritance so a carried --no-mmproj # changes the companion requirement exactly as it does for the load. @@ -4370,6 +4581,11 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre cache_type_kv = request.cache_type_kv, speculative_type = request.speculative_type, spec_draft_n_max = request.spec_draft_n_max, + gpu_memory_mode = request.gpu_memory_mode, + gpu_layers = request.gpu_layers, + n_cpu_moe = request.n_cpu_moe, + tensor_split = request.tensor_split, + gpu_ids = effective_gpu_ids, n_parallel = _n_parallel, ) if config.gguf_hf_repo: @@ -4537,6 +4753,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, tensor_parallel = llama_backend.tensor_parallel, + gpu_memory_mode = llama_backend.gpu_memory_mode, + gpu_layers = llama_backend.gpu_layers, + n_cpu_moe = llama_backend.n_cpu_moe, + tensor_split = llama_backend.tensor_split, + n_layers = llama_backend.n_layers, + n_moe_layers = llama_backend.n_moe_layers, + gpu_ids = llama_backend.gpu_ids, ) # ── Standard path: load via Unsloth/transformers ────────── @@ -4795,7 +5018,9 @@ def _requires_security_review_for_model( @router.post("/validate", response_model = ValidateModelResponse) async def validate_model( - request: ValidateModelRequest, current_subject: str = Depends(get_current_subject) + request: ValidateModelRequest, + fastapi_request: Request = None, + current_subject: str = Depends(get_current_subject), ): """ Lightweight validation endpoint for model identifiers. @@ -4823,15 +5048,39 @@ async def validate_model( detail = f"Invalid model identifier: {model_log_label}", ) - # Refuse early (before the frontend unloads to load this) if it can't fit - # alongside training, using the same settings /load uses so they agree. + # Apply the same training coexistence policy as /load before the frontend + # unloads the current model. effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # Mirror /load: reject GGUF + gpu_ids before the guard so both return 400. + # Mirror /load: GGUF supports gpu_ids, so validate the pick (a bad one is + # a clean 400) before the guard sizes the model against training VRAM. + # XPU-host picks are rejected like /load (no defined mapping from the + # picker's torch-xpu ordinals to the launcher's device spaces). if config.is_gguf and effective_gpu_ids is not None: - raise HTTPException( - status_code = 400, - detail = "gpu_ids is not supported for GGUF models yet.", - ) + from utils.hardware import DeviceType, get_device + from utils.hardware.hardware import resolve_requested_gpu_ids + + if get_device() == DeviceType.XPU: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported on Intel XPU. " + "Omit gpu_ids to use all devices." + ), + ) + if LlamaCppBackend._is_vulkan_backend(): + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported with a Vulkan " + "llama.cpp build: physical GPU ids have no defined " + "mapping to Vulkan device ordinals. Omit gpu_ids to use " + "all devices." + ), + ) + try: + resolve_requested_gpu_ids(effective_gpu_ids) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) # Both checks cover the [adapter, base] set (matching the scan route and workers): @@ -4895,16 +5144,32 @@ async def validate_model( latest_tier_active_for, config.identifier, request.hf_token ): effective_load_in_4bit = False - # Off-loop: guard does sync nvidia-smi / HF work. - await asyncio.to_thread( - _guard_chat_load_against_training, - config, - model_identifier = model_identifier, - hf_token = request.hf_token, - load_in_4bit = effective_load_in_4bit, - max_seq_length = request.max_seq_length, - requested_gpu_ids = effective_gpu_ids, - ) + # A metadata-only probe just reads the GGUF header and allocates no VRAM, + # so it must not be refused by the training guard. Real loads validate + # without include_context_length and /load applies the guard again. + if not request.include_context_length: + # Match /load's inherited llama.cpp extras and parallel slot count so + # validation cannot pass a smaller estimate than the subsequent load. + effective_extra_args = _resolve_inherited_extra_args( + request, config, model_identifier, None + ) + # Off-loop: guard does sync nvidia-smi / HF work. + await asyncio.to_thread( + _guard_chat_load_against_training, + config, + model_identifier = model_identifier, + hf_token = request.hf_token, + load_in_4bit = effective_load_in_4bit, + max_seq_length = request.max_seq_length, + requested_gpu_ids = effective_gpu_ids, + llama_extra_args = effective_extra_args, + n_parallel = ( + getattr(fastapi_request.app.state, "llama_parallel_slots", 1) + if fastapi_request is not None + else 1 + ), + gpu_memory_mode = request.gpu_memory_mode, + ) # A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a # mixed repo are inert for this load, so gating on them is a false positive. Only @@ -4918,10 +5183,15 @@ async def validate_model( # Native context length, read from the local GGUF header when present. # Lets the staged ("Load on selection" off) flow populate the context # slider before the GPU load; None until the file is downloaded. + # Staged header dims (one read): native context, total layer count, and + # MoE expert-layer count -- let the staged flow size the context, GPU- + # layers and manual --n-cpu-moe sliders before the load. context_length: Optional[int] = None + layer_count: Optional[int] = None + moe_layer_count: Optional[int] = None if request.include_context_length and is_gguf: from hub.utils.gguf import resolve_local_gguf_path - from utils.models.gguf_metadata import read_gguf_context_length + from utils.models.gguf_metadata import read_gguf_staged_dims # Best-effort: a header-read failure must never fail validation of an # otherwise-valid model (the outer except turns it into a 400). @@ -4937,9 +5207,15 @@ async def validate_model( model_identifier, request.gguf_variant ) if local_gguf: - context_length = read_gguf_context_length(local_gguf) + # Header walk reads tokenizer arrays for dense models (tens of + # ms); keep it off the event loop. + dims = await asyncio.to_thread(read_gguf_staged_dims, local_gguf) + if dims: + context_length = dims["context_length"] + layer_count = dims["layer_count"] + moe_layer_count = dims["moe_layer_count"] except Exception as e: - logger.debug("Context-length probe failed for %s: %s", model_log_label, e) + logger.debug("Header probe failed for %s: %s", model_log_label, e) return ValidateModelResponse( valid = True, @@ -4954,6 +5230,8 @@ async def validate_model( requires_trust_remote_code = requires_trust_remote_code, requires_security_review = requires_security_review, context_length = context_length, + layer_count = layer_count, + moe_layer_count = moe_layer_count, requires_transformers_upgrade = transformers_upgrade is not None, transformers_upgrade = transformers_upgrade, ) @@ -5593,6 +5871,14 @@ async def get_status(current_subject: str = Depends(get_current_subject)): speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, tensor_parallel = llama_backend.tensor_parallel, + gpu_memory_mode = llama_backend.gpu_memory_mode, + gpu_layers = llama_backend.gpu_layers, + n_cpu_moe = llama_backend.n_cpu_moe, + tensor_split = llama_backend.tensor_split, + requested_context_length = llama_backend.requested_n_ctx, + n_layers = llama_backend.n_layers, + n_moe_layers = llama_backend.n_moe_layers, + gpu_ids = llama_backend.gpu_ids, llama_cpp_supports_mtp = _supports_mtp, spec_fallback_reason = llama_backend.spec_fallback_reason, llama_cpp_prebuilt_stale = _stale, diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index bb321695cd..0806c2f513 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -2731,7 +2731,11 @@ async def get_gguf_variants( ], has_vision = response.has_vision, default_variant = response.default_variant, - context_length = _read_native_context_length(repo_id, is_local = local), + # The header walk reads tokenizer arrays on dense models (tens of + # ms per uncached file); keep it off the event loop. + context_length = await asyncio.to_thread( + _read_native_context_length, repo_id, is_local = local + ), ) except HTTPException: raise diff --git a/studio/backend/routes/training_vram.py b/studio/backend/routes/training_vram.py index fb361d3359..fd96fe2175 100644 --- a/studio/backend/routes/training_vram.py +++ b/studio/backend/routes/training_vram.py @@ -197,15 +197,18 @@ def can_load_chat_during_training( requested_gpu_ids: Optional[List[int]], is_gguf: bool = False, required_override_gb: Optional[float] = None, + single_device_gpu: Optional[str] = None, ) -> Tuple[bool, Dict[str, Any]]: """Decide if a NEW chat model can load without OOMing active training (inverse of can_keep_chat_during_training: training is already resident, so size the chat model against the free VRAM that remains). Sizes/places it the same way the loader will: HF auto reuses auto_select_gpu_ids; HF explicit requires an even-share per-GPU floor for device_map="balanced"; GGUF sizes from - required_override_gb over the visible pool. `load_in_4bit` must be effective - (LoRA can flip 4-bit -> 16-bit). Non-CUDA allows the load; default-deny on any - CUDA case it can't size, so a load never OOMs training.""" + required_override_gb over the visible pool. ``single_device_gpu`` is the + exact physical device token selected by a single-device runner. + `load_in_4bit` must be effective (LoRA can flip 4-bit -> 16-bit). Non-CUDA + allows the load; default-deny on any CUDA case it can't size, so a load never + OOMs training.""" try: from utils.hardware import ( DeviceType, @@ -251,26 +254,49 @@ def can_load_chat_during_training( } # Explicit GPUs, or GGUF: size directly and check live free VRAM. + if single_device_gpu is not None: + mode = "single_device" + elif is_gguf: + mode = "gguf" + else: + mode = "explicit" required_gb = required_override_gb if required_gb is None: required_gb, _meta = estimate_required_model_memory_gb(model_name, **est_kwargs) if required_gb is None: - mode = "explicit" if requested_gpu_ids else "gguf" return False, {"mode": mode, "reason": "estimate_unavailable"} free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", [])) - if requested_gpu_ids: + if single_device_gpu is not None: + token = str(single_device_gpu).strip() + if not token: + # Empty token = a CPU-only single-device runner (e.g. a CPU + # diffusion GGUF): it uses no GPU VRAM, so it never threatens + # active training and can always load. + return True, {"mode": "single_device", "reason": "cpu_only"} + try: + selected_gpu = int(token) + if selected_gpu < 0: + raise ValueError + except (TypeError, ValueError): + # A non-numeric device token (e.g. a CUDA UUID / MIG handle) + # can't be mapped to a free-VRAM index, but the runner still + # drives ONE device. Size against the worst-case visible device + # (min free), never the aggregate pool, so a single-device load + # is never OK'd on capacity it can't use and OOMs training. + free_vals = [min(free_by_index.values())] if free_by_index else [] + else: + free_vals = [free_by_index.get(selected_gpu, 0.0)] + elif requested_gpu_ids: # Invalid ids -> load_model 400s first, so don't block; missing id = 0. try: resolved = resolve_requested_gpu_ids(requested_gpu_ids) except ValueError: - return True, {"mode": "explicit", "reason": "invalid_gpu_ids"} + return True, {"mode": mode, "reason": "invalid_gpu_ids"} free_vals = [free_by_index.get(i, 0.0) for i in resolved] - mode = "explicit" else: # GGUF: llama.cpp picks the GPU(s); any visible GPU is a candidate. free_vals = list(free_by_index.values()) - mode = "gguf" if not free_vals: return False, {"mode": mode, "reason": "no_visible_gpus"} diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 63dba8579c..7daa4224aa 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -168,11 +168,14 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): devices, required_override = None, estimate = None, + single_device_gpu = None, + gpu_ids = None, ): with ( patch("utils.hardware.get_device", return_value = DeviceType.CUDA), patch("utils.hardware.estimate_required_model_memory_gb", return_value = (estimate, {})), patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}), + patch("utils.hardware.resolve_requested_gpu_ids", return_value = gpu_ids), patch("utils.hardware.auto_select_gpu_ids") as auto_mock, ): ok, info = tv.can_load_chat_during_training( @@ -180,9 +183,10 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): hf_token = None, load_in_4bit = True, max_seq_length = 0, - requested_gpu_ids = None, + requested_gpu_ids = gpu_ids, is_gguf = True, required_override_gb = required_override, + single_device_gpu = single_device_gpu, ) return ok, info, auto_mock @@ -198,6 +202,88 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): ok, _, _ = self._run(devices = _devices((0, 80, 35), (1, 80, 70)), required_override = 20.0) self.assertTrue(ok) + def test_no_per_gpu_floor_for_gguf_with_explicit_gpu_ids(self): + # gpu_ids narrows llama.cpp's candidate pool but does not turn its + # self-placement into HF device_map="balanced". The uneven selected + # pair therefore keeps the aggregate GGUF check without an even-share + # floor on the nearly-full card. + ok, info, _ = self._run( + devices = _devices((0, 80, 35), (1, 80, 70), (2, 80, 0)), + required_override = 20.0, + gpu_ids = [0, 1], + ) + self.assertTrue(ok) + self.assertEqual(info["mode"], "gguf") + + def test_single_device_uses_selected_gpu(self): + # The model needs 27 GB with headroom. GPU 0 has 45 GB free, while an + # unrelated training-heavy GPU 1 has only 10 GB free. + ok, info, _ = self._run( + devices = _devices((0, 80, 35), (1, 80, 70)), + required_override = 20.0, + single_device_gpu = "0", + ) + self.assertTrue(ok) + self.assertEqual(info["usable_gb"], 45.0) + + blocked, blocked_info, _ = self._run( + devices = _devices((0, 80, 35), (1, 80, 70)), + required_override = 20.0, + single_device_gpu = "1", + ) + self.assertFalse(blocked) + self.assertEqual(blocked_info["usable_gb"], 10.0) + + def test_single_device_unresolved_token_sizes_against_worst_device(self): + # A non-numeric device token (a CUDA UUID / MIG handle) can't map to a + # free-VRAM index. The runner still drives ONE device, so size against the + # worst-case visible device (min free), not the aggregate pool: one GPU + # with 80 GB free vs a 20 GB model -> allow. + ok, info, _ = self._run( + devices = _devices((0, 80, 0)), + required_override = 20.0, + single_device_gpu = "GPU-uuid", + ) + self.assertTrue(ok) + self.assertEqual(info["mode"], "single_device") + self.assertNotIn("reason", info) + + def test_single_device_unresolved_token_refuses_when_worst_device_full(self): + # Same UUID fallback, worst-case device nearly full (2 GB for a 20 GB + # model) -> refuse (default-deny), not on an unresolved-token technicality. + ok, info, _ = self._run( + devices = _devices((0, 80, 78)), + required_override = 20.0, + single_device_gpu = "GPU-uuid", + ) + self.assertFalse(ok) + self.assertNotEqual(info.get("reason"), "unresolved_gpu_id") + + def test_single_device_unresolved_token_uses_min_free_not_aggregate(self): + # The single-device runner uses ONE device but we can't tell which from a + # UUID token. Sizing against the aggregate pool would let a 20 GB model + # "fit" 160 GB of pooled free VRAM while landing on a 2 GB card and OOMing + # training. Min-free (2 GB) is the safe worst case -> refuse. + ok, info, _ = self._run( + devices = _devices((0, 80, 78), (1, 80, 0), (2, 80, 0)), + required_override = 20.0, + single_device_gpu = "GPU-uuid", + ) + self.assertFalse(ok) + self.assertEqual(info["mode"], "single_device") + + def test_single_device_cpu_token_allows(self): + # An empty device token = a CPU-only single-device runner (CPU diffusion + # GGUF): it uses no GPU VRAM, so it never threatens training -> allow + # regardless of how full the GPUs are. + ok, info, _ = self._run( + devices = _devices((0, 80, 78)), + required_override = 20.0, + single_device_gpu = "", + ) + self.assertTrue(ok) + self.assertEqual(info["reason"], "cpu_only") + def test_estimate_unavailable_refuses(self): # No override and the estimator can't size it -> default-deny. ok, info, _ = self._run(devices = _devices((0, 80, 0)), required_override = None, estimate = None) @@ -309,6 +395,8 @@ class TestChatLoadGuardRoute(unittest.TestCase): captured = None, training_active, decision, + gpu_memory_mode = "auto", + requested_gpu_ids = None, ): config = config or SimpleNamespace(is_gguf = False, is_lora = False, path = None) with _stub_guard_deps( @@ -320,7 +408,8 @@ class TestChatLoadGuardRoute(unittest.TestCase): hf_token = None, load_in_4bit = True, max_seq_length = 0, - requested_gpu_ids = None, + requested_gpu_ids = requested_gpu_ids, + gpu_memory_mode = gpu_memory_mode, ) def test_noop_when_training_inactive(self): @@ -332,6 +421,141 @@ class TestChatLoadGuardRoute(unittest.TestCase): def test_allows_when_fits(self): self._guard(training_active = True, decision = (True, {"mode": "auto"})) + def test_diffusion_detection_uses_name_before_download(self): + config = SimpleNamespace( + identifier = "unsloth/DiffusionGemma-GGUF", + gguf_hf_repo = "unsloth/DiffusionGemma-GGUF", + gguf_file = None, + ) + self.assertTrue(self.route._classify_diffusion_gguf(config)) + + def test_uncached_gguf_classification_remains_unknown(self): + config = SimpleNamespace( + identifier = "owner/renamed-model", + gguf_hf_repo = "owner/renamed-model", + gguf_variant = "Q4_K_M", + gguf_file = None, + ) + self.assertIsNone(self.route._classify_diffusion_gguf(config)) + + def test_diffusion_detection_reuses_loader_metadata_probe(self): + import tempfile + + seen = [] + + class _Probe: + is_diffusion = False + _architecture = None + + def _read_gguf_metadata(self, path): + seen.append(path) + self.is_diffusion = True + + with tempfile.TemporaryDirectory() as d: + model = Path(d) / "renamed.gguf" + model.write_bytes(b"GGUF") + config = SimpleNamespace(identifier = "local", gguf_file = str(model)) + with patch.object(self.route, "LlamaCppBackend", _Probe): + self.assertTrue(self.route._classify_diffusion_gguf(config)) + self.assertEqual(seen, [str(model)]) + + def test_local_chat_gguf_classification_is_definitive(self): + import tempfile + class _Probe: + is_diffusion = False + _architecture = "llama" + + def _read_gguf_metadata(self, _path): + pass + + with tempfile.TemporaryDirectory() as d: + model = Path(d) / "renamed.gguf" + model.write_bytes(b"GGUF") + config = SimpleNamespace(identifier = "local", gguf_file = str(model)) + with patch.object(self.route, "LlamaCppBackend", _Probe): + self.assertFalse(self.route._classify_diffusion_gguf(config)) + + def test_manual_known_normal_gguf_bypasses_training_estimate(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with patch.object(self.route, "_classify_diffusion_gguf", return_value = False): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (False, {"reason": "must not run"}), + gpu_memory_mode = "manual", + ) + self.assertEqual(captured, []) + + def test_manual_unknown_gguf_keeps_single_device_training_guard(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with ( + patch.object(self.route, "_classify_diffusion_gguf", return_value = None), + patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), + patch.object( + self.route.LlamaCppBackend, + "_diffusion_gpu_arg", + return_value = "2", + ), + ): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (True, {"mode": "single_device"}), + gpu_memory_mode = "manual", + ) + self.assertEqual(len(captured), 1) + self.assertEqual(captured[0]["single_device_gpu"], "2") + + def test_manual_diffusion_uses_single_device_guard(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with ( + patch.object(self.route, "_classify_diffusion_gguf", return_value = True), + patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), + ): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (True, {"mode": "gguf"}), + gpu_memory_mode = "manual", + requested_gpu_ids = [3, 1], + ) + self.assertEqual(len(captured), 1) + self.assertEqual(captured[0]["single_device_gpu"], "1") + self.assertEqual(captured[0]["requested_gpu_ids"], [3, 1]) + + def test_unpinned_diffusion_uses_runner_default_gpu(self): + captured = [] + config = SimpleNamespace(is_gguf = True) + with ( + patch.object(self.route, "_classify_diffusion_gguf", return_value = True), + patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), + patch.object( + self.route.LlamaCppBackend, + "_effective_gpu_count", + return_value = 2, + ), + patch.object( + self.route.LlamaCppBackend, + "_diffusion_gpu_arg", + return_value = "3", + ) as gpu_arg, + ): + self._guard( + config = config, + captured = captured, + training_active = True, + decision = (True, {"mode": "single_device"}), + gpu_memory_mode = "manual", + ) + gpu_arg.assert_called_once_with(None, cpu_only = False) + self.assertEqual(captured[0]["single_device_gpu"], "3") + def test_refuses_with_headroom_number(self): info = {"required_gb": 30.0, "usable_gb": 6.0, "needed_gb": 39.0, "mode": "auto"} with self.assertRaises(HTTPException) as exc: @@ -467,36 +691,115 @@ class TestValidateRefusesDuringTraining(unittest.TestCase): self.assertEqual(captured[0]["load_in_4bit"], False) self.assertEqual(captured[0]["max_seq_length"], 4096) - def test_rejects_gguf_with_gpu_ids_before_guard(self): - # /validate must mirror /load's GGUF + gpu_ids 400, before the VRAM guard. + def test_validate_forwards_manual_gpu_memory_mode_to_guard(self): from models.inference import ValidateModelRequest - request = ValidateModelRequest(model_path = "x.gguf", gpu_ids = [0]) + request = ValidateModelRequest( + model_path = "unsloth/model-GGUF", + gguf_variant = "Q4_K_M", + gpu_memory_mode = "manual", + ) cfg = SimpleNamespace( - identifier = "x.gguf", - display_name = "x", + identifier = "unsloth/model-GGUF", + display_name = "model-GGUF", is_gguf = True, is_lora = False, is_vision = False, path = None, base_model = None, ) - captured = [] + captured = {} with ( patch.object( self.route, "_resolve_model_identifier_for_request", - return_value = ("x.gguf", "x.gguf", False), + return_value = ("unsloth/model-GGUF", "unsloth/model-GGUF", False), ), patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), patch.object(self.route, "load_inference_config", return_value = {}), - _stub_guard_deps(training_active = True, decision = (True, {}), captured = captured), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda config, **kw: captured.update(kw), + ), ): - with self.assertRaises(HTTPException) as exc: - asyncio.run(self.route.validate_model(request, current_subject = "u")) - self.assertEqual(exc.exception.status_code, 400) - self.assertIn("gpu_ids is not supported for GGUF", exc.exception.detail) - self.assertEqual(captured, []) # guard never reached + asyncio.run(self.route.validate_model(request, current_subject = "u")) + self.assertEqual(captured.get("gpu_memory_mode"), "manual") + + def test_validate_forwards_inherited_extras_and_parallel_to_guard(self): + # Regression: /load resolves inherited same-model extras and passes the + # real slot count to the guard; validate must do the same, else it sizes + # a smaller estimate (no inherited -c/--model-draft, n_parallel=1) and + # /load then 409s after the frontend has already unloaded. + from models.inference import ValidateModelRequest + + request = ValidateModelRequest(model_path = "unsloth/Qwen3-1.7B", max_seq_length = 4096) + cfg = SimpleNamespace( + identifier = "unsloth/Qwen3-1.7B", + display_name = "Qwen3-1.7B", + is_gguf = False, + is_lora = False, + is_vision = False, + path = None, + base_model = None, + ) + captured = {} + with ( + patch.object( + self.route, + "_resolve_model_identifier_for_request", + return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False), + ), + patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), + patch.object(self.route, "load_inference_config", return_value = {}), + patch.object(self.route, "_resolve_inherited_extra_args", return_value = ["-c", "32768"]), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda config, **kw: captured.update(kw), + ), + ): + asyncio.run(self.route.validate_model(request, current_subject = "u")) + self.assertEqual(captured.get("llama_extra_args"), ["-c", "32768"]) + self.assertIn("n_parallel", captured) + + def test_metadata_probe_skips_training_guard(self): + # A header-only probe (include_context_length) allocates no VRAM, so the + # training guard must not run -- else the staging GPU-layers / MoE sliders + # it feeds are hidden exactly when a during-training user needs them. + from models.inference import ValidateModelRequest + + request = ValidateModelRequest( + model_path = "unsloth/Qwen3-1.7B", + max_seq_length = 4096, + include_context_length = True, + ) + cfg = SimpleNamespace( + identifier = "unsloth/Qwen3-1.7B", + display_name = "Qwen3-1.7B", + is_gguf = False, + is_lora = False, + is_vision = False, + path = None, + base_model = None, + ) + guard_called = [] + with ( + patch.object( + self.route, + "_resolve_model_identifier_for_request", + return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False), + ), + patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg), + patch.object(self.route, "load_inference_config", return_value = {}), + patch.object( + self.route, + "_guard_chat_load_against_training", + lambda *a, **kw: guard_called.append(True), + ), + ): + asyncio.run(self.route.validate_model(request, current_subject = "u")) + self.assertEqual(guard_called, []) # ── _estimate_gguf_required_gb (sizes the same weights the loader loads) ────── diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py index a5be07f8e3..ec0330ce05 100644 --- a/studio/backend/tests/test_gguf_metadata.py +++ b/studio/backend/tests/test_gguf_metadata.py @@ -15,6 +15,7 @@ from utils.models.gguf_metadata import ( pairing_score, read_gguf_context_length, read_gguf_general_metadata, + read_gguf_staged_dims, read_mmproj_audio_capability, ) @@ -153,6 +154,78 @@ def test_context_length_ignores_foreign_arch_key(tmp_path: Path): assert read_gguf_context_length(str(p)) is None +# --- read_gguf_staged_dims (one pass: context + layer + moe counts) ---- + + +def test_staged_dims_none_for_missing_or_non_gguf(tmp_path: Path): + assert read_gguf_staged_dims(str(tmp_path / "nope.gguf")) is None + p = tmp_path / "garbage.gguf" + p.write_bytes(b"not a gguf at all") + assert read_gguf_staged_dims(str(p)) is None + + +def test_staged_dims_moe_with_leading_dense(tmp_path: Path): + # GLM-4.7-Flash shape: context + total layers + MoE layers in one read. + p = _write_synthetic_gguf( + tmp_path / "glm.gguf", + {"general.architecture": "deepseek2"}, + extra_uint32 = { + "deepseek2.context_length": 202752, + "deepseek2.block_count": 47, + "deepseek2.expert_count": 64, + "deepseek2.leading_dense_block_count": 1, + }, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": 202752, + "layer_count": 47, + "moe_layer_count": 46, + } + + +def test_staged_dims_dense_model(tmp_path: Path): + # Dense: layer_count present, moe_layer_count 0 (slider hidden). + p = _write_synthetic_gguf( + tmp_path / "dense.gguf", + {"general.architecture": "qwen3"}, + extra_uint32 = {"qwen3.context_length": 40960, "qwen3.block_count": 36}, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": 40960, + "layer_count": 36, + "moe_layer_count": 0, + } + + +def test_staged_dims_all_moe_no_leading_dense(tmp_path: Path): + # Experts present, no leading_dense key -> every block is a MoE layer. + p = _write_synthetic_gguf( + tmp_path / "moe.gguf", + {"general.architecture": "qwen35moe"}, + extra_uint32 = {"qwen35moe.block_count": 40, "qwen35moe.expert_count": 256}, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": None, + "layer_count": 40, + "moe_layer_count": 40, + } + + +def test_staged_dims_uint64_block_count(tmp_path: Path): + # block_count stored as uint64 (vtype 10) still parses; moe == block_count. + p = _write_synthetic_gguf( + tmp_path / "moe64.gguf", + {"general.architecture": "gpt-oss"}, + extra_uint32 = {"gpt-oss.expert_count": 32}, + extra_uint64 = {"gpt-oss.block_count": 24}, + ) + assert read_gguf_staged_dims(str(p)) == { + "context_length": None, + "layer_count": 24, + "moe_layer_count": 24, + } + + def test_context_length_read_from_uint64(tmp_path: Path): # Some models store context_length as a uint64 (vtype 10). p = _write_synthetic_gguf( diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py new file mode 100644 index 0000000000..b17274197f --- /dev/null +++ b/studio/backend/tests/test_gpu_memory_mode.py @@ -0,0 +1,879 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Backend contract for the GPU Memory mode dropdown. + +The dropdown threads a single ``gpu_memory_mode`` ("auto" | "manual") from the +chat UI through the load request. "manual" lets the user own the offload: with +``gpu_layers < 0`` (Auto, the default) it hands all memory management to +llama.cpp's ``--fit on`` (no CUDA/HIP device masking, no context auto-reduce, no +gpu-layer or tensor-split planning); with ``gpu_layers >= 0`` it pins the layers +and MoE offload itself (``--fit off``). These tests pin: + + * the pydantic request/response/status contract (snake_case key, default + "auto", unknown values rejected), + * the backend ``gpu_memory_mode`` property and its reset on unload, + * the ``_already_in_target_state`` reload-detection branch, and + * that the manual + Auto-layers branch in ``load_model`` empties the probed + GPU set and drops tensor parallelism so the selection below no-ops, while + the explicit-offload branch emits ``--gpu-layers`` / ``--fit off``. +""" + +from __future__ import annotations + +import inspect +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Same external-dep stubs as the other llama_cpp unit tests so importing +# the backend doesn't drag in structlog / httpx / loggers. +_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") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +sys.modules.setdefault("structlog", _structlog_stub) + +# httpx is a real, installed backend dependency: import it so the genuine module +# is in sys.modules. A hand-rolled stub here is inevitably incomplete and, since +# setdefault installs it before real httpx loads, would poison a combined pytest +# run -- routes/inference references httpx.Response (and other attrs) at def time. +import httpx # noqa: F401 + +from core.inference import llama_cpp as llama_cpp_module +from core.inference.llama_cpp import LlamaCppBackend +from models.inference import ( + InferenceStatusResponse, + LoadRequest, + LoadResponse, +) + + +# ── Pydantic contract (snake_case key, default "auto") ─────────────── + + +def test_load_request_defaults_gpu_memory_mode_auto(): + assert LoadRequest(model_path = "owner/repo").gpu_memory_mode == "auto" + + +def test_load_request_round_trips_json_key(): + req = LoadRequest.model_validate({"model_path": "owner/repo", "gpu_memory_mode": "manual"}) + assert req.gpu_memory_mode == "manual" + assert req.model_dump()["gpu_memory_mode"] == "manual" + + +def test_load_request_rejects_unknown_mode(): + with pytest.raises(ValueError): + LoadRequest(model_path = "owner/repo", gpu_memory_mode = "bogus") + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_gpu_memory_mode(model_cls): + if model_cls is LoadResponse: + default = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + ) + manual = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + gpu_memory_mode = "manual", + ) + else: + default = model_cls() + manual = model_cls(gpu_memory_mode = "manual") + assert default.model_dump()["gpu_memory_mode"] == "auto" + assert manual.model_dump()["gpu_memory_mode"] == "manual" + + +# ── Backend property + reset ───────────────────────────────────────── + + +class _FakeProcess: + """Stand-in for subprocess.Popen so _kill_process is a no-op.""" + + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + def poll(self): + return 0 + + +def test_gpu_memory_mode_property_defaults_auto(): + assert LlamaCppBackend().gpu_memory_mode == "auto" + + +def test_gpu_memory_mode_property_reflects_field(): + backend = LlamaCppBackend() + backend._gpu_memory_mode = "manual" + assert backend.gpu_memory_mode == "manual" + + +def test_unload_resets_gpu_memory_mode(): + backend = LlamaCppBackend() + backend._process = _FakeProcess() + backend._gpu_memory_mode = "manual" + backend.unload_model() + assert backend.gpu_memory_mode == "auto" + + +# ── _already_in_target_state reload-detection branch ───────────────── + + +def _loaded_backend(gpu_memory_mode: str) -> LlamaCppBackend: + backend = LlamaCppBackend() + backend._process = _FakeProcess() # is_loaded only checks "is not None" + backend._healthy = True + backend._model_identifier = "owner/repo" + backend._hf_variant = "Q4_K_M" + backend._requested_n_ctx = 8192 + backend._cache_type_kv = None + backend._requested_spec_mode = "auto" + backend._chat_template_override = None + backend._is_vision = False + backend._extra_args = None + backend._gguf_path = None + backend._gpu_memory_mode = gpu_memory_mode + return backend + + +def _target_state(backend: LlamaCppBackend, gpu_memory_mode: str) -> bool: + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + gpu_memory_mode = gpu_memory_mode, + ) + + +@pytest.mark.parametrize("mode", ["auto", "manual"]) +def test_already_in_target_state_matches_same_mode(mode): + assert _target_state(_loaded_backend(mode), mode) is True + + +@pytest.mark.parametrize("loaded,requested", [("auto", "manual"), ("manual", "auto")]) +def test_already_in_target_state_reloads_on_mode_change(loaded, requested): + # Flipping the dropdown either direction must force a reload so the command + # is rebuilt with/without the Unsloth GPU masking. + assert _target_state(_loaded_backend(loaded), requested) is False + + +def test_already_in_target_state_ignores_mode_for_diffusion(): + # The diffusion runner is mode-agnostic (always "auto"), so a standing manual + # preference must not force a needless reload. + backend = _loaded_backend("auto") + backend._is_diffusion = True + assert _target_state(backend, "manual") is True + + +# ── load_model: manual + Auto layers bypasses Unsloth GPU management ── + + +def _load_model_source() -> str: + return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model) + + +def test_auto_layers_branch_empties_gpus_and_drops_tensor_parallel(): + # Emptying the probed set makes the selection / TP planning below no-op, so + # gpu_indices stays None and use_fit True (--fit on). + src = _load_model_source() + gate = src.find('if gpu_memory_mode == "manual" and gpu_layers < 0:') + assert gate != -1, "load_model must branch on manual + Auto layers (gpu_layers < 0)" + block = src[gate : gate + 1400] + assert "gpus = []" in block, "Auto-layers branch must empty the probed GPU set" + # --fit aborts under --split-mode tensor, so a raw-extras split-mode is stripped. + assert "strip_split_mode_only(extra_args)" in block + assert "requested_ctx if requested_ctx > 0 else 0" in block + # The branch sits before GPU selection assigns gpu_indices; --fit on is its emission. + assert gate < src.find("gpu_indices, use_fit = None, True") + assert 'cmd.extend(["--fit", "on"])' in src + # TP drops for this path, but at a guard BEFORE the quantized-KV cache-drop, so + # a requested quantized cache survives into the --fit load. + tp_drop = src.find('if tensor_parallel and gpu_memory_mode == "manual" and gpu_layers < 0:') + assert tp_drop != -1, "manual + Auto layers must drop tensor_parallel" + assert "tensor_parallel = False" in src[tp_drop : tp_drop + 400] + cache_drop = src.find("Tensor parallelism requires a non-quantized KV cache") + assert cache_drop != -1 + assert ( + tp_drop < cache_drop + ), "TP must drop before the cache-drop so a quantized KV survives --fit" + + +def test_auto_layers_never_sends_ctx_size_zero(): + # Sending "-c 0" sets fit_params_min_ctx = UINT32_MAX in llama.cpp, pinning + # the full native context and disabling --fit's reduction. So the base cmd + # must never carry -c, "-c 0" is emitted only outside the Auto-layers (--fit) + # case, and a positive context is passed through (which --fit optimizes + # layers around). + src = _load_model_source() + base_start = src.find("cmd = [") + base_end = src.find("\n ]", base_start) + base_block = src[base_start:base_end] + assert '"-c"' not in base_block, "-c must be conditional, not in the base cmd list" + assert 'cmd.extend(["-c", str(effective_ctx)])' in src, "positive ctx must pass -c" + assert 'auto_fit = gpu_memory_mode == "manual" and gpu_layers < 0' in src + zero = src.find('cmd.extend(["-c", "0"])') + assert zero != -1, '"-c 0" emission must exist outside the Auto-layers case' + guard = src.rfind("elif not auto_fit:", 0, zero) + assert guard != -1 and zero - guard < 120, '"-c 0" must sit under the not-auto_fit guard' + + +def test_manual_mode_clears_inherited_main_model_placement_env(): + env = {name: "inherited" for name in LlamaCppBackend._MANUAL_PLACEMENT_ENV_VARS} + env["LLAMA_ARG_N_GPU_LAYERS_DRAFT"] = "7" + env["UNRELATED"] = "kept" + + LlamaCppBackend._clear_manual_placement_env(env) + + assert not (set(env) & set(LlamaCppBackend._MANUAL_PLACEMENT_ENV_VARS)) + assert env["LLAMA_ARG_N_GPU_LAYERS_DRAFT"] == "7" + assert env["UNRELATED"] == "kept" + + +def test_load_model_sanitizes_manual_env_after_building_child_env(): + src = _load_model_source() + env_build = src.find("env = self._llama_server_env_for_binary(binary)") + env_clear = src.find("self._clear_manual_placement_env(env)", env_build) + launch = src.find("subprocess.Popen", env_build) + assert env_build != -1 + assert env_build < env_clear < launch + + +# ── Manual offload (--gpu-layers + --fit off + --n-cpu-moe) ─────────── + + +def test_load_request_accepts_manual(): + req = LoadRequest( + model_path = "owner/repo", + gpu_memory_mode = "manual", + gpu_layers = 20, + n_cpu_moe = 8, + tensor_split = [2, 1], + ) + assert req.gpu_memory_mode == "manual" + assert req.gpu_layers == 20 + assert req.n_cpu_moe == 8 + assert req.tensor_split == [2, 1] + + +def test_load_request_manual_defaults(): + req = LoadRequest(model_path = "owner/repo") + assert req.gpu_layers == -1 + assert req.n_cpu_moe == 0 + assert req.tensor_split is None + + +@pytest.mark.parametrize("bad", [[0, 0], [-1, 2], [float("inf"), 1], [float("nan"), 1]]) +def test_load_request_rejects_degenerate_tensor_split(bad): + # A negative/non-finite/all-zero split is dropped at launch but compared raw + # in the reload dedupe, so it would reload forever -- reject it up front. + with pytest.raises(ValueError): + LoadRequest(model_path = "owner/repo", tensor_split = bad) + + +@pytest.mark.parametrize("good", [[2, 1], [1, 1], [], None]) +def test_load_request_accepts_valid_tensor_split(good): + assert LoadRequest(model_path = "owner/repo", tensor_split = good).tensor_split == good + + +def test_route_normalizes_explicit_extras_before_reload_dedupe(): + route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + load_impl = route_src[route_src.index("async def _load_model_impl") :] + strip = load_impl.index("_stripped_explicit = strip_shadowing_flags") + normalize = load_impl.index( + 'request = request.model_copy(update = {"llama_extra_args": extra_llama_args})' + ) + dedupe = load_impl.index("and _request_matches_loaded_settings(") + assert strip < normalize < dedupe + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_manual_fields(model_cls): + if model_cls is LoadResponse: + obj = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + gpu_memory_mode = "manual", + gpu_layers = 20, + n_cpu_moe = 8, + tensor_split = [2, 1], + n_layers = 32, + n_moe_layers = 32, + ) + else: + obj = model_cls( + gpu_memory_mode = "manual", + gpu_layers = 20, + n_cpu_moe = 8, + tensor_split = [2, 1], + n_layers = 32, + n_moe_layers = 32, + ) + dumped = obj.model_dump() + assert dumped["gpu_memory_mode"] == "manual" + assert dumped["gpu_layers"] == 20 + assert dumped["n_cpu_moe"] == 8 + assert dumped["tensor_split"] == [2, 1] + assert dumped["n_layers"] == 32 + assert dumped["n_moe_layers"] == 32 + + +def test_manual_properties_default_and_reflect_and_reset(): + backend = LlamaCppBackend() + assert backend.gpu_layers == -1 and backend.n_cpu_moe == 0 + assert backend.tensor_split is None + backend._gpu_layers = 20 + backend._n_cpu_moe = 8 + backend._tensor_split = [2, 1] + assert backend.gpu_layers == 20 and backend.n_cpu_moe == 8 + assert backend.tensor_split == [2, 1] + backend._process = _FakeProcess() + backend.unload_model() + assert backend.gpu_layers == -1 and backend.n_cpu_moe == 0 + assert backend.tensor_split is None + + +def test_n_moe_layers_property(): + # 0 for a dense model (hides the slider); block_count for all-MoE; + # block_count - leading_dense otherwise (GLM-4.7-Flash: 47 - 1 -> 46). + b = LlamaCppBackend() + b._n_layers = 36 + b._n_experts = None + assert b.n_moe_layers == 0 + b._n_experts = 128 + b._leading_dense_block_count = None + assert b.n_moe_layers == 36 + b._n_layers = 47 + b._leading_dense_block_count = 1 + assert b.n_moe_layers == 46 + + +def _target_state_manual( + backend, + *, + gpu_layers, + n_cpu_moe, + tensor_split = None, +): + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + gpu_memory_mode = "manual", + gpu_layers = gpu_layers, + n_cpu_moe = n_cpu_moe, + tensor_split = tensor_split, + ) + + +def test_manual_reloads_on_gpu_layers_or_n_cpu_moe_or_split_change(): + backend = _loaded_backend("manual") + backend._gpu_layers = 20 + backend._n_cpu_moe = 0 + backend._tensor_split = None + # Same knobs -> no reload. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0) is True + # Changed layer count -> reload. + assert _target_state_manual(backend, gpu_layers = 16, n_cpu_moe = 0) is False + # Changed MoE offload -> reload. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 8) is False + # Added a GPU split -> reload. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is False + # Same GPU split -> no reload. + backend._tensor_split = [2, 1] + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is True + + +def test_auto_layers_reload_tracks_only_gpu_layers(): + # Under Auto (gpu_layers < 0) the MoE/split knobs don't apply, so a leftover + # request value must not reload -- only a gpu_layers change (Auto -> pinned) does. + backend = _loaded_backend("manual") + backend._gpu_layers = -1 + backend._n_cpu_moe = 0 + backend._tensor_split = None + # Same Auto, leftover MoE/split in the request -> still no reload. + assert _target_state_manual(backend, gpu_layers = -1, n_cpu_moe = 8, tensor_split = [2, 1]) is True + # Auto -> explicit offload reloads. + assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0) is False + + +def test_manual_offload_emits_gpu_layers_fit_off_and_n_cpu_moe(): + src = _load_model_source() + gate = src.find('elif gpu_memory_mode == "manual":') + assert gate != -1, "load_model must have an explicit-offload manual branch" + block = src[gate : gate + 700] + # Empties the probed set (skips the planner) but keeps the user's TP choice + # (only the Auto-layers branch above drops TP). + assert "gpus = []" in block + assert "tensor_parallel = False" not in block + # The cmd emits the layer count with fit disabled, gated on gpu_layers >= 0. + assert 'if gpu_memory_mode == "manual" and gpu_layers >= 0:' in src + assert 'cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"])' in src + # MoE offload uses --n-cpu-moe via _resolve_cpu_moe_flag (tested behaviorally below). + assert "_resolve_cpu_moe_flag(" in src + assert 'cmd.extend(["--n-cpu-moe", str(moe_flag)])' in src + # A count requested on a dense model is never emitted, so it must also be + # dropped from the recorded state -- else /status and /load report a count + # llama-server never received (same rule as the tensor-split drop below). + moe_emit = src.find('cmd.extend(["--n-cpu-moe", str(moe_flag)])') + assert "elif n_cpu_moe:" in src[moe_emit : moe_emit + 300] + assert "self._n_cpu_moe = 0" in src[moe_emit : moe_emit + 300] + # The offload path forces use_fit False so --fit-ctx is never added under --fit off. + emit = src.find('cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"])') + assert "use_fit = False" in src[src.rfind("\n", 0, emit) - 200 : emit + 80] + + +def test_status_reports_requested_context_length(): + # The hydration path re-seeds a Manual+Auto context pin from the REQUESTED + # n_ctx (0 = Auto); context_length only exposes the resolved value. + assert "requested_context_length" in InferenceStatusResponse.model_fields + s = InferenceStatusResponse(requested_context_length = 8192) + assert s.model_dump()["requested_context_length"] == 8192 + assert InferenceStatusResponse().model_dump()["requested_context_length"] is None + # The /status route must actually wire it from the backend (a declared-but- + # never-populated field would leave hydration silently reverting the pin). + from pathlib import Path as _P + + route_src = (_P(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + assert "requested_context_length = llama_backend.requested_n_ctx" in route_src + + +def test_manual_offload_emits_tensor_split(): + # The offload path emits --tensor-split from the per-GPU shares, only when + # provided, with >1 GPU in use, AND matching that count (a stale ratio on a + # narrowed picker or a mismatched direct-API list must not emit -- llama- + # server aborts on a split/GPU-count mismatch). + src = _load_model_source() + assert "if tensor_split and _split_gpus > 1:" in src + # Emit only on a length match AND a positive sanitized total: a mismatched + # or all-zero split aborts llama-server / assigns nothing, so it's dropped. + # The emitted list is the sanitized one (clamping tested behaviorally below). + assert "_sanitized_split = self._sanitize_tensor_split(tensor_split)" in src + assert "if len(_sanitized_split) == _split_gpus and _split_total > 0:" in src + assert '"--tensor-split"' in src + # Joined as a comma list (e.g. "2,1") within the explicit-offload cmd branch. + gate = src.find('if gpu_memory_mode == "manual" and gpu_layers >= 0:') + nxt = src.find("elif use_fit:", gate) + assert '","' in src[gate:nxt] and "tensor_split" in src[gate:nxt] + # A split with a single effective GPU is never emitted, so it must also be + # dropped from the recorded state -- else /status and /load report a ratio + # llama-server never received and the dedupe baseline preserves it. + assert "elif tensor_split:" in src[gate:nxt] + drop = src.find("elif tensor_split:", gate, nxt) + assert "self._tensor_split = None" in src[drop : drop + 250] + + +def test_sanitize_tensor_split_clamps_negative_and_non_finite(): + # Negative entries would launch a placement different from the ratio the + # UI showed; inf passes a plain > 0 total gate and would emit + # "--tensor-split inf,..." (llama.cpp normalizes shares by the running + # total, so an inf poisons the shares from that entry on). Both clamp to 0. + sanitize = LlamaCppBackend._sanitize_tensor_split + assert sanitize([2, 1]) == [2.0, 1.0] + assert sanitize([-1, 2]) == [0.0, 2.0] + assert sanitize([float("inf"), 1]) == [0.0, 1.0] + assert sanitize([float("nan"), 1]) == [0.0, 1.0] + # All-zero survives sanitization; the call site's total gate drops it. + assert sanitize([0, 0]) == [0.0, 0.0] + # Unreadable input -> []; the call site's length gate drops it. + assert sanitize(["x", 1]) == [] + assert sanitize([10**400, 1]) == [] + + +def test_zero_offload_mask_honors_device_pin_spellings(): + # A user device pin must keep the GPUs visible: llama-server aborts on a + # pin it can't see ('error: invalid device'). The pin can arrive as + # --device or its -dev alias, as the draft forms (parsed even with no + # drafter loaded), or as an inherited LLAMA_ARG_DEVICE env var. + load_src = _load_model_source() + assert "self._zero_offload_keeps_gpu_visible(cmd, env)" in load_src + block = inspect.getsource(LlamaCppBackend._cmd_has_gpu_device_pin) + for flag in ( + '"--device"', + '"-dev"', + '"--spec-draft-device"', + '"-devd"', + '"--device-draft"', + ): + assert flag in block + assert '"LLAMA_ARG_DEVICE"' in block + + +def test_resolve_cpu_moe_flag(): + # Clamp the requested MoE-layer count to the model's MoE layers, then offset + # past leading dense layers (--n-cpu-moe counts from layer 0). + R = LlamaCppBackend._resolve_cpu_moe_flag + assert R(0, 40, 0) is None # nothing requested + assert R(8, 0, 0) is None # dense model (no MoE layers) + assert R(8, 40, 0) == 8 # all-MoE: direct + assert R(100, 40, 0) == 40 # clamp to the MoE layer count + # GLM-4.7-Flash (deepseek2): block_count 47, leading_dense 1, n_moe 46. + assert R(5, 46, 1) == 6 # offset past the 1 dense layer + assert R(46, 46, 1) == 47 # all MoE on CPU == block_count + + +def test_manual_allows_tensor_parallel_via_split_mode(): + # Manual offload keeps the user's TP choice but skips the memory-based planner + # (plan_tp excludes manual, so its empty gpu set can't downgrade TP). The + # --split-mode tensor emission gates on tensor_parallel alone, so manual + # reaches it -- with tp_tensor_split None it's an even split (no + # --tensor-split). --fit off means no fit/tensor abort. + src = _load_model_source() + assert 'plan_tp = tensor_parallel and gpu_memory_mode != "manual"' in src + assert "if plan_tp:" in src + assert "if plan_tp and len(tp_gpus) < 2:" in src + sm = src.find('cmd.extend(["--split-mode", "tensor"])') + assert sm != -1, "TP must emit --split-mode tensor" + guard = src.rfind("if tensor_parallel:", 0, sm) + assert guard != -1 and sm - guard < 200, "split-mode gates on tensor_parallel" + # The tensor-split is only emitted for a planned (non-even) split, which + # manual never produces, so manual stays an even split. + assert "if tp_tensor_split and len(tp_tensor_split) > 1:" in src + + +def test_fit_sets_target_margin(): + # Manual + Auto (auto_fit) tightens the per-device VRAM margin to 512 MiB. + caps = {"supports_fit_target": True} + flags = LlamaCppBackend._ctx_integrity_flags(1, True, True, 0, 0, caps) + assert flags[flags.index("--fit-target") + 1] == "512" + # Not emitted on the legacy auto path (fit on but not auto_fit): -c 0 pins + # native there, so the tighter margin must not ride along. + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, True, False, 0, 0, caps) + # Not emitted when fit is off. + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, False, False, 0, 0, caps) + # Not emitted when the binary lacks support. + assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags( + 1, True, True, 0, 0, {"supports_fit_target": False} + ) + + +# ── GPU picker (gpu_ids -> CUDA_VISIBLE_DEVICES) ───────────────────── + + +def test_load_request_accepts_gpu_ids(): + req = LoadRequest(model_path = "owner/repo", gpu_ids = [1, 0]) + assert req.gpu_ids == [1, 0] + assert LoadRequest(model_path = "owner/repo").gpu_ids is None + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_gpu_ids(model_cls): + if model_cls is LoadResponse: + obj = model_cls(status = "loaded", model = "m", display_name = "m", inference = {}, gpu_ids = [1]) + else: + obj = model_cls(gpu_ids = [1]) + assert obj.model_dump()["gpu_ids"] == [1] + + +def test_gpu_ids_property_default_and_reset(): + backend = LlamaCppBackend() + assert backend.gpu_ids is None + backend._gpu_ids = [0, 1] + assert backend.gpu_ids == [0, 1] + backend._process = _FakeProcess() + backend.unload_model() + assert backend.gpu_ids is None + + +def _target_state_gpu_ids(backend, gpu_ids): + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + gpu_ids = gpu_ids, + ) + + +def test_gpu_ids_reload_detection_is_order_insensitive(): + backend = _loaded_backend("auto") + backend._gpu_ids = [0, 1] + # Same set, different order -> no reload. + assert _target_state_gpu_ids(backend, [1, 0]) is True + # Different set -> reload. + assert _target_state_gpu_ids(backend, [0]) is False + # Dropping the pick (auto) -> reload. + assert _target_state_gpu_ids(backend, None) is False + + +def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device(): + # The diffusion runner drives only its single lowest device, so the backend + # records [lowest]. A later multi-GPU request that still resolves to that + # same lowest device must dedupe (no needless reload); a request whose lowest + # device moves, or that drops the pick, must reload. + backend = _loaded_backend("auto") + backend._is_diffusion = True + backend._gpu_ids = [1] # loaded on the lowest of an earlier [3, 1] pick + assert _target_state_gpu_ids(backend, [3, 1]) is True + assert _target_state_gpu_ids(backend, [1]) is True + # Lowest device changes (2, not 1) -> reload. + assert _target_state_gpu_ids(backend, [3, 2]) is False + # Dropping the pick (auto) -> reload. + assert _target_state_gpu_ids(backend, None) is False + + +def test_start_diffusion_server_resets_tensor_parallel(): + # A prior tensor-parallel chat load leaves self._tensor_parallel True (load_model + # phase 1 only kills the process, it skips the unload reset). Diffusion is never + # TP, so startup must clear it -- else /status misreports TP and an identical + # diffusion re-Apply reloads against stale tensor-parallel state. + src = inspect.getsource(llama_cpp_module.LlamaCppBackend._start_diffusion_server) + assert "self._tensor_parallel = False" in src + + +def test_route_matches_loaded_settings_collapses_diffusion_gpu_ids(): + # The route-level reload dedupe mirrors the backend: for a loaded diffusion + # model it compares the request against the single recorded device, not the + # full requested list, or a same-device multi-GPU pick reloads needlessly. + route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + match_impl = route_src[route_src.index("def _request_matches_loaded_settings") :] + guard = match_impl.index("if llama_backend.is_diffusion:") + collapse = match_impl.index("[sorted(request.gpu_ids)[0]] if request.gpu_ids else None") + compare = match_impl.index("if _req_gpu_ids != llama_backend.gpu_ids:") + assert guard < collapse < compare + + +# ── Manual tensor split: child enumeration pinned to the picker's order ────── + + +def _patch_split_pin_env(monkeypatch, *, inherited, reported): + """Point the pin helper at a fake inherited mask and picker report. + ``reported`` None = enumeration unavailable (falls back to ascending).""" + import utils.hardware as hw + + monkeypatch.setattr( + LlamaCppBackend, "_resolve_visible_physical_ids", staticmethod(lambda: inherited) + ) + info = ( + {"available": False} + if reported is None + else { + "available": True, + "index_kind": "physical", + "devices": [{"index": i} for i in reported], + } + ) + monkeypatch.setattr(hw, "get_backend_visible_gpu_info", lambda: info) + + +def test_split_pin_reorders_inherited_numeric_mask(monkeypatch): + # Parent CUDA_VISIBLE_DEVICES=3,1 makes the child enumerate dev0=phys3, but + # nvidia-smi reported the picker's list ascending -- the mask must be + # re-emitted in that order or the per-GPU shares land on the wrong cards. + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3]) + env = {"CUDA_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_DEVICE_ORDER"] == "PCI_BUS_ID" + assert env["CUDA_VISIBLE_DEVICES"] == "1,3" + + +def test_split_pin_keeps_mask_order_when_picker_reported_it(monkeypatch): + # Torch-fallback enumeration (no nvidia-smi) reports devices in inherited + # mask order, so the picker's split list follows the mask -- the pin must + # keep that order, not re-sort it into a mismatch. + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [3, 1]) + env = {"CUDA_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_VISIBLE_DEVICES"] == "3,1" + + +def test_split_pin_falls_back_to_ascending_without_report(monkeypatch): + # Enumeration unavailable: ascending physical is the best guess (it matches + # the dominant nvidia-smi report order). + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = None) + env = {"CUDA_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_VISIBLE_DEVICES"] == "1,3" + + +def test_split_pin_without_mask_only_sets_pci_order(monkeypatch): + # No inherited mask (or a UUID/MIG one resolving to None): enumeration order + # is fully fixed by CUDA_DEVICE_ORDER, so no mask is written. + _patch_split_pin_env(monkeypatch, inherited = None, reported = None) + env = {} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env == {"CUDA_DEVICE_ORDER": "PCI_BUS_ID"} + + +def test_split_pin_mirrors_hip_mask_on_rocm(monkeypatch): + # ROCm: the pin must land in HIP_VISIBLE_DEVICES too, and an inherited ROCR + # mask is cleared so the mask can't apply twice (ROCR re-indexes, then HIP + # would index into the already-reduced set). + _patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3]) + torch_stub = _types.ModuleType("torch") + torch_stub.version = _types.SimpleNamespace(hip = "6.0") + monkeypatch.setitem(sys.modules, "torch", torch_stub) + env = {"CUDA_VISIBLE_DEVICES": "3,1", "ROCR_VISIBLE_DEVICES": "3,1"} + LlamaCppBackend._pin_visible_gpu_order_for_split(env) + assert env["CUDA_VISIBLE_DEVICES"] == "1,3" + assert env["HIP_VISIBLE_DEVICES"] == "1,3" + assert "ROCR_VISIBLE_DEVICES" not in env + + +# ── Diffusion single-device selection ─────────────────────────────────────── + + +def test_diffusion_gpu_arg_uses_lowest_explicit_physical_id(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,1") + monkeypatch.setenv("DG_GPU", "7") + assert LlamaCppBackend._diffusion_gpu_arg([3, 1]) == "1" + + +def test_diffusion_gpu_arg_preserves_parent_mask_order(monkeypatch): + monkeypatch.delenv("DG_GPU", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,1") + assert LlamaCppBackend._diffusion_gpu_arg(None) == "3" + + +def test_diffusion_gpu_arg_honors_override_and_cpu_mask(monkeypatch): + monkeypatch.setenv("DG_GPU", "GPU-abc") + assert LlamaCppBackend._diffusion_gpu_arg(None) == "GPU-abc" + assert LlamaCppBackend._diffusion_gpu_arg(None, cpu_only = True) == "" + + +# ── Deliberate zero-offload (manual gpu_layers=0): training-skip flag ───────── + + +def test_zero_offload_flag_false_without_companions(): + # CPU-only by construction: False lets training skip unloading a server that + # holds no VRAM. + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", "--fit", "off"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is False + + +@pytest.mark.parametrize( + "companion", + ["--mmproj", "--model-draft", "-md", "--spec-draft-model", "-hfd"], +) +def test_zero_offload_flag_true_with_companion(companion): + # mmproj / a drafter offload to GPU regardless of --gpu-layers, so the + # server still holds VRAM and training must unload it. Drafter detection + # reuses the extras parser, so pass-through aliases count too. + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", companion, "x.gguf"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_with_inline_companion_forms(): + cmd = ["llama-server", "-m", "model.gguf", "--spec-draft-model=x.gguf"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + cmd = ["llama-server", "-m", "model.gguf", "--mmproj=proj.gguf"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_with_env_drafter(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + env = {"LLAMA_ARG_SPEC_DRAFT_MODEL": "x.gguf"} + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is True + + +@pytest.mark.parametrize( + "device_args", + [ + ["--device", "CUDA0"], + ["--device=CUDA0"], + ["-dev", "CUDA0"], + ["--spec-draft-device", "CUDA0"], + ["--device-draft=CUDA0"], + ], +) +def test_zero_offload_flag_true_with_device_pin(device_args): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", *device_args] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_with_env_device_pin(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + env = {"LLAMA_ARG_DEVICE": "CUDA0"} + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is True + + +@pytest.mark.parametrize( + ("device_args", "env"), + [ + (["--device", "cpu"], {}), + (["--device=none"], {}), + (["--spec-draft-device", "cpu"], {}), + ([], {"LLAMA_ARG_DEVICE": "none"}), + (["--device", "CUDA0", "--device", "cpu"], {}), + ], +) +def test_zero_offload_flag_false_with_cpu_device_pin(device_args, env): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", *device_args] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is False + + +def test_zero_offload_flag_true_with_surviving_tensor_mode(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", "--split-mode", "tensor"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_true_for_unmasked_vulkan(monkeypatch): + monkeypatch.setattr(LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda: True)) + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True + + +def test_zero_offload_flag_none_without_gpus(): + cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"] + assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [], {}) is None + + +def test_cmd_has_gpu_companion_detection(): + # The env mask for CPU-only zero-offload loads keys off this scan: any + # --mmproj form or a drafter (flag aliases / env) keeps the GPUs visible. + has = LlamaCppBackend._cmd_has_gpu_companion + assert has(["llama-server", "-m", "m.gguf"], {}) is False + assert has(["llama-server", "--mmproj", "p.gguf"], {}) is True + assert has(["llama-server", "--mmproj=p.gguf"], {}) is True + assert has(["llama-server", "-md", "d.gguf"], {}) is True + assert has(["llama-server"], {"LLAMA_ARG_SPEC_DRAFT_MODEL": "d.gguf"}) is True + + +def test_cmd_companion_ignores_cpu_forced_drafter(): + # A CPU-pinned drafter holds no VRAM: the zero-offload mask may hide the GPUs + # and training may leave the server alone. + has = LlamaCppBackend._cmd_has_gpu_companion + cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-ngl", "0"] + assert has(cmd, {}) is False + cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-device", "cpu"] + assert has(cmd, {}) is False + # mmproj still counts even alongside a CPU drafter. + cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-ngl", "0", "--mmproj", "p.gguf"] + assert has(cmd, {}) is True diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index 69ad560788..d4f2fbe993 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -853,7 +853,13 @@ class TestRouteErrors(unittest.TestCase): self.assertIn("only supported on CUDA devices", str(exc_info.exception)) - def test_inference_route_rejects_gpu_ids_for_gguf(self): + def test_inference_route_validates_gpu_ids_for_gguf(self): + # gpu_ids is now SUPPORTED for GGUF (the GPU picker), but still + # validated: a rejected pick surfaces as a clean 400, not the old + # "not supported for GGUF" rejection. Patch the validator so the test + # is deterministic regardless of the host's (or a prior test's) GPU env. + import utils.hardware.hardware as hardware_mod + inference_route = _load_route_module( "inference_route_module_for_gguf_gpu_ids_test", "routes/inference.py", @@ -887,6 +893,11 @@ class TestRouteErrors(unittest.TestCase): ), patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), + patch.object( + hardware_mod, + "resolve_requested_gpu_ids", + side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"), + ), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( @@ -901,8 +912,11 @@ class TestRouteErrors(unittest.TestCase): ) ) + # The validator's ValueError becomes a clean 400 (not the removed + # "not supported for GGUF" rejection). self.assertEqual(exc_info.exception.status_code, 400) - self.assertIn("GGUF", exc_info.exception.detail) + self.assertIn("gpu_ids", exc_info.exception.detail.lower()) + self.assertNotIn("not supported", exc_info.exception.detail.lower()) def test_training_route_returns_400_for_invalid_gpu_ids(self): training_route = _load_route_module( diff --git a/studio/backend/tests/test_llama_cpp_no_context_shift.py b/studio/backend/tests/test_llama_cpp_no_context_shift.py index f320d29a02..662c918305 100644 --- a/studio/backend/tests/test_llama_cpp_no_context_shift.py +++ b/studio/backend/tests/test_llama_cpp_no_context_shift.py @@ -118,9 +118,17 @@ def test_flag_sits_inside_the_base_cmd_list(): "conditional branch -- otherwise some code paths would still " "run with silent context shift enabled." ) - # Pin that it sits next to -c / --ctx so the grouping makes sense. - assert '"-c"' in block assert '"--flash-attn"' in block + # -c is emitted in the conditional right after the base list, not inside + # it: auto-fit (--fit on with no pinned context) must omit -c entirely, + # because "-c 0" pins the full native context and disables --fit's + # VRAM-based sizing. Pin that it still sits next to the base block so the + # context grouping stays intact. + after = rest[end_rel : end_rel + 1000] + assert '"-c"' in after, ( + "-c must still be emitted in the conditional immediately after the " + "base cmd list (omitted only in auto-fit, where --fit sizes context)." + ) def _iter_lines_with_offset(text: str): diff --git a/studio/backend/tests/test_llama_cpp_props_readback.py b/studio/backend/tests/test_llama_cpp_props_readback.py index 488645ee5a..fe1e67edad 100644 --- a/studio/backend/tests/test_llama_cpp_props_readback.py +++ b/studio/backend/tests/test_llama_cpp_props_readback.py @@ -225,31 +225,46 @@ def test_kv_unified_added_for_multi_slot(): """Explicit --parallel N disables llama-server's auto-slots kv-unified default, splitting -c into per-slot windows of -c/N; Unsloth must restore the shared pool so one request can use the full advertised context.""" - flags = LlamaCppBackend._ctx_integrity_flags(4, False, 98304, 98304, _CAPS_ALL) + flags = LlamaCppBackend._ctx_integrity_flags(4, False, False, 98304, 98304, _CAPS_ALL) assert "--kv-unified" in flags def test_kv_unified_skipped_for_single_slot_or_old_build(): assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags( - 1, False, 98304, 98304, _CAPS_ALL + 1, False, False, 98304, 98304, _CAPS_ALL ) assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags( - 4, False, 98304, 98304, _CAPS_NONE + 4, False, False, 98304, 98304, _CAPS_NONE ) def test_fit_ctx_floors_explicit_request_under_fit(): - flags = LlamaCppBackend._ctx_integrity_flags(1, True, 98304, 98304, _CAPS_ALL) + # An explicit requested ctx floors --fit-ctx at that value on any --fit + # path, including legacy auto (auto_fit False). + flags = LlamaCppBackend._ctx_integrity_flags(1, True, False, 98304, 98304, _CAPS_ALL) assert flags[flags.index("--fit-ctx") + 1] == "98304" -def test_fit_ctx_skipped_without_fit_or_explicit_ctx_or_support(): +def test_fit_ctx_skipped_without_fit_or_support(): + # No --fit on -> no --fit-ctx. assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( - 1, False, 98304, 98304, _CAPS_ALL + 1, False, False, 98304, 98304, _CAPS_ALL ) - assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(1, True, 0, 262144, _CAPS_ALL) + # --fit on but the binary doesn't support --fit-ctx. assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( - 1, True, 98304, 98304, _CAPS_NONE + 1, True, True, 98304, 98304, _CAPS_NONE + ) + + +def test_fit_ctx_floors_auto_request_at_8192_only_under_auto_fit(): + # Manual + Auto (auto_fit) floors the auto window at 8192 so --fit can't + # shrink it to a tiny size. + flags = LlamaCppBackend._ctx_integrity_flags(1, True, True, 0, 262144, _CAPS_ALL) + assert flags[flags.index("--fit-ctx") + 1] == "8192" + # Legacy auto (fit on but not auto_fit) emits -c 0 to pin native, so the + # 8192 floor must NOT ride along and override that pin. + assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags( + 1, True, False, 0, 262144, _CAPS_ALL ) diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index ba52afad1c..c6d16363f8 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -747,6 +747,34 @@ def test_strip_shadowing_flags_defaults_strip_split_mode_too(): assert strip_shadowing_flags(["--split-mode", "tensor"]) == [] +def test_strip_offload_is_opt_in_and_covers_moe(): + base = dict( + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + ) + # Default: offload (incl. MoE) flags are NOT stripped. + assert strip_shadowing_flags(["--n-cpu-moe", "8", "--top-k", "20"], **base) == [ + "--n-cpu-moe", + "8", + "--top-k", + "20", + ] + # Opt-in strips layer AND MoE offload flags (value-aware), keeps the rest. + assert strip_shadowing_flags( + ["--n-cpu-moe", "8", "--gpu-layers", "33", "--fit", "off", "--top-k", "20"], + **base, + strip_offload = True, + ) == ["--top-k", "20"] + # Boolean --cpu-moe drops the flag only, not the following value. + assert strip_shadowing_flags(["--cpu-moe", "--seed", "-1"], **base, strip_offload = True) == [ + "--seed", + "-1", + ] + + @pytest.mark.parametrize( "args", [ @@ -796,6 +824,23 @@ def test_strip_split_mode_only_drops_tensor_split_too(): assert strip_split_mode_only(["-sm=tensor", "-ts=3,1"]) == [] +def test_strip_tensor_split_alone_preserves_split_mode(): + # Manual mode emits its own --tensor-split, so an inherited ratio is dropped + # -- but the user's --split-mode row/none/layer choice (which the manual + # ratio toggle can't express) must survive. strip_tensor_split removes only + # the ratio, unlike strip_split_mode which removes the whole group. + out = strip_shadowing_flags( + ["--split-mode", "row", "--tensor-split", "1,1", "--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + strip_tensor_split = True, + ) + assert out == ["--split-mode", "row", "--top-k", "20"] + + def test_strip_shadowing_flags_keeps_model_draft_without_spec(): out = strip_shadowing_flags( ["--model-draft", "/custom/mtp.gguf"], diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 06f72d3b9f..00c7aeac69 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -262,9 +262,12 @@ def test_proportional_tensor_split_is_emitted_in_tensor_mode(): src = _load_model_source() assert '"--tensor-split"' in src gate = src.find("if tensor_parallel:") - ts = src.find('"--tensor-split"') + # Find the TP block's emission (after the gate); manual mode emits its own + # --tensor-split earlier in the source from the user's per-GPU shares. + ts = src.find('"--tensor-split"', gate) nxt_else = src.find("self._tensor_parallel = False") assert 0 <= gate < ts < nxt_else, "--tensor-split must be emitted under `if tensor_parallel:`" + assert "tp_tensor_split" in src[gate:nxt_else] def test_mtp_decode_probe_wired_under_tensor_parallel(): diff --git a/studio/backend/tests/test_tp_vision_regression.py b/studio/backend/tests/test_tp_vision_regression.py index fb0989b306..d1372ca415 100644 --- a/studio/backend/tests/test_tp_vision_regression.py +++ b/studio/backend/tests/test_tp_vision_regression.py @@ -126,10 +126,21 @@ _ALLOWED_TP_DROP_GUARDS = { # Capability: --split-mode tensor aborted for this (binary, model) (#6415). # Self-healing -- tried by default, skipped only after a real abort (vs #6416). "tensor_parallel and self._tensor_split_aborts(binary, model_identifier)", - # Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve. - "tensor_parallel and len(tp_gpus) < 2", + # Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve. Gated + # on plan_tp (not raw tensor_parallel) so manual mode skips this planner (#6414). + "plan_tp and len(tp_gpus) < 2", # Capacity: pooled usable VRAM can't hold weights + MTP reserve -> layer split. "_tp_weight_budget_mib <= _tp_required_mib", + # Manual mode, Auto layers: --fit owns memory and is incompatible with a + # tensor split, so TP is dropped (surfaced via logger.info) before the + # cache-drop, so a quantized KV survives into the --fit load (#6414). + "tensor_parallel and gpu_memory_mode == 'manual' and (gpu_layers < 0)", + # Manual mode, explicit layers: a tensor split still needs >= 2 GPUs in use. + "tensor_parallel and gpu_memory_mode == 'manual' and (gpu_layers >= 0) and (self._effective_gpu_count(sorted(gpu_ids) if gpu_ids else None) < 2)", + # Manual mode, zero layers: nothing to split on the GPU, and a tensor-mode + # launch under the CPU-only GPU mask (no visible devices) aborts the server + # instead of the intended CPU-only load (#6414). + "gpu_memory_mode == 'manual' and gpu_layers == 0", } @@ -364,7 +375,7 @@ def test_compute_buffer_downgrade_preserves_multi_gpu_intent(): full GPU set too, so it is symmetric with the budget/geometry downgrades and doesn't collapse a multi-GPU layer load to one card (reviewer.py P1 on #6659).""" src = inspect.getsource(LlamaCppBackend.load_model) - gate = src.find("tensor_parallel and len(tp_gpus) < 2") + gate = src.find("plan_tp and len(tp_gpus) < 2") assert gate != -1 # Bound to exactly this block: from its gate to the next (budget) downgrade. nxt = src.find("_tp_weight_budget_mib <= _tp_required_mib", gate) diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py index c24ec28e1d..50b3cd3513 100644 --- a/studio/backend/utils/models/gguf_metadata.py +++ b/studio/backend/utils/models/gguf_metadata.py @@ -50,9 +50,11 @@ _CACHE_MAX_ENTRIES = 4096 # keyed by (file cache key, wanted key). None = key absent / file unreadable. _BOOL_CACHE: Dict[Tuple[_CacheKey, str], Optional[bool]] = {} -# Native training context length (``{arch}.context_length``). None = absent / -# unreadable. Lets the UI show the real context ceiling before a model loads. -_CONTEXT_CACHE: Dict[_CacheKey, Optional[int]] = {} +# GGUF header dims for the staged/deferred-load UI: context_length, layer_count +# (block_count), and moe_layer_count (block_count minus leading dense layers; 0 +# if not MoE). One cached pass fills all three so the staged sheet can size every +# slider before the model loads. None = unreadable / not a GGUF. +_DIMS_CACHE: Dict[_CacheKey, Optional[Dict[str, Optional[int]]]] = {} def _cache_key(path: str) -> Optional[_CacheKey]: @@ -142,32 +144,45 @@ def _parse_gguf_header(path: str) -> Optional[Dict[str, str]]: return out -def read_gguf_context_length(path: str) -> Optional[int]: - """Return the GGUF's native training context length (``{arch}.context_length``), - or ``None`` if missing/unreadable/not a GGUF. Cached by (path, mtime, size). - Lets the UI populate the context slider before the model is loaded.""" +def read_gguf_staged_dims(path: str) -> Optional[Dict[str, Optional[int]]]: + """GGUF header dims for the staged-load UI in one cached pass: + ``{"context_length", "layer_count", "moe_layer_count"}``. Each may be None + when absent (moe_layer_count is 0 for a dense model). Returns ``None`` if not + a GGUF / unreadable. Cached by (path, mtime, size). Lets the staged sheet size + the context, GPU-layers and MoE sliders before the model loads.""" key = _cache_key(path) if key is None: return None with _CACHE_LOCK: - if key in _CONTEXT_CACHE: - return _CONTEXT_CACHE[key] - result = _parse_gguf_context_length(path) + if key in _DIMS_CACHE: + return _DIMS_CACHE[key] + result = _parse_gguf_staged_dims(path) with _CACHE_LOCK: - while len(_CONTEXT_CACHE) >= _CACHE_MAX_ENTRIES: + while len(_DIMS_CACHE) >= _CACHE_MAX_ENTRIES: try: - _CONTEXT_CACHE.pop(next(iter(_CONTEXT_CACHE))) + _DIMS_CACHE.pop(next(iter(_DIMS_CACHE))) except StopIteration: break - _CONTEXT_CACHE[key] = result + _DIMS_CACHE[key] = result return result -def _parse_gguf_context_length(path: str) -> Optional[int]: - # The context key is architecture-namespaced (``llama.context_length`` etc.), - # so we learn the key only after reading ``general.architecture``. GGUF writes - # general.* before arch.* keys, matching the loader's own parser. - ctx_key: Optional[str] = None +def read_gguf_context_length(path: str) -> Optional[int]: + """Native training context length (``{arch}.context_length``), or ``None``. + Thin accessor over read_gguf_staged_dims.""" + dims = read_gguf_staged_dims(path) + return dims["context_length"] if dims else None + + +def _parse_gguf_arch_uints(path: str, wanted_suffixes: frozenset[str]) -> Optional[Dict[str, int]]: + """Walk a GGUF header once and return the requested architecture-namespaced + uint (vtype 4/10) keys, e.g. ``{"block_count": 32}``. Keys are + ``{arch}.``; the arch is learned from ``general.architecture`` (GGUF + writes general.* before arch.* keys, matching the loader's own parser). + Returns ``None`` if not a GGUF / unreadable, else a dict (possibly empty or + partial when some keys are absent).""" + arch: Optional[str] = None + found: Dict[str, int] = {} try: with open(path, "rb") as f: head = f.read(24) @@ -204,28 +219,68 @@ def _parse_gguf_context_length(path: str) -> Optional[int]: sbytes = f.read(slen) if len(sbytes) < slen: break - ctx_key = f"{sbytes.decode('utf-8', 'replace')}.context_length" - elif ctx_key is not None and key == ctx_key and vtype in (4, 10): + arch = sbytes.decode("utf-8", "replace") + elif ( + arch is not None + and vtype in (4, 10) + and key.startswith(f"{arch}.") + and key[len(arch) + 1 :] in wanted_suffixes + ): width = 4 if vtype == 4 else 8 n_bytes = f.read(width) if len(n_bytes) < width: break - value = struct.unpack(" 0 else None + found[key[len(arch) + 1 :]] = struct.unpack( + " Optional[Dict[str, Optional[int]]]: + vals = _parse_gguf_arch_uints( + path, + frozenset( + { + "context_length", + "block_count", + "expert_count", + "leading_dense_block_count", + } + ), + ) + if vals is None: + return None + ctx = vals.get("context_length") + block = vals.get("block_count") + # A real context/layer count is positive; treat 0/garbage as absent so the + # UI never builds a slider with max < min. + context_length = ctx if ctx and ctx > 0 else None + layer_count = block if block and block > 0 else None + # MoE layer count = block_count - leading dense layers, only when experts + # exist; else 0 (dense -> slider hidden). Mirrors n_moe_layers in + # core/inference/llama_cpp.py. + if not vals.get("expert_count") or not block: + moe_layer_count: Optional[int] = 0 + else: + moe_layer_count = max(0, block - (vals.get("leading_dense_block_count") or 0)) + return { + "context_length": context_length, + "layer_count": layer_count, + "moe_layer_count": moe_layer_count, + } # Strings (8) and arrays (9) are handled inline. diff --git a/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts b/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts index ec75b17f20..08492ab480 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts @@ -2,7 +2,9 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 // Per-model pre-load inference settings, persisted in localStorage so the load -// dialog can offer "Remember settings for ". +// dialog can offer "Remember settings for ". GGUF picks only: every +// field is a llama.cpp load knob, so all save/restore call sites gate on +// GGUF-ness (a non-GGUF blob would only snapshot leftover standing values). const KEY = "unsloth_load_settings"; @@ -12,14 +14,22 @@ export interface RememberedLoadSettings { speculativeType: string | null; specDraftNMax: number | null; tensorParallel: boolean; + // GPU Memory controls. Optional so an older blob (which lacked them) still + // parses, leaving the live knobs untouched on apply. The mode is kept with the + // manual knobs (gpuLayers/nCpuMoe are ignored outside Manual mode). A null + // selectedGpuIds is meaningful (all GPUs), so it's distinguished from absent. + // The per-GPU split ratio is deliberately NOT remembered: it's positionally + // bound to the exact GPU set/order and unvalidated, so it would mismatch. + gpuMemoryMode?: "auto" | "manual"; + gpuLayers?: number; + nCpuMoe?: number; + selectedGpuIds?: number[] | null; } -// Storage key for a pick's remembered settings. The remembered knobs are -// VRAM-budget driven (context override, KV-cache dtype, tensor-parallel), so the -// right values differ per quant. An HF repo collapses all its GGUF variants into -// one `id`, so fold the variant in to scope settings per quant. Local .gguf -// paths key by their file path (already file-specific); native drag-drop files -// key by display label, so same-named files in different folders share an entry. +// Storage key for a pick's remembered settings, scoped per quant (the VRAM-budget +// knobs differ per quant). An HF repo collapses its GGUF variants into one `id`, +// so fold the variant in. Local .gguf paths are already file-specific; native +// drag-drop files key by display label, so same-named files share an entry. export function rememberedLoadSettingsKey(selection: { id: string; ggufVariant?: string | null; diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 0bf46e7343..7083f02288 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -45,12 +45,18 @@ import { import { type PendingImageEditReference, type RagAutoInject, + GPU_LAYERS_AUTO, + loadedGpuMemoryFieldsUnlessStaged, + reconcilePersistedGpuIds, resolveLoadedSpeculativeSettings, resolveSpeculativeSettingsForLoad, + persistGpuMemoryModeOnLoad, resolveToolsEnabledOnLoad, saveSpeculativeType, useChatRuntimeStore, } from "../stores/chat-runtime-store"; +import { resolveFitMaxSeqLength, resolveManualAutoCtxPin } from "../presets/preset-policy"; +import { ensureGpuDeviceCache } from "@/hooks/use-gpu-info"; import { useExternalProvidersStore } from "../stores/external-providers-store"; import { shouldPreserveFullOutput, @@ -1489,6 +1495,13 @@ async function autoLoadSmallestModel(): Promise<{ max_seq_length: number; is_lora: boolean; gguf_variant?: string | null; + // GGUF-only: scopes the training guard to the same placement policy /load + // will use. Manual mode must match because it makes placement user-owned. + // The layer/MoE/split/KV/spec knobs are deliberately not sent: Auto mode's + // guard sizes conservatively, while Manual mode bypasses that estimate. + // The safetensors fallback omits both fields and uses HF auto-placement. + gpu_ids?: number[]; + gpu_memory_mode?: "auto" | "manual"; }): Promise { const validation = await validateModel({ ...payload, @@ -1520,12 +1533,18 @@ async function autoLoadSmallestModel(): Promise<{ return false; } const currentStore = useChatRuntimeStore.getState(); - const remembered = loadRememberedLoadSettings( - rememberedLoadSettingsKey({ - id: candidate.id, - ggufVariant: candidate.ggufVariant, - }), - ); + // Blobs are saved for GGUF picks only (the sheet gates on it), so don't + // let a legacy non-GGUF blob feed a stale context/spec choice into a + // safetensors auto-load. + const remembered = + candidate.kind === "gguf" + ? loadRememberedLoadSettings( + rememberedLoadSettingsKey({ + id: candidate.id, + ggufVariant: candidate.ggufVariant, + }), + ) + : null; const effectiveMaxSeqLength = resolveLoadMaxSeqLength({ modelId: candidate.id, ggufVariant: candidate.ggufVariant, @@ -1537,6 +1556,38 @@ async function autoLoadSmallestModel(): Promise<{ maxSeqLength: candidate.maxSeqLength, presetSource: currentStore.activePresetSource, }); + // The GPU knobs are per-model, so read them from the same remembered + // settings that fed effectiveMaxSeqLength -- on a background auto-load the + // live store holds session defaults, not the saved Manual mode / layer pin / + // GPU pick. Absent fields fall back like applyRememberedLoadSettings: the + // mode to the store (a persisted standing preference), the per-model knobs to + // their defaults. The saved GPU pick is reconciled against the GPUs present + // now, like the interactive restore. + const effectiveGpuMemoryMode = + remembered?.gpuMemoryMode ?? currentStore.gpuMemoryMode; + const effectiveGpuLayers = remembered?.gpuLayers ?? GPU_LAYERS_AUTO; + const effectiveNCpuMoe = remembered?.nCpuMoe ?? 0; + if (remembered?.selectedGpuIds != null) { + // Warm the device cache first: on a cold cache the reconcile passes the + // saved pick through unvalidated, and a stale cross-host pick then fails + // the load with the picker hidden. + await ensureGpuDeviceCache(); + } + const effectiveGpuIds = + remembered?.selectedGpuIds !== undefined + ? reconcilePersistedGpuIds(remembered.selectedGpuIds) + : null; + // Under Manual GPU memory + Auto layers, llama.cpp's --fit owns context + // sizing, so send 0 (or the pinned length). GGUF-only; a no-op otherwise. + // The context pin is per-model too, so it comes from remembered settings, + // not the live store. + const fitMaxSeqLength = resolveFitMaxSeqLength( + candidate.kind === "gguf", + effectiveGpuMemoryMode, + effectiveGpuLayers, + remembered?.contextLength ?? null, + effectiveMaxSeqLength, + ); const effectiveSpeculativeType = remembered?.speculativeType ?? specSettings.speculativeType; const effectiveSpecDraftNMax = @@ -1544,9 +1595,16 @@ async function autoLoadSmallestModel(): Promise<{ if ( !(await canAutoLoad({ model_path: candidate.id, - max_seq_length: effectiveMaxSeqLength, + max_seq_length: fitMaxSeqLength, is_lora: false, gguf_variant: candidate.ggufVariant, + // The same remembered-derived GPU pick the load below sends. + ...(candidate.kind === "gguf" + ? { + gpu_ids: effectiveGpuIds ?? undefined, + gpu_memory_mode: effectiveGpuMemoryMode, + } + : {}), })) ) { skippedAutoLoadCandidates.add( @@ -1558,7 +1616,7 @@ async function autoLoadSmallestModel(): Promise<{ const loadResp = await loadModel({ model_path: candidate.id, hf_token: hfToken, - max_seq_length: effectiveMaxSeqLength, + max_seq_length: fitMaxSeqLength, load_in_4bit: true, is_lora: false, gguf_variant: candidate.ggufVariant, @@ -1567,8 +1625,22 @@ async function autoLoadSmallestModel(): Promise<{ speculative_type: effectiveSpeculativeType, spec_draft_n_max: effectiveSpecDraftNMax, tensor_parallel: remembered?.tensorParallel ?? false, + // GGUF-only: the safetensors fallback loads via HF auto-placement (no + // explicit pins). The split ratio is deliberately never remembered + // (positionally bound to an exact GPU set), so auto-load leaves llama.cpp's + // free-VRAM default in charge rather than sending a stale store value. + ...(candidate.kind === "gguf" + ? { + gpu_memory_mode: effectiveGpuMemoryMode, + gpu_layers: effectiveGpuLayers, + n_cpu_moe: effectiveNCpuMoe, + gpu_ids: effectiveGpuIds ?? undefined, + } + : {}), }); saveSpeculativeType(effectiveSpeculativeType); + // Self-gates on is_gguf (skips diffusion), so persists only for a real GGUF load. + persistGpuMemoryModeOnLoad(loadResp, effectiveGpuMemoryMode); useChatRuntimeStore .getState() .setCheckpoint(candidate.id, candidate.ggufVariant ?? undefined); @@ -1597,6 +1669,15 @@ async function autoLoadSmallestModel(): Promise<{ store.setModels([...store.models, autoModel]); } if (candidate.kind === "gguf") { + // Keep an explicit Manual+Auto context pin the load just applied (so a + // later Apply doesn't silently revert it to auto-fit sizing), mirroring + // the interactive path's keepCustomCtx; other cases baseline on + // ggufContextLength. + const keepCustomCtx = resolveManualAutoCtxPin( + effectiveGpuMemoryMode, + effectiveGpuLayers, + remembered?.contextLength ?? null, + ); useChatRuntimeStore.setState({ ggufContextLength: loadResp.context_length ?? 131072, ggufMaxContextLength: @@ -1613,6 +1694,10 @@ async function autoLoadSmallestModel(): Promise<{ loadedKvCacheDtype: loadResp.cache_type_kv ?? null, tensorParallel: loadResp.tensor_parallel ?? false, loadedTensorParallel: loadResp.tensor_parallel ?? false, + ...loadedGpuMemoryFieldsUnlessStaged(loadResp, { + customContextLength: keepCustomCtx, + }), + loadedCustomContextLength: keepCustomCtx, defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, loadedChatTemplateOverride: null, @@ -1633,6 +1718,9 @@ async function autoLoadSmallestModel(): Promise<{ loadedKvCacheDtype: loadResp.cache_type_kv ?? null, tensorParallel: loadResp.tensor_parallel ?? false, loadedTensorParallel: loadResp.tensor_parallel ?? false, + // Non-GGUF response: clears any stale GPU baseline a prior manual-GPU + // GGUF load left, matching the interactive/status sibling load paths. + ...loadedGpuMemoryFieldsUnlessStaged(loadResp), defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, loadedChatTemplateOverride: null, @@ -1820,12 +1908,17 @@ async function autoLoadSmallestModel(): Promise<{ duration: 30000, }); try { + const rt = useChatRuntimeStore.getState(); if ( !(await canAutoLoad({ model_path: "unsloth/Qwen3.5-4B-MTP-GGUF", max_seq_length: 0, is_lora: false, gguf_variant: "UD-Q4_K_XL", + // The same live-store GPU pick the load below sends (a fresh default + // model has no remembered settings to prefer). + gpu_ids: rt.selectedGpuIds ?? undefined, + gpu_memory_mode: rt.gpuMemoryMode, })) ) { toast.dismiss(toastId); @@ -1835,6 +1928,9 @@ async function autoLoadSmallestModel(): Promise<{ const loadResp = await loadModel({ model_path: "unsloth/Qwen3.5-4B-MTP-GGUF", hf_token: hfToken, + // Model default under both modes: Auto layers + no pin means + // resolveFitMaxSeqLength returns 0 for every mode (the canAutoLoad + // preflight above sends the same). max_seq_length: 0, load_in_4bit: true, is_lora: false, @@ -1842,8 +1938,20 @@ async function autoLoadSmallestModel(): Promise<{ trust_remote_code: trustRemoteCode, speculative_type: specSettings.speculativeType, spec_draft_n_max: specSettings.specDraftNMax, + // GPU Memory mode is a standing preference, so honor it on auto-load. + // The layer/MoE/split knobs and the context pin are per-model: the live + // store may hold edits drafted for a staged pick, and a fresh default + // model has no remembered settings, so those stay at their defaults like + // the cached-candidate path. The GPU pick deliberately differs (it's the + // picker's current on-screen selection, which the canAutoLoad preflight + // above already committed to). + gpu_memory_mode: rt.gpuMemoryMode, + gpu_layers: GPU_LAYERS_AUTO, + n_cpu_moe: 0, + gpu_ids: rt.selectedGpuIds ?? undefined, }); saveSpeculativeType(specSettings.speculativeType); + persistGpuMemoryModeOnLoad(loadResp, rt.gpuMemoryMode); useChatRuntimeStore .getState() .setCheckpoint("unsloth/Qwen3.5-4B-MTP-GGUF", "UD-Q4_K_XL"); @@ -1880,6 +1988,10 @@ async function autoLoadSmallestModel(): Promise<{ loadedKvCacheDtype: loadResp.cache_type_kv ?? null, tensorParallel: loadResp.tensor_parallel ?? false, loadedTensorParallel: loadResp.tensor_parallel ?? false, + ...loadedGpuMemoryFieldsUnlessStaged(loadResp), + // Drives the GPU Memory controls' diffusion gate; set alongside the + // GPU fields on every load path so the gate can't read stale. + loadedIsDiffusion: loadResp.is_diffusion ?? false, defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, loadedIsMultimodal: isMultimodalResponse(loadResp), diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index ebf9461172..0f6af38033 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -127,28 +127,38 @@ export async function validateModel( native_path_lease: payload.nativePathLease ?? null, hf_token: payload.hf_token, gguf_variant: payload.gguf_variant ?? null, - // Send the intended load settings so validate's VRAM check matches the - // follow-up /load and doesn't unload for a load /load would then reject. + // Intended load settings so validate's preflight matches the follow-up + // /load. Default placement is sized against the selected GPUs. max_seq_length: payload.max_seq_length, load_in_4bit: payload.load_in_4bit, + gpu_ids: payload.gpu_ids, + // Manual placement is an explicit override: Auto layers use llama.cpp + // --fit, while a pinned layer count is owned by the user. Tell validate + // so it applies the same training-guard policy as /load. + gpu_memory_mode: payload.gpu_memory_mode, }), }); return parseJsonOrThrow(response); } /** - * Read a GGUF's native context length from its local header (no GPU load, no - * download). Returns null when the file isn't downloaded yet, the model isn't a - * GGUF, or it's gated. For a native (drag-drop / picked) file, pass - * `nativePathToken` so the backend reads the granted local path. Used by the - * deferred-load staging flow to fill the context slider before the single load. + * Read a GGUF's header dims (native context length, total layer count, MoE + * expert-layer count) from its local file (no GPU load, no download). All are + * null when the file isn't downloaded yet, the model isn't a GGUF, or it's + * gated. For a native (drag-drop / picked) file, pass `nativePathToken` so the + * backend reads the granted local path. Used by the deferred-load staging flow + * to size the context, GPU-layers and MoE sliders before the single load. */ -export async function fetchGgufContextLength(payload: { +export async function fetchGgufStagedMetadata(payload: { model_path: string; gguf_variant?: string | null; hf_token?: string | null; nativePathToken?: string | null; -}): Promise { +}): Promise<{ + contextLength: number | null; + layerCount: number | null; + moeLayerCount: number | null; +}> { let nativePathLease: string | null = null; if (payload.nativePathToken) { try { @@ -156,8 +166,8 @@ export async function fetchGgufContextLength(payload: { await consumeNativePathToken(payload.nativePathToken, "validate-model") ).nativePathLease; } catch { - // Lease expired / revoked: degrade to no context (the load can re-mint). - return null; + // Lease expired / revoked: degrade to no metadata (the load can re-mint). + return { contextLength: null, layerCount: null, moeLayerCount: null }; } } const response = await authFetch("/api/inference/validate", { @@ -172,7 +182,11 @@ export async function fetchGgufContextLength(payload: { }), }); const res = await parseJsonOrThrow(response); - return res.context_length ?? null; + return { + contextLength: res.context_length ?? null, + layerCount: res.layer_count ?? null, + moeLayerCount: res.moe_layer_count ?? null, + }; } export async function unloadModel(payload: UnloadModelRequest): Promise { diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index ec0ad977bf..217eaf8b6d 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -1445,9 +1445,11 @@ export function ChatPage({ // were already seeded on stage, so keepSpeculative only when a config was // saved -- otherwise the standing speculative preference should win. autoLoadStagedRef.current = (pending) => { - const remembered = loadRememberedLoadSettings( - rememberedLoadSettingsKey(pending), - ); + // Blobs are saved for GGUF picks only (the sheet gates on it), so don't + // let a legacy non-GGUF blob claim a seeded config here. + const remembered = hasGgufSource(pending) + ? loadRememberedLoadSettings(rememberedLoadSettingsKey(pending)) + : null; void selectModel({ ...pending, isDownloaded: true, @@ -2813,6 +2815,11 @@ export function ChatPage({ selectModel({ id: state.params.checkpoint, ggufVariant: state.activeGgufVariant ?? undefined, + // A native (drag-drop / picked) GGUF's checkpoint is only a display + // label, so the reload needs its path token to re-mint a lease -- + // else applying the now-exposed GPU/context controls can't resolve + // the file. Null for non-native loads, which reload by id as before. + nativePathToken: state.activeNativePathToken ?? undefined, forceReload: true, isDownloaded: true, loadingDescription: "Reloading with updated chat template.", diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index cedd298ecf..b368a811fa 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -55,6 +55,7 @@ import { Switch } from "@/components/ui/switch"; import { Textarea } from "@/components/ui/textarea"; import { InfoHint } from "@/components/ui/info-hint"; import { Tooltip, TooltipContent } from "@/components/ui/tooltip"; +import { useGpuDevices } from "@/hooks/use-gpu-info"; import { useIsMobile } from "@/hooks/use-mobile"; import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check"; import { cn } from "@/lib/utils"; @@ -99,8 +100,11 @@ import { providerSupportsFastMode, } from "./provider-capabilities"; import { + GPU_LAYERS_AUTO, + distributeByWeight, isPendingGguf, pendingSelectionMatches, + rebalanceSplit, useChatRuntimeStore, } from "./stores/chat-runtime-store"; import { RetrievalSettingsSection } from "@/features/rag/components/retrieval-settings-section"; @@ -250,6 +254,7 @@ function ParamSlider({ displayValue, info, valueSize, + disabled, }: { label: string; value: number; @@ -260,6 +265,7 @@ function ParamSlider({ displayValue?: string; info?: ReactNode; valueSize?: number; + disabled?: boolean; }) { return (
@@ -279,6 +285,7 @@ function ParamSlider({ displayValue={displayValue} ariaLabel={label} size={valueSize ?? 4} + disabled={disabled} />
onChange(snapToStep(v, step, min, max))} className="panel-slider" + disabled={disabled} /> ); @@ -540,8 +548,17 @@ export function ChatSettingsPanel({ const base = slash >= 0 ? id.slice(slash + 1) : id; return base || id; })(); + const activeNativePathToken = useChatRuntimeStore( + (s) => s.activeNativePathToken, + ); + const loadedGgufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); + // A GGUF loaded from a native path / direct .gguf has no HF variant, so key + // off the same signal the status hydration uses -- variant OR native token OR + // a GGUF context -- else the GPU Memory controls hide for a loaded local GGUF. const isLoadedGguf = - useChatRuntimeStore((s) => s.activeGgufVariant) != null; + useChatRuntimeStore((s) => s.activeGgufVariant) != null || + activeNativePathToken != null || + loadedGgufContextLength != null; // While a pick is staged the sheet configures *that* model, so its GGUF-ness // (not the currently loaded model's) decides whether the GGUF-only controls // show. Otherwise a staged non-GGUF Hub repo would inherit the loaded GGUF's @@ -607,6 +624,25 @@ export function ChatSettingsPanel({ const loadedTensorParallel = useChatRuntimeStore( (s) => s.loadedTensorParallel, ); + const gpuMemoryMode = useChatRuntimeStore((s) => s.gpuMemoryMode); + const setGpuMemoryMode = useChatRuntimeStore((s) => s.setGpuMemoryMode); + const loadedGpuMemoryMode = useChatRuntimeStore((s) => s.loadedGpuMemoryMode); + const loadedIsDiffusion = useChatRuntimeStore((s) => s.loadedIsDiffusion); + const gpuLayers = useChatRuntimeStore((s) => s.gpuLayers); + const setGpuLayers = useChatRuntimeStore((s) => s.setGpuLayers); + const loadedGpuLayers = useChatRuntimeStore((s) => s.loadedGpuLayers); + const nCpuMoe = useChatRuntimeStore((s) => s.nCpuMoe); + const setNCpuMoe = useChatRuntimeStore((s) => s.setNCpuMoe); + const loadedNCpuMoe = useChatRuntimeStore((s) => s.loadedNCpuMoe); + const splitRatio = useChatRuntimeStore((s) => s.splitRatio); + const setSplitRatio = useChatRuntimeStore((s) => s.setSplitRatio); + const loadedSplitRatio = useChatRuntimeStore((s) => s.loadedSplitRatio); + const ggufLayerCount = useChatRuntimeStore((s) => s.ggufLayerCount); + const moeLayerCount = useChatRuntimeStore((s) => s.moeLayerCount); + const selectedGpuIds = useChatRuntimeStore((s) => s.selectedGpuIds); + const setSelectedGpuIds = useChatRuntimeStore((s) => s.setSelectedGpuIds); + const loadedGpuIds = useChatRuntimeStore((s) => s.loadedGpuIds); + const gpuDevices = useGpuDevices(); const chatTemplateOverride = useChatRuntimeStore( (s) => s.chatTemplateOverride, ); @@ -614,6 +650,9 @@ export function ChatSettingsPanel({ (s) => s.loadedChatTemplateOverride, ); const customContextLength = useChatRuntimeStore((s) => s.customContextLength); + const loadedCustomContextLength = useChatRuntimeStore( + (s) => s.loadedCustomContextLength, + ); const setCustomContextLength = useChatRuntimeStore( (s) => s.setCustomContextLength, ); @@ -641,10 +680,14 @@ export function ChatSettingsPanel({ : null; useEffect(() => { if (!pendingKey) return; - const saved = loadRememberedLoadSettings(pendingKey); + // GGUF-only, like the stageOrLoad / Hub restore paths: every remembered + // field is a llama.cpp knob, so a non-GGUF pick has nothing to restore -- + // and applying its blob would clobber the standing gpuMemoryMode with a + // stale snapshot (the save on Load below is gated the same way). + const saved = pendingIsGguf ? loadRememberedLoadSettings(pendingKey) : null; setRemember(saved != null); if (saved) applyRememberedLoadSettings(saved); - }, [pendingKey, applyRememberedLoadSettings]); + }, [pendingKey, pendingIsGguf, applyRememberedLoadSettings]); // While staging, the sheet reflects the STAGED model, so its header context // takes precedence over the loaded model's (which may differ or be larger). const baseContext = pendingIsGguf ? stagedContextLength : ggufContextLength; @@ -661,15 +704,132 @@ export function ChatSettingsPanel({ const ctxDisplayValue = customContextLength ?? baseContext ?? ""; const ctxMaxValue = baseNativeContext ?? baseContext ?? null; const kvDirty = kvCacheDtype !== loadedKvCacheDtype; - const ctxDirty = customContextLength !== null; + const ctxDirty = customContextLength !== loadedCustomContextLength; const specDirty = speculativeType !== loadedSpeculativeType; const specDraftDirty = specDraftNMax !== loadedSpecDraftNMax; const tpDirty = tensorParallel !== (loadedTensorParallel ?? false); + // A loaded diffusion GGUF runs mode-agnostic (pins all layers on one GPU, + // ignores --fit/--gpu-layers), so the GPU Memory mode + manual controls don't + // apply -- hide them and don't let the preserved standing mode read as dirty. + // The GPU picker still applies (diffusion pins the chosen device). A staged pick + // keeps the controls (a pending pick's diffusion-ness isn't known until load). + const gpuModeApplies = + isGguf && (pendingSelection != null || !loadedIsDiffusion); + const gpuDirty = + gpuModeApplies && gpuMemoryMode !== (loadedGpuMemoryMode ?? "auto"); + const isManual = gpuModeApplies && gpuMemoryMode === "manual"; + // Manual with the GPU Layers slider at "Auto" (leftmost): --fit owns the whole + // layout, so the offload knobs (MoE, split, TP) don't apply. + const autoLayers = isManual && gpuLayers < 0; + // GPUs actually in use: the picked subset, or all visible when none picked. + const gpusInUse = selectedGpuIds ?? gpuDevices.map((d) => d.index); + // TP is off with fewer than 2 GPUs in use (single GPU, or the picker narrowed + // to one): tensor split is a no-op there and aborts on some archs. Mirrors the + // multi-GPU gate on the GPU picker / Split ratio. (Under Auto layers the whole + // TP control is hidden -- llama.cpp's --fit aborts under --split-mode tensor.) + const tpDisabled = gpusInUse.length <= 1; + // Manual gpu-layers ceiling = model layer count + 1 (else a safe fallback): + // llama.cpp counts the output layer as one more offloadable layer past the + // repeating blocks ("offloaded 33/33" needs -ngl 33 on a 32-block model), so + // the slider max must reach it or full offload is unreachable. While staging, + // use the staged model's layer count (read from its header). + const stagedLayerCount = pendingSelection?.layerCount ?? null; + const modelLayerCount = pendingIsGguf ? stagedLayerCount : ggufLayerCount; + const gpuLayersMax = modelLayerCount != null ? modelLayerCount + 1 : 256; + // MoE-offload slider: shown only for MoE models, capped at their MoE-layer + // count. While staging, use the staged model's count (read from its header); + // otherwise the loaded model's. + const stagedMoeLayerCount = pendingSelection?.moeLayerCount ?? null; + const moeLayersMax = pendingIsGguf + ? (stagedMoeLayerCount ?? 0) + : (moeLayerCount ?? 0); + const showMoeSlider = isManual && !autoLayers && moeLayersMax > 0; + // gpuLayers always counts; MoE only with an explicit layer count (see above). + const manualDirty = + isManual && + (gpuLayers !== loadedGpuLayers || + (!autoLayers && nCpuMoe !== (loadedNCpuMoe ?? 0))); + // GPU picker: only meaningful on multi-GPU, and only when the reported + // indices are physical (relative ordinals from a parent CUDA_VISIBLE_DEVICES + // mask can't be mapped back to pin a device). null = use all (auto). + const showGpuPicker = + isGguf && + gpuDevices.length > 1 && + gpuDevices.every((d) => d.physicalIndex); + const isGpuChecked = (index: number) => + selectedGpuIds === null || selectedGpuIds.includes(index); + const toggleGpu = (index: number) => { + const all = gpuDevices.map((d) => d.index); + const current = selectedGpuIds ?? all; + const next = current.includes(index) + ? current.filter((i) => i !== index) + : [...current, index].sort((a, b) => a - b); + if (next.length === 0) return; // keep at least one GPU selected + setSelectedGpuIds(next.length === all.length ? null : next); + // The per-GPU split is positional, so any change to the set of GPUs in use + // invalidates it: drop it (the sliders fall back to the VRAM-weighted + // default). TP needs 2+ GPUs, so disable it when only one remains. + setSplitRatio(null); + if (next.length <= 1) { + setTensorParallel(false); + } + }; + const gpuIdsKey = (ids: number[] | null) => (ids === null ? "auto" : ids.join(",")); + const gpuIdsDirty = gpuIdsKey(selectedGpuIds) !== gpuIdsKey(loadedGpuIds); + // Per-GPU layer split (--tensor-split): manual + 2+ GPUs in use. One slider + // per GPU, each a layer count; together they sum to the GPU Layers total. + const showSplitRatio = + isManual && !autoLayers && showGpuPicker && gpusInUse.length > 1; + // The total the per-GPU counts sum to (the GPU Layers slider value); 0 under + // Auto, where the split is hidden. The devices behind the GPUs in use, for + // labels + the VRAM-weighted default. + const splitTotal = Math.max(0, Math.min(gpuLayers, gpuLayersMax)); + const gpusInUseDevices = gpusInUse.map( + (i) => gpuDevices.find((d) => d.index === i) ?? null, + ); + // Displayed per-GPU counts. splitRatio is a stable reference balance (only a + // slider edit changes it), rescaled to the current total; deriving rather than + // mutating it on GPU Layers changes keeps the balance intact when the total + // passes through low values or Auto. No saved split: free-VRAM-weighted default + // (llama.cpp's unset default splits by free VRAM, so the first edit starts from + // the default's placement, not a total-VRAM ratio that can land layers on a + // busy GPU). A genuine 0 (a full GPU) is a real weight, not missing data: the + // probe's no-data case degrades to the total server-side, and an all-zero list + // falls back to an even split in distributeByWeight. Not yet sent. + const splitCounts = + splitRatio && splitRatio.length === gpusInUse.length + ? distributeByWeight(splitTotal, splitRatio) + : distributeByWeight( + splitTotal, + gpusInUseDevices.map((d) => d?.memoryFreeGb ?? d?.memoryTotalGb ?? 1), + ); + const setSplitCount = (k: number, v: number) => + setSplitRatio(rebalanceSplit(splitTotal, splitCounts, k, v)); + const splitRatioDirty = + isManual && + !autoLayers && + JSON.stringify(splitRatio ?? null) !== JSON.stringify(loadedSplitRatio ?? null); + // Auto-fit context (Manual + Auto layers): <= 0 means "Auto" (--fit sizes it); + // a positive value pins it. Surface the length --fit chose once it's loaded. + const fitCtxAuto = autoLayers && (customContextLength ?? 0) <= 0; + const loadedAutoLayers = + loadedGpuMemoryMode === "manual" && (loadedGpuLayers ?? GPU_LAYERS_AUTO) < 0; + const fitResolvedCtx = + fitCtxAuto && loadedAutoLayers ? ggufContextLength : null; // A saved chat-template override is a reload-time setting too, so surface // Apply for a template-only edit (otherwise it could never be applied). const templateDirty = chatTemplateOverride !== loadedChatTemplateOverride; const modelSettingsDirty = - kvDirty || ctxDirty || specDirty || specDraftDirty || tpDirty || templateDirty; + kvDirty || + ctxDirty || + specDirty || + specDraftDirty || + tpDirty || + gpuDirty || + manualDirty || + gpuIdsDirty || + splitRatioDirty || + templateDirty; const [presetNameInput, setPresetNameInput] = useState(activePreset); const [systemPromptEditorOpen, setSystemPromptEditorOpen] = useState(false); const [systemPromptDraft, setSystemPromptDraft] = useState(""); @@ -980,7 +1140,64 @@ export function ChatSettingsPanel({ )} {isGguf && ( <> - {showContextControl && ( + {showContextControl && (autoLayers ? ( +
+
+
+ + Context Length + + + Auto: llama.cpp's --fit sizes the context to fit VRAM. + Set a length to pin it instead -- --fit then optimizes + GPU layer offload around it. The length --fit chose + shows here after loading. + +
+ { + setCustomContextLength(v > 0 ? v : null); + }} + ariaLabel="Context Length" + size={8} + disabled={modelControlsDisabled} + /> +
+ { + // Far-left snaps to Auto; otherwise to the nearest 1024. + if (v < 512) { + setCustomContextLength(null); + } else { + setCustomContextLength(Math.round(v / 1024) * 1024); + } + }} + className="panel-slider" + disabled={modelControlsDisabled} + /> + {fitResolvedCtx != null && ( +

+ llama.cpp loaded {fitResolvedCtx.toLocaleString()} tokens. +

+ )} +
+ ) : (
@@ -1036,7 +1253,7 @@ export function ChatSettingsPanel({

)}
- )} + ))}
@@ -1191,6 +1408,163 @@ export function ChatSettingsPanel({ )} )} + {gpuModeApplies && ( +
+
+ + GPU Memory + + +
+
+ Default: Unsloth + fits the model and context to your GPUs. +
+
+ Manual: set GPU + Layers yourself. Leave it on Auto to let llama.cpp size + the context and offload overflow (including MoE experts) + to RAM. +
+
+ +
+
+ +
+
+ )} + {isManual && ( + <> + + Layers to keep on the GPU (--gpu-layers); the rest run + on CPU. Auto lets llama.cpp size the split (and the + context) to fit VRAM. At the maximum, the whole model + is on the GPU. + + } + /> + {showMoeSlider && ( + + Keep the experts of this many MoE layers on the CPU + (--n-cpu-moe) to save VRAM. 0 = all experts on the + GPU; at the maximum, all are on the CPU. + + } + /> + )} + {showSplitRatio && ( +
+
+ + Layers per GPU + + + Splits GPU Layers across GPUs (--tensor-split). + Without Tensor Parallelism each value is the layer + count on that GPU; with it, every GPU holds a slice + of each layer, so the values are only a ratio. + +
+ {gpusInUseDevices.map((d, k) => ( + setSplitCount(k, v)} + valueSize={6} + disabled={modelControlsDisabled} + /> + ))} +
+ )} + + )} + {showGpuPicker && ( +
+
+ + GPUs + + + Which GPUs this model may use. Unchecked GPUs are hidden + from llama.cpp (CUDA_VISIBLE_DEVICES, or + HIP_VISIBLE_DEVICES on ROCm). Leave all checked to use + every GPU. + +
+
+ {gpuDevices.map((d) => ( +
+ + GPU {d.index}: {d.name} + {d.memoryTotalGb + ? ` · ${Math.round(d.memoryTotalGb)} GB` + : ""} + + toggleGpu(d.index)} + data-test-id={`gpu-pick-${d.index}`} + disabled={modelControlsDisabled} + /> +
+ ))} +
+
+ )} + {gpuModeApplies && !autoLayers && (
@@ -1206,10 +1580,11 @@ export function ChatSettingsPanel({ className="panel-switch shrink-0" checked={tensorParallel} onCheckedChange={setTensorParallel} - disabled={modelControlsDisabled} + disabled={tpDisabled || modelControlsDisabled} data-test-id="tensor-parallel-switch" />
+ )} )} {/* No persistent "enable custom code" toggle: it is consented per model @@ -1228,14 +1603,21 @@ export function ChatSettingsPanel({ {Math.round((stagedDownloadFraction ?? 0) * 100)}%

)} - + {/* GGUF picks only: a non-GGUF pick shows none of the load + knobs the blob captures, so there is nothing to remember. */} + {pendingIsGguf && ( + + )} {stagedLoading ? ( // Mid-load: nothing to load or abandon until it settles, so disable.
) : null} - + {/* The template override is a load-time knob too (applied on the next + reload) and the in-flight load already snapshotted it, so lock its + editors like the sibling controls -- a mid-load save would be + silently clobbered by the load response despite its toast. */} +
)} @@ -2086,7 +2477,7 @@ function BypassPermissionsToggle() { ); } -function ChatTemplateFields() { +function ChatTemplateFields({ disabled = false }: { disabled?: boolean }) { const defaultTemplate = useChatRuntimeStore((s) => s.defaultChatTemplate); const override = useChatRuntimeStore((s) => s.chatTemplateOverride); const setOverride = useChatRuntimeStore((s) => s.setChatTemplateOverride); @@ -2120,7 +2511,8 @@ function ChatTemplateFields() { @@ -2131,7 +2523,8 @@ function ChatTemplateFields() { -
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 10e0904e4f..3003b52230 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -29,9 +29,14 @@ import { } from "../api/chat-api"; import { formatEta, formatRate } from "../utils/format-transfer"; import { + GPU_LAYERS_AUTO, isLocalModelPath, + loadedGpuMemoryFields, + loadedGpuMemoryFieldsUnlessStaged, pendingSelectionMatches, + persistGpuMemoryModeOnLoad, readPersistedSpeculativeType, + reconcilePersistedGpuIds, resolveToolsEnabledOnLoad, saveSpeculativeType, useChatRuntimeStore, @@ -46,9 +51,12 @@ import { } from "../lib/apply-inference-status-to-store"; import { mergeBackendRecommendedInference, + resolveFitMaxSeqLength, resolveLoadMaxSeqLength, + resolveManualAutoCtxPin, } from "../presets/preset-policy"; import { recordLastLocalModelLoad } from "../utils/last-local-model-load"; +import { ensureGpuDeviceCache } from "@/hooks/use-gpu-info"; import { isMultimodalResponse, } from "../types/api"; @@ -291,9 +299,12 @@ async function syncInferenceStatusToStore(options?: { if (statusRes.active_model && !isExternalSelectionActive) { const checkpointId = resolveInferenceCheckpointId(statusRes); if (checkpointId) { + const previousGgufVariant = + useChatRuntimeStore.getState().activeGgufVariant; setCheckpoint(checkpointId, statusRes.gguf_variant); applyActiveModelStatusToStore(statusRes, { previousCheckpoint: selectedCheckpoint, + previousGgufVariant, }); // setModels(listRes...) above used catalog data, which omits audio // capability. Re-apply live status so attach gates survive a refresh. @@ -511,7 +522,11 @@ export function useChatModelRuntime() { typeof selection === "string" ? false : selection.isDownloaded ?? false; const model = models.find((entry) => entry.id === modelId); const lora = loras.find((entry) => entry.id === modelId); - const isGguf = explicitIsGguf ?? model?.isGguf ?? false; + // A native path-token selection is a local GGUF by construction (the + // native model intents only grant .gguf files), but its id is a display + // label that need not end in ".gguf" -- without this, Manual + Auto + // layers would pin the UI context instead of letting --fit size it. + const isGguf = explicitIsGguf ?? model?.isGguf ?? nativePathToken != null; const loraIsAdapter = lora?.exportType === "lora"; const isLora = explicitIsLora ?? model?.isLora ?? loraIsAdapter ?? false; @@ -578,18 +593,27 @@ export function useChatModelRuntime() { let trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false; let approvedRemoteCodeFingerprint: string | null = null; const maxSeqLength = stateBeforeUnload.params.maxSeqLength; + const previousActiveNativePathToken = + stateBeforeUnload.activeNativePathToken; const previousIsGguf = previousModel?.isGguf === true || previousVariant != null + || previousActiveNativePathToken != null || (previousCheckpoint?.toLowerCase().endsWith(".gguf") ?? false); - const rollbackMaxSeqLength = previousIsGguf - ? (stateBeforeUnload.ggufContextLength ?? 0) - : maxSeqLength; + // Respect the rolled-back model's auto-layers mode: a Manual+Auto model + // with an unpinned (auto) context must reload with 0 (so --fit + // re-auto-sizes), not the positive context it happened to pick (which + // the backend would treat as a pin). + const rollbackMaxSeqLength = resolveFitMaxSeqLength( + previousIsGguf, + stateBeforeUnload.loadedGpuMemoryMode ?? "auto", + stateBeforeUnload.loadedGpuLayers ?? GPU_LAYERS_AUTO, + stateBeforeUnload.loadedCustomContextLength, + previousIsGguf ? (stateBeforeUnload.ggufContextLength ?? 0) : maxSeqLength, + ); const hfToken = stateBeforeUnload.hfToken || null; const previousModelRequiresTrustRemoteCode = stateBeforeUnload.modelRequiresTrustRemoteCode; - const previousActiveNativePathToken = - stateBeforeUnload.activeNativePathToken; // Snapshot the load settings at click time, before the awaits below // (validation, the trust dialog, unload). For a staged Load these knobs // stay editable and a sheet-close revert (abandonStagedModel) can fire @@ -598,11 +622,29 @@ export function useChatModelRuntime() { // updates this snapshot in lock-step so non-staged loads are unchanged. const loadChatTemplateOverride = stateBeforeUnload.chatTemplateOverride; const loadKvCacheDtype = stateBeforeUnload.kvCacheDtype; - const loadCustomContextLength = stateBeforeUnload.customContextLength; + // gpuMemoryMode is a standing preference (kept across a model switch); + // the rest are per-model knobs the reset below clears, so they are + // re-baselined there in lock-step with the store. + let loadCustomContextLength = stateBeforeUnload.customContextLength; const loadGgufContextLength = stateBeforeUnload.ggufContextLength; const loadTensorParallel = stateBeforeUnload.tensorParallel; const loadActivePresetSource = stateBeforeUnload.activePresetSource; const loadActiveGgufVariant = stateBeforeUnload.activeGgufVariant; + const loadGpuMemoryMode = stateBeforeUnload.gpuMemoryMode; + let loadGpuLayers = stateBeforeUnload.gpuLayers; + let loadNCpuMoe = stateBeforeUnload.nCpuMoe; + let loadSplitRatio = stateBeforeUnload.splitRatio; + // Reconcile the persisted pick against the GPUs present now, so a stale + // cross-host / now-hidden pick is dropped before /load rather than + // rejected there. Warm the device cache first: load-on-selection can + // run before any GPU hook mounted, and a cold cache would pass the + // pick through unvalidated. validateGpuIds derives from this too. + if (stateBeforeUnload.selectedGpuIds != null) { + await ensureGpuDeviceCache(); + } + let loadSelectedGpuIds = reconcilePersistedGpuIds( + stateBeforeUnload.selectedGpuIds, + ); let loadSpeculativeType = stateBeforeUnload.speculativeType; let loadSpecDraftNMax = stateBeforeUnload.specDraftNMax; try { @@ -615,16 +657,47 @@ export function useChatModelRuntime() { // context can exceed maxSeqLength, so sizing on raw maxSeqLength could // pass, unload, then have /load refuse it. Uses the click-time // snapshot (same values loadModel uses below), so the two agree. - const validateMaxSeqLength = resolveLoadMaxSeqLength({ - modelId, - ggufVariant, - customContextLength: loadCustomContextLength, - ggufContextLength: loadGgufContextLength, - currentCheckpoint, - activeGgufVariant: loadActiveGgufVariant, - maxSeqLength, - presetSource: loadActivePresetSource, - }); + // Mirror what /load does on a cross-model switch: the reset below + // clears the per-model Auto-layers context pin + GPU pick, and + // Manual+Auto sizes context through resolveFitMaxSeqLength. + // gpuMemoryMode is a standing preference, kept across the switch. + // A same-repo quant switch (same checkpoint, different gguf_variant) + // is a different model for per-model knobs: the pinned context, + // gpuLayers, GPU pick, and MoE offload are scoped per variant, so + // treat a variant change like a model switch and re-baseline them. + const switchingModelOrVariant = + currentCheckpoint !== modelId || + (loadActiveGgufVariant ?? null) !== (ggufVariant ?? null); + const resetsPerModelSettings = Boolean( + currentCheckpoint && switchingModelOrVariant && !keepSpeculative, + ); + const validateCustomContextLength = resetsPerModelSettings + ? null + : loadCustomContextLength; + const validateGpuIds = resetsPerModelSettings + ? null + : loadSelectedGpuIds; + // The reset below re-baselines gpuLayers to Auto; mirror it here. + const validateGpuLayers = resetsPerModelSettings + ? GPU_LAYERS_AUTO + : loadGpuLayers; + const validateMaxSeqLength = resolveFitMaxSeqLength( + isGguf, + loadGpuMemoryMode, + validateGpuLayers, + validateCustomContextLength, + resolveLoadMaxSeqLength({ + modelId, + ggufVariant, + isGguf, + customContextLength: validateCustomContextLength, + ggufContextLength: loadGgufContextLength, + currentCheckpoint, + activeGgufVariant: loadActiveGgufVariant, + maxSeqLength, + presetSource: loadActivePresetSource, + }), + ); const validation = await validateModel({ model_path: modelId, nativePathLease: validateNativePathLease, @@ -633,6 +706,8 @@ export function useChatModelRuntime() { load_in_4bit: true, is_lora: isLora, gguf_variant: ggufVariant ?? null, + gpu_ids: validateGpuIds ?? undefined, + ...(isGguf ? { gpu_memory_mode: loadGpuMemoryMode } : {}), }); // Upgrade consent runs before the security dialogs; Accept installs and the load continues. if (validation.requires_transformers_upgrade) { @@ -697,18 +772,52 @@ export function useChatModelRuntime() { // keepSpeculative skips this for a staged Load: the user picked the // mode for this model on the sidebar, so honor it (the backend still // falls back at runtime if the model has no MTP head). - if (currentCheckpoint && currentCheckpoint !== modelId && !keepSpeculative) { + if (resetsPerModelSettings) { const persistedSpeculativeType = readPersistedSpeculativeType(); useChatRuntimeStore.setState({ speculativeType: persistedSpeculativeType, loadedSpeculativeType: persistedSpeculativeType, specDraftNMax: null, loadedSpecDraftNMax: null, + // Per-model GPU knobs must not follow onto a different model + // (gpuMemoryMode is a standing preference and is kept). + selectedGpuIds: null, + gpuLayers: GPU_LAYERS_AUTO, + nCpuMoe: 0, + splitRatio: null, + // A Manual+Auto context pin is per-model; clear it so a different + // model loads at Auto/native, not the previous model's pin. + customContextLength: null, }); loadSpeculativeType = persistedSpeculativeType; loadSpecDraftNMax = null; + // Keep the click-time snapshot in lock-step with the store reset so + // the load below sizes against the cleared per-model knobs, not the + // previous model's (gpuMemoryMode is standing, so left as captured). + loadCustomContextLength = null; + loadSelectedGpuIds = null; + loadGpuLayers = GPU_LAYERS_AUTO; + loadNCpuMoe = 0; + loadSplitRatio = null; } + // Pinning layers on the SAME model keeps the currently resolved + // context: with no explicit pin, a manual+pinned reload would send 0, + // which the backend's --fit off branch treats as the NATIVE context -- + // far larger than the sheet shows when the load was fit-sized (Default + // or Manual + Auto layers may auto-reduce context to fit VRAM), a + // likely OOM. ggufContextLength is that resolved value; a model already + // at native reloads unchanged, so this is safe for any prior mode. + if ( + isGguf && + !switchingModelOrVariant && + loadGpuMemoryMode === "manual" && + loadGpuLayers >= 0 && + loadCustomContextLength == null && + (loadGgufContextLength ?? 0) > 0 + ) { + loadCustomContextLength = loadGgufContextLength; + } const effectiveMaxSeqLength = resolveLoadMaxSeqLength({ modelId, ggufVariant, @@ -720,13 +829,20 @@ export function useChatModelRuntime() { maxSeqLength, presetSource: loadActivePresetSource, }); + const loadMaxSeqLength = resolveFitMaxSeqLength( + isGguf, + loadGpuMemoryMode, + loadGpuLayers, + loadCustomContextLength, + effectiveMaxSeqLength, + ); const effectiveChatTemplateOverride = loadChatTemplateOverride?.trim() ? loadChatTemplateOverride : null; const loadResponse = await loadModel({ model_path: modelId, nativePathLease: loadNativePathLease, hf_token: hfToken, - max_seq_length: effectiveMaxSeqLength, + max_seq_length: loadMaxSeqLength, load_in_4bit: true, is_lora: isLora, gguf_variant: ggufVariant ?? null, @@ -737,6 +853,11 @@ export function useChatModelRuntime() { speculative_type: loadSpeculativeType, spec_draft_n_max: loadSpecDraftNMax, tensor_parallel: loadTensorParallel, + gpu_memory_mode: loadGpuMemoryMode, + gpu_layers: loadGpuLayers, + n_cpu_moe: loadNCpuMoe, + tensor_split: loadSplitRatio ?? undefined, + gpu_ids: loadSelectedGpuIds ?? undefined, }); // If cancelled while loading, don't update UI to show @@ -747,6 +868,9 @@ export function useChatModelRuntime() { // preference now (the requested intent, not the resolved echo; // saveSpeculativeType keeps only the universal auto/ngram/off). saveSpeculativeType(loadSpeculativeType); + // Persist the GPU Memory mode only on a successful load (not on + // dropdown change), so an abandoned selection doesn't stick. + persistGpuMemoryModeOnLoad(loadResponse, loadGpuMemoryMode); const currentParams = useChatRuntimeStore.getState().params; setParams( @@ -782,9 +906,13 @@ export function useChatModelRuntime() { const reportedNativeCtx = loadResponse.is_gguf ? (loadResponse.native_context_length ?? null) : null; - // A successful reload has applied settings, so clear pending custom - // context state and display the backend-reported effective context. - const keepCustomCtx = null; + // Keep an explicit Manual+Auto context pin (so a later Apply doesn't + // revert it to Auto); other cases baseline on ggufContextLength. + const keepCustomCtx = resolveManualAutoCtxPin( + loadGpuMemoryMode, + loadGpuLayers, + loadCustomContextLength, + ); const reasoningAlwaysOn = loadResponse.reasoning_always_on ?? false; const reasoningStyle = loadResponse.reasoning_style ?? "enable_thinking"; const supportsReasoning = loadResponse.supports_reasoning ?? false; @@ -837,11 +965,13 @@ export function useChatModelRuntime() { loadedKvCacheDtype: loadedKv, tensorParallel: loadedTp, loadedTensorParallel: loadedTp, + ...loadedGpuMemoryFields(loadResponse), speculativeType: loadedSpec, loadedSpeculativeType: loadedSpec, specDraftNMax: loadResponse.spec_draft_n_max ?? null, loadedSpecDraftNMax: loadResponse.spec_draft_n_max ?? null, customContextLength: keepCustomCtx, + loadedCustomContextLength: keepCustomCtx, defaultChatTemplate: loadResponse.chat_template ?? null, chatTemplateOverride: effectiveChatTemplateOverride, loadedChatTemplateOverride: effectiveChatTemplateOverride, @@ -938,7 +1068,7 @@ export function useChatModelRuntime() { } } try { - await loadModel({ + const rollbackResponse = await loadModel({ model_path: previousCheckpoint, nativePathLease: rollbackNativePathLease, hf_token: hfToken, @@ -951,14 +1081,51 @@ export function useChatModelRuntime() { // Resend the previous model's pinned approval so restoring it is not re-blocked. approved_remote_code_fingerprint: approvedRemoteCodeFingerprints.get(previousCheckpoint) ?? null, + chat_template_override: + stateBeforeUnload.loadedChatTemplateOverride, + cache_type_kv: stateBeforeUnload.loadedKvCacheDtype, + speculative_type: + stateBeforeUnload.loadedSpeculativeType, + spec_draft_n_max: + stateBeforeUnload.loadedSpecDraftNMax, // Restore the previous model in the split mode it was running, // not the default layer split. tensor_parallel: stateBeforeUnload.loadedTensorParallel ?? false, + gpu_memory_mode: stateBeforeUnload.loadedGpuMemoryMode ?? "auto", + gpu_layers: stateBeforeUnload.loadedGpuLayers ?? -1, + n_cpu_moe: stateBeforeUnload.loadedNCpuMoe ?? 0, + tensor_split: stateBeforeUnload.loadedSplitRatio ?? undefined, + gpu_ids: stateBeforeUnload.loadedGpuIds ?? undefined, }); + const rollbackSpeculativeType = normalizeSpeculativeType( + rollbackResponse.speculative_type, + ); useChatRuntimeStore.setState({ activeNativePathToken: previousActiveNativePathToken ?? null, - loadedSpeculativeType: null, - loadedSpecDraftNMax: null, + loadedSpeculativeType: rollbackSpeculativeType, + loadedSpecDraftNMax: + rollbackResponse.spec_draft_n_max ?? null, + loadedKvCacheDtype: rollbackResponse.cache_type_kv ?? null, + loadedChatTemplateOverride: + stateBeforeUnload.loadedChatTemplateOverride, + // Re-baseline the GPU knobs from the rolled-back load's own + // response (the shared seeding every load path uses): the + // refresh() below can't do it, since the status reseed is + // gated off while modelLoading is still true. A failed staged + // Load stays staged for retry, so the staged hold applies. + ...loadedGpuMemoryFieldsUnlessStaged(rollbackResponse, { + tensorParallel: rollbackResponse.tensor_parallel ?? false, + loadedTensorParallel: + rollbackResponse.tensor_parallel ?? false, + // refresh() is held while modelLoading remains true, so + // restore the rolled-back model's context pin directly. + customContextLength: + stateBeforeUnload.loadedCustomContextLength, + }), + loadedTensorParallel: + rollbackResponse.tensor_parallel ?? false, + loadedCustomContextLength: + stateBeforeUnload.loadedCustomContextLength, }); await refresh(); } catch { diff --git a/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts b/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts index d8076c720b..a3e7a2d264 100644 --- a/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts +++ b/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts @@ -7,7 +7,7 @@ import { useRepoDownload } from "@/features/hub/download-manager/use-repo-downlo import type { DownloadJob } from "@/features/hub/download-manager/use-repo-download"; import { useLatestRef } from "@/features/hub/hooks/use-latest-ref"; -import { fetchGgufContextLength } from "../api/chat-api"; +import { fetchGgufStagedMetadata } from "../api/chat-api"; import { isPendingGguf, pendingSelectionMatches, @@ -46,8 +46,16 @@ export function useStagedModelPreparation(opts?: { const pendingDownloaded = useChatRuntimeStore( (s) => s.pendingSelection?.isDownloaded ?? false, ); - const pendingHasContext = useChatRuntimeStore( - (s) => s.pendingSelection?.contextLength != null, + // "Already probed" must key off layerCount / moeLayerCount, which only the + // full header probe fills (it sets all three together, so either is a + // reliable marker). contextLength alone can be list-seeded from + // /gguf-variants, which returns no layer/MoE counts -- treating it as + // complete would skip the probe and leave the GPU Layers slider at its 256 + // fallback and the MoE slider hidden until the model loads. + const pendingHasMetadata = useChatRuntimeStore( + (s) => + s.pendingSelection?.layerCount != null || + s.pendingSelection?.moeLayerCount != null, ); const setPendingSelection = useChatRuntimeStore((s) => s.setPendingSelection); const onAutoLoadRef = useLatestRef(opts?.onAutoLoad); @@ -69,25 +77,31 @@ export function useStagedModelPreparation(opts?: { if (!current?.id || !isPendingGguf(current)) return; const { id, ggufVariant, nativePathToken } = current; try { - const contextLength = await fetchGgufContextLength({ - model_path: id, - gguf_variant: ggufVariant, - hf_token: useChatRuntimeStore.getState().hfToken || null, - nativePathToken, - }); + const { contextLength, layerCount, moeLayerCount } = + await fetchGgufStagedMetadata({ + model_path: id, + gguf_variant: ggufVariant, + hf_token: useChatRuntimeStore.getState().hfToken || null, + nativePathToken, + }); // Apply only if the same model is still staged (the user may have switched // picks or loaded/cancelled while the request was in flight). const latest = useChatRuntimeStore.getState().pendingSelection; if ( latest && - contextLength != null && - pendingSelectionMatches(latest, { id, ggufVariant, nativePathToken }) + pendingSelectionMatches(latest, { id, ggufVariant, nativePathToken }) && + (contextLength != null || layerCount != null || moeLayerCount != null) ) { - setPendingSelection({ ...latest, contextLength }); + setPendingSelection({ + ...latest, + contextLength, + layerCount, + moeLayerCount, + }); } } catch { - // Leave contextLength null: the context slider stays hidden and the user - // can still load (context fills in from the load response afterwards). + // Leave metadata null: the context/MoE sliders stay hidden and the user + // can still load (they fill in from the load response afterwards). } }, [setPendingSelection]); @@ -125,7 +139,7 @@ export function useStagedModelPreparation(opts?: { if ( !pendingId || (!pendingIsGguf && !pendingIsHubRepo) || - pendingHasContext + pendingHasMetadata ) { return; } @@ -146,7 +160,7 @@ export function useStagedModelPreparation(opts?: { pendingIsGguf, pendingIsHubRepo, pendingDownloaded, - pendingHasContext, + pendingHasMetadata, startDownloadRef, fetchMetadataRef, ]); diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts index a4b5f848e2..69bb38bbbe 100644 --- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -2,13 +2,17 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { getInferenceStatus } from "../api/chat-api"; -import { mergeBackendRecommendedInference } from "../presets/preset-policy"; +import { + mergeBackendRecommendedInference, + resolveManualAutoCtxPin, +} from "../presets/preset-policy"; import { clampReasoningEffortToLevels } from "../provider-capabilities"; import { CHAT_REASONING_ENABLED_KEY, type ReasoningEffort, type ReasoningStyle, loadOptionalBool, + loadedGpuMemoryFields, resolveToolsEnabledOnLoad, useChatRuntimeStore, } from "../stores/chat-runtime-store"; @@ -20,6 +24,10 @@ import type { ChatModelSummary } from "../types/runtime"; type LocalReasoningEffort = Extract; +function sameArray(a: T[] | null, b: T[] | null): boolean { + return JSON.stringify(a) === JSON.stringify(b); +} + // Canonicalises backend / persisted speculative mode values onto the UI modes. export function normalizeSpeculativeType( v: string | null | undefined, @@ -119,6 +127,10 @@ function ensureActiveModelInStoreList( export type ApplyInferenceStatusOptions = { previousCheckpoint?: string; + /** activeGgufVariant BEFORE the caller's setCheckpoint synced it to the + * status -- without it a variant-only switch underneath the tab reads as + * steady state and the hydration reseed keeps the old quant's baselines. */ + previousGgufVariant?: string | null; }; /** Mirror refresh() hydration so adopted CLI models get reasoning/tools flags. */ @@ -144,9 +156,13 @@ export function applyActiveModelStatusToStore( ); } + const previousGgufVariant = + options.previousGgufVariant !== undefined + ? options.previousGgufVariant + : store.activeGgufVariant; const hydratingExistingModel = previousCheckpoint !== checkpointId || - store.activeGgufVariant !== (status.gguf_variant ?? null); + previousGgufVariant !== (status.gguf_variant ?? null); const supportsReasoning = status.supports_reasoning ?? false; const reasoningAlwaysOn = status.reasoning_always_on ?? false; const reasoningStyle = status.reasoning_style ?? "enable_thinking"; @@ -185,6 +201,66 @@ export function applyActiveModelStatusToStore( // While a load is in flight, performLoad owns the load params. Seeding them // from a stale poll here would clobber the values the load dialog just set. const seedLoadParams = !prevState.modelLoading; + // A Manual + Auto-layers load sent its positive context pin as max_seq_length, + // and status only exposes the RESOLVED context; re-seed the pin from the + // requested value (parity with the load paths' keepCustomCtx). Baselines + // unconditionally: anything but an applicable pin is null, so a previous + // model's pin can't survive a model change underneath and reload at the old length. + const gpuPin = status.is_gguf + ? resolveManualAutoCtxPin( + status.gpu_memory_mode ?? "auto", + status.gpu_layers ?? -1, + status.requested_context_length ?? null, + ) + : null; + const incomingGpuMode = status.is_gguf + ? (status.gpu_memory_mode ?? "auto") + : null; + const incomingGpuLayers = + incomingGpuMode === "manual" ? (status.gpu_layers ?? null) : null; + const incomingNCpuMoe = + incomingGpuMode === "manual" ? (status.n_cpu_moe ?? null) : null; + const incomingSplit = + incomingGpuMode === "manual" ? (status.tensor_split ?? null) : null; + const incomingGpuIds = status.is_gguf ? (status.gpu_ids ?? null) : null; + const gpuStatusChanged = + prevState.loadedGpuMemoryMode !== incomingGpuMode || + prevState.loadedGpuLayers !== incomingGpuLayers || + prevState.loadedNCpuMoe !== incomingNCpuMoe || + !sameArray(prevState.loadedSplitRatio, incomingSplit) || + !sameArray(prevState.loadedGpuIds, incomingGpuIds) || + prevState.loadedCustomContextLength !== gpuPin; + const gpuMemoryEditsPending = + (prevState.loadedGpuMemoryMode !== null && + prevState.gpuMemoryMode !== prevState.loadedGpuMemoryMode) || + (prevState.loadedGpuMemoryMode === "manual" && + (prevState.gpuLayers !== prevState.loadedGpuLayers || + prevState.nCpuMoe !== prevState.loadedNCpuMoe || + !sameArray(prevState.splitRatio, prevState.loadedSplitRatio))) || + prevState.customContextLength !== prevState.loadedCustomContextLength; + const gpuIdsEditPending = !sameArray( + prevState.selectedGpuIds, + prevState.loadedGpuIds, + ); + const incomingGpuFields = loadedGpuMemoryFields(status); + // A same-model reload from another client advances every loaded baseline. + // Preserve each editable group only when this tab has an unapplied change. + const preserveSameModelEdits = gpuStatusChanged && !hydratingExistingModel; + const gpuStatusFields = { + ...incomingGpuFields, + customContextLength: gpuPin, + loadedCustomContextLength: gpuPin, + ...(preserveSameModelEdits && + gpuMemoryEditsPending && { + gpuMemoryMode: prevState.gpuMemoryMode, + gpuLayers: prevState.gpuLayers, + nCpuMoe: prevState.nCpuMoe, + splitRatio: prevState.splitRatio, + customContextLength: prevState.customContextLength, + }), + ...(preserveSameModelEdits && + gpuIdsEditPending && { selectedGpuIds: prevState.selectedGpuIds }), + }; useChatRuntimeStore.setState({ supportsReasoning, @@ -215,30 +291,51 @@ export function applyActiveModelStatusToStore( loadedIsMultimodal: isMultimodalResponse(status), loadedIsDiffusion: status.is_diffusion ?? false, specFallbackReason: status.spec_fallback_reason ?? null, + // The spec / KV seeds share the GPU-fields reseed mechanism below: a + // non-GGUF status leaves their loaded baselines null, so the "unseeded" + // guard re-fires every refresh -- hold them too while a staged pick's + // settings are being edited, or the refresh resets the staged edit. + // hydratingExistingModel reopens every load-param seed: when the active + // model changed underneath this tab (auto-switch, another client), the + // old model's baselines are stale and must adopt the new status. ...(seedLoadParams && - prevState.loadedSpeculativeType === null && { + prevState.pendingSelection == null && + (prevState.loadedSpeculativeType === null || hydratingExistingModel) && { speculativeType: currentSpecType, loadedSpeculativeType: currentSpecType, }), ...(seedLoadParams && + prevState.pendingSelection == null && status.spec_draft_n_max !== undefined && - prevState.loadedSpecDraftNMax === null && - prevState.specDraftNMax === null && { + (hydratingExistingModel || + (prevState.loadedSpecDraftNMax === null && + prevState.specDraftNMax === null)) && { specDraftNMax: status.spec_draft_n_max ?? null, loadedSpecDraftNMax: status.spec_draft_n_max ?? null, }), ...(seedLoadParams && + prevState.pendingSelection == null && status.cache_type_kv !== undefined && - prevState.loadedKvCacheDtype === null && { + (prevState.loadedKvCacheDtype === null || hydratingExistingModel) && { kvCacheDtype: status.cache_type_kv, loadedKvCacheDtype: status.cache_type_kv, }), ...(seedLoadParams && + prevState.pendingSelection == null && status.tensor_parallel !== undefined && - prevState.loadedTensorParallel === null && { + (prevState.loadedTensorParallel === null || hydratingExistingModel) && { tensorParallel: status.tensor_parallel, loadedTensorParallel: status.tensor_parallel, }), + // Re-seed on first hydration, model/variant changes, or a same-model backend + // placement change. gpuStatusFields preserves dirty local edits in the last + // case while advancing their loaded baselines. + ...(seedLoadParams && + prevState.pendingSelection == null && + (prevState.loadedGpuMemoryMode === null || + hydratingExistingModel || + gpuStatusChanged) && + gpuStatusFields), ...(status.chat_template_override !== undefined && prevState.loadedChatTemplateOverride === null && prevState.chatTemplateOverride === null && { @@ -298,7 +395,11 @@ export async function tryAdoptServerActiveModel(): Promise { if (previousCheckpoint) { return true; } + const previousGgufVariant = useChatRuntimeStore.getState().activeGgufVariant; store.setCheckpoint(checkpointId, status.gguf_variant); - applyActiveModelStatusToStore(status, { previousCheckpoint }); + applyActiveModelStatusToStore(status, { + previousCheckpoint, + previousGgufVariant, + }); return true; } diff --git a/studio/frontend/src/features/chat/presets/preset-policy.ts b/studio/frontend/src/features/chat/presets/preset-policy.ts index f96ee91f1b..23d79a35e1 100644 --- a/studio/frontend/src/features/chat/presets/preset-policy.ts +++ b/studio/frontend/src/features/chat/presets/preset-policy.ts @@ -339,3 +339,34 @@ export function resolveLoadMaxSeqLength({ } return maxSeqLength; } + +/** + * Adjust a resolved max-seq-length for the GPU Memory mode. Under Manual + Auto + * layers (GGUF, gpuLayers < 0) llama.cpp's --fit owns context sizing, so send 0 + * (the backend omits -c) unless the user pinned a length; every other case keeps + * the resolved fallback. Shared by every GGUF load path so they can't drift. + */ +export function resolveFitMaxSeqLength( + isGguf: boolean | null | undefined, + gpuMemoryMode: "auto" | "manual", + gpuLayers: number, + customContextLength: number | null, + fallback: number, +): number { + if (!isGguf || gpuMemoryMode !== "manual" || gpuLayers >= 0) return fallback; + return customContextLength && customContextLength > 0 ? customContextLength : 0; +} + +// A Manual + Auto-layers load sends its positive context pin as max_seq_length; +// keep it across a status reseed/Apply so the model isn't reverted to auto-fit +// sizing. Anything else (Auto mode, pinned layers, no pin) baselines to null. +// The caller keeps its own isGguf/targetIsGguf guard inline. +export function resolveManualAutoCtxPin( + gpuMemoryMode: "auto" | "manual", + gpuLayers: number, + customContextLength: number | null, +): number | null { + return gpuMemoryMode === "manual" && gpuLayers < 0 && (customContextLength ?? 0) > 0 + ? customContextLength + : null; +} diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index a0813fe27b..31e50ee60c 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -84,6 +84,8 @@ import { useTransformersUpgradeDialogStore, } from "@/features/transformers-upgrade"; import { loadModel, validateModel } from "./api/chat-api"; +import { resolveFitMaxSeqLength, resolveManualAutoCtxPin } from "./presets/preset-policy"; +import { ensureGpuDeviceCache } from "@/hooks/use-gpu-info"; import { parseExternalModelId, providerTypeSupportsVision, @@ -95,8 +97,11 @@ import { usePlusMenuPrefsStore, } from "./stores/plus-menu-prefs-store"; import { + loadedGpuMemoryFieldsUnlessStaged, type ReasoningEffort, + reconcilePersistedGpuIds, resolveLoadedSpeculativeSettings, + persistGpuMemoryModeOnLoad, resolveSpeculativeSettingsForLoad, saveSpeculativeType, useChatRuntimeStore, @@ -1037,10 +1042,32 @@ export function SharedComposer({ return parts[parts.length - 1] || id; } + // Warm the device cache before the snapshot below reconciles the GPU + // pick: on a cold cache the reconcile passes a stale pick through. + if (store.selectedGpuIds != null) { + await ensureGpuDeviceCache(); + } + // The GPU/offload knobs both compare loads must use, snapshotted at Send. + // ensureModelLoaded runs sequentially and the first load's response echo + // (loadedGpuMemoryFields) rewrites the live store -- a non-GGUF or Auto + // first model resets gpuLayers/nCpuMoe/split/pick to defaults -- so + // reading the store per load would hand model 2 the first model's echoed + // defaults instead of the settings the user pressed Send with. + const compareLoadKnobs = { + gpuMemoryMode: store.gpuMemoryMode, + gpuLayers: store.gpuLayers, + nCpuMoe: store.nCpuMoe, + splitRatio: store.splitRatio, + // Reconcile the pick against the GPUs present now, like the model-switch + // path: an early remember-restore can hold a stale cross-host pick that + // /load would reject (the device cache is populated by send time). + selectedGpuIds: reconcilePersistedGpuIds(store.selectedGpuIds), + tensorParallel: store.tensorParallel, + customContextLength: store.customContextLength, + }; // Set when an accepted transformers install unloaded the active model // server-side; a later failure must then clear the stale checkpoint. let upgradeUnloadedActive = false; - // Helper: load a model and update store checkpoint async function ensureModelLoaded( sel: CompareModelSelection, @@ -1057,15 +1084,35 @@ export function SharedComposer({ if (isAlreadyActive) { return "ready"; } + const targetIsGguf = + sel.id.toLowerCase().endsWith(".gguf") || sel.ggufVariant != null; + // Size validation exactly as the load below, so the training-guard + // preflight checks the footprint that actually loads (under Manual + Auto + // layers the load sends 0 / the pinned context, not raw maxSeqLength). + const compareMaxSeqLength = resolveFitMaxSeqLength( + targetIsGguf, + compareLoadKnobs.gpuMemoryMode, + compareLoadKnobs.gpuLayers, + compareLoadKnobs.customContextLength, + maxSeqLength, + ); const validation = await validateModel({ model_path: sel.id, hf_token: currentStore.hfToken || null, - max_seq_length: maxSeqLength, + max_seq_length: compareMaxSeqLength, load_in_4bit: true, is_lora: sel.isLora, gguf_variant: sel.ggufVariant ?? null, trust_remote_code: loadTrustRemoteCode, chat_template_override: effectiveChatTemplateOverride, + // Scope the validate to the picked GPUs. GGUF-only, like the load + // below: a non-GGUF target must not inherit a hidden GGUF GPU pick. + ...(targetIsGguf + ? { + gpu_ids: compareLoadKnobs.selectedGpuIds ?? undefined, + gpu_memory_mode: compareLoadKnobs.gpuMemoryMode, + } + : {}), }); // Upgrade dialog first (mirrors the primary load path). if (validation.requires_transformers_upgrade) { @@ -1114,7 +1161,7 @@ export function SharedComposer({ const resp = await loadModel({ model_path: sel.id, hf_token: useChatRuntimeStore.getState().hfToken || null, - max_seq_length: maxSeqLength, + max_seq_length: compareMaxSeqLength, load_in_4bit: true, is_lora: sel.isLora, gguf_variant: sel.ggufVariant ?? null, @@ -1123,10 +1170,25 @@ export function SharedComposer({ chat_template_override: effectiveChatTemplateOverride, speculative_type: specSettings.speculativeType, spec_draft_n_max: specSettings.specDraftNMax, - // Honor the Tensor Parallelism toggle on compare loads too. - tensor_parallel: currentStore.tensorParallel, + // Honor the Tensor Parallelism + GPU Memory choices on compare loads. + // GGUF-only, like the auto-load path: the picker is a GGUF control, + // so a non-GGUF target loads via HF auto-placement instead of being + // pinned to a leftover GGUF pick it can't even show. + tensor_parallel: compareLoadKnobs.tensorParallel, + ...(targetIsGguf + ? { + gpu_memory_mode: compareLoadKnobs.gpuMemoryMode, + gpu_layers: compareLoadKnobs.gpuLayers, + n_cpu_moe: compareLoadKnobs.nCpuMoe, + tensor_split: compareLoadKnobs.splitRatio ?? undefined, + gpu_ids: compareLoadKnobs.selectedGpuIds ?? undefined, + } + : {}), }); saveSpeculativeType(specSettings.speculativeType); + // Persist the GPU Memory mode on a non-diffusion GGUF compare-load too, + // so an applied manual choice survives a restart. + persistGpuMemoryModeOnLoad(resp, compareLoadKnobs.gpuMemoryMode); upgradeUnloadedActive = false; const store = useChatRuntimeStore.getState(); store.setCheckpoint( @@ -1136,6 +1198,17 @@ export function SharedComposer({ store.setModelRequiresTrustRemoteCode( resp.requires_trust_remote_code ?? false, ); + // Keep an explicit Manual+Auto context pin the load just applied (so a + // later Apply/Reset doesn't silently revert the model to auto-fit + // sizing), mirroring the interactive path's keepCustomCtx. Non-GGUF + // compare loads don't send the pin, so their baseline clears. + const keepCustomCtx = targetIsGguf + ? resolveManualAutoCtxPin( + compareLoadKnobs.gpuMemoryMode, + compareLoadKnobs.gpuLayers, + compareLoadKnobs.customContextLength, + ) + : null; useChatRuntimeStore.setState({ supportsReasoning: resp.supports_reasoning ?? false, reasoningAlwaysOn: resp.reasoning_always_on ?? false, @@ -1144,6 +1217,32 @@ export function SharedComposer({ supportsTools: resp.supports_tools ?? false, tensorParallel: resp.tensor_parallel ?? false, loadedTensorParallel: resp.tensor_parallel ?? false, + customContextLength: keepCustomCtx, + loadedCustomContextLength: keepCustomCtx, + // Seed the loaded GGUF context (interactive/auto-load parity): the + // settings sheet keys the GGUF GPU controls off it for a direct .gguf + // with no variant, and a later Apply reads it as the resolved context. + ...(targetIsGguf + ? { + ggufContextLength: resp.context_length ?? 131072, + ggufMaxContextLength: + resp.max_context_length ?? resp.context_length ?? 131072, + ggufNativeContextLength: resp.native_context_length ?? null, + } + : { ggufContextLength: null }), + // Compare loads resolve by id (HF repo / local path), never through a + // native-path lease, so a token left by a previously loaded native + // GGUF is stale here -- isLoadedGguf keys off it, and a stale token + // would dress a non-GGUF compare load in GGUF controls. Mirror the + // interactive path, which writes it on every load success. + activeNativePathToken: null, + // Held under an open staged pick: setCheckpoint preserves a stage on + // the empty->active transition, so a compare load can complete with + // staged GPU edits still on screen. + ...loadedGpuMemoryFieldsUnlessStaged(resp), + // Drives the GPU Memory controls' diffusion gate; set alongside the + // GPU fields on every load path so the gate can't read stale. + loadedIsDiffusion: resp.is_diffusion ?? false, loadedIsMultimodal: isMultimodalResponse(resp), ...resolveLoadedSpeculativeSettings(resp), }); diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 192ce1ec69..5786947118 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -7,6 +7,10 @@ import { mirrorHfTokenInto, useHfTokenStore, } from "@/features/hub"; +import { + cachedPinnableGpuIndices, + ensureGpuDeviceCache, +} from "@/hooks/use-gpu-info"; import { toast } from "@/lib/toast"; import { create } from "zustand"; import { isExternalModelId, parseExternalModelId } from "../external-providers"; @@ -74,6 +78,7 @@ export const CHAT_RAG_AUTOINJECT_MIN_SCORE_KEY = export const CHAT_RAG_OCR_KEY = "unsloth_chat_rag_ocr_scanned"; export const CHAT_RAG_CAPTION_KEY = "unsloth_chat_rag_caption_figures"; export const CHAT_SPECULATIVE_TYPE_KEY = "unsloth_chat_speculative_type"; +export const CHAT_GPU_MEMORY_MODE_KEY = "unsloth_chat_gpu_memory_mode"; // Persist only the model-agnostic intents (auto/ngram/off). MTP modes // (mtp/mtp+ngram) and spec_draft_n_max stay session-only: a persisted MTP @@ -497,6 +502,213 @@ export function saveSpeculativeType(value: string | null): void { } } +// GPU Memory strategy is a standing preference (like speculative type), not a +// per-model setting: a "manual" choice persists across model switches and reloads. +export function readPersistedGpuMemoryMode(): "auto" | "manual" { + return loadString(CHAT_GPU_MEMORY_MODE_KEY, "auto") === "manual" ? "manual" : "auto"; +} + +export function saveGpuMemoryMode(value: "auto" | "manual"): void { + saveString(CHAT_GPU_MEMORY_MODE_KEY, value); +} + +/** Persist the GPU Memory mode after a load, but only for a non-diffusion GGUF: + * non-GGUF has no such mode, and diffusion runs mode-agnostic (reports "auto"), + * so neither must clobber the standing manual preference. */ +export function persistGpuMemoryModeOnLoad( + resp: { is_gguf?: boolean; is_diffusion?: boolean }, + mode: "auto" | "manual", +): void { + if (resp.is_gguf && !resp.is_diffusion) saveGpuMemoryMode(mode); +} + +// Manual-mode gpu_layers sentinel: -1 = Auto (hand layer + context sizing to +// llama.cpp's --fit). The Manual default; "all on GPU" is the slider's max. +export const GPU_LAYERS_AUTO = -1; + +// Round real-valued shares to integers summing exactly to `total`, giving the +// leftover units to the largest fractional parts (largest-remainder method). +function largestRemainder(shares: number[], total: number): number[] { + const out = shares.map((x) => Math.floor(x)); + let rem = total - out.reduce((a, b) => a + b, 0); + const byFrac = shares + .map((x, i) => ({ i, frac: x - Math.floor(x) })) + .sort((a, b) => b.frac - a.frac); + for (let k = 0; rem > 0 && k < byFrac.length; k++, rem--) out[byFrac[k].i] += 1; + return out; +} + +// Spread `total` layers across GPUs in proportion to `weights` (e.g. per-GPU +// VRAM), as integers summing exactly to `total`; even split for all-zero/empty +// weights. Default per-GPU layer split before the user edits it (mirrors +// llama.cpp's free-VRAM default). +export function distributeByWeight(total: number, weights: number[]): number[] { + if (weights.length === 0) return []; + const t = Math.max(0, Math.floor(total)); + const sum = weights.reduce((a, b) => a + b, 0); + const w = sum > 0 ? weights : weights.map(() => 1); + const wSum = w.reduce((a, b) => a + b, 0); + return largestRemainder( + w.map((x) => (t * x) / wSum), + t, + ); +} + +// Set GPU `index` to `value` and rebalance the rest so per-GPU counts still sum +// to `total`; others absorb the remainder in proportion to their counts (evenly +// if all zero). The --tensor-split editor: counts are sent verbatim, and +// llama.cpp gives each GPU exactly its count when gpu_layers == sum(counts). +export function rebalanceSplit( + total: number, + counts: number[], + index: number, + value: number, +): number[] { + const v = Math.max(0, Math.min(value, total)); + const out = counts.slice(); + const otherIdx = counts.map((_, i) => i).filter((i) => i !== index); + // No other GPU to absorb the remainder: this one holds everything. + if (otherIdx.length === 0) { + out[index] = total; + return out; + } + out[index] = v; + const dist = distributeByWeight( + total - v, + otherIdx.map((i) => counts[i]), + ); + otherIdx.forEach((i, k) => (out[i] = dist[k])); + return out; +} + +// Validate a persisted gpu_ids pick against the GPUs present right now, before +// restoring it from remembered settings. Returns null (= automatic) when the +// pick is stale (none of the saved ids exist, or the host can't pin a multi-GPU +// set), so a saved [1] on a now-1-GPU host doesn't get sent and rejected with no +// way to clear it. A null pick (= automatic) passes through unchanged, and an +// unpopulated device cache leaves the pick alone (the backend still guards). +export function reconcilePersistedGpuIds( + ids: number[] | null, +): number[] | null { + if (ids == null) return ids; + const pinnable = cachedPinnableGpuIndices(); + if (pinnable === null) return ids; // cache not ready: can't validate, keep it + const kept = ids.filter((i) => pinnable.includes(i)); + return kept.length > 0 ? kept : null; +} + +// Store fields derived from a load/status response's GPU-memory settings. +// Shared by every load path so the manual-knob round-trip can't drift. +export function loadedGpuMemoryFields(resp: { + is_gguf?: boolean; + is_diffusion?: boolean; + gpu_memory_mode?: "auto" | "manual"; + gpu_layers?: number; + n_cpu_moe?: number; + tensor_split?: number[] | null; + n_layers?: number | null; + n_moe_layers?: number; + gpu_ids?: number[] | null; +}) { + // GPU-memory state is meaningful only for a GGUF chat load. A non-GGUF response + // still carries gpu_memory_mode (its default "auto" is serialized), so gate on + // the authoritative is_gguf flag, not the field's presence -- otherwise loading + // a transformers model would reset the standing manual preference. + if (!resp.is_gguf) { + // Clear the GPU pick / offload baseline a prior GGUF load may have left, so it + // reflects the non-GGUF model (no pin) -- else a stale loadedGpuIds reads as + // dirty (gpuIdsDirty is ungated) and Reset restores it while the picker is + // hidden. gpuMemoryMode (the standing preference) is kept, but its loaded + // baseline clears to null so Reset preserves the preference, not a stale mode. + return { + selectedGpuIds: null, + loadedGpuIds: null, + loadedGpuMemoryMode: null, + gpuLayers: GPU_LAYERS_AUTO, + loadedGpuLayers: null, + nCpuMoe: 0, + loadedNCpuMoe: null, + splitRatio: null, + loadedSplitRatio: null, + ggufLayerCount: null, + moeLayerCount: null, + }; + } + const mode = resp.gpu_memory_mode ?? "auto"; + const gpuIds = resp.gpu_ids ?? null; + // Layer/MoE/split knobs apply (and are reported) only in manual mode; in auto + // the server ignores them, so don't seed the loaded baseline or the editable + // knobs with values it never applied. In manual, the server reports gpu_layers + // = -1 under Auto, which round-trips the slider back to its Auto position. + const manualKnobs = + mode === "manual" + ? { + loadedGpuLayers: resp.gpu_layers ?? null, + loadedNCpuMoe: resp.n_cpu_moe ?? null, + loadedSplitRatio: resp.tensor_split ?? null, + gpuLayers: resp.gpu_layers ?? GPU_LAYERS_AUTO, + nCpuMoe: resp.n_cpu_moe ?? 0, + splitRatio: resp.tensor_split ?? null, + } + : { + loadedGpuLayers: null, + loadedNCpuMoe: null, + loadedSplitRatio: null, + // Auto ignores these, so reset the editable knobs too (not just the + // loaded baseline) -- else a later switch back to Manual would snapshot + // and send a previous model's stale gpuLayers/nCpuMoe/split that this + // load never applied. Mirrors the non-GGUF branch above. + gpuLayers: GPU_LAYERS_AUTO, + nCpuMoe: 0, + splitRatio: null, + }; + return { + // A diffusion GGUF runs mode-agnostic (pins all layers on one GPU, reports + // "auto"), so adopt everything a chat GGUF does EXCEPT the live standing + // preference -- the next chat load must still honor the user's manual choice. + // The loaded baseline is still "auto", but the UI hides mode controls for a + // loaded diffusion model so it can't read as dirty against the preference. + ...(resp.is_diffusion ? {} : { gpuMemoryMode: mode }), + loadedGpuMemoryMode: mode, + ggufLayerCount: resp.n_layers ?? null, + // MoE expert-layer count: the n_cpu_moe slider max, and 0 hides the slider. + moeLayerCount: resp.n_moe_layers ?? null, + // The picker reflects what loaded (the request sent the user's pick). + selectedGpuIds: gpuIds, + loadedGpuIds: gpuIds, + ...manualKnobs, + }; +} + +/** loadedGpuMemoryFields (plus any seedExtras), unless a staged pick is open. + * + * With a staged pick open (the load fired mid-staging), preserve its editable + * GPU knobs and seedExtras, but still advance every loaded baseline. Otherwise + * cancelling the stage restores its edits onto the newly loaded model. The + * status reseed cannot repair that while pendingSelection holds it off. + */ +export function loadedGpuMemoryFieldsUnlessStaged( + resp: Parameters[0], + seedExtras?: T, +) { + const fields = loadedGpuMemoryFields(resp); + if (useChatRuntimeStore.getState().pendingSelection != null) { + return { + loadedGpuMemoryMode: fields.loadedGpuMemoryMode, + loadedGpuLayers: fields.loadedGpuLayers, + loadedNCpuMoe: fields.loadedNCpuMoe, + loadedSplitRatio: fields.loadedSplitRatio, + loadedGpuIds: fields.loadedGpuIds, + // These are metadata ceilings for the model that actually loaded, not + // editable values from the open stage. Advance them with the baselines + // so abandoning the stage cannot expose the previous model's limits. + ggufLayerCount: fields.ggufLayerCount, + moeLayerCount: fields.moeLayerCount, + }; + } + return { ...fields, ...seedExtras }; +} + /** A local model staged for a deferred load (see `pendingSelection`). Shape is * a subset of the load hook's `SelectedModelInput`, structurally assignable. */ export type PendingModelSelection = { @@ -515,6 +727,13 @@ export type PendingModelSelection = { * Scoped here (not the shared `ggufContextLength`) so a staged model's * metadata never pollutes the currently-loaded model's context display. */ contextLength?: number | null; + /** Total layer count (GGUF block_count); the manual gpu-layers ceiling is + * this + 1 (llama.cpp counts the output layer as offloadable too); + * scoped here like contextLength. */ + layerCount?: number | null; + /** MoE expert-layer count from the GGUF header (manual --n-cpu-moe ceiling); + * 0 for dense models, scoped here like contextLength. */ + moeLayerCount?: number | null; /** "Load on selection" on + un-cached GGUF: download via the manager (global * indicator) without opening the sheet, then load once the download finishes. */ autoLoad?: boolean; @@ -743,6 +962,32 @@ type ChatRuntimeStore = { tensorParallel: boolean; /** Backend-reported tensor-parallel state; null until first hydrated. */ loadedTensorParallel: boolean | null; + /** GPU memory strategy for GGUF loads. "auto" = Unsloth picks GPUs and context + * to fit; "manual" = you own the offload (gpuLayers < 0 = Auto/--fit, >= 0 + * pins layers + nCpuMoe). */ + gpuMemoryMode: "auto" | "manual"; + /** Backend-reported gpu memory mode; null until first hydrated. */ + loadedGpuMemoryMode: "auto" | "manual" | null; + /** Manual mode: layers to offload to GPU. -1 = Auto (--fit); >= model layer + * count = all. */ + gpuLayers: number; + loadedGpuLayers: number | null; + /** Manual mode: MoE expert layers to keep on CPU (--n-cpu-moe); 0 = none. */ + nCpuMoe: number; + loadedNCpuMoe: number | null; + /** Manual mode: per-GPU layer counts (--tensor-split), in GPU-in-use order; + * null = unset (llama.cpp splits by free VRAM). */ + splitRatio: number[] | null; + /** Backend-reported per-GPU split ratio (--tensor-split); null = unset. */ + loadedSplitRatio: number[] | null; + /** Model layer count (GGUF block_count); the manual gpu-layers ceiling is + * this + 1 (the output layer is offloadable too). */ + ggufLayerCount: number | null; + /** MoE expert-layer count: the nCpuMoe slider max; 0/null hides the slider. */ + moeLayerCount: number | null; + /** Picked physical GPU indices (null = use all / automatic). */ + selectedGpuIds: number[] | null; + loadedGpuIds: number[] | null; /** Persisted: when false, picking a local model stages it as * `pendingSelection` (and opens settings) instead of loading immediately, * so load settings can be set before the single load. */ @@ -766,6 +1011,9 @@ type ChatRuntimeStore = { * per step, cleared when the run ends, never persisted into the transcript. */ activeDiffusionCanvas: DiffusionCanvasFrame | null; customContextLength: number | null; + /** The pinned context the loaded model used (null = Auto), so dirty-tracking + * and a later fit Apply can tell an explicit pin apart from Auto. */ + loadedCustomContextLength: number | null; defaultChatTemplate: string | null; chatTemplateOverride: string | null; loadedChatTemplateOverride: string | null; @@ -884,6 +1132,11 @@ type ChatRuntimeStore = { * which skip the sheet but must still honor a saved config. */ applyRememberedLoadSettings: (settings: RememberedLoadSettings) => void; setTensorParallel: (value: boolean) => void; + setGpuMemoryMode: (mode: "auto" | "manual") => void; + setGpuLayers: (value: number) => void; + setNCpuMoe: (value: number) => void; + setSplitRatio: (value: number[] | null) => void; + setSelectedGpuIds: (ids: number[] | null) => void; setLoadOnSelection: (value: boolean) => void; setExpandQuantizations: (value: boolean) => void; setShowAllQuantizations: (value: boolean) => void; @@ -1101,11 +1354,12 @@ function setScalarSettingVersion( /** The "revert to the loaded model" baseline for the editable load knobs. * Shared by resetModelSettingsToLoaded (full revert) and stageModel (which - * overrides speculative to start a fresh pick from the standing default). */ + * overrides speculative and the per-model GPU knobs to start a fresh pick). */ function loadedBaselineSettings(s: ChatRuntimeStore) { const hasLoadedModel = Boolean(s.params.checkpoint); return { - customContextLength: null, + // Revert to the loaded model's pin (null = Auto), not a blanket Auto. + customContextLength: s.loadedCustomContextLength, kvCacheDtype: s.loadedKvCacheDtype, tensorParallel: s.loadedTensorParallel ?? false, speculativeType: hasLoadedModel @@ -1113,6 +1367,20 @@ function loadedBaselineSettings(s: ChatRuntimeStore) { : readPersistedSpeculativeType(), specDraftNMax: hasLoadedModel ? s.loadedSpecDraftNMax : null, chatTemplateOverride: s.loadedChatTemplateOverride, + // GPU memory mode is a standing preference; revert to the loaded model's + // mode (or the persisted default when nothing is loaded). Manual knobs and + // the GPU pick are per-model and revert to their loaded baseline. A loaded + // model with no applicable mode -- diffusion ("auto" baseline) or non-GGUF + // (null baseline) -- keeps the live preference so Reset can't drop it. + gpuMemoryMode: !hasLoadedModel + ? readPersistedGpuMemoryMode() + : s.loadedIsDiffusion + ? s.gpuMemoryMode + : (s.loadedGpuMemoryMode ?? s.gpuMemoryMode), + gpuLayers: s.loadedGpuLayers ?? GPU_LAYERS_AUTO, + nCpuMoe: s.loadedNCpuMoe ?? 0, + splitRatio: s.loadedSplitRatio ?? null, + selectedGpuIds: s.loadedGpuIds, }; } @@ -1213,6 +1481,18 @@ export const useChatRuntimeStore = create((set, get) => ({ loadedSpecDraftNMax: null, tensorParallel: false, loadedTensorParallel: null, + gpuMemoryMode: readPersistedGpuMemoryMode(), + loadedGpuMemoryMode: null, + gpuLayers: GPU_LAYERS_AUTO, + loadedGpuLayers: null, + nCpuMoe: 0, + loadedNCpuMoe: null, + splitRatio: null, + loadedSplitRatio: null, + ggufLayerCount: null, + moeLayerCount: null, + selectedGpuIds: null, + loadedGpuIds: null, loadOnSelection: loadBool(CHAT_LOAD_ON_SELECTION_KEY, true), expandQuantizations: loadBool(CHAT_EXPAND_QUANTIZATIONS_KEY, false), showAllQuantizations: loadBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, true), @@ -1221,6 +1501,7 @@ export const useChatRuntimeStore = create((set, get) => ({ loadedIsMultimodal: false, loadedIsDiffusion: false, customContextLength: null, + loadedCustomContextLength: null, defaultChatTemplate: null, chatTemplateOverride: null, loadedChatTemplateOverride: null, @@ -1455,9 +1736,23 @@ export const useChatRuntimeStore = create((set, get) => ({ loadedSpecDraftNMax: null, tensorParallel: false, loadedTensorParallel: null, + // Standing preference: survives unload, unlike the per-model knobs above. + gpuMemoryMode: readPersistedGpuMemoryMode(), + loadedGpuMemoryMode: null, + gpuLayers: GPU_LAYERS_AUTO, + loadedGpuLayers: null, + nCpuMoe: 0, + loadedNCpuMoe: null, + splitRatio: null, + loadedSplitRatio: null, + ggufLayerCount: null, + moeLayerCount: null, + selectedGpuIds: null, + loadedGpuIds: null, loadedIsMultimodal: false, loadedIsDiffusion: false, customContextLength: null, + loadedCustomContextLength: null, defaultChatTemplate: null, chatTemplateOverride: null, loadedChatTemplateOverride: null, @@ -1753,17 +2048,67 @@ export const useChatRuntimeStore = create((set, get) => ({ setSpeculativeType: (speculativeType) => set({ speculativeType }), setSpecDraftNMax: (specDraftNMax) => set({ specDraftNMax }), setTensorParallel: (tensorParallel) => set({ tensorParallel }), + // Standing preference, but persisted only on a successful load (see + // use-chat-model-runtime), not on selection -- so an unapplied pick the user + // resets/abandons doesn't stick to the next session. + setGpuMemoryMode: (gpuMemoryMode) => set({ gpuMemoryMode }), + setGpuLayers: (gpuLayers) => set({ gpuLayers }), + setNCpuMoe: (nCpuMoe) => set({ nCpuMoe }), + setSplitRatio: (splitRatio) => set({ splitRatio }), + setSelectedGpuIds: (selectedGpuIds) => set({ selectedGpuIds }), resetModelSettingsToLoaded: () => set((s) => loadedBaselineSettings(s)), - applyRememberedLoadSettings: (settings) => + applyRememberedLoadSettings: (settings) => { + const gpuCacheWasCold = cachedPinnableGpuIndices() === null; + const restoredGpuIds = + settings.selectedGpuIds !== undefined + ? reconcilePersistedGpuIds(settings.selectedGpuIds) + : undefined; // Coalesce every field: a blob persisted by an older/newer build can omit // keys, and a raw spread would push `undefined` into fields typed non-null. + // The GPU knobs are spread only when present, but first reset the per-model + // ones to defaults: this path (load-on-selection) starts from the loaded + // model's baseline and skips the model-switch reset, so a blob omitting + // gpuLayers/nCpuMoe/selectedGpuIds (older build) or splitRatio (never + // remembered) must not inherit the previous model's placement. gpuMemoryMode + // (standing preference) is NOT reset, only applied when the blob carries it; + // selectedGpuIds keeps a meaningful null (all GPUs), so it keys off undefined. set({ + gpuLayers: GPU_LAYERS_AUTO, + nCpuMoe: 0, + splitRatio: null, + selectedGpuIds: null, customContextLength: settings.contextLength ?? null, kvCacheDtype: settings.kvCacheDtype ?? null, speculativeType: settings.speculativeType ?? "auto", specDraftNMax: settings.specDraftNMax ?? null, tensorParallel: settings.tensorParallel ?? false, - }), + ...(settings.gpuMemoryMode != null && { + gpuMemoryMode: settings.gpuMemoryMode, + }), + ...(settings.gpuLayers != null && { gpuLayers: settings.gpuLayers }), + ...(settings.nCpuMoe != null && { nCpuMoe: settings.nCpuMoe }), + ...(restoredGpuIds !== undefined && { + // Reconcile against the GPUs present now (see reconcilePersistedGpuIds): + // a saved [1] on a 1-GPU host (or under relative/UUID visibility) would + // hide the picker yet still send gpu_ids, which the backend rejects. + selectedGpuIds: restoredGpuIds, + }), + }); + // A cold cache makes the synchronous restore provisional. Reconcile again + // when the shared fetch completes, but only if this exact restored array is + // still current so a user edit, stage change, or load cannot be overwritten. + if (gpuCacheWasCold && restoredGpuIds != null) { + void ensureGpuDeviceCache().then(() => { + set((state) => { + if (state.selectedGpuIds !== restoredGpuIds) return state; + const reconciled = reconcilePersistedGpuIds(restoredGpuIds); + return reconciled === restoredGpuIds + ? state + : { selectedGpuIds: reconciled }; + }); + }); + } + }, setLoadOnSelection: (loadOnSelection) => { saveBool(CHAT_LOAD_ON_SELECTION_KEY, loadOnSelection); set({ loadOnSelection }); @@ -1798,6 +2143,22 @@ export const useChatRuntimeStore = create((set, get) => ({ // Load's keepSpeculative) a forced MTP mode onto a model that may lack it. speculativeType: readPersistedSpeculativeType(), specDraftNMax: null, + // Keep the on-screen GPU Memory selection (loadedBaselineSettings would + // otherwise revert it to the loaded model's mode, dropping a Manual choice + // just made). Use the live store value, not the persisted one, which can + // lag a mode hydrated from an out-of-band load. + gpuMemoryMode: s.gpuMemoryMode, + // Per-model GPU knobs start from defaults too so a fresh pick doesn't + // inherit the loaded model's layer/MoE/split/GPU choices, matching the + // immediate-switch reset. + gpuLayers: GPU_LAYERS_AUTO, + nCpuMoe: 0, + splitRatio: null, + selectedGpuIds: null, + // Fresh pick starts at Auto context (loadedBaselineSettings would + // otherwise restore the current model's pin). Leaves the baseline + // intact, like the GPU knobs, so abandoning restores the loaded pin. + customContextLength: null, }; }); }, diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index d72c406fdd..c24ddde5f5 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -65,6 +65,18 @@ export interface LoadModelRequest { * of by layer for GGUF models. Multi-GPU only; no effect on a single GPU. */ tensor_parallel?: boolean | null; + /** GPU memory strategy for GGUF models. "auto" (default): Unsloth selects GPUs + * and caps context to fit VRAM. "manual": you own the offload -- gpu_layers + * -1 (Auto) hands sizing to llama.cpp's --fit, >= 0 pins layers/n_cpu_moe. */ + gpu_memory_mode?: "auto" | "manual"; + /** Manual mode: layers to offload to GPU (--gpu-layers, --fit off); -1 = Auto (--fit). */ + gpu_layers?: number; + /** Manual mode: MoE expert layers to keep on CPU (--n-cpu-moe); 0 = none. */ + n_cpu_moe?: number; + /** Manual mode: relative model share per GPU (--tensor-split), in GPU order. */ + tensor_split?: number[] | null; + /** Picked physical GPU indices (omit/empty = automatic). */ + gpu_ids?: number[]; } export interface ValidateModelResponse { @@ -80,6 +92,13 @@ export interface ValidateModelResponse { requires_security_review?: boolean; /** Native context length from the local GGUF header; null until downloaded. */ context_length?: number | null; + /** Total layer count (GGUF block_count); the manual gpu-layers ceiling is + * this + 1 (llama.cpp counts the output layer as offloadable too); null + * until downloaded. */ + layer_count?: number | null; + /** MoE expert-layer count from the GGUF header (manual --n-cpu-moe ceiling); + * 0 for dense models, null until downloaded. */ + moe_layer_count?: number | null; /** Architecture only shipped by a newer transformers; UI pauses on the upgrade dialog. */ requires_transformers_upgrade?: boolean; /** Set only when requires_transformers_upgrade. */ @@ -159,6 +178,14 @@ export interface LoadModelResponse { spec_draft_n_max?: number | null; /** Whether tensor-parallel split (--split-mode tensor) is active. */ tensor_parallel?: boolean; + gpu_memory_mode?: "auto" | "manual"; + gpu_layers?: number; + n_cpu_moe?: number; + tensor_split?: number[] | null; + n_layers?: number | null; + /** Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not MoE. */ + n_moe_layers?: number; + gpu_ids?: number[] | null; } export interface UnloadModelRequest { @@ -203,6 +230,17 @@ export interface InferenceStatusResponse { spec_draft_n_max?: number | null; /** Whether tensor-parallel split (--split-mode tensor) is active. */ tensor_parallel?: boolean; + gpu_memory_mode?: "auto" | "manual"; + gpu_layers?: number; + n_cpu_moe?: number; + tensor_split?: number[] | null; + /** n_ctx the active GGUF load was invoked with (0 = Auto); re-seeds a + * Manual + Auto-layers context pin on hydration. Null for non-GGUF. */ + requested_context_length?: number | null; + gpu_ids?: number[] | null; + n_layers?: number | null; + /** Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not MoE. */ + n_moe_layers?: number; /** * Why MTP was disabled on the loaded model despite being requested. * "binary_no_mtp" / "binary_outdated" -> updating llama.cpp would re-enable diff --git a/studio/frontend/src/hooks/use-gpu-info.ts b/studio/frontend/src/hooks/use-gpu-info.ts index 1e313acdf3..db2cc021be 100644 --- a/studio/frontend/src/hooks/use-gpu-info.ts +++ b/studio/frontend/src/hooks/use-gpu-info.ts @@ -15,6 +15,19 @@ export interface GpuInfo { systemRamTotalGb: number } +export interface SystemGpuDevice { + index: number; + name: string; + memoryTotalGb: number; + /** Free VRAM at fetch time. Degrades to the total when the utilization + * probe had no usage data; 0 only when the total is unknown too. */ + memoryFreeGb: number; + /** "physical" = `index` is a stable physical/PCI id safe to pin via gpu_ids; + * "relative" = an ordinal into a parent CUDA_VISIBLE_DEVICES mask, which the + * backend can't map back, so the picker must not offer it. */ + physicalIndex: boolean; +} + const DEFAULT_GPU: GpuInfo = { available: false, name: "Unknown", @@ -25,70 +38,135 @@ const DEFAULT_GPU: GpuInfo = { systemRamTotalGb: 0 }; -// Module-level cache so multiple components share one fetch. -let cachedGpu: GpuInfo | null = null; -let fetchPromise: Promise | null = null; +// One module-level cache so every GPU hook shares a single /api/system fetch. +let cachedSystem: SystemInfoResponse | null = null; +let systemPromise: Promise | null = null; -async function fetchGpuOnce(): Promise { - if (cachedGpu) return cachedGpu; - if (fetchPromise) return fetchPromise; - - fetchPromise = (async () => { +async function fetchSystemOnce(): Promise { + if (cachedSystem) return cachedSystem; + if (systemPromise) return systemPromise; + systemPromise = (async () => { try { const res = await authFetch("/api/system"); if (!res.ok) throw new Error(`HTTP ${res.status}`); - - const data = await res.json() as SystemInfoResponse; - const gpuData = data?.gpu; - - // CPU/RAM exist even on hosts without a GPU, so populate them on every path. - // No discrete GPU (e.g. Mac): still surface system RAM so memory math - // (unified memory) has a budget to work with. - const base = { - cpuCore: data?.cpu?.physical_count ?? 0, - cpuThread: data?.cpu?.logical_count ?? 0, - systemRamAvailableGb: data?.memory?.available_gb ?? 0, - systemRamTotalGb: data?.memory?.total_gb ?? 0, - }; - - const devices = gpuData?.devices ?? []; - const info: GpuInfo = - gpuData?.available && devices.length - ? { - ...base, - available: true, - name: devices[0]?.name ?? "Unknown", - memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0), - } - : { ...DEFAULT_GPU, ...base }; - cachedGpu = info; - return info; + cachedSystem = (await res.json()) as SystemInfoResponse; + return cachedSystem; } catch { - // Reset promise so subsequent calls retry (e.g. backend wasn't ready) - fetchPromise = null; - return DEFAULT_GPU; + systemPromise = null; // reset so a later call retries (backend not ready) + return null; } })(); + return systemPromise; +} - return fetchPromise; +function toGpuInfo(data: SystemInfoResponse | null): GpuInfo { + // CPU/RAM exist even on GPU-less hosts (e.g. Mac), so populate them on every + // path: unified-memory math still needs a RAM budget to work with. + const base = { + cpuCore: data?.cpu?.physical_count ?? 0, + cpuThread: data?.cpu?.logical_count ?? 0, + systemRamAvailableGb: data?.memory?.available_gb ?? 0, + systemRamTotalGb: data?.memory?.total_gb ?? 0, + }; + const gpuData = data?.gpu; + const devices = gpuData?.devices ?? []; + if (!gpuData?.available || !devices.length) { + return { ...DEFAULT_GPU, ...base }; + } + return { + ...base, + available: true, + name: devices[0]?.name ?? "Unknown", + memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0), + }; +} + +function toGpuDevices(data: SystemInfoResponse | null): SystemGpuDevice[] { + // Unpinnable configurations must hide every pick surface: XPU indices are + // torch-xpu ordinals no applicator speaks, and Vulkan-only builds pin ggml's + // own ordinals -- /load and /validate 400 picks on both, so the backend + // reports gpu.gguf_gpu_ids_supported and every gate keyed on physicalIndex + // (picker, persisted-pick reconcile) follows it. The device flavor lives on + // the TOP-LEVEL device_backend field; absent support info defaults to + // pinnable (older backend). + const pinnableBackend = + data?.device_backend !== "xpu" && + data?.gpu?.gguf_gpu_ids_supported !== false; + return (data?.gpu?.devices ?? []) + .filter((d) => typeof d.index === "number") + .map((d) => ({ + index: d.index as number, + name: d.name ?? `GPU ${d.index}`, + memoryTotalGb: d.memory_total_gb ?? 0, + memoryFreeGb: d.vram_free_gb ?? 0, + physicalIndex: pinnableBackend && d.index_kind === "physical", + })); +} + +/** Aggregate GPU info from /api/system; shares one module-level fetch across all GPU hooks. */ +export function useGpuInfo(): GpuInfo { + const [gpu, setGpu] = useState( + cachedSystem ? toGpuInfo(cachedSystem) : DEFAULT_GPU, + ); + useEffect(() => { + // No early return on cachedSystem: a consumer mounting as the cache fills + // (between render and effect) would otherwise stay stuck at the default. + let cancelled = false; + fetchSystemOnce().then((d) => { + if (!cancelled) setGpu(toGpuInfo(d)); + }); + return () => { + cancelled = true; + }; + }, []); + return gpu; +} + +/** All backend-visible GPUs (index, name, total VRAM); shares the same fetch. */ +export function useGpuDevices(): SystemGpuDevice[] { + const [devices, setDevices] = useState( + cachedSystem ? toGpuDevices(cachedSystem) : [], + ); + useEffect(() => { + // No early return on cachedSystem: a consumer mounting as the cache fills + // (between render and effect) would otherwise stay stuck at the default. + let cancelled = false; + fetchSystemOnce().then((d) => { + if (!cancelled) setDevices(toGpuDevices(d)); + }); + return () => { + cancelled = true; + }; + }, []); + return devices; } /** - * Fetch GPU info from /api/system. Cached at module level, so only one request - * is made no matter how many components call this hook. + * Await the shared /api/system fetch so cachedPinnableGpuIndices (and the + * store's reconcilePersistedGpuIds) can validate a persisted pick before a + * load path sends it -- on a cold cache the reconcile passes ids through + * unvalidated, and a stale cross-host pick then fails /load with the picker + * hidden. Resolves immediately once the module cache is warm; a failed fetch + * keeps the cache cold, preserving the "can't validate, backend guards" + * degradation. */ -export function useGpuInfo(): GpuInfo { - const [gpu, setGpu] = useState(cachedGpu ?? DEFAULT_GPU); +export async function ensureGpuDeviceCache(): Promise { + await fetchSystemOnce(); +} - useEffect(() => { - if (cachedGpu) return; - - let cancelled = false; - fetchGpuOnce().then((info) => { - if (!cancelled) setGpu(info); - }); - return () => { cancelled = true; }; - }, []); - - return gpu; -} \ No newline at end of file +/** + * Pinnable physical GPU indices from the already-fetched /api/system cache, for + * non-React code (the store) that needs to validate a persisted `gpu_ids` pick + * without triggering a fetch. Returns: + * - `null` when the cache isn't populated yet (caller can't validate, so keep + * the pick and let the backend guard reject a truly bad one); + * - `[]` when the host has no pinnable multi-GPU set (single GPU, or relative/ + * UUID-masked indices) -- the picker is hidden, so any saved pick is stale; + * - the physical indices otherwise. + */ +export function cachedPinnableGpuIndices(): number[] | null { + if (!cachedSystem) return null; + const physical = toGpuDevices(cachedSystem).filter((d) => d.physicalIndex); + // Mirrors the sheet's showGpuPicker gate: only a 2+ physical-GPU host can pin. + return physical.length > 1 ? physical.map((d) => d.index) : []; +} diff --git a/studio/frontend/src/hooks/use-system.ts b/studio/frontend/src/hooks/use-system.ts index a135cce86e..8cfe2bace4 100644 --- a/studio/frontend/src/hooks/use-system.ts +++ b/studio/frontend/src/hooks/use-system.ts @@ -40,6 +40,9 @@ export interface SystemInfoResponse { gpu: { available: boolean; backend?: string; + /** Whether GGUF loads accept an explicit gpu_ids pick (false on XPU hosts + * and Vulkan-only builds, where /load and /validate 400 picks). */ + gguf_gpu_ids_supported?: boolean; backend_cuda_visible_devices?: string | null; parent_visible_gpu_ids?: number[]; index_kind?: string;