ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: GPU memory configuration for GGUF models (#6414)

Browse source 
* 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 <unslothshared@gmail.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
This commit is contained in:
oobabooga 2026-07-19 09:46:22 -03:00 • committed by GitHub
commit 5f1f30ec82
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
31 changed files with 4352 additions and 397 deletions
Download patch file Download diff file Expand all files Collapse all files

715
studio/backend/core/inference/llama_cpp.py
View file

File diff suppressed because it is too large Load diff

27
studio/backend/core/inference/llama_server_args.py
View file

@ -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] = []

14
studio/backend/main.py
View file

@ -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

144
studio/backend/models/inference.py
View file

@ -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 = (

518
studio/backend/routes/inference.py
View file

@ -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,

6
studio/backend/routes/models.py
View file

@ -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

42
studio/backend/routes/training_vram.py
View file

@ -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"}

333
studio/backend/tests/test_chat_load_during_training.py
View file

@ -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) ──────

73
studio/backend/tests/test_gguf_metadata.py
View file

@ -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(

879
studio/backend/tests/test_gpu_memory_mode.py Normal file
View file

@ -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

18
studio/backend/tests/test_gpu_selection.py
View file

@ -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(

12
studio/backend/tests/test_llama_cpp_no_context_shift.py
View file

@ -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):

31
studio/backend/tests/test_llama_cpp_props_readback.py
View file

@ -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
)

45
studio/backend/tests/test_llama_server_args.py
View file

@ -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"],

5
studio/backend/tests/test_tensor_parallel.py
View file

@ -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():

17
studio/backend/tests/test_tp_vision_regression.py
View file

@ -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)

109
studio/backend/utils/models/gguf_metadata.py
View file

@ -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}.<suffix>``; 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("<I" if vtype == 4 else "<Q", n_bytes)[0]
# A real context length is positive; treat 0/garbage as
# absent so the UI never builds a slider with max < min.
return value if value > 0 else None
found[key[len(arch) + 1 :]] = struct.unpack(
"<I" if vtype == 4 else "<Q", n_bytes
)[0]
if len(found) == len(wanted_suffixes):
break
else:
if not _skip_gguf_value(f, vtype):
break
except (struct.error, UnicodeDecodeError):
break
except OSError as e:
logger.debug(f"read_gguf_context_length: cannot open {path}: {e}")
logger.debug(f"_parse_gguf_arch_uints: cannot open {path}: {e}")
return None
except Exception as e:
logger.debug(f"read_gguf_context_length: parse failure on {path}: {e}")
logger.debug(f"_parse_gguf_arch_uints: parse failure on {path}: {e}")
return None
return None
return found
def _parse_gguf_staged_dims(path: str) -> 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.

24
studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts
View file

@ -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 <model>".
// dialog can offer "Remember settings for <model>". 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;

128
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -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<boolean> {
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),

38
studio/frontend/src/features/chat/api/chat-api.ts
View file

@ -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<ValidateModelResponse>(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<number | null> {
}): 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<ValidateModelResponse>(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<void> {

13
studio/frontend/src/features/chat/chat-page.tsx
View file

@ -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.",

450
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -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 (
<div className="space-y-3.5">
@ -279,6 +285,7 @@ function ParamSlider({
displayValue={displayValue}
ariaLabel={label}
size={valueSize ?? 4}
disabled={disabled}
/>
</div>
<Slider
@ -288,6 +295,7 @@ function ParamSlider({
value={[value]}
onValueChange={([v]) => onChange(snapToStep(v, step, min, max))}
className="panel-slider"
disabled={disabled}
/>
</div>
);
@ -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 ? (
<div className="space-y-3.5">
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Context Length
</span>
<InfoHint>
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.
</InfoHint>
</div>
<NumericValueInput
value={fitCtxAuto ? 0 : (customContextLength ?? 0)}
displayValue={fitCtxAuto ? "Auto" : undefined}
min={0}
max={ctxMaxValue ?? undefined}
step={1}
onChange={(v) => {
setCustomContextLength(v > 0 ? v : null);
}}
ariaLabel="Context Length"
size={8}
disabled={modelControlsDisabled}
/>
</div>
<Slider
min={0}
max={ctxMaxValue ?? 4096}
step={1024}
value={[
fitCtxAuto
? 0
: Math.min(
customContextLength ?? 0,
ctxMaxValue ?? 4096,
),
]}
onValueChange={([v]) => {
// 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 && (
<p className="text-[11px] text-nav-fg/40">
llama.cpp loaded {fitResolvedCtx.toLocaleString()} tokens.
</p>
)}
</div>
) : (
<div className="space-y-3.5">
<div className="flex items-center justify-between gap-3">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
@ -1036,7 +1253,7 @@ export function ChatSettingsPanel({
</p>
)}
</div>
)}
))}
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
@ -1191,6 +1408,163 @@ export function ChatSettingsPanel({
)}
</>
)}
{gpuModeApplies && (
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
GPU Memory
</span>
<InfoHint>
<div className="flex flex-col gap-1.5">
<div>
<span className="font-medium">Default:</span> Unsloth
fits the model and context to your GPUs.
</div>
<div>
<span className="font-medium">Manual:</span> set GPU
Layers yourself. Leave it on Auto to let llama.cpp size
the context and offload overflow (including MoE experts)
to RAM.
</div>
</div>
</InfoHint>
</div>
<div className="flex shrink-0 items-center gap-1.5">
<Select
value={gpuMemoryMode}
onValueChange={(v) => {
setGpuMemoryMode(v as "auto" | "manual");
}}
// An in-flight staged load already snapshotted its
// settings, so edits here could not apply -- disable like
// the sibling context/KV/spec controls.
disabled={modelControlsDisabled}
>
<SelectTrigger
animateRadius={false}
icon={ChevronDownStandardIcon}
iconClassName="size-3.5"
className="grid h-7 w-[136px] min-w-0 grid-cols-[minmax(0,1fr)_auto] items-center gap-1 rounded-full border-transparent bg-black/[0.04] dark:bg-white/[0.05] hover:bg-black/[0.06] dark:hover:bg-white/[0.1] pl-3 pr-2 py-0 text-[13px]! font-medium text-nav-fg focus-visible:ring-0 focus-visible:border-transparent [&_[data-slot=select-value]]:min-w-0 [&_[data-slot=select-value]]:truncate [&>svg]:shrink-0"
data-test-id="gpu-memory-mode-select"
>
<SelectValue />
</SelectTrigger>
<SelectContent className="menu-soft-surface ring-0 border-0 rounded-lg">
<SelectItem value="auto">Default</SelectItem>
<SelectItem value="manual">Manual</SelectItem>
</SelectContent>
</Select>
</div>
</div>
)}
{isManual && (
<>
<ParamSlider
label="GPU Layers"
value={Math.max(GPU_LAYERS_AUTO, Math.min(gpuLayers, gpuLayersMax))}
min={GPU_LAYERS_AUTO}
max={gpuLayersMax}
step={1}
onChange={setGpuLayers}
disabled={modelControlsDisabled}
displayValue={autoLayers ? "Auto" : undefined}
valueSize={6}
info={
<>
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 && (
<ParamSlider
label="MoE Layers on CPU"
value={Math.min(nCpuMoe, moeLayersMax)}
min={0}
max={moeLayersMax}
step={1}
onChange={setNCpuMoe}
disabled={modelControlsDisabled}
valueSize={6}
info={
<>
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 && (
<div className="space-y-3.5">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Layers per GPU
</span>
<InfoHint>
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.
</InfoHint>
</div>
{gpusInUseDevices.map((d, k) => (
<ParamSlider
key={d?.index ?? k}
label={`GPU ${d?.index ?? k}`}
value={Math.min(splitCounts[k] ?? 0, splitTotal)}
min={0}
max={splitTotal}
step={1}
onChange={(v) => setSplitCount(k, v)}
valueSize={6}
disabled={modelControlsDisabled}
/>
))}
</div>
)}
</>
)}
{showGpuPicker && (
<div className="space-y-2">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
GPUs
</span>
<InfoHint>
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.
</InfoHint>
</div>
<div className="flex flex-col gap-2">
{gpuDevices.map((d) => (
<div
key={d.index}
className="flex items-center justify-between gap-3"
>
<span className="min-w-0 truncate text-[12px] text-nav-fg/80">
GPU {d.index}: {d.name}
{d.memoryTotalGb
? ` · ${Math.round(d.memoryTotalGb)} GB`
: ""}
</span>
<Switch
className="panel-switch shrink-0"
checked={isGpuChecked(d.index)}
onCheckedChange={() => toggleGpu(d.index)}
data-test-id={`gpu-pick-${d.index}`}
disabled={modelControlsDisabled}
/>
</div>
))}
</div>
</div>
)}
{gpuModeApplies && !autoLayers && (
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
@ -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"
/>
</div>
)}
</>
)}
{/* No persistent "enable custom code" toggle: it is consented per model
@ -1228,14 +1603,21 @@ export function ChatSettingsPanel({
{Math.round((stagedDownloadFraction ?? 0) * 100)}%
</p>
)}
<label className="flex cursor-pointer items-center gap-2 pb-1.5 text-[12px] text-muted-foreground">
<Checkbox
className="size-3.5 rounded-full [&_[data-slot=checkbox-indicator]_svg]:size-2.5"
checked={remember}
onCheckedChange={(v) => setRemember(v === true)}
/>
Remember settings next time
</label>
{/* GGUF picks only: a non-GGUF pick shows none of the load
knobs the blob captures, so there is nothing to remember. */}
{pendingIsGguf && (
<label className="flex cursor-pointer items-center gap-2 pb-1.5 text-[12px] text-muted-foreground">
<Checkbox
className="size-3.5 rounded-full [&_[data-slot=checkbox-indicator]_svg]:size-2.5"
checked={remember}
onCheckedChange={(v) => setRemember(v === true)}
// The save/clear already ran in the Load click handler, so
// a mid-load toggle could not apply -- lock it like the knobs.
disabled={modelControlsDisabled}
/>
Remember settings next time
</label>
)}
{stagedLoading ? (
// Mid-load: nothing to load or abandon until it settles, so disable.
<Button
@ -1255,9 +1637,10 @@ export function ChatSettingsPanel({
// Persist (or clear) this model's load knobs before loading.
// Context is stored as the override (null = auto), never the
// resolved native value, so restoring can't force an OOM.
const pid = pendingSelection
? rememberedLoadSettingsKey(pendingSelection)
: null;
// GGUF-only, like the restore effect: saving for a
// non-GGUF pick would snapshot leftover standing values
// its hidden controls never showed.
const pid = pendingIsGguf ? pendingKey : null;
if (pid) {
if (remember) {
saveRememberedLoadSettings(pid, {
@ -1266,6 +1649,10 @@ export function ChatSettingsPanel({
speculativeType,
specDraftNMax,
tensorParallel,
gpuMemoryMode,
gpuLayers,
nCpuMoe,
selectedGpuIds,
});
} else {
clearRememberedLoadSettings(pid);
@ -1319,7 +1706,11 @@ export function ChatSettingsPanel({
</Button>
</div>
) : null}
<ChatTemplateFields />
{/* 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. */}
<ChatTemplateFields disabled={modelControlsDisabled} />
</div>
</CollapsibleSection>
)}
@ -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() {
<button
type="button"
onClick={openEditor}
className="cursor-pointer text-left text-[13px] font-medium tracking-nav text-nav-fg"
disabled={disabled}
className="cursor-pointer text-left text-[13px] font-medium tracking-nav text-nav-fg disabled:pointer-events-none disabled:opacity-50"
>
Chat Template
</button>
@ -2131,7 +2523,8 @@ function ChatTemplateFields() {
<button
type="button"
onClick={() => setOverride(null)}
className="nav-icon-btn text-nav-icon-idle hover:bg-panel-surface-hover hover:text-black dark:hover:text-white"
disabled={disabled}
className="nav-icon-btn text-nav-icon-idle hover:bg-panel-surface-hover hover:text-black dark:hover:text-white disabled:pointer-events-none disabled:opacity-50"
aria-label="Revert chat template"
>
<HugeiconsIcon
@ -2155,7 +2548,8 @@ function ChatTemplateFields() {
<button
type="button"
onClick={openEditor}
className="nav-icon-btn text-nav-icon-idle hover:bg-panel-surface-hover hover:text-black dark:hover:text-white"
disabled={disabled}
className="nav-icon-btn text-nav-icon-idle hover:bg-panel-surface-hover hover:text-black dark:hover:text-white disabled:pointer-events-none disabled:opacity-50"
aria-label="Edit chat template"
>
<HugeiconsIcon
@ -2218,7 +2612,13 @@ function ChatTemplateFields() {
>
Cancel
</Button>
<Button type="button" onClick={saveEditor} disabled={!draftDirty}>
{/* Also locked mid-load: an autoLoad can start with this dialog
already open, and a save then would be silently clobbered. */}
<Button
type="button"
onClick={saveEditor}
disabled={!draftDirty || disabled}
>
Save
</Button>
</div>

217
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -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 {

46
studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts
View file

@ -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,
]);

117
studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts
View file

@ -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<ReasoningEffort, "low" | "medium" | "high">;
function sameArray<T>(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<boolean> {
if (previousCheckpoint) {
return true;
}
const previousGgufVariant = useChatRuntimeStore.getState().activeGgufVariant;
store.setCheckpoint(checkpointId, status.gguf_variant);
applyActiveModelStatusToStore(status, { previousCheckpoint });
applyActiveModelStatusToStore(status, {
previousCheckpoint,
previousGgufVariant,
});
return true;
}

31
studio/frontend/src/features/chat/presets/preset-policy.ts
View file

@ -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;
}

109
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -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),
});

369
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -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<T extends object>(
resp: Parameters<typeof loadedGpuMemoryFields>[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<K extends ScalarSettingKey>(
/** 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<ChatRuntimeStore>((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<ChatRuntimeStore>((set, get) => ({
loadedIsMultimodal: false,
loadedIsDiffusion: false,
customContextLength: null,
loadedCustomContextLength: null,
defaultChatTemplate: null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
@ -1455,9 +1736,23 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((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<ChatRuntimeStore>((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<ChatRuntimeStore>((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,
};
});
},

38
studio/frontend/src/features/chat/types/api.ts
View file

@ -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

186
studio/frontend/src/hooks/use-gpu-info.ts
View file

@ -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<GpuInfo> | null = null;
// One module-level cache so every GPU hook shares a single /api/system fetch.
let cachedSystem: SystemInfoResponse | null = null;
let systemPromise: Promise<SystemInfoResponse | null> | null = null;
async function fetchGpuOnce(): Promise<GpuInfo> {
if (cachedGpu) return cachedGpu;
if (fetchPromise) return fetchPromise;
fetchPromise = (async () => {
async function fetchSystemOnce(): Promise<SystemInfoResponse | null> {
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<GpuInfo>(
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<SystemGpuDevice[]>(
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<GpuInfo>(cachedGpu ?? DEFAULT_GPU);
export async function ensureGpuDeviceCache(): Promise<void> {
await fetchSystemOnce();
}
useEffect(() => {
if (cachedGpu) return;
let cancelled = false;
fetchGpuOnce().then((info) => {
if (!cancelled) setGpu(info);
});
return () => { cancelled = true; };
}, []);
return gpu;
}
/**
* 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) : [];
}

3
studio/frontend/src/hooks/use-system.ts
View file

@ -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;