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

Merge 0573a0e0f0 into 3212710a4a

Browse source
This commit is contained in:
Ayushman 2026-07-29 09:09:54 +00:00 • committed by GitHub
commit 2435d18cbc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
53 changed files with 4011 additions and 770 deletions
Download patch file Download diff file Expand all files Collapse all files

89
studio/backend/core/inference/_vulkan_probe.py
View file

@ -37,19 +37,6 @@ def _igpu_flags_and_names(base, lib, count: int) -> tuple[list[bool], list[str]]
"""
flags = [False] * count
names = [""] * count
# The name lookup is bound OUTSIDE the type-detection try: a ggml-base
# without ggml_backend_dev_description (older/custom build) must degrade to
# unnamed devices, not abort before the iGPU flags are read (which would
# count an iGPU's shared RAM as VRAM).
describe = None
try:
base.ggml_backend_dev_description.restype = ctypes.c_char_p
base.ggml_backend_dev_description.argtypes = [ctypes.c_void_p]
describe = base.ggml_backend_dev_description
except Exception:
pass
try:
lib.ggml_backend_vk_reg.restype = ctypes.c_void_p
lib.ggml_backend_vk_reg.argtypes = []
@ -59,33 +46,50 @@ def _igpu_flags_and_names(base, lib, count: int) -> tuple[list[bool], list[str]]
base.ggml_backend_reg_dev_get.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
base.ggml_backend_dev_type.restype = ctypes.c_int
base.ggml_backend_dev_type.argtypes = [ctypes.c_void_p]
reg = lib.ggml_backend_vk_reg()
if not reg:
return flags, names
dev_count = base.ggml_backend_reg_dev_count(reg)
for i in range(min(count, dev_count)):
dev = base.ggml_backend_reg_dev_get(reg, i)
if dev:
flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU
if describe is not None:
try:
desc = describe(dev)
if desc:
# Tabs/newlines would corrupt the line protocol;
# spaces are safe.
names[i] = (
desc.decode("utf-8", errors = "replace")
.replace("\t", " ")
.replace("\n", " ")
.strip()
)
except Exception:
pass
except Exception:
# Best-effort: any failure degrades to "discrete"/"unnamed" so the
# memory readings still get through instead of crashing the probe.
pass
return flags, names
# Bound outside the type-detection try above: a ggml-base without the
# description symbol (older/custom build) must degrade to unnamed devices,
# not abort before the iGPU flags are read (which would count an iGPU's
# shared RAM as VRAM).
name_functions = []
for symbol in ("ggml_backend_dev_description", "ggml_backend_dev_name"):
try:
function = getattr(base, symbol)
function.restype = ctypes.c_char_p
function.argtypes = [ctypes.c_void_p]
name_functions.append(function)
except Exception:
pass
for i in range(min(count, dev_count)):
try:
dev = base.ggml_backend_reg_dev_get(reg, i)
except Exception:
continue
if not dev:
continue
try:
flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU
except Exception:
pass
for function in name_functions:
try:
raw_name = function(dev)
except Exception:
continue
if raw_name:
# Tabs/newlines would corrupt the line protocol; spaces are safe.
name = raw_name.decode("utf-8", errors = "replace")
names[i] = name.replace("\t", " ").replace("\r", " ").replace("\n", " ").strip()
break
return flags, names
@ -115,12 +119,27 @@ def main() -> int:
else:
base_name, vk_name = "libggml-base.so", "libggml-vulkan.so"
def _find_lib(directory, stem):
"""Find an unversioned library or a versioned runtime soname."""
path = os.path.join(directory, stem)
if os.path.isfile(path):
return path
for entry in sorted(os.listdir(directory)):
if entry.startswith(stem + "."):
return os.path.join(directory, entry)
return None
# RTLD_GLOBAL exposes ggml-base's symbols to ggml-vulkan on POSIX. getattr
# falls back to 0 where the flag doesn't exist (Windows CDLL ignores mode).
_rtld_global = getattr(ctypes, "RTLD_GLOBAL", 0)
base_path = _find_lib(bindir, base_name)
vk_path = _find_lib(bindir, vk_name)
if not base_path or not vk_path:
print(f"ggml-vulkan load failed: library not found in {bindir}", file = sys.stderr)
return 1
try:
base = ctypes.CDLL(os.path.join(bindir, base_name), mode = _rtld_global)
lib = ctypes.CDLL(os.path.join(bindir, vk_name), mode = _rtld_global)
base = ctypes.CDLL(base_path, mode = _rtld_global)
lib = ctypes.CDLL(vk_path, mode = _rtld_global)
except OSError as e:
print(f"ggml-vulkan load failed: {e}", file = sys.stderr)
return 1

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

@ -131,7 +131,7 @@ LLAMA_SERVER_NOT_FOUND_DETAIL = (
"then try again. (Advanced: set LLAMA_SERVER_PATH to an existing binary.)"
)
# Shared by the route, pre-teardown and post-metadata rejections (#7205).
# Shared by the pre-teardown and post-metadata rejections (#7205).
_VULKAN_DIFFUSION_GPU_IDS_ERROR = (
"GPU selection (gpu_ids) is not supported for a DiffusionGemma "
"GGUF on a Vulkan llama.cpp build: the diffusion runner selects "
@ -1868,18 +1868,14 @@ def _extra_args_draft_offloaded_to_cpu(
device flag has no env). An embedded MTP head follows the main -ngl, so these
draft-only flags don't move it. Last-wins, so only each flag's final value counts."""
ngl_flags = {"--spec-draft-ngl", "-ngld", "--gpu-layers-draft", "--n-gpu-layers-draft"}
dev_flags = {"--spec-draft-device", "-devd", "--device-draft"}
args = [str(a) for a in extra_args] if extra_args else []
last_ngl: Optional[str] = None
last_dev: Optional[str] = None
for i, raw in enumerate(args):
flag = _flag_name(raw)
_, eq, inline = raw.partition("=")
value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
if flag in ngl_flags:
last_ngl = value
elif flag in dev_flags:
last_dev = value
if last_ngl is None:
last_ngl = (os.environ if env is None else env).get("LLAMA_ARG_N_GPU_LAYERS_DRAFT")
if last_ngl is not None:
@ -1888,6 +1884,7 @@ def _extra_args_draft_offloaded_to_cpu(
return True
except (TypeError, ValueError):
pass
last_dev = _extra_args_draft_device(extra_args)
if last_dev is not None:
devs = [d.strip().lower() for d in last_dev.split(",") if d.strip()]
if devs and all(d in ("cpu", "none") for d in devs):
@ -1895,6 +1892,41 @@ def _extra_args_draft_offloaded_to_cpu(
return False
def _extra_args_device(
extra_args: Optional[Iterable[str]], flags: Collection[str]
) -> Optional[str]:
"""Return the last value for any device flag in ``flags``."""
args = [str(a) for a in extra_args] if extra_args else []
value: Optional[str] = None
for i, raw in enumerate(args):
flag = _flag_name(raw)
_, eq, inline = raw.partition("=")
if flag in flags:
value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
return value
def _extra_args_main_device(extra_args: Optional[Iterable[str]]) -> Optional[str]:
"""Return the last explicit main-device value, if any."""
return _extra_args_device(extra_args, {"--device", "-dev"})
def _extra_args_draft_device(extra_args: Optional[Iterable[str]]) -> Optional[str]:
"""Return the last explicit draft-device value, if any."""
return _extra_args_device(extra_args, {"--spec-draft-device", "-devd", "--device-draft"})
def _extra_args_draft_device_pin(extra_args: Optional[Iterable[str]]) -> Optional[str]:
"""Return a GPU draft-device override; cpu/none do not conflict with a pin."""
last_dev = _extra_args_draft_device(extra_args)
if last_dev is None:
return None
devs = [d.strip() for d in last_dev.split(",") if d.strip()]
if not devs or all(d.lower() in ("cpu", "none") for d in devs):
return None
return last_dev
def _extra_args_n_ubatch(
extra_args: Optional[Iterable[str]],
env: Optional[Mapping[str, str]] = None,
@ -2067,10 +2099,6 @@ def _backfill_usage_from_timings(usage, timings):
return out
def _vulkan_lib_filename() -> str:
return "ggml-vulkan.dll" if sys.platform == "win32" else "libggml-vulkan.so"
# Host RAM to leave free on an integrated GPU, matching llama.cpp's own --fit
# margin (default 1024 MiB per device). ggml reports an iGPU's "VRAM" as shared
# system RAM, so hold back the same margin rather than inventing a larger one.
@ -2111,6 +2139,19 @@ def _llama_lib_dir(binary: str) -> Path:
return resolved.parent
def _lib_dir_has_ggml_backend(lib_dir: Path, backend: str) -> bool:
"""Match an exact or versioned ggml backend library soname."""
stem = f"ggml-{backend}.dll" if sys.platform == "win32" else f"libggml-{backend}.so"
try:
return any(
f.name == stem or f.name.startswith(stem + ".")
for f in lib_dir.iterdir()
if f.is_file()
)
except OSError:
return False
def _is_external_link(path: Path) -> bool:
"""True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink
or a Windows directory junction / reparse point. Such a link resolves into
@ -3235,13 +3276,10 @@ class LlamaCppBackend:
if not binary:
return False
lib_dir = _llama_lib_dir(binary)
if not (lib_dir / _vulkan_lib_filename()).is_file():
if not _lib_dir_has_ggml_backend(lib_dir, "vulkan"):
return False
for _backend in ("cuda", "hip"):
sibling = (
f"ggml-{_backend}.dll" if sys.platform == "win32" else f"libggml-{_backend}.so"
)
if (lib_dir / sibling).is_file():
if _lib_dir_has_ggml_backend(lib_dir, _backend):
return False
return True
@ -3550,6 +3588,42 @@ class LlamaCppBackend:
return []
return ["--device", ",".join(f"Vulkan{i}" for i in gpu_indices)]
@staticmethod
def _backend_lacks_gpu_lib(binary: Optional[str] = None) -> bool:
"""Detect a split-library build proven to be CPU-only.
Unknown and static layouts return False to avoid rejecting custom GPU builds.
"""
binary = binary or LlamaCppBackend._find_llama_server_binary()
if not binary:
return False
lib_dir = _llama_lib_dir(binary)
if not lib_dir or not lib_dir.is_dir():
return False
if any(_lib_dir_has_ggml_backend(lib_dir, b) for b in ("vulkan", "cuda", "hip")):
return False
return any(_lib_dir_has_ggml_backend(lib_dir, b) for b in ("cpu", "base"))
def is_vulkan_build(self) -> bool:
"""Whether the resolved llama-server uses Vulkan ordinals."""
try:
return self._is_vulkan_backend(self._find_llama_server_binary())
except Exception:
return False
@staticmethod
def _strip_device_extra_args(extra_args):
"""Drop device and main-GPU flags so gpu_ids remains authoritative."""
return strip_shadowing_flags(
extra_args,
strip_context = False,
strip_cache = False,
strip_spec = False,
strip_template = False,
strip_split_mode = False,
strip_device = True,
)
@staticmethod
def _get_gpu_free_memory(binary: Optional[str] = None) -> list[tuple[int, int]]:
"""Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by
@ -3692,6 +3766,12 @@ class LlamaCppBackend:
logger.debug(f"torch GPU probe failed: {e}")
return []
@staticmethod
def _vulkan_auto_gpu_memory(gpus: list[tuple[int, int, int]]) -> list[tuple[int, int, int]]:
"""Prefer discrete devices for automatic Vulkan placement."""
discrete = [gpu for gpu in gpus if gpu[2] > 0]
return discrete or gpus
@staticmethod
def _run_vulkan_probe(binary: Optional[str] = None) -> list[dict]:
"""Run ``_vulkan_probe.py`` and parse its per-device lines.
@ -3709,17 +3789,11 @@ class LlamaCppBackend:
if not binary:
return []
binary_dir = _llama_lib_dir(binary)
if not (binary_dir / _vulkan_lib_filename()).is_file():
if not _lib_dir_has_ggml_backend(binary_dir, "vulkan"):
return []
env = child_env_without_native_path_secret()
# Pass any inherited GGML_VK_VISIBLE_DEVICES through to ggml unchanged so
# the probe enumerates the same device list the launch will, named
# Vulkan0..N in the compact order reported here and pinned by that name
# via --device -- probe, mask, and pin stay in one index space. Do NOT
# filter the mask in Python: ggml parses the env var in raw
# vkEnumeratePhysicalDevices space while this probe reports the compact
# post-filter ordinal, so a Python filter would compare mismatched spaces.
# Let ggml apply GGML_VK_VISIBLE_DEVICES; Python sees compact ordinals.
if sys.platform != "win32":
# Let the loader resolve sibling ggml libs next to the binary.
existing_ld = env.get("LD_LIBRARY_PATH", "")
@ -6748,6 +6822,7 @@ class LlamaCppBackend:
# Explicit GPU placement pool (issue #7164). None/[] = auto-select;
# the fitter may pin the smallest subset of this pool that fits.
gpu_ids: Optional[List[int]] = None,
gpu_ids_are_vulkan_ordinals: Optional[bool] = None,
n_threads: Optional[int] = None,
n_gpu_layers: Optional[int] = None, # caller compat, unused
n_parallel: int = 1,
@ -6788,6 +6863,7 @@ class LlamaCppBackend:
"n_cpu_moe": n_cpu_moe,
"tensor_split": list(tensor_split) if tensor_split is not None else None,
"gpu_ids": list(gpu_ids) if gpu_ids is not None else None,
"gpu_ids_are_vulkan_ordinals": gpu_ids_are_vulkan_ordinals,
"n_threads": n_threads,
"n_gpu_layers": n_gpu_layers,
"n_parallel": n_parallel,
@ -6849,6 +6925,12 @@ class LlamaCppBackend:
# GGUF uses the diffusion runner, and its arch is only known after the header.
binary = self._find_llama_server_binary()
is_vulkan_backend = self._is_vulkan_backend(binary)
_vulkan_ordinal_pin = (
is_vulkan_backend and bool(gpu_ids) and gpu_ids_are_vulkan_ordinals is not False
)
_cpu_only_pin = (
bool(gpu_ids) and not is_vulkan_backend and self._backend_lacks_gpu_lib(binary)
)
# Without --kv-unified an explicit --parallel N splits -c into windows of -c/N, so on a
# build lacking the flag the default of 4 would quarter every context window for a
@ -6876,8 +6958,10 @@ class LlamaCppBackend:
# not, so a stale gpu_ids=[99] used to kill the server then 400, leaving
# nothing running (#7239). _get_gpu_memory needs only the binary (safe pre-
# download) and reuses the later fit's issubset logic. Guarded on a found
# Vulkan build + a pin so a deferred not-found stays deferred for diffusion.
if is_vulkan_backend and gpu_ids and binary:
# Vulkan build + a pin so a deferred not-found stays deferred for diffusion,
# and on the ordinals being unconfirmed: a caller-confirmed physical CUDA pin
# is not in this namespace and must not be range-checked against it.
if _vulkan_ordinal_pin and binary:
_pf_wanted = {int(x) for x in gpu_ids}
_pf_probed = {g[0] for g in self._get_gpu_memory(binary)}
if not _pf_wanted.issubset(_pf_probed):
@ -6886,9 +6970,11 @@ class LlamaCppBackend:
f"present. Available Vulkan devices: {sorted(_pf_probed)}."
)
# Classify before killing the healthy server (#7205); Phase 2 reuses this path.
# Classify before killing the healthy server (#7205); Phase 2 reuses this
# path. Diffusion changes the meaning or validity of requested placement,
# so both the remote and the local branch settle it above the teardown.
_preflight_model_path = None
if is_vulkan_backend and gpu_ids and hf_repo:
if hf_repo and (_vulkan_ordinal_pin or _cpu_only_pin):
_resolved_repo = _resolve_repo_id_casing(hf_repo)
if _resolved_repo != hf_repo:
logger.info(
@ -6903,13 +6989,25 @@ class LlamaCppBackend:
hf_variant = hf_variant,
hf_token = hf_token,
)
self._reject_vulkan_diffusion_gpu_ids_before_teardown(
_preflight_model_path,
model_identifier,
_preflight_is_diffusion = self._gguf_path_is_diffusion(
_preflight_model_path, model_identifier
)
if _preflight_is_diffusion:
if _vulkan_ordinal_pin:
raise ValueError(_VULKAN_DIFFUSION_GPU_IDS_ERROR)
elif _cpu_only_pin:
raise ValueError(
f"Requested gpu_ids {list(gpu_ids)} but the llama.cpp build has "
"no GPU backend (CPU-only build); it would ignore the pin and run "
"on CPU. Omit gpu_ids to run on CPU."
)
elif is_vulkan_backend and gpu_ids and gguf_path and not hf_repo:
if not Path(gguf_path).is_file():
raise FileNotFoundError(f"GGUF file not found: {gguf_path}")
# Only an ordinal pin is ambiguous for the diffusion runner; a
# caller-confirmed physical pin keeps its CUDA placement path.
if _vulkan_ordinal_pin and gguf_path and not hf_repo:
self._reject_vulkan_diffusion_gpu_ids_before_teardown(
gguf_path,
model_identifier,
@ -6990,23 +7088,14 @@ class LlamaCppBackend:
# Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server;
# serve them with the diffusion runner (same OpenAI-compat interface).
if self._is_diffusion:
# Final defense: route and pre-teardown preflights reject before Phase 1.
# The diffusion runner cannot translate Vulkan ordinals to CUDA IDs, so
# this fires unless the caller confirmed the pin is physical CUDA IDs.
if is_vulkan_backend and gpu_ids and gpu_ids_are_vulkan_ordinals is not False:
raise ValueError(_VULKAN_DIFFUSION_GPU_IDS_ERROR)
# Not a tensor/layer GGUF: clear any preserved-fallback flag from a
# prior load (this path skips the command builder that clears it).
self._layer_preserves_tensor_intent = False
# On a Vulkan build gpu_ids are ggml Vulkan ordinals, but the diffusion
# runner selects its device by CUDA physical index (_diffusion_gpu_arg
# forwards gpu_ids[0] as a CUDA/DG_GPU token) with no mapping to them.
# The route rejects a CONFIRMED-diffusion pick up front; an uncached GGUF
# only classified as diffusion post-download still reaches here with a
# pin, so drop it and serve on the default device (like an unpinned load).
if gpu_ids and is_vulkan_backend:
logger.warning(
"Ignoring gpu_ids %s for diffusion GGUF on a Vulkan build: "
"the diffusion runner cannot map ggml Vulkan ordinals; "
"serving on the default device.",
gpu_ids,
)
gpu_ids = None
with self._lock:
if self._cancel_event.is_set():
logger.info("Load cancelled before diffusion server start")
@ -7122,7 +7211,8 @@ class LlamaCppBackend:
self._gpu_layers = -1
self._n_cpu_moe = 0
self._tensor_split = None
self._gpu_ids = sorted(gpu_ids) if gpu_ids else None
self._requested_gpu_ids = sorted(gpu_ids) if gpu_ids else None
self._gpu_ids = list(self._requested_gpu_ids) if self._requested_gpu_ids else None
# Manual offload skips the TP planner but still emits --split-mode
# tensor at launch; drop it when fewer than 2 GPUs are in use --
# tensor split is a no-op there and aborts on some architectures.
@ -7234,6 +7324,7 @@ class LlamaCppBackend:
# empty `gpus` so the speculative defaults stay GPU-aware and the
# CPU-fallback check still knows GPUs were present.
_detected_gpus: list[tuple[int, int]] = []
total_by_idx: dict[int, int] = {}
model_size = None # set in the fit try; used by the APU RAM guard
# Layer-fallback min GPUs; raised below on a tensor downgrade. Bound
# before the try so the --fit-on except path still has it (no UnboundLocal).
@ -7255,6 +7346,8 @@ class LlamaCppBackend:
# per-GPU headroom (correct when the GPU is already partly used).
# Pass binary so a Vulkan build probes ggml's Vulkan ordinals.
_gpu_mem = self._get_gpu_memory(binary)
if is_vulkan_backend and not gpu_ids and gpu_memory_mode != "manual":
_gpu_mem = self._vulkan_auto_gpu_memory(_gpu_mem)
gpus = [(idx, free) for idx, free, _t in _gpu_mem]
# Restrict the fit (and thus the layer plan + pin env) to the
# selected GPUs; fail-open if none match so a stale UI choice
@ -8142,6 +8235,24 @@ class LlamaCppBackend:
# model can't spill onto an unpicked GPU.
if gpu_ids and gpu_indices is None:
gpu_indices = sorted(gpu_ids)
# Auto Vulkan fit prefers discrete GPUs and keeps that pool pinned.
elif (
is_vulkan_backend
and gpu_memory_mode != "manual"
and use_fit
and gpu_indices is None
and _detected_gpus
):
gpu_indices = sorted(idx for idx, _free in _detected_gpus)
# Status reports the explicit subset the launch actually uses,
# while reload dedupe compares requested_gpu_ids so repeating a
# wider request that fit narrowed does not restart the server.
if gpu_ids:
effective_pin = gpu_indices if gpu_indices is not None else gpu_ids
self._gpu_ids = sorted(int(idx) for idx in effective_pin)
else:
self._gpu_ids = None
# Unified-memory APUs load weights into system RAM (under WSL the VM
# cap, not the ROCm-reported VRAM, is the real ceiling); refuse an
@ -8198,6 +8309,8 @@ class LlamaCppBackend:
elif not auto_fit:
cmd.extend(["-c", "0"])
server_caps = self.probe_server_capabilities(binary)
# Report a clean public model id (matching GET /v1/models) rather
# than the raw -m path in llama-server's own /v1/models and the
# "model" field of its chat/completions responses.
@ -8397,6 +8510,10 @@ class LlamaCppBackend:
# Speculative decoding. See _build_speculative_flags for the
# mode resolution, benchmarks, and llama.cpp references.
_vulkan_pin_ids = gpu_indices if gpu_indices is not None else (gpu_ids or None)
_draft_device = _extra_args_main_device(extra_args) if gpu_ids is None else None
if _draft_device is None and is_vulkan_backend and _vulkan_pin_ids:
_draft_device = ",".join(f"Vulkan{i}" for i in _vulkan_pin_ids)
launch_mtp_draft_path = self._resolve_launch_mtp_path(
mtp_draft_path = mtp_draft_path,
)
@ -8409,6 +8526,7 @@ class LlamaCppBackend:
gpus = bool(_detected_gpus),
binary = binary,
mtp_draft_path = launch_mtp_draft_path,
draft_device = _draft_device,
)
# Remember where the spec block sits so a drafter-load failure
# can be retried with these flags swapped out (see below).
@ -8504,10 +8622,6 @@ class LlamaCppBackend:
", ".join(unsupported_cache_flags),
)
# Vulkan pins via --device (a cmd arg), before user extras so a user
# --device wins. Fall back to raw ids when the fit did not narrow.
_vulkan_pin_ids = gpu_indices if gpu_indices is not None else (gpu_ids or None)
# Record the pin actually applied (fit-narrowed gpu_indices, else the raw
# request) for the keep-warm loop, dedupe, and /status, so an explicit
# [0, 1] narrowed to [0] records [0] and /status never echoes an ordinal
@ -8544,12 +8658,22 @@ class LlamaCppBackend:
if is_vulkan_backend and _vulkan_pin_ids is not None:
cmd += LlamaCppBackend._vulkan_pin_args(_vulkan_pin_ids)
# User pass-through args go last so llama.cpp's last-wins parsing
# lets the user override Unsloth's auto-set flags. Already
# validated by the route via validate_extra_args().
# User pass-through args go last. Placement flags are removed
# below when the Studio picker owns the GPU selection.
if extra_args:
cmd.extend(str(a) for a in extra_args)
logger.info(f"Appending user extra args to llama-server: {list(extra_args)}")
_emit_extra_args = list(extra_args)
if gpu_ids is not None:
# gpu_ids owns placement, so remove competing device flags.
_emit_extra_args = self._strip_device_extra_args(extra_args)
if _emit_extra_args != list(extra_args):
logger.info(
"Dropped user device placement flags from extra args: "
"explicit gpu_ids owns placement."
)
cmd.extend(str(a) for a in _emit_extra_args)
logger.info(
f"Appending user extra args to llama-server: {list(_emit_extra_args)}"
)
kv_cache_unified = _kv_unified_from_args(cmd)
@ -8590,6 +8714,11 @@ class LlamaCppBackend:
if _ct_raw and _ct_raw not in self._TENSOR_PARALLEL_KV_TYPES:
env.pop(_ct_var, None)
# A gpu_ids pin also owns inherited device placement.
if gpu_ids is not None:
env.pop("LLAMA_ARG_DEVICE", None)
env.pop("LLAMA_ARG_MAIN_GPU", None)
# Windows + full offload: PASSIVE OMP + 2 threads stop
# spin-wait burning CPU. CPU/partial offload keeps default
# OMP parallelism. #5692.
@ -9102,7 +9231,14 @@ class LlamaCppBackend:
# clears, list sets. Source records hf_variant for the route's
# same_source check.
if extra_args is not None:
self._extra_args = list(extra_args)
# Persist the same device-stripped extras the command used when gpu_ids
# owns placement, so a later inheriting reload (after gpu_ids clears)
# can't resurrect the dropped --device (#7188).
self._extra_args = (
self._strip_device_extra_args(extra_args)
if gpu_ids is not None
else list(extra_args)
)
self._extra_args_source = (model_identifier, hf_variant)
self._requested_n_ctx = int(n_ctx)
# Local n_parallel may have been reduced above; the snapshot has the ask.
@ -9195,6 +9331,7 @@ class LlamaCppBackend:
gpus: bool,
binary: Optional[str],
mtp_draft_path: Optional[str] = None,
draft_device: Optional[str] = None,
) -> List[str]:
"""Return the llama-server flag list for the requested spec mode.
@ -9337,6 +9474,8 @@ class LlamaCppBackend:
# main GGUF.
if mtp_draft_path:
flags.extend(["--model-draft", mtp_draft_path])
if draft_device:
flags.extend(["--spec-draft-device", draft_device])
logger.info(f"Using separate MTP drafter: {mtp_draft_path}")
spec_value = mtp_token
ngram_knobs: list[str] = []
@ -9660,7 +9799,11 @@ class LlamaCppBackend:
# layer); only an explicit list forces equality.
if extra_args is not None:
current = list(self._extra_args) if self._extra_args is not None else []
if list(extra_args) != current:
candidate = list(extra_args)
# Compare the same device-stripped extras that load_model persisted.
if gpu_ids is not None:
candidate = self._strip_device_extra_args(candidate)
if candidate != current:
return False
self._record_matching_gpu_request(gpu_ids)
return True
@ -10768,7 +10911,7 @@ class LlamaCppBackend:
windows of ``-c / N``; restore the shared pool so one request can use
the full context. With ``--fit on``, ``--fit-ctx`` floors the fit step
at an explicitly requested ctx so it offloads or fails instead of
silently shrinking the window. The 8192 auto-floor and the tighter
silently shrinking the window. The 8192 auto-floor and tighter
``--fit-target`` margin apply only under Manual + Auto (``auto_fit``),
which omits ``-c``: on the legacy auto path ``-c 0`` already pins the
native window and ``--fit-ctx 8192`` would override it down to 8192.
@ -10785,9 +10928,8 @@ class LlamaCppBackend:
# shrink the window below a usable size.
flags.extend(["--fit-ctx", "8192"])
if use_fit and auto_fit and caps.get("supports_fit_target"):
# llama.cpp's --fit leaves 1 GiB free per device by default;
# tighten that to 512 MiB so it packs more of the model onto
# the GPU before spilling to system RAM.
# Leave 512 MiB free per device so llama.cpp packs more of the
# model onto the GPU before spilling to system RAM.
flags.extend(["--fit-target", "512"])
return flags

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

@ -198,6 +198,10 @@ _TEMPLATE_FLAGS: frozenset[str] = frozenset(
_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
# llama.cpp placement flags. Opt-in (users may pass them under auto-select):
# stripped only when gpu_ids is set, so they cannot override the selected pool
# or choose a main GPU outside it (#7188).
_DEVICE_FLAGS: frozenset[str] = frozenset({"--device", "-dev", "--main-gpu", "-mg"})
# 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
@ -215,7 +219,13 @@ _SHADOWING_FLAGS: frozenset[str] = (
# 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", "-cmoe", "--cpu-moe"}
{
"--spec-default",
"--jinja",
"--no-jinja",
"-cmoe",
"--cpu-moe",
}
)
@ -471,6 +481,7 @@ def strip_shadowing_flags(
strip_split_mode: bool = True,
strip_tensor_split: bool = False,
strip_offload: bool = False,
strip_device: bool = False,
) -> list[str]:
"""Strip flags that shadow first-class Unsloth settings.
@ -484,7 +495,8 @@ def strip_shadowing_flags(
``--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.
row/none/layer choice intact. ``strip_device`` is enabled when ``gpu_ids``
owns placement.
"""
shadowing: set[str] = set()
if strip_context:
@ -501,6 +513,8 @@ def strip_shadowing_flags(
shadowing |= _TENSOR_SPLIT_FLAGS
if strip_offload:
shadowing |= _OFFLOAD_SHADOWING_FLAGS
if strip_device:
shadowing |= _DEVICE_FLAGS
tokens = [str(a) for a in (args or [])]
out: list[str] = []

24
studio/backend/main.py
View file

@ -1265,20 +1265,26 @@ def _get_cached_system_gpu_info(logger) -> tuple[dict[str, Any], dict[str, Any]]
)
enriched_devices.append(enriched_dev)
# Whether GGUF loads accept an explicit gpu_ids pick. /load and /validate
# 400 picks on XPU hosts, where no visibility mask speaks torch-xpu
# ordinals. A Vulkan build IS pinnable: its picks are ggml ordinals, the
# same space `--device Vulkan<i>` uses, so check it first and let it
# through even on an XPU host (the XPU ban is about torch ordinals).
is_vulkan_build = False
try:
from core.inference.llama_cpp import LlamaCppBackend
from utils.hardware import DeviceType, get_device
is_vulkan_build = LlamaCppBackend._is_vulkan_backend()
gpu_ids_supported = is_vulkan_build or get_device() != DeviceType.XPU
llama_uses_vulkan = LlamaCppBackend._is_vulkan_backend()
if llama_uses_vulkan:
# The separate inference inventory below owns Vulkan ordinals.
# Keep this false so a failed Vulkan probe cannot expose torch
# indices that llama.cpp would interpret in another namespace.
gpu_ids_supported = False
else:
# XPU indices cannot yet be applied safely across Level Zero's
# FLAT and COMPOSITE hierarchy modes. A proven CPU-only
# llama.cpp build cannot apply a CUDA pin either.
gpu_ids_supported = (
get_device() != DeviceType.XPU and not LlamaCppBackend._backend_lacks_gpu_lib()
)
except Exception as e:
logger.debug(f"Could not resolve gpu_ids support: {e}")
llama_uses_vulkan = False
gpu_ids_supported = True
# Preserve backend/index metadata from the visibility probe. In
# particular, a CPU training host can expose a Vulkan inference GPU and
@ -1288,6 +1294,7 @@ def _get_cached_system_gpu_info(logger) -> tuple[dict[str, Any], dict[str, Any]]
**visibility_info,
"available": visibility_info.get("available", False),
"devices": enriched_devices,
"backend": visibility_info.get("backend"),
"gguf_gpu_ids_supported": gpu_ids_supported,
}
@ -1296,6 +1303,7 @@ def _get_cached_system_gpu_info(logger) -> tuple[dict[str, Any], dict[str, Any]]
# If Vulkan is installed but its probe fails, retain the unavailable
# Vulkan shape instead of budgeting training GPUs that llama.cpp cannot use.
if visibility_info.get("backend") == "vulkan":
gpu_info["gguf_gpu_ids_supported"] = bool(enriched_devices)
inference_gpu_info = gpu_info
else:
vulkan_info = get_vulkan_inference_gpu_info()

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

@ -80,9 +80,10 @@ class LoadRequest(BaseModel):
None,
description = (
"GPU placement pool, for example [0, 1]. Omit or pass [] to use "
"automatic selection. CUDA/ROCm and Intel XPU values are physical "
"GPU indices; Vulkan values are ggml device ordinals. Explicit "
"physical IDs are unsupported when the parent visibility mask uses "
"automatic selection. CUDA/ROCm values are physical GPU indices; "
"Vulkan values are ggml device ordinals. Explicit selection is not "
"supported on XPU, and physical IDs are unsupported when the parent "
"visibility mask uses "
"non-numeric or subdevice entries, including CUDA_VISIBLE_DEVICES "
"with UUID/MIG entries and ZE_AFFINITY_MASK with subdevice tokens "
"(for example '0.0,0.1') or FLAT-hierarchy tile handles. For GGUF "
@ -334,6 +335,9 @@ class ValidateModelResponse(BaseModel):
identifier: Optional[str] = Field(None, description = "Resolved model identifier")
display_name: Optional[str] = Field(None, description = "Display name derived from identifier")
is_gguf: bool = Field(False, description = "Whether this is a GGUF model (llama.cpp)")
is_diffusion: bool = Field(
False, description = "Whether this is a block-diffusion model (DiffusionGemma)"
)
is_lora: bool = Field(False, description = "Whether this is a LoRA adapter")
is_vision: bool = Field(False, description = "Whether this is a vision-capable model")
requires_trust_remote_code: bool = Field(

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

@ -1005,6 +1005,7 @@ try:
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_STREAM_STALL_TIMEOUT_S,
_canonicalize_spec_mode,
_extra_args_draft_device_pin,
_extra_args_n_ubatch,
_extra_args_set_spec_type,
_hf_offline_if_dns_dead,
@ -1049,6 +1050,7 @@ except ImportError:
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_STREAM_STALL_TIMEOUT_S,
_canonicalize_spec_mode,
_extra_args_draft_device_pin,
_extra_args_n_ubatch,
_extra_args_set_spec_type,
_hf_offline_if_dns_dead,
@ -3454,7 +3456,22 @@ def _request_matches_loaded_settings(
):
return False
else:
if list(request.llama_extra_args) != backend_extra:
# Compare against the managed flags that load_model actually persisted.
_strip_dev = bool(request.gpu_ids)
_request_extra = (
strip_shadowing_flags(
request.llama_extra_args,
strip_context = False,
strip_cache = False,
strip_spec = False,
strip_template = False,
strip_split_mode = False,
strip_device = _strip_dev,
)
if _strip_dev
else list(request.llama_extra_args)
)
if _request_extra != backend_extra:
return False
# A separate drafter (Gemma's root mtp-*.gguf) appearing or disappearing
# next to the loaded weights changes the launch command (--model-draft),
@ -4598,21 +4615,11 @@ def _estimate_gguf_required_gb(
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. Default
placement keeps that unknown case guarded until the header is available.
"""
"""Classify a GGUF as diffusion, normal, or unknown before loading."""
identity = " ".join(
str(getattr(config, attr, "") or "") for attr in ("identifier", "gguf_hf_repo", "gguf_file")
).lower()
# Name-only hint, used ONLY as a pre-download fallback, scoped to the
# DiffusionGemma runner family: a bare "diffusion" substring is common in
# ordinary text-model names/paths (e.g. "stable-diffusion-prompt"), and treating
# those as diffusion falsely rejects a valid Vulkan+gpu_ids GGUF (#7239). Normalize
# non-alphanumerics so "DiffusionGemma"/"diffusion-gemma" collapse to one token.
# The local header below stays authoritative.
# Only use the specific DiffusionGemma family name as a header fallback.
name_says_diffusion = "diffusiongemma" in _re.sub(r"[^a-z0-9]+", "", identity)
try:
@ -4624,42 +4631,92 @@ def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]:
from hub.utils.gguf import resolve_local_gguf_path
main = resolve_local_gguf_path(repo, variant)
if main and Path(main).is_file():
# The local GGUF header is authoritative (same probe the loader uses), so
# it can't be fooled by a "diffusion"-flavored name/path.
probe = LlamaCppBackend()
probe._read_gguf_metadata(str(main))
if probe.is_diffusion:
return True
# A decoded architecture proves a normal llama-server GGUF; no architecture
# means the probe was inconclusive, so fall through to the name hint below.
if getattr(probe, "_architecture", None):
return False
except Exception as e:
logger.debug("Could not identify diffusion GGUF for training guard: %s", e)
# Header unavailable (remote uncached) or inconclusive: True only for the
# DiffusionGemma name family; otherwise None keeps an unknown remote GGUF guarded
# as potentially diffusion until its header proves otherwise.
return True if name_says_diffusion else None
async def _resolve_gguf_gpu_ids_for_request(
config: ModelConfig, gpu_ids: Optional[List[int]]
) -> Optional[List[int]]:
"""Resolve and fully validate an explicit GGUF GPU placement pool.
def _reject_draft_device_with_gpu_ids(
gpu_ids: Optional[List[int]],
extra_args: Optional[list[str]],
*,
gpu_ids_are_vulkan_ordinals: bool,
) -> None:
"""Reject a physical drafter pin beside Vulkan-ordinal main placement."""
if not gpu_ids or not gpu_ids_are_vulkan_ordinals:
return
draft_device = _extra_args_draft_device_pin(extra_args)
if draft_device is not None:
raise HTTPException(
status_code = 400,
detail = (
f"A draft-model device override ('{draft_device}') cannot be combined "
"with explicit gpu_ids: it would place the speculative drafter outside "
"the pinned GPUs the training guard budgeted. Remove the draft-device "
"flag to follow gpu_ids, or set it to none."
),
)
CUDA and ROCm use physical IDs. Vulkan uses ggml ordinals, so its device
existence check comes from the same ggml probe used by the loader. Both
/load and /validate call this before their training guard or any teardown.
"""
_DIFFUSION_KIND_UNSET = object()
async def _resolve_gguf_gpu_ids_for_request(
config: ModelConfig,
gpu_ids: Optional[List[int]],
*,
diffusion_kind: Optional[bool] | object = _DIFFUSION_KIND_UNSET,
) -> tuple[Optional[List[int]], bool]:
"""Validate GGUF GPU IDs and report whether they are Vulkan ordinals."""
if not gpu_ids:
return None
return None, False
from utils.hardware import DeviceType, get_device
from utils.hardware.hardware import resolve_requested_gpu_ids
is_vulkan = LlamaCppBackend._is_vulkan_backend()
if get_device() == DeviceType.XPU and not is_vulkan:
llama_backend = get_llama_cpp_backend()
is_vulkan_build = await asyncio.to_thread(llama_backend.is_vulkan_build)
if diffusion_kind is _DIFFUSION_KIND_UNSET:
diffusion_kind = _classify_diffusion_gguf(config)
confirmed_diffusion = diffusion_kind is True
definitively_non_diffusion = diffusion_kind is False
device = get_device()
lacks_gpu_lib = getattr(llama_backend, "_backend_lacks_gpu_lib", None)
# ROCm is deliberately DeviceType.CUDA internally because it uses
# torch.cuda.*. Only the API label changes to "rocm", so this accepts both
# CUDA and ROCm physical IDs while rejecting device namespaces the
# diffusion runner cannot apply.
diffusion_physical_ids_supported = device == DeviceType.CUDA
if confirmed_diffusion and not diffusion_physical_ids_supported:
raise HTTPException(
status_code = 400,
detail = (
"GPU selection (gpu_ids) for DiffusionGemma requires CUDA or ROCm. "
"Omit gpu_ids on this host."
),
)
if confirmed_diffusion and is_vulkan_build:
raise HTTPException(
status_code = 400,
detail = (
"GPU selection (gpu_ids) is not supported for a DiffusionGemma "
"GGUF on a Vulkan llama.cpp build: the picker uses Vulkan ordinals, "
"which have no defined mapping to CUDA physical indices. Omit gpu_ids "
"to use the default device."
),
)
ids_are_vulkan_ordinals = is_vulkan_build
if device == DeviceType.XPU and not ids_are_vulkan_ordinals:
raise HTTPException(
status_code = 400,
detail = (
@ -4668,24 +4725,31 @@ async def _resolve_gguf_gpu_ids_for_request(
),
)
if is_vulkan and _classify_diffusion_gguf(config) is True:
if (
device == DeviceType.CUDA
and not ids_are_vulkan_ordinals
and definitively_non_diffusion
and callable(lacks_gpu_lib)
and await asyncio.to_thread(lacks_gpu_lib)
):
raise HTTPException(
status_code = 400,
detail = (
"GPU selection (gpu_ids) is not supported for a DiffusionGemma "
"GGUF on a Vulkan llama.cpp build: the diffusion runner selects "
"its device by CUDA physical index, which has no defined mapping "
"to ggml Vulkan device ordinals. Omit gpu_ids to use the default "
"device."
f"Requested gpu_ids {list(gpu_ids)} but the llama.cpp build has "
"no GPU backend (CPU-only build); it would ignore the pin and run "
"on CPU. Omit gpu_ids to run on CPU."
),
)
try:
resolved = resolve_requested_gpu_ids(gpu_ids, is_vulkan = is_vulkan)
resolved = resolve_requested_gpu_ids(
gpu_ids,
is_vulkan = ids_are_vulkan_ordinals,
)
except ValueError as exc:
raise HTTPException(status_code = 400, detail = str(exc)) from exc
if is_vulkan and resolved:
if ids_are_vulkan_ordinals and resolved:
binary = LlamaCppBackend._find_llama_server_binary()
if binary:
probed = {
@ -4701,7 +4765,7 @@ async def _resolve_gguf_gpu_ids_for_request(
),
)
return resolved
return resolved, ids_are_vulkan_ordinals
def _guard_chat_load_against_training(
@ -4717,6 +4781,8 @@ def _guard_chat_load_against_training(
cache_type_kv: Optional[str] = None,
tensor_parallel: bool = False,
gpu_memory_mode: Literal["auto", "manual"] = "auto",
gpu_ids_are_vulkan_ordinals: bool = False,
diffusion_kind: Optional[bool] | object = _DIFFUSION_KIND_UNSET,
) -> None:
"""Protect active training from automatically placed chat-model loads.
@ -4740,29 +4806,31 @@ 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 diffusion_kind is _DIFFUSION_KIND_UNSET:
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
# Vulkan GGUF pins are ggml ordinals, not CUDA physical IDs. Detect this
# before deriving a possible diffusion fallback device so an unknown remote
# GGUF never sends its ordinal through the CUDA single-device path.
is_vulkan = False
if is_gguf:
try:
is_vulkan = LlamaCppBackend._is_vulkan_backend()
except Exception as e:
logger.warning("Could not detect Vulkan backend for chat-load guard: %s", e)
diffusion_gpu = None
if is_gguf and diffusion_kind is not False and not (is_vulkan and requested_gpu_ids):
if is_gguf and diffusion_kind is not False and not gpu_ids_are_vulkan_ordinals:
# Use the same token selection as the runner: an explicit pick wins,
# followed by DG_GPU, the first parent-visible token, then GPU 0.
# followed by DG_GPU, the first parent-visible token, then GPU 0. Suppressed
# for a Vulkan-ordinal pin so single-device CUDA budgeting can't override the
# Vulkan-ordinal path (single_device_gpu wins in can_load_chat_during_training).
diffusion_gpu = LlamaCppBackend._diffusion_gpu_arg(
requested_gpu_ids,
cpu_only = LlamaCppBackend._effective_gpu_count() == 0,
)
# Detected once: both the tensor-parallel KV sizing below and the Vulkan
# free-VRAM view need the same answer. An ordinal pin only exists on a
# Vulkan build, so it settles the question without probing the binary.
binary = LlamaCppBackend._find_llama_server_binary() if is_gguf else None
is_vulkan_backend = bool(
is_gguf
and (gpu_ids_are_vulkan_ordinals or (binary and LlamaCppBackend._is_vulkan_backend(binary)))
)
# Size with the count that will actually launch, or a load that fits gets a
# 409: diffusion never receives --parallel, and load_model clamps to 1 on an
# llama-server without --kv-unified. An unclassified GGUF keeps the ask.
@ -4787,13 +4855,34 @@ def _guard_chat_load_against_training(
cache_type_kv = cache_type_kv,
tensor_parallel = (
_effective_tensor_parallel(llama_extra_args, tensor_parallel)
and (is_vulkan or LlamaCppBackend._effective_gpu_count(requested_gpu_ids) >= 2)
and (
is_vulkan_backend
or LlamaCppBackend._effective_gpu_count(requested_gpu_ids) >= 2
)
),
)
if is_gguf
else None
)
vulkan_free_vram_gb = None
if is_gguf:
if is_vulkan_backend and (gpu_ids_are_vulkan_ordinals or diffusion_kind is False):
gpu_memory = LlamaCppBackend._get_gpu_memory(binary)
if not requested_gpu_ids:
gpu_memory = LlamaCppBackend._vulkan_auto_gpu_memory(gpu_memory)
vulkan_free_vram_gb = {
index: free_mib / 1024.0 for index, free_mib, _total_mib in gpu_memory
}
elif is_vulkan_backend and diffusion_kind is None and requested_gpu_ids:
# Until the header is available, the model may use either the Vulkan
# llama-server or the CUDA-only diffusion runner, so an explicit pin
# cannot be budgeted: neither device namespace can stand in for the
# other. Automatic placement has no ordinal to mis-map, so it keeps
# the torch view below rather than refusing every uncached remote
# GGUF while training runs.
vulkan_free_vram_gb = {}
ok, info = can_load_chat_during_training(
model_name = model_identifier,
hf_token = hf_token,
@ -4801,7 +4890,8 @@ def _guard_chat_load_against_training(
max_seq_length = max_seq_length,
requested_gpu_ids = requested_gpu_ids,
is_gguf = is_gguf,
is_vulkan = is_vulkan,
gpu_ids_are_vulkan_ordinals = gpu_ids_are_vulkan_ordinals,
vulkan_free_vram_gb = vulkan_free_vram_gb,
required_override_gb = required_override_gb,
single_device_gpu = diffusion_gpu,
)
@ -4846,14 +4936,15 @@ def _resolve_inherited_extra_args(
if not getattr(config, "is_gguf", False):
return extra_llama_args
llama_backend = get_llama_cpp_backend()
if not llama_backend.extra_args:
stored_args = getattr(llama_backend, "extra_args", None)
if not stored_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
source = getattr(llama_backend, "extra_args_source", None)
# 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
@ -4885,7 +4976,7 @@ def _resolve_inherited_extra_args(
# (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,
stored_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),
@ -4893,7 +4984,7 @@ def _resolve_inherited_extra_args(
"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),
strip_split_mode = _should_strip_split_mode(request, stored_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.
@ -5480,15 +5571,36 @@ async def _load_model_impl(
detail = f"Invalid model identifier: {model_log_label}",
)
# Resolve inherited extras once before command-dependent preflights.
extra_llama_args = _resolve_inherited_extra_args(
request,
config,
model_identifier,
extra_llama_args,
effective_chat_template_override,
)
# Normalize gpu_ids: empty list means auto-selection, same as None
effective_gpu_ids = request.gpu_ids if request.gpu_ids else None
gpu_ids_are_vulkan_ordinals = False
diffusion_kind = _classify_diffusion_gguf(config) if config.is_gguf else False
# Validate the full GGUF placement pool before the training guard so an
# invalid physical ID or Vulkan ordinal is a clean 400, not a masked VRAM
# 409. The same helper is used by /validate.
# Invalid GPU IDs must fail before the training coexistence guard.
gguf_gpu_ids: Optional[List[int]] = None
if config.is_gguf:
gguf_gpu_ids = await _resolve_gguf_gpu_ids_for_request(config, effective_gpu_ids)
(
gguf_gpu_ids,
gpu_ids_are_vulkan_ordinals,
) = await _resolve_gguf_gpu_ids_for_request(
config,
effective_gpu_ids,
diffusion_kind = diffusion_kind,
)
_reject_draft_device_with_gpu_ids(
gguf_gpu_ids,
extra_llama_args,
gpu_ids_are_vulkan_ordinals = gpu_ids_are_vulkan_ordinals,
)
if not config.is_gguf and _mlx_distributed_launch_detected():
raise HTTPException(
status_code = 400,
@ -5518,18 +5630,6 @@ async def _load_model_impl(
"architectures)"
)
# 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(
@ -5545,6 +5645,8 @@ async def _load_model_impl(
cache_type_kv = request.cache_type_kv,
tensor_parallel = bool(request.tensor_parallel),
gpu_memory_mode = request.gpu_memory_mode,
gpu_ids_are_vulkan_ordinals = gpu_ids_are_vulkan_ordinals,
diffusion_kind = diffusion_kind,
)
# ── GGUF path: load via llama-server ──────────────────────
@ -5633,6 +5735,7 @@ async def _load_model_impl(
gpu_layers = request.gpu_layers,
n_cpu_moe = request.n_cpu_moe,
tensor_split = request.tensor_split,
gpu_ids_are_vulkan_ordinals = gpu_ids_are_vulkan_ordinals,
n_parallel = _n_parallel,
# Issue #7164: explicit GPU pin resolved to physical ids above.
gpu_ids = gguf_gpu_ids,
@ -6128,11 +6231,29 @@ async def validate_model(
detail = f"Invalid model identifier: {model_log_label}",
)
effective_extra_args = _resolve_inherited_extra_args(
request, config, model_identifier, None
)
# 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
gpu_ids_are_vulkan_ordinals = False
diffusion_kind = _classify_diffusion_gguf(config) if config.is_gguf else False
if config.is_gguf:
await _resolve_gguf_gpu_ids_for_request(config, effective_gpu_ids)
(
validated_gpu_ids,
gpu_ids_are_vulkan_ordinals,
) = await _resolve_gguf_gpu_ids_for_request(
config,
effective_gpu_ids,
diffusion_kind = diffusion_kind,
)
_reject_draft_device_with_gpu_ids(
validated_gpu_ids,
effective_extra_args,
gpu_ids_are_vulkan_ordinals = gpu_ids_are_vulkan_ordinals,
)
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):
@ -6200,11 +6321,6 @@ async def validate_model(
# training guard must not refuse it. Real loads omit include_context_length /
# include_chat_template, and /load applies the guard again.
if not (request.include_context_length or request.include_chat_template):
# 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,
@ -6228,6 +6344,8 @@ async def validate_model(
cache_type_kv = request.cache_type_kv,
tensor_parallel = request.tensor_parallel,
gpu_memory_mode = request.gpu_memory_mode,
gpu_ids_are_vulkan_ordinals = gpu_ids_are_vulkan_ordinals,
diffusion_kind = diffusion_kind,
)
# A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a
@ -6301,6 +6419,7 @@ async def validate_model(
if native_grant_backed
else getattr(config, "display_name", config.identifier),
is_gguf = is_gguf,
is_diffusion = is_gguf and diffusion_kind is True,
is_lora = getattr(config, "is_lora", False),
is_vision = getattr(config, "is_vision", False),
requires_trust_remote_code = requires_trust_remote_code,

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

@ -225,7 +225,8 @@ def can_load_chat_during_training(
max_seq_length: int,
requested_gpu_ids: Optional[List[int]],
is_gguf: bool = False,
is_vulkan: bool = False,
gpu_ids_are_vulkan_ordinals: bool = False,
vulkan_free_vram_gb: Optional[Dict[int, float]] = None,
required_override_gb: Optional[float] = None,
single_device_gpu: Optional[str] = None,
) -> Tuple[bool, Dict[str, Any]]:
@ -234,15 +235,11 @@ def can_load_chat_during_training(
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. A Vulkan GGUF selection picks by ggml
Vulkan ordinal (separate index space from CUDA ids), so its requested_gpu_ids is
NOT resolved against the CUDA set (which would raise -> invalid_gpu_ids -> bypass
the OOM check); conservatively size an N-device request against the least-free
N visible GPUs instead.
``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). CPU/MLX allows the load; default-deny on any CUDA/XPU 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). CPU/MLX
allows the load; default-deny on any CUDA/XPU case it can't size, so a load
never OOMs training."""
try:
from utils.hardware import (
DeviceType,
@ -263,11 +260,6 @@ def can_load_chat_during_training(
max_seq_length = max_seq_length or 2048,
)
# A Vulkan GGUF selection uses ggml Vulkan ordinals, not CUDA physical ids;
# size it against the full visible pool (GGUF self-placement) rather than
# resolving ordinals against the CUDA parent-visible set.
vulkan_gguf = is_gguf and is_vulkan
# HF auto: reuse the loader's selector; fits iff its pick clears the margin.
if not requested_gpu_ids and not is_gguf:
_selected, meta = auto_select_gpu_ids(model_name, **est_kwargs)
@ -293,7 +285,8 @@ def can_load_chat_during_training(
}
# Explicit GPUs, or GGUF: size directly and check live free VRAM.
if requested_gpu_ids and vulkan_gguf:
uses_vulkan_memory = vulkan_free_vram_gb is not None
if is_gguf and uses_vulkan_memory:
mode = "gguf_vulkan"
elif single_device_gpu is not None:
mode = "single_device"
@ -307,17 +300,13 @@ def can_load_chat_during_training(
if required_gb is None:
return False, {"mode": mode, "reason": "estimate_unavailable"}
free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", []))
if requested_gpu_ids and vulkan_gguf:
# Vulkan ordinals cannot be mapped to CUDA physical indices. Budget
# the least-free N visible cards for an N-device request. If that
# conservative subset fits, any physical mapping of the ordinals
# fits, without collapsing a multi-GPU request to one card.
visible_free = list(free_by_index.values())
if not visible_free:
return False, {"mode": "gguf_vulkan", "reason": "no_visible_gpus"}
n_pins = min(len(requested_gpu_ids), len(visible_free))
free_vals = sorted(visible_free)[:n_pins]
free_by_index = (
vulkan_free_vram_gb
if uses_vulkan_memory
else _free_vram_by_index(get_visible_gpu_utilization().get("devices", []))
)
if requested_gpu_ids and gpu_ids_are_vulkan_ordinals:
free_vals = [free_by_index.get(int(gpu_id), 0.0) for gpu_id in requested_gpu_ids]
elif single_device_gpu is not None:
token = str(single_device_gpu).strip()
if not token:
@ -358,20 +347,27 @@ def can_load_chat_during_training(
needed_gb = required_gb * SAFETY_MARGIN + KEEP_FLOOR_GB
aggregate_fits = usable_gb >= needed_gb
# device_map="balanced" shards across GPUs: an even-share floor stops one
# near-full GPU hiding behind aggregate capacity. GGUF self-places, no floor.
# Explicit HF placement uses balanced sharding across a known number of
# GPUs. GGUF pins are candidate pools: llama.cpp may narrow an uneven
# pool to the smallest fitting subset, so only their aggregate matters.
min_free_gb = min(free_vals)
per_gpu_fits = True
per_gpu_needed_gb = None
if mode == "explicit" and len(free_vals) > 1:
per_gpu_fits = min_free_gb >= needed_gb / len(free_vals)
per_gpu_needed_gb = needed_gb / len(free_vals)
if per_gpu_needed_gb is not None:
per_gpu_fits = min_free_gb >= per_gpu_needed_gb
return aggregate_fits and per_gpu_fits, {
info = {
"mode": mode,
"required_gb": round(required_gb, 3),
"usable_gb": round(usable_gb, 3),
"needed_gb": round(needed_gb, 3),
"min_free_gb": round(min_free_gb, 3),
}
if per_gpu_needed_gb is not None:
info["per_gpu_needed_gb"] = round(per_gpu_needed_gb, 3)
return aggregate_fits and per_gpu_fits, info
except Exception as e:
# Never let a sizing failure load a chat model into a training OOM.
logger.warning("Chat-load coexistence probe failed; will refuse: %s", e)

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

@ -128,8 +128,7 @@ class TestCanLoadExplicitHF(_GpuCacheResetMixin, unittest.TestCase):
auto_mock.assert_not_called() # explicit never calls the auto selector
def test_per_gpu_floor_blocks_uneven_split(self):
# free [45, 10]; aggregate 45 + 10*0.85 = 53.5 >= needed 27, but the 10 GB
# GPU is below the even-share floor 27/2 = 13.5 -> refuse (would OOM it).
# Aggregate capacity passes, but the 10 GB shard fails its 13.5 GB floor.
ok, info, _ = self._run(
required = 20.0, devices = _devices((0, 80, 35), (1, 80, 70)), gpu_ids = [0, 1]
)
@ -170,7 +169,8 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
estimate = None,
single_device_gpu = None,
gpu_ids = None,
is_vulkan = False,
gpu_ids_are_vulkan_ordinals = False,
vulkan_free_vram_gb = None,
):
with (
patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
@ -186,7 +186,8 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
max_seq_length = 0,
requested_gpu_ids = gpu_ids,
is_gguf = True,
is_vulkan = is_vulkan,
gpu_ids_are_vulkan_ordinals = gpu_ids_are_vulkan_ordinals,
vulkan_free_vram_gb = vulkan_free_vram_gb,
required_override_gb = required_override,
single_device_gpu = single_device_gpu,
)
@ -205,10 +206,7 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
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.
# llama.cpp self-placement uses aggregate capacity within the pinned pool.
ok, info, _ = self._run(
devices = _devices((0, 80, 35), (1, 80, 70), (2, 80, 0)),
required_override = 20.0,
@ -237,39 +235,77 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
self.assertEqual(blocked_info["usable_gb"], 10.0)
def test_vulkan_pin_takes_precedence_over_unknown_diffusion_fallback(self):
# An uncached GGUF can carry a speculative single-device fallback while
# its explicit pin is actually a ggml Vulkan ordinal. Never interpret
# that ordinal as the same-numbered CUDA physical device.
# Use the Vulkan probe, not unrelated CUDA telemetry or a speculative
# CUDA diffusion device.
ok, info, _ = self._run(
devices = _devices((0, 80, 0), (1, 80, 78)),
devices = _devices((0, 80, 0)),
required_override = 20.0,
single_device_gpu = "0",
gpu_ids = [0],
is_vulkan = True,
gpu_ids_are_vulkan_ordinals = True,
vulkan_free_vram_gb = {0: 2.0, 1: 80.0},
)
self.assertFalse(ok)
self.assertEqual(info["mode"], "gguf_vulkan")
self.assertEqual(info["usable_gb"], 2.0)
def test_vulkan_multi_gpu_guard_counts_requested_devices(self):
# The ordinal mapping is unknown, so use the least-free two visible
# cards for a two-device request. Their aggregate capacity is still
# available instead of collapsing the request to one card.
# Vulkan ordinals select the same entries reported by the Vulkan probe.
ok, info, _ = self._run(
devices = _devices((0, 80, 70), (1, 80, 70), (2, 80, 0)),
devices = _devices((0, 80, 0)),
required_override = 10.0,
gpu_ids = [0, 1],
is_vulkan = True,
gpu_ids_are_vulkan_ordinals = True,
vulkan_free_vram_gb = {0: 10.0, 1: 10.0, 2: 80.0},
)
self.assertTrue(ok)
self.assertEqual(info["mode"], "gguf_vulkan")
self.assertEqual(info["usable_gb"], 18.5)
def test_vulkan_uneven_pool_uses_fitting_subset(self):
# The loader may choose the 20 GB device alone from this candidate pool.
ok, info, _ = self._run(
devices = _devices((0, 80, 0), (1, 80, 0)),
required_override = 10.0,
gpu_ids = [0, 1],
gpu_ids_are_vulkan_ordinals = True,
vulkan_free_vram_gb = {0: 20.0, 1: 2.0},
)
self.assertTrue(ok)
self.assertNotIn("per_gpu_needed_gb", info)
def test_vulkan_uneven_pool_is_order_independent(self):
ok, _, _ = self._run(
devices = _devices((0, 80, 0), (1, 80, 0)),
required_override = 10.0,
gpu_ids = [1, 0],
gpu_ids_are_vulkan_ordinals = True,
vulkan_free_vram_gb = {0: 20.0, 1: 2.0},
)
self.assertTrue(ok)
def test_vulkan_pool_with_insufficient_aggregate_refuses(self):
ok, _, _ = self._run(
devices = _devices((0, 80, 0), (1, 80, 0)),
required_override = 10.0,
gpu_ids = [0, 1],
gpu_ids_are_vulkan_ordinals = True,
vulkan_free_vram_gb = {0: 8.0, 1: 2.0},
)
self.assertFalse(ok)
def test_vulkan_auto_uses_full_vulkan_pool(self):
ok, info, _ = self._run(
devices = _devices((0, 80, 79)),
required_override = 20.0,
vulkan_free_vram_gb = {0: 30.0, 1: 30.0},
)
self.assertTrue(ok)
self.assertEqual(info["mode"], "gguf_vulkan")
self.assertEqual(info["usable_gb"], 55.5)
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.
# Unknown single-device tokens use the least-free visible GPU.
ok, info, _ = self._run(
devices = _devices((0, 80, 0)),
required_override = 20.0,
@ -291,10 +327,7 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
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.
# Aggregate capacity cannot justify an unresolved single-device load.
ok, info, _ = self._run(
devices = _devices((0, 80, 78), (1, 80, 0), (2, 80, 0)),
required_override = 20.0,
@ -304,9 +337,7 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
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.
# An empty device token is CPU-only and consumes no VRAM.
ok, info, _ = self._run(
devices = _devices((0, 80, 78)),
required_override = 20.0,
@ -451,6 +482,7 @@ class TestChatLoadGuardRoute(unittest.TestCase):
decision,
gpu_memory_mode = "auto",
requested_gpu_ids = None,
gpu_ids_are_vulkan_ordinals = False,
llama_extra_args = None,
cache_type_kv = None,
tensor_parallel = False,
@ -470,6 +502,7 @@ class TestChatLoadGuardRoute(unittest.TestCase):
cache_type_kv = cache_type_kv,
tensor_parallel = tensor_parallel,
gpu_memory_mode = gpu_memory_mode,
gpu_ids_are_vulkan_ordinals = gpu_ids_are_vulkan_ordinals,
)
def test_noop_when_training_inactive(self):
@ -574,6 +607,98 @@ class TestChatLoadGuardRoute(unittest.TestCase):
self.assertEqual(captured[0]["single_device_gpu"], "1")
self.assertEqual(captured[0]["requested_gpu_ids"], [3, 1])
def test_unclassified_gguf_on_vulkan_build_budgets_as_ordinals(self):
# Unknown GGUFs still use the Vulkan ordinal namespace selected by the build.
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,
"_get_gpu_memory",
return_value = [(0, 2048, 8192), (1, 4096, 8192)],
),
):
self._guard(
config = config,
captured = captured,
training_active = True,
decision = (True, {"mode": "gguf_vulkan"}),
requested_gpu_ids = [0, 1],
gpu_ids_are_vulkan_ordinals = True,
)
self.assertEqual(len(captured), 1)
self.assertTrue(captured[0]["gpu_ids_are_vulkan_ordinals"])
self.assertEqual(captured[0]["vulkan_free_vram_gb"], {0: 2.0, 1: 4.0})
self.assertIsNone(captured[0]["single_device_gpu"])
def test_auto_vulkan_uses_vulkan_probe_for_training_budget(self):
captured = []
config = SimpleNamespace(is_gguf = True)
with (
patch.object(self.route, "_classify_diffusion_gguf", return_value = False),
patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5),
patch.object(
self.route.LlamaCppBackend,
"_find_llama_server_binary",
return_value = "/tmp/llama-server",
),
patch.object(self.route.LlamaCppBackend, "_is_vulkan_backend", return_value = True),
patch.object(
self.route.LlamaCppBackend,
"_get_gpu_memory",
return_value = [(0, 3072, 0), (1, 2048, 8192)],
),
):
self._guard(
config = config,
captured = captured,
training_active = True,
decision = (True, {"mode": "gguf_vulkan"}),
)
self.assertEqual(captured[0]["vulkan_free_vram_gb"], {1: 2.0})
def test_unclassified_gguf_budget_depends_on_the_pin(self):
# An uncached remote GGUF that neither the header nor the name can classify
# may still turn out to be the CUDA-only diffusion runner, so the Vulkan
# free-VRAM map cannot stand in for it. The pin decides what that means:
cases = (
# No gpu_ids -> no ordinal to mis-map, so fall back to the torch view
# (None) rather than an empty map, which would read as "no free VRAM
# anywhere" and 409 every such load during training.
(None, "single_device", None),
# An explicit pin is the opposite case: the ordinal belongs to exactly
# one of the two device namespaces and neither can stand in for the
# other, so the guard must fail closed with an empty map.
([1], "gguf_vulkan", {}),
)
for requested_gpu_ids, mode, expected in cases:
with self.subTest(requested_gpu_ids = requested_gpu_ids):
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,
"_find_llama_server_binary",
return_value = "/tmp/llama-server",
),
patch.object(
self.route.LlamaCppBackend, "_is_vulkan_backend", return_value = True
),
):
self._guard(
config = config,
captured = captured,
training_active = True,
decision = (True, {"mode": mode}),
requested_gpu_ids = requested_gpu_ids,
)
self.assertEqual(len(captured), 1)
self.assertEqual(captured[0]["vulkan_free_vram_gb"], expected)
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:
@ -617,6 +742,11 @@ class TestChatLoadGuardRoute(unittest.TestCase):
"_effective_gpu_count",
return_value = 0,
),
patch.object(
self.route.LlamaCppBackend,
"_find_llama_server_binary",
return_value = "/fake/llama-server",
),
patch.object(self.route.LlamaCppBackend, "_is_vulkan_backend", return_value = True),
):
self._guard(
@ -771,10 +901,7 @@ class TestValidateRefusesDuringTraining(unittest.TestCase):
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.
# Validate and load must size the same inherited command.
from models.inference import ValidateModelRequest
request = ValidateModelRequest(
@ -815,9 +942,7 @@ class TestValidateRefusesDuringTraining(unittest.TestCase):
self.assertTrue(captured.get("tensor_parallel"))
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.
# Header-only probes allocate no VRAM.
from models.inference import ValidateModelRequest
request = ValidateModelRequest(
@ -852,6 +977,41 @@ class TestValidateRefusesDuringTraining(unittest.TestCase):
asyncio.run(self.route.validate_model(request, current_subject = "u"))
self.assertEqual(guard_called, [])
def test_metadata_probe_reports_diffusion(self):
from models.inference import ValidateModelRequest
request = ValidateModelRequest(
model_path = "unsloth/DiffusionGemma-GGUF",
include_context_length = True,
)
cfg = SimpleNamespace(
identifier = "unsloth/DiffusionGemma-GGUF",
display_name = "DiffusionGemma-GGUF",
is_gguf = True,
is_lora = False,
is_vision = False,
gguf_file = None,
gguf_hf_repo = "unsloth/DiffusionGemma-GGUF",
gguf_variant = None,
path = None,
base_model = None,
)
with (
patch.object(
self.route,
"_resolve_model_identifier_for_request",
return_value = (
"unsloth/DiffusionGemma-GGUF",
"unsloth/DiffusionGemma-GGUF",
False,
),
),
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
patch.object(self.route, "load_inference_config", return_value = {}),
):
response = asyncio.run(self.route.validate_model(request, current_subject = "u"))
self.assertTrue(response.is_diffusion)
def _validate_gguf_template(
self,
*,
@ -1131,6 +1291,24 @@ class TestLoadModelGuardIntegration(unittest.TestCase):
def setUpClass(cls):
cls.route = _load_inference_route()
def test_cuda_gpu_ids_allow_matching_draft_device(self):
self.route._reject_draft_device_with_gpu_ids(
[1],
["--spec-draft-device", "CUDA0"],
gpu_ids_are_vulkan_ordinals = False,
)
def test_vulkan_gpu_ids_reject_underscore_draft_device_alias(self):
with self.assertRaises(HTTPException) as exc:
self.route._reject_draft_device_with_gpu_ids(
[0],
["--spec_draft_device", "Vulkan1"],
gpu_ids_are_vulkan_ordinals = True,
)
self.assertEqual(exc.exception.status_code, 400)
self.assertIn("Vulkan1", exc.exception.detail)
def test_refusal_409_and_no_unload(self):
import contextlib
from unittest.mock import MagicMock
@ -1178,6 +1356,73 @@ class TestLoadModelGuardIntegration(unittest.TestCase):
inf._shutdown_subprocess.assert_not_called()
llama.unload_model.assert_not_called()
def test_gguf_inherited_draft_device_rejected_under_vulkan_gpu_ids(self):
# Check effective inherited extras, not only the raw request.
import contextlib
from unittest.mock import MagicMock
from models.inference import LoadRequest
inf = SimpleNamespace(active_model_name = None)
inf.unload_model = MagicMock()
inf._shutdown_subprocess = MagicMock()
llama = SimpleNamespace(
is_loaded = True,
model_identifier = "x.gguf",
hf_variant = None,
extra_args = ["--spec-draft-device", "CUDA1"],
extra_args_source = ("x.gguf", None),
is_vulkan_build = lambda: True,
)
llama.unload_model = MagicMock()
cfg = SimpleNamespace(
is_gguf = True,
is_lora = False,
is_vision = False,
path = None,
base_model = None,
identifier = "x.gguf",
display_name = "x",
gguf_variant = None,
)
request = LoadRequest(model_path = "x.gguf", gpu_ids = [0], max_seq_length = 4096)
captured = []
with (
patch("utils.transformers_version.latest_tier_active_for", return_value = False),
patch.object(
self.route,
"validate_extra_args",
side_effect = lambda args: list(args) if args else None,
),
patch.object(
self.route,
"_resolve_model_identifier_for_request",
return_value = ("x.gguf", "x.gguf", False),
),
patch.object(self.route, "resolve_effective_chat_template_override", return_value = None),
patch.object(self.route, "_request_matches_loaded_settings", return_value = False),
patch.object(
self.route,
"_resolve_gguf_gpu_ids_for_request",
return_value = ([0], True),
),
patch.object(self.route, "get_inference_backend", return_value = inf),
patch.object(self.route, "get_llama_cpp_backend", return_value = llama),
patch.object(self.route, "_hf_offline_if_dns_dead", lambda: contextlib.nullcontext()),
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
_stub_guard_deps(training_active = True, decision = (True, {}), captured = captured),
):
with self.assertRaises(HTTPException) as exc:
asyncio.run(
self.route.load_model(request, fastapi_request = MagicMock(), current_subject = "u")
)
self.assertEqual(exc.exception.status_code, 400)
self.assertIn("draft-model device", exc.exception.detail)
self.assertIn("set it to none", exc.exception.detail)
self.assertEqual(captured, [])
inf.unload_model.assert_not_called()
llama.unload_model.assert_not_called()
if __name__ == "__main__":
unittest.main()

3
studio/backend/tests/test_gguf_load_cache_reuse.py
View file

@ -142,7 +142,8 @@ async def _inline_to_thread(func, /, *args, **kwargs):
async def _no_gguf_gpu_ids(*_args, **_kwargs):
return None
# Mirrors the resolver's no-gpu_ids early return: no ids, not Vulkan ordinals.
return None, False
class TestLoadReusesCachedCopy:

4
studio/backend/tests/test_gpu_memory_mode.py
View file

@ -619,11 +619,15 @@ def test_gguf_load_and_status_responses_include_requested_gpu_pool():
def test_gpu_ids_property_default_and_reset():
backend = LlamaCppBackend()
assert backend.gpu_ids is None
assert backend.requested_gpu_ids is None
backend._gpu_ids = [0, 1]
backend._requested_gpu_ids = [0, 1, 2]
assert backend.gpu_ids == [0, 1]
assert backend.requested_gpu_ids == [0, 1, 2]
backend._process = _FakeProcess()
backend.unload_model()
assert backend.gpu_ids is None
assert backend.requested_gpu_ids is None
def _target_state_gpu_ids(backend, gpu_ids):

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

@ -419,8 +419,16 @@ class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase):
return_value = True,
),
patch(
"core.inference.llama_cpp.LlamaCppBackend._get_gpu_memory",
return_value = [(0, 7402, 8192)],
"core.inference.llama_cpp.LlamaCppBackend.vulkan_device_inventory",
return_value = [
{
"index": 0,
"name": "Vulkan0",
"free_mib": 7402,
"total_mib": 8192,
"is_igpu": False,
}
],
),
):
result = get_vulkan_inference_gpu_info()
@ -455,8 +463,20 @@ class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase):
return_value = True,
),
patch(
"core.inference.llama_cpp.LlamaCppBackend._get_gpu_memory",
return_value = [(0, 12288, 0)],
"core.inference.llama_cpp.LlamaCppBackend.vulkan_device_inventory",
return_value = [
{
"index": 0,
"name": "Vulkan0",
"free_mib": 12288,
"total_mib": 32768,
"is_igpu": True,
}
],
),
patch(
"core.inference.llama_cpp._apply_igpu_host_reserve_mib",
return_value = 12288,
),
):
result = get_vulkan_inference_gpu_info()
@ -476,8 +496,16 @@ class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase):
return_value = True,
),
patch(
"core.inference.llama_cpp.LlamaCppBackend._get_gpu_memory",
return_value = [(1, 6144, 8192)],
"core.inference.llama_cpp.LlamaCppBackend.vulkan_device_inventory",
return_value = [
{
"index": 1,
"name": "Vulkan1",
"free_mib": 6144,
"total_mib": 8192,
"is_igpu": False,
}
],
),
patch(
"utils.hardware.nvidia.get_backend_visible_gpu_info",
@ -506,7 +534,7 @@ class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase):
return_value = True,
),
patch(
"core.inference.llama_cpp.LlamaCppBackend._get_gpu_memory",
"core.inference.llama_cpp.LlamaCppBackend.vulkan_device_inventory",
return_value = [],
),
):
@ -980,10 +1008,7 @@ class TestRouteErrors(unittest.TestCase):
self.assertIn("only supported on CUDA and Intel XPU", str(exc_info.exception))
def test_inference_route_resolves_gguf_gpu_ids(self):
# GGUF gpu_ids are now supported: /load routes them through the same
# resolution as non-GGUF loads (rejecting only genuinely invalid ids with
# the resolver's actionable message) rather than a blanket "not supported"
# reject, so /validate can stay consistent with /load (#7239).
# GGUF IDs use the normal resolver instead of a blanket rejection.
import utils.hardware.hardware as hardware_mod
inference_route = _load_route_module(
@ -1015,8 +1040,6 @@ class TestRouteErrors(unittest.TestCase):
"ModelConfig",
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
),
# Patch both the package re-export and the defining module so the stub
# fires no matter which import path the route uses.
patch("utils.hardware.resolve_requested_gpu_ids", _fake_resolve),
patch.object(hardware_mod, "resolve_requested_gpu_ids", _fake_resolve),
patch.object(
@ -1040,7 +1063,6 @@ class TestRouteErrors(unittest.TestCase):
)
)
# The selection was routed through resolution (not the old blanket reject).
self.assertEqual(exc_info.exception.status_code, 400)
self.assertIn("SENTINEL", exc_info.exception.detail)
self.assertNotIn("not supported for GGUF", exc_info.exception.detail)
@ -1138,18 +1160,57 @@ class TestRouteErrors(unittest.TestCase):
return_value = None,
),
):
resolved = asyncio.run(
resolved, uses_vulkan_ordinals = asyncio.run(
inference_route._resolve_gguf_gpu_ids_for_request(config, [1, 0])
)
self.assertEqual(resolved, [0, 1])
self.assertTrue(uses_vulkan_ordinals)
def test_diffusion_gpu_ids_accept_rocm_physical_path(self):
import utils.hardware as hardware_pkg
import utils.hardware.hardware as hardware_mod
inference_route = _load_route_module(
"inference_route_module_for_diffusion_rocm_path_test",
"routes/inference.py",
)
config = SimpleNamespace(is_gguf = True)
fake_backend = SimpleNamespace(
is_vulkan_build = lambda: False,
_backend_lacks_gpu_lib = lambda *a, **k: False,
)
with (
patch.object(hardware_mod, "IS_ROCM", True),
patch.object(hardware_pkg, "get_device", return_value = DeviceType.CUDA),
patch.object(
hardware_mod,
"resolve_requested_gpu_ids",
return_value = [0, 1],
),
patch.object(
inference_route,
"get_llama_cpp_backend",
return_value = fake_backend,
),
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
):
self.assertEqual(hardware_mod._backend_label(DeviceType.CUDA), "rocm")
resolved, uses_vulkan_ordinals = asyncio.run(
inference_route._resolve_gguf_gpu_ids_for_request(
config,
[1, 0],
diffusion_kind = True,
)
)
self.assertEqual(resolved, [0, 1])
self.assertFalse(uses_vulkan_ordinals)
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
import utils.hardware as hardware_pkg
inference_route = _load_route_module(
"inference_route_module_for_gguf_gpu_ids_test2",
@ -1177,8 +1238,6 @@ class TestRouteErrors(unittest.TestCase):
"ModelConfig",
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
),
# Patch both the package re-export and the defining module so the stub
# fires no matter which import path the route uses.
patch(
"utils.hardware.resolve_requested_gpu_ids",
side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"),
@ -1209,11 +1268,124 @@ 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("gpu_ids", exc_info.exception.detail.lower())
self.assertNotIn("not supported", exc_info.exception.detail.lower())
def test_inference_route_rejects_gpu_ids_on_cpu_only_llama_build(self):
# A CPU-only llama.cpp build cannot honor a CUDA visibility pin.
import utils.hardware as hardware_pkg
inference_route = _load_route_module(
"inference_route_module_for_cpu_only_gpu_ids_test",
"routes/inference.py",
)
request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [0, 1])
model_config = SimpleNamespace(
is_gguf = True,
is_lora = False,
gguf_hf_repo = None,
gguf_file = "/tmp/test.gguf",
gguf_mmproj_file = None,
gguf_variant = None,
identifier = "unsloth/test.gguf",
display_name = "unsloth/test.gguf",
is_vision = False,
is_audio = False,
audio_type = None,
has_audio_input = False,
)
fake_backend = SimpleNamespace(
is_loaded = False,
model_identifier = None,
is_vulkan_build = lambda: False,
_backend_lacks_gpu_lib = lambda *a, **k: True,
)
with (
patch.object(
inference_route,
"ModelConfig",
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
),
patch.object(inference_route, "_classify_diffusion_gguf", return_value = False),
patch.object(inference_route, "_guard_chat_load_against_training", return_value = None),
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
patch.object(inference_route, "get_llama_cpp_backend", return_value = fake_backend),
patch.object(hardware_pkg, "get_device", return_value = hardware_pkg.DeviceType.CUDA),
):
with self.assertRaises(HTTPException) as exc_info:
asyncio.run(
inference_route._load_model_impl(
request,
SimpleNamespace(
app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)),
),
current_subject = "test-user",
)
)
self.assertEqual(exc_info.exception.status_code, 400)
self.assertIn("cpu-only build", exc_info.exception.detail.lower())
def test_diffusion_gguf_on_vulkan_build_rejects_ordinal_pin(self):
# The GGUF picker supplies Vulkan ordinals, which cannot be reinterpreted
# as the CUDA physical IDs used by the diffusion runner.
import utils.hardware as hardware_pkg
inference_route = _load_route_module(
"inference_route_module_for_diffusion_cuda_path_test",
"routes/inference.py",
)
request = LoadRequest(model_path = "unsloth/diffusion.gguf", gpu_ids = [0, 1])
model_config = SimpleNamespace(
is_gguf = True,
is_lora = False,
gguf_hf_repo = None,
gguf_file = "/tmp/diffusion.gguf",
gguf_mmproj_file = None,
gguf_variant = None,
identifier = "unsloth/diffusion.gguf",
display_name = "unsloth/diffusion.gguf",
is_vision = False,
is_audio = False,
audio_type = None,
has_audio_input = False,
)
fake_backend = SimpleNamespace(
is_loaded = False,
model_identifier = None,
is_vulkan_build = lambda: True,
_backend_lacks_gpu_lib = lambda *a, **k: False,
)
with (
patch.object(
inference_route,
"ModelConfig",
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
),
patch.object(inference_route, "_classify_diffusion_gguf", return_value = True),
patch.object(
_hw_module,
"resolve_requested_gpu_ids",
side_effect = AssertionError("Vulkan ordinal reached the CUDA resolver"),
),
patch.object(inference_route, "_guard_chat_load_against_training", return_value = None),
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
patch.object(inference_route, "get_llama_cpp_backend", return_value = fake_backend),
patch.object(hardware_pkg, "get_device", return_value = hardware_pkg.DeviceType.CUDA),
):
with self.assertRaises(HTTPException) as exc_info:
asyncio.run(
inference_route._load_model_impl(
request,
SimpleNamespace(
app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)),
),
current_subject = "test-user",
)
)
self.assertEqual(exc_info.exception.status_code, 400)
self.assertIn("no defined mapping", exc_info.exception.detail)
def test_training_route_returns_400_for_invalid_gpu_ids(self):
training_route = _load_route_module(

499
studio/backend/tests/test_install_resolve_prebuilt.py
View file

@ -110,11 +110,18 @@ def test_macos_upstream_pin_only_for_explicit_pre26_upstream():
assert ilp.pinned_macos_release_tag(tahoe, UPSTREAM) is None
def _run_resolve(monkeypatch, capsys, plans_or_exc):
def _run_resolve(
monkeypatch,
capsys,
plans_or_exc,
*,
host = None,
extra_args = (),
):
monkeypatch.setattr(
ilp,
"detect_host",
lambda: _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64"),
lambda: host or _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64"),
)
def _resolver(tag, host, repo, published_release_tag):
@ -126,7 +133,14 @@ def _run_resolve(monkeypatch, capsys, plans_or_exc):
monkeypatch.setattr(
sys,
"argv",
["install_llama_prebuilt.py", "--resolve-prebuilt", "latest", "--output-format", "json"],
[
"install_llama_prebuilt.py",
"--resolve-prebuilt",
"latest",
"--output-format",
"json",
*extra_args,
],
)
rc = ilp.main()
assert rc == ilp.EXIT_SUCCESS
@ -155,6 +169,33 @@ def test_resolve_prebuilt_unavailable(monkeypatch, capsys):
assert out["repo"] == FORK
def test_resolve_prebuilt_backend_flag_filters_cpu_fallback(monkeypatch, capsys):
cpu_plan = ilp.InstallReleasePlan(
requested_tag = "latest",
llama_tag = "b9585",
release_tag = "release",
attempts = [
ilp.AssetChoice(
repo = FORK,
tag = "release",
name = "app-release-linux-x64-cpu",
url = "https://example/cpu",
source_label = "published",
install_kind = "linux-cpu",
)
],
approved_checksums = SimpleNamespace(),
)
out = _run_resolve(
monkeypatch,
capsys,
[cpu_plan],
host = _host(is_linux = True, is_x86_64 = True),
extra_args = ("--llama-backend", "vulkan"),
)
assert out["prebuilt_available"] is False
def _run_resolve_capture_host(monkeypatch, capsys):
"""Drive --resolve-prebuilt and return the host the resolver was handed."""
seen = {}
@ -401,6 +442,99 @@ def test_direct_upstream_x86_intel_prefers_vulkan():
assert "linux-cpu" in kinds
def _published_vulkan_bundle(*install_kinds):
profiles = {
"linux-vulkan": "linux-vulkan-x64",
"windows-vulkan": "windows-vulkan-x64",
"linux-cpu": "linux-cpu-x64",
"linux-arm64": "linux-cpu-arm64",
"windows-cpu": "windows-cpu-x64",
}
artifacts = [
ilp.PublishedLlamaArtifact(
asset_name = f"app-release-{install_kind}",
install_kind = install_kind,
runtime_line = None,
coverage_class = None,
supported_sms = [],
min_sm = None,
max_sm = None,
bundle_profile = profiles.get(install_kind),
rank = 60,
)
for install_kind in install_kinds
]
return ilp.PublishedReleaseBundle(
repo = FORK,
release_tag = "release",
upstream_tag = "b9925",
assets = {
artifact.asset_name: f"https://example/{artifact.asset_name}" for artifact in artifacts
},
artifacts = artifacts,
)
def test_fork_linux_intel_prefers_published_vulkan_bundle():
bundle = _published_vulkan_bundle("linux-vulkan", "linux-cpu")
attempts = ilp._linux_published_attempts(
_host(is_linux = True, is_x86_64 = True, has_intel_gpu = True),
bundle,
)
assert [attempt.install_kind for attempt in attempts] == ["linux-vulkan", "linux-cpu"]
assert attempts[0].source_label == "published"
assert ["llama-diffusion-gemma-visual-server"] in ilp.runtime_payload_health_groups(attempts[0])
def test_fork_linux_arm64_does_not_select_x64_vulkan_bundle():
bundle = _published_vulkan_bundle("linux-vulkan", "linux-arm64")
attempts = ilp._linux_published_attempts(
_host(
machine = "aarch64",
is_linux = True,
is_arm64 = True,
has_intel_gpu = True,
),
bundle,
)
assert [attempt.install_kind for attempt in attempts] == ["linux-arm64"]
def test_fork_windows_intel_prefers_published_vulkan_bundle():
bundle = _published_vulkan_bundle("windows-vulkan", "windows-cpu")
checksums = ilp.ApprovedReleaseChecksums(
repo = FORK,
release_tag = bundle.release_tag,
upstream_tag = bundle.upstream_tag,
artifacts = {
artifact.asset_name: ilp.ApprovedArtifactHash(
asset_name = artifact.asset_name,
sha256 = "a" * 64,
repo = FORK,
kind = artifact.install_kind,
)
for artifact in bundle.artifacts
},
)
attempts = ilp.resolve_release_asset_choice(
_host(
system = "Windows",
is_windows = True,
is_x86_64 = True,
has_intel_gpu = True,
),
bundle.upstream_tag,
bundle,
checksums,
)
assert [attempt.install_kind for attempt in attempts] == ["windows-vulkan", "windows-cpu"]
assert attempts[0].source_label == "published"
assert ["llama-diffusion-gemma-visual-server.exe"] in ilp.runtime_payload_health_groups(
attempts[0]
)
def test_linux_vulkan_health_glob_matches_bare_cpu_lib():
# The widened glob must cover both arch-suffixed (x64) and bare (arm64) CPU
# libs so a valid Vulkan install is not re-flagged unhealthy every check.
@ -417,15 +551,98 @@ def test_linux_vulkan_health_glob_matches_bare_cpu_lib():
assert ["libggml-cpu-*.so*"] not in groups
def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin():
# Routing fork -> upstream also drops the fork release pin, which is in a
# different tag namespace and would make the upstream resolver miss.
@pytest.mark.parametrize(
"backend, legacy, expected_backend, expected",
[
(None, None, None, False),
("vulkan", None, "vulkan", True),
(" VULKAN ", None, "vulkan", True),
("auto", "1", None, True),
(None, "true", None, True),
(None, "on", None, True),
("auto", None, None, False),
("cpu", "0", "cpu", False),
("hip", "1", "hip", False),
("rocm", "on", "hip", False),
(" HIP ", "1", "hip", False),
(" ROCM ", "1", "hip", False),
],
)
def test_force_vulkan_requested_accepts_public_selector_and_legacy_alias(
monkeypatch, backend, legacy, expected_backend, expected
):
# UNSLOTH_LLAMA_CPP_BACKEND is the public selector; a recognized non-vulkan
# value is authoritative, so it opts out even against a stale legacy alias.
# auto/unset/unknown leave the backend unpinned so selection stays automatic.
for name, value in (
("UNSLOTH_LLAMA_CPP_BACKEND", backend),
("UNSLOTH_FORCE_VULKAN", legacy),
):
if value is None:
monkeypatch.delenv(name, raising = False)
else:
monkeypatch.setenv(name, value)
assert ilp.llama_backend_from_env() == expected_backend
assert ilp.force_vulkan_requested() is expected
def test_forced_vulkan_filters_cpu_fallback_before_validation():
vulkan = ilp.AssetChoice(
repo = UPSTREAM,
tag = "b9925",
name = "llama-b9925-bin-ubuntu-vulkan-x64.tar.gz",
url = "https://example/vulkan",
source_label = "upstream",
install_kind = "linux-vulkan",
)
cpu = ilp.AssetChoice(
repo = UPSTREAM,
tag = "b9925",
name = "llama-b9925-bin-ubuntu-x64.tar.gz",
url = "https://example/cpu",
source_label = "upstream",
install_kind = "linux-cpu",
)
assert ilp._vulkan_only_attempts([vulkan, cpu]) == [vulkan]
def test_forced_vulkan_skips_cpu_only_release_plans():
vulkan = ilp.AssetChoice(
repo = UPSTREAM,
tag = "b9924",
name = "llama-b9924-bin-ubuntu-vulkan-x64.tar.gz",
url = "https://example/vulkan",
source_label = "upstream",
install_kind = "linux-vulkan",
)
cpu = ilp.AssetChoice(
repo = UPSTREAM,
tag = "b9925",
name = "llama-b9925-bin-ubuntu-x64.tar.gz",
url = "https://example/cpu",
source_label = "upstream",
install_kind = "linux-cpu",
)
plans = [
ilp.InstallReleasePlan("latest", "b9925", "b9925", [cpu], SimpleNamespace()),
ilp.InstallReleasePlan("latest", "b9924", "b9924", [vulkan], SimpleNamespace()),
]
filtered = ilp._vulkan_only_release_plans(plans)
assert [plan.release_tag for plan in filtered] == ["b9924"]
assert filtered[0].attempts == [vulkan]
def test_route_to_vulkan_prebuilt_auto_intel_keeps_fork_release():
# The fork now publishes a verified Vulkan app bundle that includes the
# DiffusionGemma runner, so keep its release namespace and pin.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
routed, repo, tag, _persist = ilp._route_to_vulkan_prebuilt(
host, FORK, "b9596-mix-abc", force_cpu = False
)
assert repo == UPSTREAM
assert tag == ""
assert repo == FORK
assert tag == "b9596-mix-abc"
assert routed.has_intel_gpu is True
@ -439,6 +656,94 @@ def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin():
assert tag == "b9596"
def _linux_arm64_vulkan_host():
return _host(
system = "Linux",
machine = "aarch64",
is_linux = True,
is_arm64 = True,
has_intel_gpu = True,
)
def test_route_to_vulkan_prebuilt_linux_arm64_falls_back_to_upstream():
# The fork ships no ARM64 Vulkan bundle, so a forced-Vulkan Linux ARM64 host
# must keep planning against upstream (which publishes
# llama-<tag>-bin-ubuntu-vulkan-arm64.tar.gz). The fork pin is dropped because
# the two repos use different tag namespaces.
routed, repo, tag, persist = ilp._route_to_vulkan_prebuilt(
_linux_arm64_vulkan_host(),
FORK,
"b9596-mix-abc",
force_cpu = False,
llama_backend = "vulkan",
)
assert repo == UPSTREAM
assert tag == ""
assert persist == "vulkan"
assert routed.has_intel_gpu is True
def test_forced_vulkan_linux_arm64_still_resolves_a_vulkan_bundle():
# End to end for the routing above: without it the fork planner yields only the
# ARM64 CPU attempt, and the strict-Vulkan filter then leaves nothing to install.
routed, repo, _tag, _persist = ilp._route_to_vulkan_prebuilt(
_linux_arm64_vulkan_host(),
FORK,
"",
force_cpu = False,
llama_backend = "vulkan",
)
fork_attempts = ilp._linux_published_attempts(
routed, _published_vulkan_bundle("linux-vulkan", "linux-arm64")
)
assert [attempt.install_kind for attempt in fork_attempts] == ["linux-arm64"]
release = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(release, routed, repo, "latest")
filtered = ilp._vulkan_only_release_plans([plan])
assert [attempt.name for attempt in filtered[0].attempts] == [
"llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz"
]
assert filtered[0].attempts[0].repo == UPSTREAM
def test_route_to_vulkan_prebuilt_linux_arm64_keeps_an_explicit_repo_override():
# A hand-picked --published-repo is the user's call; only the default fork
# repo is rerouted.
other = "someone/llama.cpp"
_routed, repo, tag, _persist = ilp._route_to_vulkan_prebuilt(
_linux_arm64_vulkan_host(),
other,
"b9596",
force_cpu = False,
llama_backend = "vulkan",
)
assert repo == other
assert tag == "b9596"
def test_route_to_vulkan_prebuilt_linux_x86_64_keeps_the_fork_bundle():
# Only ARM64 lacks a fork Vulkan bundle; x86_64 must stay on the fork release
# so it keeps the DiffusionGemma visual server.
_routed, repo, tag, _persist = ilp._route_to_vulkan_prebuilt(
_host(is_linux = True, is_x86_64 = True, has_intel_gpu = True),
FORK,
"b9596-mix-abc",
force_cpu = False,
llama_backend = "vulkan",
)
assert repo == FORK
assert tag == "b9596-mix-abc"
def test_route_to_vulkan_prebuilt_cpu_fallback_wins():
# --cpu-fallback suppresses Vulkan routing even for an Intel host.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
@ -450,6 +755,33 @@ def test_route_to_vulkan_prebuilt_cpu_fallback_wins():
assert routed is host
@pytest.mark.parametrize("via", ["env", "argument"])
@pytest.mark.parametrize("backend", ["hip", "rocm", "cpu"])
def test_non_vulkan_backend_suppresses_intel_auto_route(monkeypatch, backend, via):
# The selector suppresses the Intel auto-route the same way whether it arrives
# from the environment or from --llama-backend. Pinned to Linux ARM64, the one
# host where the suppression shows up in a returned VALUE: without it the Intel
# auto-route fires, falls through to the ARM64 fork -> upstream reroute and
# hands back UPSTREAM with the fork pin dropped. On x86_64 the auto-route
# rewrites neither host nor repo, so only a log line differs and the same case
# there cannot fail.
kwargs = {}
if via == "env":
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", backend)
else:
kwargs["llama_backend"] = backend
host = _linux_arm64_vulkan_host()
routed, repo, tag, persist = ilp._route_to_vulkan_prebuilt(
host, FORK, "b9596-mix-abc", force_cpu = False, **kwargs
)
assert routed is host
assert repo == FORK, repo
assert tag == "b9596-mix-abc", tag
assert persist is None
@pytest.mark.parametrize("cpu_flag", ["--cpu-fallback", "--force-cpu"])
def test_resolve_prebuilt_cpu_fallback_overrides_intel_vulkan(monkeypatch, capsys, cpu_flag):
"""Either CPU flag via CLI must suppress Vulkan even on an Intel GPU host: both
@ -485,6 +817,47 @@ def test_resolve_prebuilt_cpu_fallback_overrides_intel_vulkan(monkeypatch, capsy
assert seen["repo"] == FORK
def test_cpu_backend_env_forces_cpu_in_resolver(monkeypatch, capsys):
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "cpu")
monkeypatch.setattr(
ilp,
"detect_host",
lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True),
)
seen, _ = _run_resolve_capture_host(monkeypatch, capsys)
assert seen["host"].has_intel_gpu is False
assert seen["repo"] == FORK
def test_cpu_backend_env_forces_cpu_in_direct_install(monkeypatch, tmp_path):
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "cpu")
monkeypatch.setattr(
ilp,
"detect_host",
lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True),
)
seen = {}
def _route(
host,
repo,
tag,
*,
force_cpu,
llama_backend = None,
):
seen["host"] = host
seen["force_cpu"] = force_cpu
raise RuntimeError("stop after routing")
monkeypatch.setattr(ilp, "_route_to_vulkan_prebuilt", _route)
with pytest.raises(RuntimeError, match = "stop after routing"):
ilp.install_prebuilt(tmp_path, "latest", FORK, "")
assert seen["host"].has_intel_gpu is False
assert seen["force_cpu"] is True
@pytest.mark.parametrize(
"flags, expect_force, expect_persist",
[
@ -574,15 +947,15 @@ def test_route_to_vulkan_prebuilt_non_intel_unchanged():
assert routed is host
def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys):
def test_resolve_prebuilt_intel_host_keeps_fork_manifest(monkeypatch, capsys):
# The --resolve-prebuilt probe must agree with the install path: an
# auto-detected Intel host resolves against upstream (Vulkan), not the fork.
# auto-detected Intel host resolves the fork's Vulkan app bundle.
monkeypatch.setattr(
ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
)
seen, out = _run_resolve_capture_host(monkeypatch, capsys)
assert seen["repo"] == UPSTREAM
assert out["repo"] == UPSTREAM
assert seen["repo"] == FORK
assert out["repo"] == FORK
# ---------------------------------------------------------------------------
@ -844,7 +1217,7 @@ def _windows_amd_host(**overrides):
def test_route_to_vulkan_prebuilt_auto_fallback_for_legacy_amd_gfx():
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx803"])
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert repo == FORK
assert persist == "vulkan"
assert routed.has_intel_gpu is True
assert routed.has_rocm is False
@ -883,7 +1256,7 @@ def test_route_to_vulkan_prebuilt_auto_fallback_when_no_amd_gpu_reaches_floor():
rocm_gfx_targets = ["gfx803", "gfx900"],
)
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert repo == FORK
assert persist == "vulkan"
assert routed.has_rocm is False
@ -951,7 +1324,7 @@ def test_masked_probe_suppression_does_not_touch_non_amd_auto_paths(monkeypatch)
_routed, repo, _tag, _persist = ilp._route_to_vulkan_prebuilt(
host, FORK, "pin", force_cpu = False
)
assert repo == UPSTREAM
assert repo == FORK
def test_route_to_vulkan_prebuilt_hip_masked_host_still_honours_explicit_optin(monkeypatch):
@ -965,7 +1338,7 @@ def test_route_to_vulkan_prebuilt_hip_masked_host_still_honours_explicit_optin(m
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(
host, FORK, "pin", force_cpu = False, llama_backend = "vulkan"
)
assert repo == UPSTREAM
assert repo == FORK
assert persist == "vulkan"
@ -1060,13 +1433,13 @@ def test_route_to_vulkan_prebuilt_gfx1034_keeps_rocm():
def test_route_to_vulkan_prebuilt_explicit_opt_in_on_mixed_amd(monkeypatch):
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "vulkan")
host = _windows_amd_host(
rocm_gfx_target = "gfx1201",
rocm_gfx_targets = ["gfx1201", "gfx803"],
)
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert repo == FORK
assert persist == "vulkan"
assert routed.has_rocm is False
@ -1087,21 +1460,16 @@ def test_direct_upstream_windows_amd_legacy_gfx_routes_to_vulkan():
assert plan.attempts[0].install_kind == "windows-vulkan"
def test_llama_backend_env_requests_vulkan(monkeypatch):
assert ilp.llama_backend_from_env() is None
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
assert ilp.llama_backend_from_env() == "vulkan"
assert ilp.force_vulkan_requested() is True
def test_llama_cpp_backend_env_does_not_trigger_vulkan(monkeypatch):
# UNSLOTH_LLAMA_CPP_BACKEND is a separate setup variable (auto/cpu) whose other values
# setup warns about and ignores, so reading it here would opt in behind that warning.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
def test_hip_backend_env_suppresses_auto_vulkan_fallback_on_unsupported_gfx(monkeypatch):
# An explicit HIP choice keeps HIP on an unsupported gfx target instead of
# silently substituting Vulkan.
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "vulkan")
assert ilp.llama_backend_from_env() is None
assert ilp.force_vulkan_requested() is False
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "hip")
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx803"])
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert routed is host
assert repo == FORK
assert persist is None
def test_route_to_vulkan_prebuilt_hidden_physical_nvidia_amd_not_rerouted():
@ -1121,7 +1489,7 @@ def test_route_to_vulkan_prebuilt_hidden_physical_nvidia_amd_not_rerouted():
def test_route_to_vulkan_prebuilt_explicit_opt_in_overrides_hidden_nvidia(monkeypatch):
# The physical-NVIDIA guard only gates the AMD auto path; an explicit opt-in wins.
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "vulkan")
host = _windows_amd_host(
rocm_gfx_target = "gfx803",
rocm_gfx_targets = ["gfx803"],
@ -1129,7 +1497,7 @@ def test_route_to_vulkan_prebuilt_explicit_opt_in_overrides_hidden_nvidia(monkey
has_usable_nvidia = False,
)
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert repo == FORK
assert persist == "vulkan"
@ -1221,7 +1589,7 @@ def test_windows_hip_gfx_floor_covers_every_fork_windows_rocm_bundle():
@pytest.mark.parametrize("gfx", _FORK_WINDOWS_ROCM_GFX)
def test_route_to_vulkan_prebuilt_keeps_every_fork_windows_rocm_arch(gfx, monkeypatch):
# No ambient opt-in: this asserts the AUTO path leaves covered archs alone.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
host = _windows_amd_host(rocm_gfx_target = gfx, rocm_gfx_targets = [gfx])
routed, repo, tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
@ -1235,7 +1603,7 @@ def test_forwarded_gfx_does_not_undo_visible_device_auto_vulkan(monkeypatch):
# Under CUDA_VISIBLE_DEVICES=1 setup.ps1 still resolves GPU 0 and forwards gfx1100, but
# detect_host() resolved the visible gfx1010, so folding the forward in must not
# reinstate gfx1100 and install a HIP bundle the visible GPU cannot run.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = _windows_amd_host(rocm_gfx_target = "gfx1010", rocm_gfx_targets = ["gfx1100", "gfx1010"])
@ -1258,7 +1626,7 @@ def test_forwarded_gfx_absent_from_probe_keeps_the_physical_hip_card(monkeypatch
# the HIP target but must not delete the probe's inventory, or the floor check concludes
# no AMD GPU here reaches HIP and auto-routes to Vulkan, which ignores the HIP mask and
# enumerates the reserved gfx1100.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx1100", "gfx803"])
@ -1274,7 +1642,7 @@ def test_forwarded_gfx_absent_from_probe_keeps_the_physical_hip_card(monkeypatch
def test_forwarded_gfx_absent_from_probe_keeps_a_single_probed_hip_card(monkeypatch):
# Same rule on a single-GPU box: a stale below-floor forward over a probe-confirmed
# gfx1100 must not auto-route that machine to Vulkan.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
@ -1286,12 +1654,12 @@ def test_forwarded_gfx_absent_from_probe_keeps_a_single_probed_hip_card(monkeypa
def test_forwarded_gfx_absent_from_probe_still_allows_explicit_vulkan(monkeypatch):
# The physical-inventory rule gates the AUTO path only; naming the backend wins.
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "vulkan")
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
host = ilp._apply_host_overrides(host, override_rocm_gfx = "gfx803")
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert repo == FORK
assert persist == "vulkan"
@ -1299,7 +1667,7 @@ def test_forwarded_gfx_on_unprobed_host_still_auto_vulkans(monkeypatch):
# Negative control: a driver-only AMD host runs no successful probe (no hipinfo, amd-smi
# suppressed), so --rocm-gfx is the ONLY source of the arch and there is no inventory to
# preserve. This is the #7357 path the feature exists for; it must still reach Vulkan.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = _windows_amd_host(rocm_gfx_target = None, rocm_gfx_targets = [])
@ -1307,14 +1675,14 @@ def test_forwarded_gfx_on_unprobed_host_still_auto_vulkans(monkeypatch):
assert host.rocm_gfx_targets == ["gfx803"]
assert ilp._should_auto_vulkan_for_amd_windows(host, FORK) is True
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert repo == FORK
assert persist == "vulkan"
def test_forwarded_gfx_still_fills_an_unprobed_arch(monkeypatch):
# Negative control: on an amd-smi-only host detect_host() reports no arch, so the
# forward is the only source and must still apply.
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False)
host = _windows_amd_host(rocm_gfx_target = None, rocm_gfx_targets = [])
@ -1326,9 +1694,9 @@ def test_forwarded_gfx_still_fills_an_unprobed_arch(monkeypatch):
assert persist is None
def test_llama_backend_hip_opts_out_of_auto_vulkan(monkeypatch):
# hip names a backend, so it keeps the fork path even on an auto-fallback arch.
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "hip")
def test_llama_cpp_backend_cpu_opts_out_of_auto_vulkan(monkeypatch):
# cpu is an explicit non-Vulkan choice, so it suppresses the auto-fallback.
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "cpu")
host = _windows_amd_host(rocm_gfx_target = "gfx803", rocm_gfx_targets = ["gfx803"])
routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert routed is host
@ -1338,10 +1706,10 @@ def test_llama_backend_hip_opts_out_of_auto_vulkan(monkeypatch):
def test_explicit_backend_beats_legacy_force_vulkan(monkeypatch):
# A stale UNSLOTH_FORCE_VULKAN must not overrule UNSLOTH_LLAMA_BACKEND=rocm (== hip).
# A stale UNSLOTH_FORCE_VULKAN must not overrule the canonical CPU choice.
monkeypatch.setenv("UNSLOTH_FORCE_VULKAN", "1")
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "rocm")
assert ilp.resolved_llama_backend() == "hip"
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "cpu")
assert ilp.resolved_llama_backend() == "cpu"
assert ilp.force_vulkan_requested() is False
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
@ -1351,7 +1719,7 @@ def test_explicit_backend_beats_legacy_force_vulkan(monkeypatch):
def test_unknown_llama_backend_value_falls_through_to_legacy_flag(monkeypatch):
# An unrecognised value is ignored, not an error, so the legacy flag still works.
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "banana")
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "banana")
assert ilp.resolved_llama_backend() is None
assert ilp.force_vulkan_requested() is False
monkeypatch.setenv("UNSLOTH_FORCE_VULKAN", "1")
@ -1360,13 +1728,13 @@ def test_unknown_llama_backend_value_falls_through_to_legacy_flag(monkeypatch):
def test_llama_backend_flag_beats_conflicting_env(monkeypatch):
# --llama-backend is the caller's explicit request and outranks the env.
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "hip")
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "cpu")
assert ilp.force_vulkan_requested("vulkan") is True
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(
host, FORK, "pin", force_cpu = False, llama_backend = "vulkan"
)
assert repo == UPSTREAM
assert repo == FORK
assert persist == "vulkan"
@ -1395,14 +1763,14 @@ def _windows_arm64_host(**overrides):
@pytest.mark.parametrize(
"env, flag",
[
({"UNSLOTH_LLAMA_BACKEND": "vulkan"}, None),
({"UNSLOTH_LLAMA_CPP_BACKEND": "vulkan"}, None),
({"UNSLOTH_FORCE_VULKAN": "1"}, None),
({}, "vulkan"),
],
)
def test_vulkan_opt_in_ignored_on_windows_arm64(monkeypatch, env, flag):
# Upstream builds win-vulkan for x64 only (arm64 gets CPU + opencl-adreno), so rewriting
# the host would only swap the published arm64 bundle for the upstream CPU one.
def test_vulkan_opt_in_keeps_windows_arm64_host_for_strict_rejection(monkeypatch, env, flag):
# Windows arm64 has no compatible Vulkan bundle. Keep the real host shape so
# strict filtering rejects the CPU plan instead of routing through x64.
for name, value in env.items():
monkeypatch.setenv(name, value)
host = _windows_arm64_host()
@ -1416,10 +1784,10 @@ def test_vulkan_opt_in_ignored_on_windows_arm64(monkeypatch, env, flag):
def test_vulkan_opt_in_still_routes_on_windows_x64(monkeypatch):
# Negative control for the arm64 guard: x64 keeps its Vulkan routing.
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", "vulkan")
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", "vulkan")
host = _windows_amd_host(rocm_gfx_target = "gfx1100", rocm_gfx_targets = ["gfx1100"])
_routed, repo, _tag, persist = ilp._route_to_vulkan_prebuilt(host, FORK, "pin", force_cpu = False)
assert repo == UPSTREAM
assert repo == FORK
assert persist == "vulkan"
@ -1490,12 +1858,9 @@ def test_marker_records_no_backend_when_vulkan_fell_back_to_cpu(tmp_path):
assert marker["llama_backend"] == "vulkan"
# UNSLOTH_LLAMA_CPP_BACKEND (setup.sh/setup.ps1, "auto"|"cpu") and
# UNSLOTH_LLAMA_BACKEND (this module, a backend name) are different variables at
# different layers, and both accept "cpu". setup translates its own =cpu into
# --force-cpu to pin the CPU-only bundle on a GPU host, which is what keeps Intel
# iGPU Vulkan crashes away (#7213). Vulkan is opt-in here, so no trigger it adds
# may outrank that flag on any host.
# setup translates UNSLOTH_LLAMA_CPP_BACKEND=cpu into --force-cpu to pin the
# CPU-only bundle on a GPU host, which is what keeps Intel iGPU Vulkan crashes
# away (#7213). Vulkan is opt-in, so no trigger may outrank that flag on any host.
_SIM_PLATFORMS = {
# WSL presents as Linux to this resolver, so it rides the Linux row.
"Linux": dict(
@ -1577,16 +1942,16 @@ def _sim_host(platform_name, gpu_name):
@pytest.mark.parametrize("platform_name", sorted(_SIM_PLATFORMS))
@pytest.mark.parametrize("gpu_name", sorted(_SIM_GPUS))
@pytest.mark.parametrize("backend_env", [None, "vulkan", "hip", "rocm", "cpu"])
@pytest.mark.parametrize("backend_env", [None, "auto", "vulkan", "cpu"])
def test_forced_cpu_outranks_every_vulkan_trigger(
monkeypatch, platform_name, gpu_name, backend_env
):
"""A deliberate CPU install stays CPU on every host, whatever asks for Vulkan."""
monkeypatch.delenv("UNSLOTH_FORCE_VULKAN", raising = False)
if backend_env is None:
monkeypatch.delenv("UNSLOTH_LLAMA_BACKEND", raising = False)
monkeypatch.delenv("UNSLOTH_LLAMA_CPP_BACKEND", raising = False)
else:
monkeypatch.setenv("UNSLOTH_LLAMA_BACKEND", backend_env)
monkeypatch.setenv("UNSLOTH_LLAMA_CPP_BACKEND", backend_env)
# The legacy switch too, so a stale one cannot smuggle Vulkan past --force-cpu.
monkeypatch.setenv("UNSLOTH_FORCE_VULKAN", "1")

273
studio/backend/tests/test_llama_cpp_placement.py Normal file
View file

@ -0,0 +1,273 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Focused integration tests for explicit GGUF GPU placement."""
from __future__ import annotations
import os
import struct
import subprocess
import sys
import types
from pathlib import Path
from unittest.mock import patch
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
def _stub_module(name: str, **attrs):
if name in sys.modules:
return
try:
__import__(name)
return
except Exception:
module = types.ModuleType(name)
for key, value in attrs.items():
setattr(module, key, value)
sys.modules[name] = module
_stub_module("loggers", get_logger = lambda name: __import__("logging").getLogger(name))
_stub_module("structlog", get_logger = lambda *a, **k: __import__("logging").getLogger("stub"))
_stub_module(
"jwt",
decode = lambda *a, **k: {},
ExpiredSignatureError = type("ExpiredSignatureError", (Exception,), {}),
InvalidTokenError = type("InvalidTokenError", (Exception,), {}),
)
if "httpx" not in sys.modules:
try:
import httpx # noqa: F401
except Exception:
module = types.ModuleType("httpx")
for name in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
):
setattr(module, name, type(name, (Exception,), {}))
module.Timeout = type("Timeout", (), {"__init__": lambda self, *a, **k: None})
module.Client = type(
"Client",
(),
{
"__init__": lambda self, **kwargs: None,
"__enter__": lambda self: self,
"__exit__": lambda self, *args: None,
},
)
sys.modules["httpx"] = module
from core.inference.llama_cpp import LlamaCppBackend
_REAL_POPEN = subprocess.Popen
def _write_gguf(path: Path, architecture: str = "llama") -> Path:
def string(value: str) -> bytes:
data = value.encode()
return struct.pack("<Q", len(data)) + data
metadata = string("general.architecture") + struct.pack("<I", 8) + string(architecture)
path.write_bytes(struct.pack("<IIQQ", 0x46554747, 3, 0, 1) + metadata)
return path
def _backend(tmp_path: Path, *, vulkan: bool, memory):
backend = LlamaCppBackend()
gguf = _write_gguf(tmp_path / "model.gguf")
backend._get_gpu_memory = lambda _binary = None: list(memory)
backend._get_gpu_free_memory = lambda _binary = None: [
(index, free) for index, free, _total in memory
]
backend._read_gguf_metadata = lambda _path: None
backend._can_estimate_kv = lambda: False
backend._get_gguf_size_bytes = lambda _path: 1024
backend._mmproj_vram_bytes = lambda _path: 0
backend._resolve_launch_mmproj_path = lambda **kwargs: None
backend._apu_ram_shortfall_message = lambda *args, **kwargs: None
backend._amd_apu_wants_unified_memory = lambda *args, **kwargs: False
backend._find_llama_server_binary = lambda include_denied = False: "/fake/llama-server"
backend._is_vulkan_backend = lambda _binary = None: vulkan
backend._wait_for_health = lambda timeout: True
backend._detect_audio_type_strict = lambda: None
backend._apply_detected_audio = lambda _detected: True
return backend, gguf
def _launch(backend, gguf, **load_kwargs):
captured = {}
def fake_popen(cmd, **kwargs):
if not cmd or str(cmd[0]) != "/fake/llama-server":
return _REAL_POPEN(cmd, **kwargs)
captured["cmd"] = list(cmd)
captured["env"] = kwargs.get("env") or dict(os.environ)
return type(
"Process",
(),
{
"pid": 123,
"stdout": (),
"poll": lambda self: None,
"terminate": lambda self: None,
"wait": lambda self, timeout = None: 0,
"kill": lambda self: None,
},
)()
with patch.object(subprocess, "Popen", side_effect = fake_popen):
assert backend.load_model(
gguf_path = str(gguf),
model_identifier = "test",
**load_kwargs,
)
return captured
def test_vulkan_selection_uses_ordinals_and_owns_device_flags(tmp_path):
backend, gguf = _backend(
tmp_path,
vulkan = True,
memory = [(0, 10_000, 16_000), (1, 8_000, 16_000)],
)
backend._select_gpus = lambda *args, **kwargs: ([1], False)
result = _launch(
backend,
gguf,
gpu_ids = [0, 1],
extra_args = ["--device", "Vulkan0", "--main-gpu", "0", "--top-k", "5"],
)
cmd = result["cmd"]
assert cmd[cmd.index("--device") + 1] == "Vulkan1"
assert cmd.count("--device") == 1
assert "--main-gpu" not in cmd
assert cmd[cmd.index("--top-k") + 1] == "5"
assert backend.requested_gpu_ids == [0, 1]
assert backend.gpu_ids == [1]
@pytest.mark.parametrize(
"gpu_ids,extra_args,expected_draft,user_device_survives",
[
(None, None, "Vulkan1", False),
(None, ["--device", "Vulkan1", "-dev=Vulkan0"], "Vulkan0", True),
([1], ["--device", "Vulkan1", "-dev=Vulkan0"], "Vulkan1", False),
],
)
def test_vulkan_fit_and_mtp_drafter_follow_placement_owner(
tmp_path, gpu_ids, extra_args, expected_draft, user_device_survives
):
backend, gguf = _backend(
tmp_path,
vulkan = True,
memory = [(0, 24_000, 0), (1, 8_000, 16_000)],
)
planned = []
def fallback(_model_size, gpus, *args, **kwargs):
planned.append(list(gpus))
return None, True
backend._select_gpus = fallback
backend.probe_server_capabilities = lambda _binary = None: {
"mtp_token": "draft-mtp",
"spec_draft_n_max_flag": "--spec-draft-n-max",
}
backend._resolve_launch_mtp_path = lambda **_kwargs: "/fake/mtp.gguf"
result = _launch(
backend,
gguf,
mtp_draft_path = "/fake/mtp.gguf",
speculative_type = "mtp",
gpu_ids = gpu_ids,
extra_args = extra_args,
)
assert planned
assert all(gpus == [(1, 8_000)] for gpus in planned)
cmd = result["cmd"]
assert cmd[cmd.index("--device") + 1] == "Vulkan1"
assert cmd[cmd.index("--spec-draft-device") + 1] == expected_draft
assert ("-dev=Vulkan0" in cmd) is user_device_survives
def test_cuda_selection_uses_visibility_and_removes_environment_placement(tmp_path, monkeypatch):
monkeypatch.setenv("LLAMA_ARG_DEVICE", "CUDA0")
monkeypatch.setenv("LLAMA_ARG_MAIN_GPU", "0")
backend, gguf = _backend(
tmp_path,
vulkan = False,
memory = [(0, 10_000, 16_000), (1, 8_000, 16_000)],
)
backend._select_gpus = lambda *args, **kwargs: ([1], False)
result = _launch(backend, gguf, gpu_ids = [1])
assert result["env"]["CUDA_VISIBLE_DEVICES"] == "1"
assert "LLAMA_ARG_DEVICE" not in result["env"]
assert "LLAMA_ARG_MAIN_GPU" not in result["env"]
def test_backend_detection_accepts_versioned_vulkan_soname(tmp_path):
binary = tmp_path / "llama-server"
binary.write_bytes(b"x")
lib_dir = tmp_path / "lib"
lib_dir.mkdir()
prefix = "" if sys.platform == "win32" else "lib"
extension = "dll" if sys.platform == "win32" else "so"
(lib_dir / f"{prefix}ggml-vulkan.{extension}.0").write_bytes(b"x")
with patch("core.inference.llama_cpp._llama_lib_dir", return_value = lib_dir):
assert LlamaCppBackend._is_vulkan_backend(str(binary)) is True
assert LlamaCppBackend._backend_lacks_gpu_lib(str(binary)) is False
def test_cpu_only_detection_requires_a_proven_split_library_layout(tmp_path):
binary = tmp_path / "llama-server"
binary.write_bytes(b"x")
lib_dir = tmp_path / "lib"
lib_dir.mkdir()
prefix = "" if sys.platform == "win32" else "lib"
extension = "dll" if sys.platform == "win32" else "so"
(lib_dir / f"{prefix}ggml-cpu.{extension}").write_bytes(b"x")
with patch("core.inference.llama_cpp._llama_lib_dir", return_value = lib_dir):
assert LlamaCppBackend._backend_lacks_gpu_lib(str(binary)) is True
(lib_dir / f"{prefix}ggml-vulkan.{extension}").write_bytes(b"x")
with patch("core.inference.llama_cpp._llama_lib_dir", return_value = lib_dir):
assert LlamaCppBackend._backend_lacks_gpu_lib(str(binary)) is False
def test_diffusion_does_not_reinterpret_vulkan_ordinals(tmp_path):
gguf = _write_gguf(tmp_path / "diffusion.gguf", "diffusion-gemma")
backend = LlamaCppBackend()
backend._find_llama_server_binary = lambda include_denied = False: "/fake/llama-server"
backend._is_vulkan_backend = lambda _binary = None: True
backend._get_gpu_memory = lambda _binary = None: [(1, 8_000, 8_000)]
backend._download_gguf = lambda **kwargs: str(gguf)
backend._read_gguf_metadata = lambda _path: setattr(backend, "_is_diffusion", True)
backend._start_diffusion_server = lambda **kwargs: pytest.fail(
"Vulkan ordinal reached the CUDA diffusion runner"
)
with pytest.raises(ValueError, match = "no defined mapping"):
backend.load_model(
hf_repo = "renamed/model",
hf_variant = "Q4_K_M",
model_identifier = "renamed/model",
speculative_type = "off",
gpu_ids = [1],
)

2
studio/backend/tests/test_llama_cpp_update.py
View file

@ -499,7 +499,7 @@ def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
time.sleep(0.05)
assert job["state"] == "success", job
assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1"
assert popen_kwargs["env"]["UNSLOTH_LLAMA_BACKEND"] == "vulkan"
assert popen_kwargs["env"]["UNSLOTH_LLAMA_CPP_BACKEND"] == "vulkan"
assert "--llama-backend" in captured["cmd"] and "vulkan" in captured["cmd"]

48
studio/backend/tests/test_llama_cpp_vulkan_probe.py
View file

@ -53,16 +53,39 @@ _maybe_stub("loggers", _build_loggers_stub)
_maybe_stub("structlog", lambda: _types.ModuleType("structlog"))
from core.inference import llama_cpp as _llama_mod # noqa: E402
from core.inference._vulkan_probe import _igpu_flags_and_names # noqa: E402
from core.inference.llama_cpp import ( # noqa: E402
LlamaCppBackend,
_llama_lib_dir,
_vulkan_lib_filename,
)
MIB = 1024 * 1024
GIB = 1024 * MIB
class _FakeCFunction:
def __init__(self, result):
self.result = result
def __call__(self, *_args):
return self.result
def test_missing_description_symbol_keeps_igpu_detection():
base = _types.SimpleNamespace(
ggml_backend_reg_dev_count = _FakeCFunction(1),
ggml_backend_reg_dev_get = _FakeCFunction(1),
ggml_backend_dev_type = _FakeCFunction(2),
ggml_backend_dev_name = _FakeCFunction(b"Legacy Vulkan iGPU"),
)
lib = _types.SimpleNamespace(ggml_backend_vk_reg = _FakeCFunction(1))
flags, names = _igpu_flags_and_names(base, lib, 1)
assert flags == [True]
assert names == ["Legacy Vulkan iGPU"]
def _make_vulkan_install(tmp_path: Path) -> str:
"""A binary whose sibling dir holds the Vulkan ggml lib, so the
reader's ``is_vulkan_backend`` sibling-file check passes."""
@ -70,7 +93,8 @@ def _make_vulkan_install(tmp_path: Path) -> str:
bindir.mkdir(parents = True)
binary = bindir / ("llama-server.exe" if sys.platform == "win32" else "llama-server")
binary.write_bytes(b"stub")
(bindir / _vulkan_lib_filename()).write_bytes(b"stub")
vulkan_lib = "ggml-vulkan.dll" if sys.platform == "win32" else "libggml-vulkan.so"
(bindir / vulkan_lib).write_bytes(b"stub")
return str(binary)
@ -97,8 +121,10 @@ def _row(
free_bytes: int,
is_igpu: int,
total_bytes: int = 0,
name: str | None = None,
) -> str:
return f"{idx}\t{free_bytes}\t{is_igpu}\t{total_bytes}"
row = f"{idx}\t{free_bytes}\t{is_igpu}\t{total_bytes}"
return f"{row}\t{name}" if name is not None else row
def test_integrated_gpu_leaves_host_margin(tmp_path):
@ -189,5 +215,21 @@ def test_shell_wrapper_entrypoint_resolves_to_real_lib_dir(tmp_path):
assert LlamaCppBackend._is_vulkan_backend(str(wrapper)) is True
@pytest.mark.skipif(sys.platform == "win32", reason = "soname versioning is POSIX")
def test_versioned_only_vulkan_soname_is_probed(tmp_path):
# Split-library install: only the versioned soname libggml-vulkan.so.0 exists
# (no unversioned dev symlink). The reader must still classify Vulkan and run
# the probe instead of returning [] and rejecting gpu_ids before launch (#7188).
bindir = tmp_path / "build" / "bin"
bindir.mkdir(parents = True)
binary = bindir / "llama-server"
binary.write_bytes(b"stub")
(bindir / "libggml-vulkan.so.0").write_bytes(b"stub")
rows = [_row(0, 23 * GIB, is_igpu = 0, total_bytes = 24 * GIB)]
with _mock_probe(rows):
gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(str(binary))
assert gpus == [(0, 23 * 1024, 24 * 1024)], gpus
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))

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

@ -93,6 +93,9 @@ validate_extra_args = _lsa.validate_extra_args
["-fit", "off"],
["--fit", "on"],
["--fit-ctx", "8192"],
# Memory placement flags (soft-managed; shadowed on inherit)
["--mlock"],
["--no-mmap", "--mlock"],
],
)
def test_pass_through_allowed(args):
@ -319,6 +322,9 @@ def test_is_managed_flag_false_for_pass_through():
assert is_managed_flag("--flash-attn") is False
assert is_managed_flag("-ngl") is False
assert is_managed_flag("--threads") is False
# Memory placement flags are pass-through (shadowed on inherit only).
assert is_managed_flag("--mlock") is False
assert is_managed_flag("--no-mmap") is False
# ── strip_shadowing_flags ─────────────────────────────────────────────
@ -383,6 +389,34 @@ def test_strip_shadowing_flags_keeps_spec_when_spec_disabled():
assert out == ["--spec-type", "ngram-mod", "--draft-min", "48", "--top-k", "20"]
def test_strip_shadowing_flags_keeps_device_by_default():
# --device is pass-through by default (users may pin when Unsloth auto-selects).
out = strip_shadowing_flags(
["--device", "Vulkan1", "--top-k", "20"],
strip_context = False,
strip_cache = False,
strip_spec = False,
strip_template = False,
strip_split_mode = False,
)
assert out == ["--device", "Vulkan1", "--top-k", "20"]
def test_strip_shadowing_flags_drops_device_when_requested():
# strip_device drops device placement flags when gpu_ids owns placement.
for flag in ("--device", "-dev", "--main-gpu", "-mg"):
out = strip_shadowing_flags(
[flag, "Vulkan1", "--top-k", "20"],
strip_context = False,
strip_cache = False,
strip_spec = False,
strip_template = False,
strip_split_mode = False,
strip_device = True,
)
assert out == ["--top-k", "20"], flag
def test_strip_shadowing_flags_drops_mtp_flags_when_requested():
# MTP / draft-mtp flags must drop when speculative_type re-applies.
out = strip_shadowing_flags(

28
studio/backend/tests/test_mtp_vram_budget.py
View file

@ -68,6 +68,7 @@ from core.inference.llama_cpp import ( # noqa: E402
_CTX_FIT_VRAM_FRACTION,
LlamaCppBackend,
_extra_args_draft_cache_types,
_extra_args_draft_device_pin,
_extra_args_draft_offloaded_to_cpu,
_extra_args_mtp_draft_path,
_extra_args_n_ubatch,
@ -664,6 +665,33 @@ class TestExtraArgsMtpDetection:
)
assert _extra_args_draft_offloaded_to_cpu([], env = {}) is False
@pytest.mark.parametrize(
"args,expected",
[
# No draft-device flag -> no pin to reject.
(None, None),
([], None),
(["-c", "4096"], None),
# cpu / none offload is a supported placement, not a GPU escape.
(["--spec-draft-device", "cpu"], None),
(["--spec-draft-device", "CPU,none"], None),
(["-devd", "none"], None),
# A real GPU device escapes the gpu_ids pin -> return the offending value.
(["--spec-draft-device", "CUDA1"], "CUDA1"),
(["--device-draft", "Vulkan2"], "Vulkan2"),
(["--spec_draft_device", "Vulkan1"], "Vulkan1"),
(["-devd", "CUDA0,CPU"], "CUDA0,CPU"), # any GPU in the list conflicts
# inline flag=value form.
(["--spec-draft-device=Vulkan3"], "Vulkan3"),
(["--device_draft=Vulkan4"], "Vulkan4"),
# last-wins: only the final draft-device value counts.
(["--spec-draft-device", "CUDA1", "--spec-draft-device", "cpu"], None),
(["--spec-draft-device", "cpu", "--spec-draft-device", "CUDA1"], "CUDA1"),
],
)
def test_draft_device_pin(self, args, expected):
assert _extra_args_draft_device_pin(args) == expected
@pytest.mark.parametrize(
"args,expected",
[

274
studio/backend/tests/test_setup_llama_cpp_backend.py
View file

@ -1,12 +1,12 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""setup.sh and setup.ps1 must map UNSLOTH_LLAMA_CPP_BACKEND=cpu to
install_llama_prebuilt.py's --force-cpu so users can force the CPU-only prebuilt
on GPU hosts (#7213). The match is case-insensitive and whitespace-trimmed, an
unrecognized value warns instead of silently falling back, and macOS warns (no
CPU-only bundle). Runs the real block extracted from each script so the tests
track the shipped logic.
"""Backend selector coverage for setup.sh and setup.ps1.
cpu maps to install_llama_prebuilt.py's persisted --force-cpu option. vulkan is
accepted and passed through in the environment for the installer to consume.
The match is case-insensitive and whitespace-trimmed, unknown values warn, and
macOS warns for the CPU-only choice.
"""
import os
@ -35,11 +35,19 @@ def _run(value: str | None, system: str = "Linux") -> tuple[list[str], str]:
# Pass the value through env (not the script text) so whitespace survives, and
# stub the setup.sh logging helpers the unknown-value branch calls. system sets
# _HOST_SYSTEM so the macOS (Darwin) no-op branch can be exercised.
env = {k: v for k, v in os.environ.items() if k != "UNSLOTH_LLAMA_CPP_BACKEND"}
env = {
k: v
for k, v in os.environ.items()
if k not in ("UNSLOTH_LLAMA_CPP_BACKEND", "UNSLOTH_FORCE_VULKAN")
}
if value is not None:
env["UNSLOTH_LLAMA_CPP_BACKEND"] = value
harness = (
f'_PREBUILT_CMD=()\nC_WARN=""\n_HOST_SYSTEM="{system}"\n'
f'set -u\n_PREBUILT_CMD=()\nC_WARN=""\nC_OK=""\n_HOST_SYSTEM="{system}"\n'
'_source_backend_choice="$(printf \'%s\' "${UNSLOTH_LLAMA_CPP_BACKEND:-auto}" '
"| awk '{$1=$1; print tolower($0)}')\"\n"
'_source_legacy_force_vulkan="$(printf \'%s\' "${UNSLOTH_FORCE_VULKAN:-}" '
"| awk '{$1=$1; print tolower($0)}')\"\n"
'step() { printf "STEP: %s\\n" "$*" >&2; }\n'
f"{_backend_block()}\n"
'printf "%s\\n" "${_PREBUILT_CMD[@]}"'
@ -81,7 +89,25 @@ def test_backend_auto_no_flag_no_warn(value):
@_SKIP_NO_BASH
@pytest.mark.parametrize("value", ["vulkan", "gpu", "cuda"])
@pytest.mark.parametrize("value", ["vulkan", "VULKAN", " vulkan "])
def test_backend_vulkan_is_accepted(value):
args, stderr = _run(value)
assert "--force-cpu" not in args
assert args[-2:] == ["--llama-backend", "vulkan"]
assert "Ignoring" not in stderr
@_SKIP_NO_BASH
@pytest.mark.parametrize("value", ["hip", "HIP", "rocm", " ROCM "])
def test_backend_hip_opt_out_is_accepted(value):
args, stderr = _run(value)
assert "--force-cpu" not in args
assert "--llama-backend" not in args
assert "Ignoring" not in stderr
@_SKIP_NO_BASH
@pytest.mark.parametrize("value", ["gpu", "cuda"])
def test_backend_unknown_warns_and_no_flag(value):
args, stderr = _run(value)
assert "--force-cpu" not in args
@ -100,26 +126,202 @@ def test_arm64_recovery_uses_transient_cpu_fallback():
assert "--force-cpu" not in block
def test_explicit_vulkan_prebuilt_failure_does_not_change_backend():
sh = _SETUP_SH.read_text(encoding = "utf-8")
failure = sh.index('step "llama.cpp" "prebuilt install failed"')
source_build = sh.index("_NEED_LLAMA_SOURCE_BUILD=true", failure)
guard = sh.index('if [ "$_explicit_vulkan_backend" = true ]', failure)
assert guard < source_build
guarded = sh[guard:source_build]
assert "will not substitute a ROCm or CPU source build" in guarded
assert "exit 1" in guarded
ps1 = _SETUP_PS1.read_text(encoding = "utf-8")
failure = ps1.index('step "llama.cpp" "prebuilt install failed"')
source_build = ps1.index("$NeedLlamaSourceBuild = $true", failure)
guard = ps1.index("if ($explicitVulkanBackend)", failure)
assert guard < source_build
guarded = ps1[guard:source_build]
assert "will not substitute a CUDA, ROCm, or CPU source build" in guarded
assert "exit 1" in guarded
def test_explicit_vulkan_source_build_fails_closed():
sh = _SETUP_SH.read_text(encoding = "utf-8")
local_branch = sh.index('if [ "$_LOCAL_LLAMA_CPP_LINKED" = true ]; then')
prebuilt_branch = sh.index('elif [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then', local_branch)
guarded = sh[local_branch:prebuilt_branch]
assert (
'elif [ "$_explicit_vulkan_source_build" = true ] && '
'[ "$_NEED_LLAMA_SOURCE_BUILD" = true ]; then'
) in guarded
assert "Vulkan source builds are not supported" in guarded
assert "exit 1" in guarded
ps1 = _SETUP_PS1.read_text(encoding = "utf-8")
local_branch = ps1.index("if ($LocalLlamaCppLinked) {")
prebuilt_branch = ps1.index(
'} elseif ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {', local_branch
)
guarded = ps1[local_branch:prebuilt_branch]
assert "} elseif ($explicitVulkanSourceBuild -and $NeedLlamaSourceBuild) {" in guarded
assert "Vulkan source builds are not supported" in guarded
assert "exit 1" in guarded
def test_force_compile_sets_need_source_build_before_vulkan_guard():
# UNSLOTH_LLAMA_FORCE_COMPILE=1 combined with an explicit Vulkan backend must
# hit the "Vulkan source builds are not supported" rejection, not silently
# fall through to a CUDA/ROCm/CPU source build. That only holds if
# _NEED_LLAMA_SOURCE_BUILD/$NeedLlamaSourceBuild is already true by the time
# the explicit-Vulkan elif guard runs, i.e. set earlier in the script.
sh = _SETUP_SH.read_text(encoding = "utf-8")
force_compile_set = sh.index(
'if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then\n _NEED_LLAMA_SOURCE_BUILD=true'
)
vulkan_guard = sh.index('elif [ "$_explicit_vulkan_source_build" = true ] && ')
assert force_compile_set < vulkan_guard
ps1 = _SETUP_PS1.read_text(encoding = "utf-8")
force_compile_set = ps1.index(
'if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {\n $NeedLlamaSourceBuild = $true'
)
vulkan_guard = ps1.index("} elseif ($explicitVulkanSourceBuild -and $NeedLlamaSourceBuild) {")
assert force_compile_set < vulkan_guard
def test_legacy_force_vulkan_gets_the_same_strict_fallback():
sh = _SETUP_SH.read_text(encoding = "utf-8")
assert "_legacy_force_vulkan=" in sh
assert "1|true|yes|on) _explicit_vulkan_backend=true" in sh
ps1 = _SETUP_PS1.read_text(encoding = "utf-8")
assert "$legacyForceVulkan = $sourceLegacyForceVulkan" in ps1
assert '$legacyForceVulkan -in @("1", "true", "yes", "on")' in ps1
def _source_backend_choice_block() -> str:
text = _SETUP_SH.read_text(encoding = "utf-8")
m = re.search(
r'_source_backend_choice="\$\(printf.*?\n.*?_explicit_vulkan_source_build=false\n'
r".*?\nfi\n",
text,
re.DOTALL,
)
assert m, "_source_backend_choice block not found in setup.sh"
return m.group(0)
@_SKIP_NO_BASH
@pytest.mark.parametrize(
"backend, force_vulkan, expected_backend, expected_explicit",
[
(None, None, "auto", "false"),
("vulkan", None, "vulkan", "true"),
("auto", "on", "auto", "true"),
("banana", "1", "banana", "true"),
("cpu", "1", "cpu", "false"),
("hip", "1", "hip", "false"),
("rocm", "1", "rocm", "false"),
],
)
def test_llama_backend_source_choice_in_setup_sh(
backend, force_vulkan, expected_backend, expected_explicit
):
env = {
k: v
for k, v in os.environ.items()
if k not in ("UNSLOTH_LLAMA_CPP_BACKEND", "UNSLOTH_FORCE_VULKAN")
}
if backend is not None:
env["UNSLOTH_LLAMA_CPP_BACKEND"] = backend
if force_vulkan is not None:
env["UNSLOTH_FORCE_VULKAN"] = force_vulkan
harness = (
"set -u\n_HOST_SYSTEM=Linux\n"
f"{_source_backend_choice_block()}\n"
'printf "%s:%s" "$_source_backend_choice" "$_explicit_vulkan_source_build"'
)
out = subprocess.run(
["bash", "-c", harness], capture_output = True, text = True, env = env, check = True
)
assert out.stdout == f"{expected_backend}:{expected_explicit}"
def _ps1_search(pattern: str, flags = 0) -> str:
m = re.search(pattern, _SETUP_PS1.read_text(encoding = "utf-8"), flags)
assert m, f"setup.ps1 block not found: {pattern}"
return m.group(0)
def _run_ps1(value: str | None) -> str:
# The override is normalized (assign + warn) at the top of the prebuilt block and
# applied to $prebuiltArgs lower down; compose both real snippets.
@_SKIP_NO_PWSH
@pytest.mark.parametrize(
"backend, force_vulkan, expected_explicit",
[
(None, None, "False"),
("vulkan", None, "True"),
("auto", "on", "True"),
("cpu", "1", "False"),
("hip", "1", "False"),
("rocm", "1", "False"),
],
)
def test_llama_backend_source_choice_in_setup_ps1(backend, force_vulkan, expected_explicit):
normalize = _ps1_search(
r'\$llamaBackend = "\$\(\$env:UNSLOTH_LLAMA_CPP_BACKEND\)".*?Write-Host.*?\n\s*\}',
r'\$sourceLlamaBackend = "\$\(\$env:UNSLOTH_LLAMA_CPP_BACKEND\)".*?'
r"\$explicitVulkanSourceBuild = \(.*?\n\)\n",
re.DOTALL,
)
apply_flag = _ps1_search(
r'if \(\$llamaBackend -eq "cpu"\) \{\s*\$prebuiltArgs \+= "--force-cpu"\s*\}'
env = {
k: v
for k, v in os.environ.items()
if k not in ("UNSLOTH_LLAMA_CPP_BACKEND", "UNSLOTH_FORCE_VULKAN")
}
if backend is not None:
env["UNSLOTH_LLAMA_CPP_BACKEND"] = backend
if force_vulkan is not None:
env["UNSLOTH_FORCE_VULKAN"] = force_vulkan
out = subprocess.run(
[
"pwsh",
"-NoProfile",
"-Command",
# Brace the name: PowerShell reads "$var:" as a scope qualifier and
# fails to parse, so the probe never ran on a host that has pwsh.
f'{normalize}\n"RESULT:${{sourceLlamaBackend}}:$explicitVulkanSourceBuild"',
],
capture_output = True,
text = True,
env = env,
check = True,
)
env = {k: v for k, v in os.environ.items() if k != "UNSLOTH_LLAMA_CPP_BACKEND"}
expected_backend = (backend or "").strip().lower()
assert out.stdout.strip() == f"RESULT:{expected_backend}:{expected_explicit}"
def _run_ps1(value: str | None) -> str:
# One snippet, not two: NORMALIZE already spans the whole if/elseif chain up to
# and including the "Ignoring UNSLOTH_LLAMA_CPP_BACKEND" warn, so the --force-cpu
# append is inside it. Composing the apply snippet on top of it made the harness
# emit the flag twice while setup.ps1 appends it once.
normalize = _ps1_search(
r"\$llamaBackend = \$sourceLlamaBackend.*?Ignoring UNSLOTH_LLAMA_CPP_BACKEND.*?\n\s*\}",
re.DOTALL,
)
assert normalize.count('$prebuiltArgs += "--force-cpu"') == 1, normalize
env = {
k: v
for k, v in os.environ.items()
if k not in ("UNSLOTH_LLAMA_CPP_BACKEND", "UNSLOTH_FORCE_VULKAN")
}
if value is not None:
env["UNSLOTH_LLAMA_CPP_BACKEND"] = value
harness = f'$prebuiltArgs = @()\n{normalize}\n{apply_flag}\n"ARGS:" + ($prebuiltArgs -join ",")'
harness = (
"$prebuiltArgs = @()\n"
'$sourceLlamaBackend = "$($env:UNSLOTH_LLAMA_CPP_BACKEND)".Trim().ToLowerInvariant()\n'
'$sourceLegacyForceVulkan = "$($env:UNSLOTH_FORCE_VULKAN)".Trim().ToLowerInvariant()\n'
f'{normalize}\n"ARGS:" + ($prebuiltArgs -join ",")'
)
out = subprocess.run(
["pwsh", "-NoProfile", "-Command", harness],
capture_output = True,
@ -134,8 +336,10 @@ def _run_ps1(value: str | None) -> str:
@pytest.mark.parametrize("value", ["cpu", "CPU", "Cpu", " cpu ", "CPU\t"])
def test_ps1_backend_cpu_appends_flag(value):
out = _run_ps1(value)
assert "--force-cpu" in out
assert "Ignoring" not in out
# The whole argv, not a substring: an `in` check passed while the harness was
# emitting --force-cpu twice. setup.ps1 appends it exactly once, and nothing
# else, for a cpu override.
assert out.strip() == "ARGS:--force-cpu"
@_SKIP_NO_PWSH
@ -147,7 +351,37 @@ def test_ps1_backend_auto_no_flag_no_warn(value):
@_SKIP_NO_PWSH
@pytest.mark.parametrize("value", ["vulkan", "gpu", "cuda"])
@pytest.mark.parametrize("value", ["vulkan", "VULKAN", " vulkan "])
def test_ps1_backend_vulkan_is_accepted(value):
out = _run_ps1(value)
assert "--force-cpu" not in out
assert "--llama-backend,vulkan" in out
assert "Ignoring" not in out
@_SKIP_NO_PWSH
@pytest.mark.parametrize("value", ["hip", "HIP", "rocm", " ROCM "])
def test_ps1_backend_hip_opt_out_is_accepted(value):
out = _run_ps1(value)
assert "--force-cpu" not in out
assert "--llama-backend" not in out
assert "Ignoring" not in out
def test_ps1_forced_vulkan_fails_closed_on_windows_arm64():
ps1 = _SETUP_PS1.read_text(encoding = "utf-8")
arm64_guard = ps1.index("elseif ($windowsArm64)")
vulkan_flag = ps1.index(
'$prebuiltArgs += @("--llama-backend", "vulkan")',
arm64_guard,
)
assert arm64_guard < vulkan_flag
assert "no Windows ARM64 Vulkan bundle is published" in ps1
assert "Unset UNSLOTH_FORCE_VULKAN" in ps1
@_SKIP_NO_PWSH
@pytest.mark.parametrize("value", ["gpu", "cuda"])
def test_ps1_backend_unknown_warns_and_no_flag(value):
out = _run_ps1(value)
assert "--force-cpu" not in out

65
studio/backend/tests/test_system_vulkan_gpu_info.py
View file

@ -56,14 +56,43 @@ def test_system_gpu_info_preserves_vulkan_visibility_metrics(monkeypatch):
assert gpu["available"] is False
assert gpu["backend"] == "cpu"
assert gpu["index_kind"] == "relative"
# A Vulkan llama.cpp build accepts gpu_ids even when torch training is
# CPU-only: the pick is a ggml ordinal, not a torch device index.
assert gpu["gguf_gpu_ids_supported"] is True
# The training inventory must not advertise physical pins for a Vulkan
# llama.cpp build. Its ordinals live in inference_gpu below.
assert gpu["gguf_gpu_ids_supported"] is False
# Torch's view stays empty; the ggml ordinals stay in inference_gpu.
assert gpu["devices"] == []
assert inference_gpu["backend"] == "vulkan"
assert inference_gpu["devices"] == [vulkan_device]
def test_system_gpu_info_withholds_gguf_pin_when_the_vulkan_probe_enumerates_nothing(monkeypatch):
"""A Vulkan build whose probe returns no ordinals has nothing valid to pin,
so the picker must be told pins are unsupported rather than offered an empty
namespace it would 400 on."""
import utils.hardware as hardware
monkeypatch.setattr(
hardware,
"get_backend_visible_gpu_info",
lambda: {"available": False, "backend": "cpu", "devices": [], "index_kind": "relative"},
)
monkeypatch.setattr(
hardware,
"get_visible_gpu_utilization",
lambda: {"available": False, "backend": "cpu", "devices": []},
)
monkeypatch.setattr(hardware, "get_vulkan_inference_gpu_info", lambda: None)
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.setattr(LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda: True))
monkeypatch.setattr(main, "_system_gpu_cache", None)
gpu, _ = main._get_cached_system_gpu_info(SimpleNamespace(debug = lambda *args: None))
assert gpu["gguf_gpu_ids_supported"] is False
def test_system_gpu_info_keeps_forced_vulkan_separate_from_training_metrics(monkeypatch):
import utils.hardware as hardware
@ -181,17 +210,17 @@ def test_vulkan_inference_gpu_uses_real_device_names_and_igpu_flag(monkeypatch):
reserve is applied; budgeting off the raw shared total would hand out the
whole machine's RAM with no OS headroom.
"""
from core.inference import llama_cpp
from core.inference.llama_cpp import LlamaCppBackend
from utils.hardware.hardware import get_vulkan_inference_gpu_info
monkeypatch.setattr(
LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda binary = None: True)
)
# Fit view: discrete card keeps its total, iGPU reports 0 with capped free.
monkeypatch.setattr(
LlamaCppBackend,
"_get_gpu_memory",
staticmethod(lambda binary = None: [(0, 15 * 1024, 16 * 1024), (1, 12 * 1024, 0)]),
llama_cpp,
"_apply_igpu_host_reserve_mib",
lambda free_mib, is_igpu: 12 * 1024 if is_igpu else free_mib,
)
monkeypatch.setattr(
LlamaCppBackend,
@ -231,24 +260,28 @@ def test_vulkan_inference_gpu_uses_real_device_names_and_igpu_flag(monkeypatch):
assert igpu["memory_total_gb"] == 12.0
def test_vulkan_inference_gpu_falls_back_to_ordinal_names(monkeypatch):
"""A probe that cannot resolve descriptions must not lose the device list:
names degrade to Vulkan<i> and the memory readings still get through."""
def test_vulkan_inference_gpu_uses_inventory_fallback_names(monkeypatch):
"""The inventory's fallback name must flow through unchanged."""
from core.inference.llama_cpp import LlamaCppBackend
from utils.hardware.hardware import get_vulkan_inference_gpu_info
monkeypatch.setattr(
LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda binary = None: True)
)
monkeypatch.setattr(
LlamaCppBackend,
"_get_gpu_memory",
staticmethod(lambda binary = None: [(0, 15 * 1024, 16 * 1024)]),
)
monkeypatch.setattr(
LlamaCppBackend,
"vulkan_device_inventory",
staticmethod(lambda binary = None: (_ for _ in ()).throw(RuntimeError("probe failed"))),
staticmethod(
lambda binary = None: [
{
"index": 0,
"name": "Vulkan0",
"free_mib": 15 * 1024,
"total_mib": 16 * 1024,
"is_igpu": False,
}
]
),
)
info = get_vulkan_inference_gpu_info()

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

@ -843,3 +843,42 @@ def test_already_in_target_state_reloads_on_tensor_off_after_fallback():
)
# A genuine layer load (no preserved intent) -> dedupe, no churn.
assert _backend(False)._already_in_target_state(**kwargs) is True
# ── route dedup: gpu_ids device strip (#7164/#7188) ───────────────────────────
def _dedup_loaded_backend(*, extra_args):
"""A loaded GGUF backend for route dedup tests."""
b = LlamaCppBackend()
b._model_identifier = "owner/repo"
b._requested_n_ctx = 0
b._cache_type_kv = None
b._tensor_parallel = False
b._layer_preserves_tensor_intent = False
b._extra_args = list(extra_args) if extra_args else None
b._requested_spec_mode = "auto"
b._chat_template_override = None
b._gguf_path = None
b._gpu_ids = None
return b
def test_explicit_gpu_ids_dedupes_when_device_already_stripped():
"""A GGUF loaded with explicit gpu_ids had a user --device stripped from its stored
extras. A repeat identical request re-sending --device must still dedupe: the request-
side strip (gated on gpu_ids) compares equal to the stripped backend extras, so the
load hits the fast path instead of a needless reload / training 409 (#7188)."""
from models.inference import LoadRequest
inference_routes = _load_inference_routes_module()
req = LoadRequest(
model_path = "owner/repo",
gpu_ids = [0],
llama_extra_args = ["--device", "Vulkan3", "--top-k", "5"],
)
backend = _dedup_loaded_backend(extra_args = ["--top-k", "5"])
backend._gpu_ids = [0]
backend._requested_gpu_ids = [0]
assert inference_routes._request_matches_loaded_settings(req, backend) is True

28
studio/backend/utils/hardware/hardware.py
View file

@ -2586,7 +2586,10 @@ def get_vulkan_inference_gpu_info() -> Optional[Dict[str, Any]]:
# Vulkan is a llama.cpp inference backend, not a PyTorch training device, so
# keep it separate from the PyTorch/MLX training-device report.
try:
from core.inference.llama_cpp import LlamaCppBackend
from core.inference.llama_cpp import (
LlamaCppBackend,
_apply_igpu_host_reserve_mib,
)
except Exception as e:
logger.debug("Could not inspect the llama.cpp Vulkan backend: %s", e)
return None
@ -2606,23 +2609,12 @@ def get_vulkan_inference_gpu_info() -> Optional[Dict[str, Any]]:
"devices": [],
"index_kind": "vulkan",
}
# Identity (real device description, explicit iGPU flag) comes from the
# inventory; the memory numbers stay on _get_gpu_memory, which applies the
# iGPU host reserve and zeroes a shared total. Budgeting an APU off the raw
# shared total instead would hand out the whole machine's RAM with no OS
# headroom. Join by ordinal; a probe failure just leaves names unresolved.
identity: Dict[int, Dict[str, Any]] = {}
try:
identity = {row["index"]: row for row in LlamaCppBackend.vulkan_device_inventory()}
except Exception as e:
logger.debug("Vulkan device inventory failed, falling back to ordinals: %s", e)
try:
for ordinal, free_mib, total_mib in LlamaCppBackend._get_gpu_memory():
info = identity.get(ordinal, {})
# _get_gpu_memory reports total 0 for a shared pool; prefer the
# explicit flag when the inventory resolved this ordinal.
shared_memory = bool(info["is_igpu"]) if "is_igpu" in info else total_mib == 0
for row in LlamaCppBackend.vulkan_device_inventory():
ordinal = row["index"]
shared_memory = bool(row["is_igpu"])
free_mib = _apply_igpu_host_reserve_mib(row["free_mib"], shared_memory)
total_mib = 0 if shared_memory else row["total_mib"]
budget_mib = total_mib or free_mib
used_mib = max(0, total_mib - free_mib) if total_mib else None
result["devices"].append(
@ -2632,7 +2624,7 @@ def get_vulkan_inference_gpu_info() -> Optional[Dict[str, Any]]:
# so unlike a torch-xpu relative ordinal these are selectable.
"index_kind": "vulkan",
"visible_ordinal": ordinal,
"name": info.get("name") or f"Vulkan{ordinal}",
"name": row["name"],
"memory_total_gb": round(budget_mib / 1024, 2),
"vram_used_gb": round(used_mib / 1024, 2) if used_mib is not None else None,
"vram_free_gb": round(free_mib / 1024, 2),

2
studio/backend/utils/llama_cpp_update.py
View file

@ -470,7 +470,7 @@ def _run_llama_phase(
# otherwise re-route and silently replace it. Re-assert via setup's env/CLI flags.
if llama_backend == "vulkan" or (asset and "vulkan" in asset.lower()):
env["UNSLOTH_FORCE_VULKAN"] = "1"
env["UNSLOTH_LLAMA_BACKEND"] = "vulkan"
env["UNSLOTH_LLAMA_CPP_BACKEND"] = "vulkan"
_flow.stream_installer(
cmd,
env,

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

@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { getAuthToken } from "@/features/auth";
import { prepareHfTokenForUse } from "@/features/hf-auth";
import { resolveInitialConfig } from "@/features/model-picker";
import { projectHasSources } from "@/features/rag/api/rag-api";
import { apiUrl } from "@/lib/api-base";
@ -99,6 +100,7 @@ import { resolveLoadMaxSeqLength } from "../presets/preset-policy";
import {
generateAudio,
GenerationLengthError,
fetchGgufStagedMetadata,
listCachedGguf,
listCachedModels,
listGgufVariants,
@ -1569,9 +1571,38 @@ async function autoLoadSmallestModel(): Promise<{
// the load with the picker hidden.
await ensureGpuDeviceCache();
}
let isDiffusion = false;
if (
candidate.kind === "gguf" &&
(config.selectedGpuIds != null || config.tensorParallel === true)
) {
// Prepare the token before this probe, exactly as validateModel/loadModel
// (and the interactive performLoad path) do: the Hub 401s an invalid
// Authorization header even for a public repo, so sending the raw stored
// token here would abort this candidate before the existing anonymous/
// replacement-token recovery flow could run.
const preparedToken = await prepareHfTokenForUse(hfToken);
if (!preparedToken.proceed) {
throw new Error("Model load cancelled.");
}
isDiffusion = (
await fetchGgufStagedMetadata({
model_path: modelPath,
gguf_variant: candidate.ggufVariant,
hf_token: preparedToken.token,
})
).isDiffusion;
}
const effectiveTensorParallel = isDiffusion
? false
: config.tensorParallel;
const effectiveGpuIds =
config.selectedGpuIds !== undefined
? reconcilePersistedGpuIds(config.selectedGpuIds)
? reconcilePersistedGpuIds(
config.selectedGpuIds,
config.selectedGpuIndexKind,
isDiffusion,
)
: 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.
@ -1598,7 +1629,7 @@ async function autoLoadSmallestModel(): Promise<{
is_lora: false,
gguf_variant: candidate.ggufVariant,
cache_type_kv: config.kvCacheDtype,
tensor_parallel: config.tensorParallel,
tensor_parallel: effectiveTensorParallel,
// The same remembered-derived GPU pick the load below sends.
...(candidate.kind === "gguf"
? {
@ -1627,7 +1658,7 @@ async function autoLoadSmallestModel(): Promise<{
cache_type_kv: config.kvCacheDtype,
speculative_type: effectiveSpeculativeType,
spec_draft_n_max: effectiveSpecDraftNMax,
tensor_parallel: config.tensorParallel,
tensor_parallel: effectiveTensorParallel,
// 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
@ -1939,6 +1970,13 @@ async function autoLoadSmallestModel(): Promise<{
);
try {
const rt = useChatRuntimeStore.getState();
if (rt.selectedGpuIds != null) {
await ensureGpuDeviceCache();
}
const defaultGpuIds = reconcilePersistedGpuIds(
rt.selectedGpuIds,
rt.selectedGpuIndexKind,
);
if (
!(await canAutoLoad({
model_path: "unsloth/Qwen3.5-4B-MTP-GGUF",
@ -1947,7 +1985,7 @@ async function autoLoadSmallestModel(): Promise<{
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_ids: defaultGpuIds ?? undefined,
gpu_memory_mode: rt.gpuMemoryMode,
}))
) {
@ -1978,7 +2016,7 @@ async function autoLoadSmallestModel(): Promise<{
gpu_memory_mode: rt.gpuMemoryMode,
gpu_layers: GPU_LAYERS_AUTO,
n_cpu_moe: 0,
gpu_ids: rt.selectedGpuIds ?? undefined,
gpu_ids: defaultGpuIds ?? undefined,
});
saveSpeculativeType(specSettings.speculativeType);
persistGpuMemoryModeOnLoad(loadResp, rt.gpuMemoryMode);

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

@ -216,6 +216,7 @@ export async function fetchGgufStagedMetadata(payload: {
contextLength: number | null;
layerCount: number | null;
moeLayerCount: number | null;
isDiffusion: boolean;
}> {
let nativePathLease: string | null = null;
if (payload.nativePathToken) {
@ -225,7 +226,12 @@ export async function fetchGgufStagedMetadata(payload: {
).nativePathLease;
} catch {
// Lease expired / revoked: degrade to no metadata (the load can re-mint).
return { contextLength: null, layerCount: null, moeLayerCount: null };
return {
contextLength: null,
layerCount: null,
moeLayerCount: null,
isDiffusion: false,
};
}
}
const response = await authFetch("/api/inference/validate", {
@ -244,6 +250,7 @@ export async function fetchGgufStagedMetadata(payload: {
contextLength: res.context_length ?? null,
layerCount: res.layer_count ?? null,
moeLayerCount: res.moe_layer_count ?? null,
isDiffusion: res.is_diffusion ?? false,
};
}

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

@ -494,6 +494,7 @@ type CompareModelSelection = {
id: string;
isLora: boolean;
ggufVariant?: string;
isDiffusion?: boolean;
config?: PerModelConfig;
};
@ -842,6 +843,7 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
const globalCheckpoint = useChatRuntimeStore((s) => s.params.checkpoint);
const globalGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant);
const globalIsDiffusion = useChatRuntimeStore((s) => s.loadedIsDiffusion);
const active = useChatActive();
const compareRunning = useChatRuntimeStore(
(s) => Object.keys(s.runningByThreadId).length > 0,
@ -850,6 +852,7 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
id: globalCheckpoint || "",
isLora: false,
ggufVariant: globalGgufVariant ?? undefined,
isDiffusion: globalIsDiffusion,
});
const [model2, setModel2] = useState<CompareModelSelection>({
id: "",
@ -935,6 +938,7 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
id,
isLora: meta.isLora,
ggufVariant: meta.ggufVariant,
isDiffusion: meta.isDiffusion,
config: meta.config,
})
}
@ -965,6 +969,7 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
id,
isLora: meta.isLora,
ggufVariant: meta.ggufVariant,
isDiffusion: meta.isDiffusion,
config: meta.config,
})
}
@ -2016,6 +2021,9 @@ export function ChatPage({
} = useActiveModelConfig();
const activeModelIsGguf =
runtimeCheckpoint != null && !isExternalModel && runtimeModelIsGguf;
const activeModelIsDiffusion = useChatRuntimeStore(
(s) => s.loadedIsDiffusion,
);
const activeModelIsLora = useMemo(() => {
const checkpoint = inferenceParams.checkpoint;
if (!checkpoint || isExternalModel) return false;
@ -2416,12 +2424,11 @@ export function ChatPage({
const previousConfig = currentRuntimePerModelConfig({
includeMaxSeqLength: true,
});
const hasAppliedConfig = applyModelLoadConfigToRuntime(
selection.config ?? rememberedConfigFor(selection),
);
const loadConfig =
selection.config ?? rememberedConfigFor(selection);
await selectModel({
...selection,
...(hasAppliedConfig ? { keepSpeculative: true } : {}),
...(loadConfig ? { config: loadConfig, keepSpeculative: true } : {}),
previousConfig,
});
},
@ -2758,6 +2765,7 @@ export function ChatPage({
isDownloaded: meta?.isDownloaded || isSameLoadedModel,
expectedBytes: meta?.expectedBytes,
isGguf: meta?.isGguf,
isDiffusion: meta?.isDiffusion,
config: meta?.config,
nativePathToken: meta?.nativePathToken,
nativePathExpiresAtMs: meta?.nativePathExpiresAtMs,
@ -2800,6 +2808,7 @@ export function ChatPage({
nativePathToken: nativeToken ?? undefined,
nativePathExpiresAtMs: nativeExpiry,
isGguf: activeModelIsGguf,
isDiffusion: activeModelIsDiffusion,
isDownloaded: true,
config,
forceReload: true,
@ -2810,6 +2819,7 @@ export function ChatPage({
activeGgufVariant,
activeModelIsLora,
activeModelIsGguf,
activeModelIsDiffusion,
handleCheckpointChange,
],
);
@ -3495,6 +3505,7 @@ export function ChatPage({
modelId={inferenceParams.checkpoint}
ggufVariant={activeGgufVariant ?? null}
isGguf={activeModelIsGguf}
isDiffusion={activeModelIsDiffusion}
nativeContextLength={ggufNativeContextLength}
loadedContextLength={ggufContextLength}
loadedConfig={activeModelConfig}

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

@ -9,6 +9,7 @@ import {
useTransformersUpgradeDialogStore,
} from "@/features/transformers-upgrade";
import { consumeNativePathToken } from "@/features/native-intents/api";
import { prepareHfTokenForUse } from "@/features/hf-auth";
import {
notifyNative,
primeNativeNotificationPermission,
@ -21,6 +22,7 @@ import {
getGgufDownloadProgress,
getInferenceStatus,
getLoadProgress,
fetchGgufStagedMetadata,
listLoras,
listModels,
loadModel,
@ -89,6 +91,8 @@ export type SelectedModelInput = {
/** Direct local .gguf file (no HF variant / native token) — still a GGUF
* source, so the staging flow treats it as one. */
isGguf?: boolean;
/** Staged metadata confirmed the separate DiffusionGemma runner. */
isDiffusion?: boolean;
throwOnError?: boolean;
/** Keep the current speculative-decoding choice across the model switch
* instead of resetting it to the standing preference. */
@ -496,15 +500,23 @@ export function useChatModelRuntime() {
: selection.nativePathExpiresAtMs ?? null;
const explicitIsGguf =
typeof selection === "string" ? undefined : selection.isGguf;
let isDiffusion =
typeof selection === "string" ? undefined : selection.isDiffusion;
const restorePreviousConfig = () => {
if (typeof selection !== "string" && selection.previousConfig) {
applyPerModelConfigToRuntime(selection.previousConfig, {
isDiffusion:
useChatRuntimeStore.getState().loadedIsDiffusion,
});
}
};
const throwOnError =
typeof selection === "string" ? false : selection.throwOnError ?? false;
const keepSpeculative =
typeof selection === "string" ? false : selection.keepSpeculative ?? false;
const currentVariant = useChatRuntimeStore.getState().activeGgufVariant;
if (!forceReload && (!modelId || (params.checkpoint === modelId && (ggufVariant ?? null) === (currentVariant ?? null)))) {
if (typeof selection !== "string" && selection.previousConfig) {
applyPerModelConfigToRuntime(selection.previousConfig);
}
restorePreviousConfig();
return;
}
// A load is already in flight. If it's this exact pick (id + variant + token),
@ -518,9 +530,9 @@ export function useChatModelRuntime() {
loadingModelRef.current ??
useChatRuntimeStore.getState().loadingModelPick;
if (!inFlightLoad) return false;
if (typeof selection !== "string" && selection.previousConfig) {
applyPerModelConfigToRuntime(selection.previousConfig);
}
// The helper form, not an inline apply: it also carries the loaded
// diffusion flag, which the restored config needs to stay correct.
restorePreviousConfig();
const loadingSamePick =
inFlightLoad.id === modelId &&
(inFlightLoad.ggufVariant ?? null) === (ggufVariant ?? null) &&
@ -675,12 +687,38 @@ export function useChatModelRuntime() {
// default), which the resolved baseline cannot express. previousConfig
// is the snapshot the picker took before pre-applying the target's
// config, so the live control is only the outgoing one without it.
// Read before the staged-metadata await below so an in-flight change
// cannot perturb the outgoing snapshot.
const previousNParallel =
typeof selection !== "string" && selection.previousConfig
? (selection.previousConfig.nParallel ?? null)
: useChatRuntimeStore.getState().nParallel;
if (isGguf && isDiffusion === undefined) {
// Prepare the token exactly as validateModel/loadModel do (and as
// the compare path does): the Hub rejects an invalid Authorization
// header with 401 even for a public repo, so sending the raw stored
// token here would abort the load before the existing "continue
// anonymously / replace token" recovery flow could run.
const preparedToken = await prepareHfTokenForUse(
useChatRuntimeStore.getState().hfToken || null,
);
if (!preparedToken.proceed) {
throw new Error("Model load cancelled.");
}
isDiffusion = (
await fetchGgufStagedMetadata({
model_path: modelId,
gguf_variant: ggufVariant ?? null,
hf_token: preparedToken.token,
nativePathToken: nativePathToken ?? null,
})
).isDiffusion;
}
const targetIsDiffusion = isDiffusion === true;
if (pendingLoadConfig) {
applyPerModelConfigToRuntime(pendingLoadConfig);
applyPerModelConfigToRuntime(pendingLoadConfig, {
isDiffusion: targetIsDiffusion,
});
}
const currentCheckpoint =
useChatRuntimeStore.getState().params.checkpoint;
@ -739,8 +777,10 @@ export function useChatModelRuntime() {
pendingLoadConfig?.customContextLength ??
stateBeforeUnload.customContextLength;
const loadGgufContextLength = stateBeforeUnload.ggufContextLength;
const loadTensorParallel =
pendingLoadConfig?.tensorParallel ?? stateBeforeUnload.tensorParallel;
const loadTensorParallel = targetIsDiffusion
? false
: (pendingLoadConfig?.tensorParallel ??
stateBeforeUnload.tensorParallel);
const loadActivePresetSource = stateBeforeUnload.activePresetSource;
const loadActiveGgufVariant = stateBeforeUnload.activeGgufVariant;
const loadGpuMemoryMode =
@ -763,8 +803,16 @@ export function useChatModelRuntime() {
}
let loadSelectedGpuIds =
pendingLoadConfig?.selectedGpuIds !== undefined
? reconcilePersistedGpuIds(pendingLoadConfig.selectedGpuIds)
: reconcilePersistedGpuIds(stateBeforeUnload.selectedGpuIds);
? reconcilePersistedGpuIds(
pendingLoadConfig.selectedGpuIds,
pendingLoadConfig.selectedGpuIndexKind,
targetIsDiffusion,
)
: reconcilePersistedGpuIds(
stateBeforeUnload.selectedGpuIds,
stateBeforeUnload.selectedGpuIndexKind,
targetIsDiffusion,
);
let loadSpeculativeType =
pendingLoadConfig?.speculativeType != null
? normalizeSpeculativeType(pendingLoadConfig.speculativeType)
@ -931,6 +979,7 @@ export function useChatModelRuntime() {
// Per-model GPU knobs must not follow onto a different model
// (gpuMemoryMode is a standing preference and is kept).
selectedGpuIds: null,
selectedGpuIndexKind: null,
gpuLayers: GPU_LAYERS_AUTO,
nCpuMoe: 0,
splitRatio: null,
@ -952,7 +1001,11 @@ export function useChatModelRuntime() {
pendingLoadConfig?.customContextLength ?? null;
loadSelectedGpuIds =
pendingLoadConfig?.selectedGpuIds !== undefined
? reconcilePersistedGpuIds(pendingLoadConfig.selectedGpuIds)
? reconcilePersistedGpuIds(
pendingLoadConfig.selectedGpuIds,
pendingLoadConfig.selectedGpuIndexKind,
targetIsDiffusion,
)
: null;
loadGpuLayers = pendingLoadConfig?.gpuLayers ?? GPU_LAYERS_AUTO;
loadNCpuMoe = pendingLoadConfig?.nCpuMoe ?? 0;
@ -1629,9 +1682,7 @@ export function useChatModelRuntime() {
resetLoadingUi();
}
} catch (error) {
if (typeof selection !== "string" && selection.previousConfig) {
applyPerModelConfigToRuntime(selection.previousConfig);
}
restorePreviousConfig();
if (abortCtrl.signal.aborted) return; // User cancelled, nothing to report
resetLoadingUi();
const message =

1
studio/frontend/src/features/chat/index.ts
View file

@ -50,6 +50,7 @@ export {
readPersistedSpeculativeType,
readPersistedGpuMemoryMode,
reconcilePersistedGpuIds,
reconcilePersistedGpuSelection,
GPU_LAYERS_AUTO,
} from "./stores/chat-runtime-store";
export {

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

@ -24,6 +24,7 @@ import {
isMultimodalResponse,
} from "../types/api";
import type { ChatModelSummary } from "../types/runtime";
import { sameGpuSelection } from "@/hooks/gpu-selection";
type LocalReasoningEffort = Extract<ReasoningEffort, "low" | "medium" | "high">;
@ -244,15 +245,21 @@ export function applyActiveModelStatusToStore(
incomingGpuMode === "manual" ? (status.n_cpu_moe ?? null) : null;
const incomingSplit =
incomingGpuMode === "manual" ? (status.tensor_split ?? null) : null;
const incomingGpuIds = status.is_gguf
? (status.requested_gpu_ids ?? status.gpu_ids ?? null)
: null;
const incomingGpuFields = loadedGpuMemoryFields(status);
const incomingGpuIds = incomingGpuFields.loadedGpuIds;
const incomingGpuIndexKind = incomingGpuFields.loadedGpuIndexKind;
const gpuStatusChanged =
prevState.loadedGpuMemoryMode !== incomingGpuMode ||
prevState.loadedGpuLayers !== incomingGpuLayers ||
prevState.loadedNCpuMoe !== incomingNCpuMoe ||
!sameArray(prevState.loadedSplitRatio, incomingSplit) ||
!sameArray(prevState.loadedGpuIds, incomingGpuIds) ||
!sameGpuSelection(
{
ids: prevState.loadedGpuIds,
indexKind: prevState.loadedGpuIndexKind,
},
{ ids: incomingGpuIds, indexKind: incomingGpuIndexKind },
) ||
prevState.loadedCustomContextLength !== gpuPin;
const gpuMemoryEditsPending =
(prevState.loadedGpuMemoryMode !== null &&
@ -262,11 +269,16 @@ export function applyActiveModelStatusToStore(
prevState.nCpuMoe !== prevState.loadedNCpuMoe ||
!sameArray(prevState.splitRatio, prevState.loadedSplitRatio))) ||
prevState.customContextLength !== prevState.loadedCustomContextLength;
const gpuIdsEditPending = !sameArray(
prevState.selectedGpuIds,
prevState.loadedGpuIds,
const gpuIdsEditPending = !sameGpuSelection(
{
ids: prevState.selectedGpuIds,
indexKind: prevState.selectedGpuIndexKind,
},
{
ids: prevState.loadedGpuIds,
indexKind: prevState.loadedGpuIndexKind,
},
);
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;
@ -283,7 +295,10 @@ export function applyActiveModelStatusToStore(
customContextLength: prevState.customContextLength,
}),
...(preserveSameModelEdits &&
gpuIdsEditPending && { selectedGpuIds: prevState.selectedGpuIds }),
gpuIdsEditPending && {
selectedGpuIds: prevState.selectedGpuIds,
selectedGpuIndexKind: prevState.selectedGpuIndexKind,
}),
};
useChatRuntimeStore.setState({

9
studio/frontend/src/features/chat/presets/preset-load-config.ts
View file

@ -128,7 +128,8 @@ export function normalizePresetLoadConfig(
...(nCpuMoe !== undefined ? { nCpuMoe } : {}),
};
return hasPresetLoadConfig(normalized) ? normalized : undefined;
const coalesced = coalesceDefaultLoadKnobs(normalized);
return hasPresetLoadConfig(coalesced) ? coalesced : undefined;
}
export function hasPresetLoadConfig(
@ -177,9 +178,8 @@ export function capturePresetLoadConfig(): PresetLoadConfig | undefined {
? { nCpuMoe: snapshot.nCpuMoe }
: {}),
};
return hasPresetLoadConfig(coalesceDefaultLoadKnobs(captured))
? coalesceDefaultLoadKnobs(captured)
: undefined;
const coalesced = coalesceDefaultLoadKnobs(captured);
return hasPresetLoadConfig(coalesced) ? coalesced : undefined;
}
function coalesceDefaultLoadKnobs(
@ -226,6 +226,7 @@ export function applyPresetLoadConfig(
gpuLayers: config.gpuLayers,
nCpuMoe: config.nCpuMoe,
selectedGpuIds: store.selectedGpuIds,
selectedGpuIndexKind: store.selectedGpuIndexKind,
});
}

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

@ -87,7 +87,12 @@ import {
confirmTransformersUpgradeIfNeeded,
useTransformersUpgradeDialogStore,
} from "@/features/transformers-upgrade";
import { loadModel, validateModel } from "./api/chat-api";
import { prepareHfTokenForUse } from "@/features/hf-auth";
import {
fetchGgufStagedMetadata,
loadModel,
validateModel,
} from "./api/chat-api";
import { resolveFitMaxSeqLength, resolveManualAutoCtxPin } from "./presets/preset-policy";
import { ensureGpuDeviceCache } from "@/hooks/use-gpu-info";
import {
@ -102,6 +107,7 @@ import {
usePlusMenuPrefsStore,
} from "./stores/plus-menu-prefs-store";
import {
GPU_LAYERS_AUTO,
loadedGpuMemoryFields,
type ReasoningEffort,
reconcilePersistedGpuIds,
@ -423,6 +429,7 @@ type CompareModelSelection = {
id: string;
isLora: boolean;
ggufVariant?: string;
isDiffusion?: boolean;
config?: PerModelConfig;
};
@ -1020,11 +1027,8 @@ export function SharedComposer({
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),
customContextLength: store.customContextLength,
selectedGpuIds: store.selectedGpuIds,
selectedGpuIndexKind: store.selectedGpuIndexKind,
};
// Set when an accepted transformers install unloaded the active model
// server-side; a later failure must then clear the stale checkpoint.
@ -1044,19 +1048,42 @@ export function SharedComposer({
: resolveInitialConfig(sel.id, sel.ggufVariant ?? null);
const ownConfig = resolved.config;
const ownRemembered = resolved.remembered;
const isAlreadyActive =
currentStore.params.checkpoint === sel.id &&
(currentStore.activeGgufVariant ?? null) ===
(sel.ggufVariant ?? null);
if (isAlreadyActive && !config && !loadedFromConfig) {
return "ready";
}
const targetIsGguf =
(sel.ggufVariant ?? null) != null ||
sel.id.toLowerCase().endsWith(".gguf");
let resolvedIsDiffusion = sel.isDiffusion;
if (targetIsGguf && resolvedIsDiffusion === undefined) {
const preparedToken = await prepareHfTokenForUse(
currentStore.hfToken,
);
if (!preparedToken.proceed) {
throw new Error("Model load cancelled.");
}
resolvedIsDiffusion = (
await fetchGgufStagedMetadata({
model_path: sel.id,
gguf_variant: sel.ggufVariant ?? null,
hf_token: preparedToken.token,
})
).isDiffusion;
}
// Mirror single-view resolveLoadMaxSeqLength: a GGUF pane with no explicit
// context loads at native (0 -> n_ctx_train), not the session maxSeqLength,
// which would silently shrink the shown context.
const isGgufLoad =
(sel.ggufVariant ?? null) != null ||
sel.id.toLowerCase().endsWith(".gguf");
// A non-GGUF pane with no saved maxSeqLength falls back to the app default,
// not the active model's shared runtime snapshot: else comparing a saved
// 128K model against an unconfigured one loads the latter at 128K and OOMs.
const effectiveMaxSeqLength =
ownConfig.customContextLength ??
normalizeMaxSeqLength(ownConfig.maxSeqLength) ??
(isGgufLoad ? 0 : DEFAULT_MAX_SEQ_LENGTH);
(targetIsGguf ? 0 : DEFAULT_MAX_SEQ_LENGTH);
const effectiveChatTemplateOverride = cleanCompareChatTemplate(
ownConfig.chatTemplateOverride,
);
@ -1068,22 +1095,38 @@ export function SharedComposer({
ownConfig.specDraftNMax,
)
: specSettings.specDraftNMax;
const effectiveTensorParallel = ownRemembered
? ownConfig.tensorParallel
: fallbackTensorParallel;
const effectiveTensorParallel = resolvedIsDiffusion
? false
: ownRemembered
? ownConfig.tensorParallel
: fallbackTensorParallel;
if (ownConfig.selectedGpuIds != null) {
await ensureGpuDeviceCache();
}
const effectiveGpuMemoryMode =
ownConfig.gpuMemoryMode ?? compareLoadKnobs.gpuMemoryMode;
resolvedIsDiffusion
? "auto"
: (ownConfig.gpuMemoryMode ?? compareLoadKnobs.gpuMemoryMode);
const effectiveGpuLayers =
ownConfig.gpuLayers ?? compareLoadKnobs.gpuLayers;
resolvedIsDiffusion
? GPU_LAYERS_AUTO
: (ownConfig.gpuLayers ?? compareLoadKnobs.gpuLayers);
const effectiveNCpuMoe =
ownConfig.nCpuMoe ?? compareLoadKnobs.nCpuMoe;
resolvedIsDiffusion
? 0
: (ownConfig.nCpuMoe ?? compareLoadKnobs.nCpuMoe);
const effectiveSelectedGpuIds =
ownConfig.selectedGpuIds !== undefined
? reconcilePersistedGpuIds(ownConfig.selectedGpuIds)
: compareLoadKnobs.selectedGpuIds;
? reconcilePersistedGpuIds(
ownConfig.selectedGpuIds,
ownConfig.selectedGpuIndexKind,
resolvedIsDiffusion === true,
)
: reconcilePersistedGpuIds(
compareLoadKnobs.selectedGpuIds,
compareLoadKnobs.selectedGpuIndexKind,
resolvedIsDiffusion === true,
);
// A pane's context comes from its own config only: a saved pin, or null
// (Auto/native). It must not inherit the active model's shared snapshot --
// resolveFitMaxSeqLength would treat that as a pin and load this pane at
@ -1091,15 +1134,6 @@ export function SharedComposer({
const effectiveCustomContextLength = ownConfig.customContextLength;
let loadTrustRemoteCode = trustRemoteCode;
let approvedRemoteCodeFingerprint: string | null = null;
const isAlreadyActive =
currentStore.params.checkpoint === sel.id &&
(currentStore.activeGgufVariant ?? null) ===
(sel.ggufVariant ?? null);
if (isAlreadyActive && !config && !loadedFromConfig) {
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).
@ -1268,7 +1302,7 @@ export function SharedComposer({
// Record the context this pane loaded with (like the single-model path)
// so when it becomes the active model, the UI and later reload/save use
// its context, not the previous/default one.
customContextLength: isGgufLoad
customContextLength: targetIsGguf
? (ownConfig.customContextLength ?? keepCustomCtx)
: null,
ggufContextLength: resp.is_gguf ? (resp.context_length ?? null) : null,
@ -1285,7 +1319,7 @@ export function SharedComposer({
activeNativePathExpiresAtMs: null,
...resolveLoadedSpeculativeSettings(resp),
});
if (!isGgufLoad) {
if (!targetIsGguf) {
// Non-GGUF panes carry their context in params.maxSeqLength.
store.setParams({
...useChatRuntimeStore.getState().params,

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

@ -2,7 +2,12 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { mirrorHfTokenInto, useHfTokenStore } from "@/features/hub";
import { cachedPinnableGpuIndices } from "@/hooks/use-gpu-info";
import {
cachedPinnableGpuIndexKind,
reconcileCachedGpuSelection,
type ReconciledGpuSelection,
type GpuIndexKind,
} from "@/hooks/use-gpu-info";
import { toast } from "@/lib/toast";
import { create } from "zustand";
import { isExternalModelId, parseExternalModelId } from "../external-providers";
@ -628,14 +633,35 @@ export function rebalanceSplit(
// 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).
// An explicit null namespace means discovery had not completed when the live
// state was captured, while an absent namespace is a legacy physical-ID pick.
export function reconcilePersistedGpuIds(
ids: number[] | null,
savedIndexKind?: GpuIndexKind | null,
forDiffusion = false,
): 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;
return reconcilePersistedGpuSelection(
ids,
savedIndexKind,
forDiffusion,
).ids;
}
export function reconcilePersistedGpuSelection(
ids: number[] | null,
savedIndexKind?: GpuIndexKind | null,
forDiffusion = false,
): ReconciledGpuSelection {
return reconcileCachedGpuSelection(ids, savedIndexKind, forDiffusion);
}
export function requestedGpuIdsFromResponse(resp: {
gpu_ids?: number[] | null;
requested_gpu_ids?: number[] | null;
}): number[] | null {
return Object.prototype.hasOwnProperty.call(resp, "requested_gpu_ids")
? (resp.requested_gpu_ids ?? null)
: (resp.gpu_ids ?? null);
}
// Store fields derived from a load/status response's GPU-memory settings.
@ -664,7 +690,9 @@ export function loadedGpuMemoryFields(resp: {
// baseline clears to null so Reset preserves the preference, not a stale mode.
return {
selectedGpuIds: null,
selectedGpuIndexKind: null,
loadedGpuIds: null,
loadedGpuIndexKind: null,
loadedGpuMemoryMode: null,
gpuLayers: GPU_LAYERS_AUTO,
loadedGpuLayers: null,
@ -679,7 +707,20 @@ export function loadedGpuMemoryFields(resp: {
const mode = resp.gpu_memory_mode ?? "auto";
// Keep the user's placement pool editable across status/load hydration.
// gpu_ids remains the effective fitted subset for diagnostics.
const gpuIds = resp.requested_gpu_ids ?? resp.gpu_ids ?? null;
const reportedGpuIds = requestedGpuIdsFromResponse(resp);
const gpuIndexKind =
reportedGpuIds == null
? null
: cachedPinnableGpuIndexKind(resp.is_diffusion === true);
// A numeric ID is unsafe to adopt or persist once discovery says it is NOT a
// physical CUDA/ROCm ID or a Vulkan ordinal (gpuIndexKind === null, cache
// warm). But while discovery is still cold (gpuIndexKind === undefined) the
// namespace is merely deferred, not rejected -- keep the just-applied pin so
// a reload/rollback in that window doesn't omit gpu_ids and let llama.cpp
// fall back to every device. A later status refresh resolves the namespace
// once the shared system cache warms.
const gpuIds =
reportedGpuIds != null && gpuIndexKind !== null ? reportedGpuIds : 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
@ -719,7 +760,12 @@ export function loadedGpuMemoryFields(resp: {
moeLayerCount: resp.n_moe_layers ?? null,
// The picker reflects the requested placement pool, not a fitted subset.
selectedGpuIds: gpuIds,
// gpuIndexKind is `undefined` only in the deferred (cache-cold) case here,
// since a `null` kind already forced gpuIds to null above. Normalize to the
// explicit-null-namespace convention (discovery not complete yet).
selectedGpuIndexKind: gpuIds == null ? null : (gpuIndexKind ?? null),
loadedGpuIds: gpuIds,
loadedGpuIndexKind: gpuIds == null ? null : (gpuIndexKind ?? null),
...manualKnobs,
};
}
@ -1001,9 +1047,13 @@ type ChatRuntimeStore = {
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). */
/** Picked IDs in the backend-declared GPU namespace (null = automatic). */
selectedGpuIds: number[] | null;
/** Namespace used by selectedGpuIds; kept with deferred persisted picks. */
selectedGpuIndexKind: GpuIndexKind | null;
loadedGpuIds: number[] | null;
/** Backend-reported namespace paired with loadedGpuIds. */
loadedGpuIndexKind: GpuIndexKind | null;
/** Persisted: expand every On Device GGUF repo's quantizations by default
* instead of waiting for a click. */
expandQuantizations: boolean;
@ -1183,7 +1233,10 @@ type ChatRuntimeStore = {
setGpuLayers: (value: number) => void;
setNCpuMoe: (value: number) => void;
setSplitRatio: (value: number[] | null) => void;
setSelectedGpuIds: (ids: number[] | null) => void;
setSelectedGpuIds: (
ids: number[] | null,
indexKind?: GpuIndexKind | null,
) => void;
setExpandQuantizations: (value: boolean) => void;
setShowAllQuantizations: (value: boolean) => void;
setFitOnDeviceOnly: (value: boolean) => void;
@ -1512,7 +1565,9 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
ggufLayerCount: null,
moeLayerCount: null,
selectedGpuIds: null,
selectedGpuIndexKind: null,
loadedGpuIds: null,
loadedGpuIndexKind: null,
expandQuantizations: loadBool(CHAT_EXPAND_QUANTIZATIONS_KEY, false),
showAllQuantizations: loadBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, true),
fitOnDeviceOnly: loadBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, false),
@ -1898,7 +1953,9 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
ggufLayerCount: null,
moeLayerCount: null,
selectedGpuIds: null,
selectedGpuIndexKind: null,
loadedGpuIds: null,
loadedGpuIndexKind: null,
loadedIsMultimodal: false,
loadedIsDiffusion: false,
customContextLength: null,
@ -2292,7 +2349,12 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
setGpuLayers: (gpuLayers) => set({ gpuLayers }),
setNCpuMoe: (nCpuMoe) => set({ nCpuMoe }),
setSplitRatio: (splitRatio) => set({ splitRatio }),
setSelectedGpuIds: (selectedGpuIds) => set({ selectedGpuIds }),
setSelectedGpuIds: (selectedGpuIds, selectedGpuIndexKind = null) =>
set({
selectedGpuIds,
selectedGpuIndexKind:
selectedGpuIds == null ? null : selectedGpuIndexKind,
}),
setExpandQuantizations: (expandQuantizations) => {
saveBool(CHAT_EXPAND_QUANTIZATIONS_KEY, expandQuantizations);
set({ expandQuantizations });

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

@ -85,7 +85,7 @@ export interface LoadModelRequest {
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). */
/** Picked CUDA/ROCm physical IDs or Vulkan ordinals (omit/empty = automatic). */
gpu_ids?: number[];
}
@ -95,6 +95,7 @@ export interface ValidateModelResponse {
identifier?: string | null;
display_name?: string | null;
is_gguf?: boolean;
is_diffusion?: boolean;
is_lora?: boolean;
is_vision?: boolean;
requires_trust_remote_code?: boolean;

11
studio/frontend/src/features/hub/catalog/sampling-settings-dialog.tsx
View file

@ -103,6 +103,9 @@ export function SamplingSettingsButton({ className }: { className?: string }) {
const { selectModel } = useChatModelRuntime();
const modelLoading = useChatRuntimeStore((s) => s.modelLoading);
const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant);
const activeModelIsDiffusion = useChatRuntimeStore(
(s) => s.loadedIsDiffusion,
);
const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
const ggufNativeContextLength = useChatRuntimeStore(
(s) => s.ggufNativeContextLength,
@ -131,7 +134,9 @@ export function SamplingSettingsButton({ className }: { className?: string }) {
const previousConfig = currentRuntimePerModelConfig({
includeMaxSeqLength: true,
});
applyPerModelConfigToRuntime(config);
applyPerModelConfigToRuntime(config, {
isDiffusion: activeModelIsDiffusion,
});
void selectModel({
id: activeCheckpoint,
source: "local",
@ -139,13 +144,14 @@ export function SamplingSettingsButton({ className }: { className?: string }) {
nativePathToken: nativeToken ?? undefined,
nativePathExpiresAtMs: nativeExpiry,
isGguf: activeModelIsGguf,
isDiffusion: activeModelIsDiffusion,
isDownloaded: true,
keepSpeculative: true,
previousConfig,
forceReload: true,
});
},
[selectModel, activeModelIsGguf],
[selectModel, activeModelIsGguf, activeModelIsDiffusion],
);
// Reasoning + tools: same store bindings as the chat page.
@ -235,6 +241,7 @@ export function SamplingSettingsButton({ className }: { className?: string }) {
modelId={checkpoint}
ggufVariant={activeGgufVariant ?? null}
isGguf={activeModelIsGguf}
isDiffusion={activeModelIsDiffusion}
nativeContextLength={ggufNativeContextLength}
loadedContextLength={ggufContextLength}
loadedConfig={activeModelConfig}

252
studio/frontend/src/features/model-picker/components/model-config-page.tsx
View file

@ -19,7 +19,15 @@ import {
readPersistedSpeculativeType,
useChatRuntimeStore,
} from "@/features/chat";
import { useGpuDevices } from "@/hooks/use-gpu-info";
import { prepareHfTokenForUse } from "@/features/hf-auth";
import {
type GpuIndexKind,
type SystemGpuDevice,
cachedPinnableGpuContext,
pinnableGpuContext,
reconcileGpuSelection,
useGpuDevices,
} from "@/hooks/use-gpu-info";
import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
import { toast } from "@/lib/toast";
import { ArrowLeft01Icon } from "@hugeicons/core-free-icons";
@ -75,14 +83,16 @@ const SELECT_TRIGGER_CLASS = `grid h-8 min-w-0 grid-cols-[minmax(0,1fr)_auto] it
const NUMBER_INPUT_CLASS = `h-8 w-[92px] ${CONTROL_SURFACE} pl-3 pr-2 py-0 text-right text-ui-13 font-medium text-nav-fg outline-none focus-visible:ring-0`;
const KV_CACHE_DTYPE_DEFAULT = "f16";
const SPECULATIVE_TYPE_LABELS: Record<(typeof SPECULATIVE_TYPES)[number], string> =
{
auto: "Auto",
mtp: "MTP",
ngram: "Ngram",
"mtp+ngram": "MTP+Ngram",
off: "Off",
};
const SPECULATIVE_TYPE_LABELS: Record<
(typeof SPECULATIVE_TYPES)[number],
string
> = {
auto: "Auto",
mtp: "MTP",
ngram: "Ngram",
"mtp+ngram": "MTP+Ngram",
off: "Off",
};
function hasNonDefaultAdvanced(config: PerModelConfig): boolean {
return (
@ -99,6 +109,65 @@ function hasNonDefaultAdvanced(config: PerModelConfig): boolean {
);
}
function withoutUnsupportedDiffusionSettings(
config: PerModelConfig,
currentGpuIndexKind: GpuIndexKind | null = null,
): PerModelConfig {
const hasUnsupportedGpuPick =
config.selectedGpuIds != null &&
(config.selectedGpuIndexKind === "vulkan" ||
currentGpuIndexKind === "vulkan");
if (
(config.gpuMemoryMode ?? "auto") === "auto" &&
config.gpuLayers == null &&
config.nCpuMoe == null &&
!config.tensorParallel &&
!hasUnsupportedGpuPick
) {
return config;
}
return {
...config,
gpuMemoryMode: "auto",
gpuLayers: undefined,
nCpuMoe: undefined,
tensorParallel: false,
...(hasUnsupportedGpuPick
? {
selectedGpuIds: undefined,
selectedGpuIndexKind: undefined,
}
: {}),
};
}
function reconcileConfigGpuSelection(
config: PerModelConfig,
isDiffusion: boolean,
gpuDevices?: SystemGpuDevice[],
): PerModelConfig {
const context = cachedPinnableGpuContext(isDiffusion, gpuDevices);
const supported = isDiffusion
? withoutUnsupportedDiffusionSettings(config, context.indexKind ?? null)
: config;
if (supported.selectedGpuIds == null) {
return supported;
}
const reconciled = reconcileGpuSelection(
supported.selectedGpuIds,
supported.selectedGpuIndexKind,
context.indexKind,
context.ids,
);
const next = {
...supported,
selectedGpuIds: reconciled.ids ?? undefined,
selectedGpuIndexKind:
reconciled.ids === null ? undefined : reconciled.indexKind,
};
return perModelConfigsEqual(next, supported) ? supported : next;
}
function ChatTemplateSetting({
config,
onEditTemplate,
@ -250,6 +319,8 @@ function GpuMemorySettings({
update,
layerCount,
moeLayerCount,
isDiffusion,
gpuDevices,
gpuLayersInputRef,
moeLayersInputRef,
}: {
@ -257,10 +328,11 @@ function GpuMemorySettings({
update: (patch: Partial<PerModelConfig>) => void;
layerCount: number | null;
moeLayerCount: number | null;
isDiffusion: boolean;
gpuDevices: SystemGpuDevice[];
gpuLayersInputRef?: Ref<NumericValueInputHandle>;
moeLayersInputRef?: Ref<NumericValueInputHandle>;
}) {
const gpuDevices = useGpuDevices();
const mode = config.gpuMemoryMode ?? "auto";
const isManual = mode === "manual";
const gpuLayers = config.gpuLayers ?? GPU_LAYERS_AUTO;
@ -273,26 +345,29 @@ function GpuMemorySettings({
const moeLayersMax = moeLayerCount ?? 0;
const showMoeSlider = isManual && !autoLayers && moeLayersMax > 0;
const selectedGpuIds = config.selectedGpuIds ?? null;
const singleGpuInUse =
(selectedGpuIds ?? gpuDevices.map((device) => device.index)).length <= 1;
// Multi-GPU only, and only with physical indices (relative ordinals from a
// CUDA_VISIBLE_DEVICES mask can't be mapped back to pin a device). null = all (auto).
const showGpuPicker =
gpuDevices.length > 1 && gpuDevices.every((d) => d.physicalIndex);
const gpuContext = pinnableGpuContext(gpuDevices, isDiffusion);
const pinnableDevices = gpuContext.devices ?? [];
const gpuIndexKind = gpuContext.indexKind ?? null;
const singleGpuInUse = (selectedGpuIds ?? gpuContext.ids ?? []).length <= 1;
// Multi-GPU only, with one backend-declared index namespace. null = automatic.
const showGpuPicker = (gpuContext.ids?.length ?? 0) > 1;
const isGpuChecked = (index: number) =>
selectedGpuIds === null || selectedGpuIds.includes(index);
const toggleGpu = (index: number) => {
const all = gpuDevices.map((d) => d.index);
const all = gpuContext.ids ?? [];
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
update({ selectedGpuIds: next.length === all.length ? null : next });
update({
selectedGpuIds: next,
selectedGpuIndexKind: gpuIndexKind,
});
};
return (
<>
<div className={ROW_CLASS}>
<div className={isDiffusion ? "hidden" : ROW_CLASS}>
<div className="flex min-w-0 items-center gap-1.5">
<span className={LABEL_CLASS}>GPU Memory</span>
<InfoHint>
@ -312,9 +387,7 @@ function GpuMemorySettings({
<Select
value={mode}
onValueChange={(v) =>
// Returning to Default must clear the Manual-only knobs, else a
// remembered config keeps stale gpuLayers/nCpuMoe/GPU pick that a
// later load re-applies while the page shows Default.
// Default clears every Manual-only setting.
update(
v === "manual"
? { gpuMemoryMode: "manual" }
@ -323,6 +396,7 @@ function GpuMemorySettings({
gpuLayers: undefined,
nCpuMoe: undefined,
selectedGpuIds: undefined,
selectedGpuIndexKind: undefined,
},
)
}
@ -341,7 +415,7 @@ function GpuMemorySettings({
</SelectContent>
</Select>
</div>
{isManual && (
{!isDiffusion && isManual && (
<>
<AdvancedGpuSlider
label="GPU Layers"
@ -354,8 +428,8 @@ function GpuMemorySettings({
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.
Auto lets llama.cpp size the split (and the context) to fit
VRAM. At the maximum, the whole model is on the GPU.
</>
}
/>
@ -383,14 +457,13 @@ function GpuMemorySettings({
<div className="flex min-w-0 items-center gap-1.5">
<span className={LABEL_CLASS}>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. At least one GPU must stay
selected.
By default, Unsloth chooses GPUs automatically. Editing this list
makes the checked GPUs the explicit candidate pool. At least one
GPU must stay selected.
</InfoHint>
</div>
<div className="flex flex-col gap-2">
{gpuDevices.map((d) => (
{pinnableDevices.map((d) => (
<div
key={d.index}
className="flex items-center justify-between gap-3"
@ -424,6 +497,8 @@ function GgufAdvancedSettings({
onEditTemplate,
layerCount,
moeLayerCount,
isDiffusion,
gpuDevices,
gpuLayersInputRef,
moeLayersInputRef,
}: {
@ -434,6 +509,8 @@ function GgufAdvancedSettings({
onEditTemplate: () => void;
layerCount: number | null;
moeLayerCount: number | null;
isDiffusion: boolean;
gpuDevices: SystemGpuDevice[];
gpuLayersInputRef?: Ref<NumericValueInputHandle>;
moeLayersInputRef?: Ref<NumericValueInputHandle>;
}) {
@ -602,6 +679,8 @@ function GgufAdvancedSettings({
update={update}
layerCount={layerCount}
moeLayerCount={moeLayerCount}
isDiffusion={isDiffusion}
gpuDevices={gpuDevices}
gpuLayersInputRef={gpuLayersInputRef}
moeLayersInputRef={moeLayersInputRef}
/>
@ -614,10 +693,11 @@ function GgufAdvancedSettings({
interface ModelConfigPageProps {
target: ModelPickTarget;
onBack?: () => void;
onRun: (config: PerModelConfig) => void;
onRun: (config: PerModelConfig, isDiffusion?: boolean) => void;
loadedConfig?: PerModelConfig | null;
loadedContextLength?: number | null;
initialConfig?: PerModelConfig | null;
isDiffusion?: boolean;
variant?: "page" | "sidebar";
}
@ -628,6 +708,7 @@ export function ModelConfigPage({
loadedConfig = null,
loadedContextLength = null,
initialConfig = null,
isDiffusion = false,
variant = "page",
}: ModelConfigPageProps) {
const rememberId = useId();
@ -642,6 +723,7 @@ export function ModelConfigPage({
const loadedMaxContextLength = useChatRuntimeStore(
(s) => s.ggufMaxContextLength,
);
const gpuDevices = useGpuDevices();
const resolveInitial = () => {
const resolved = resolveInitialConfig(target.id, target.ggufVariant);
if (loadedConfig) {
@ -658,7 +740,9 @@ export function ModelConfigPage({
return resolved;
};
const [initial] = useState(resolveInitial);
const [config, setConfig] = useState<PerModelConfig>(() => initial.config);
const [config, setConfig] = useState<PerModelConfig>(() =>
reconcileConfigGpuSelection(initial.config, isDiffusion, gpuDevices),
);
const [remember, setRemember] = useState(() => initial.remembered);
const [savedRemember, setSavedRemember] = useState(() => initial.remembered);
const [speculativeFallback] = useState(readPersistedSpeculativeType);
@ -705,33 +789,45 @@ export function ModelConfigPage({
contextLength: number | null;
layerCount: number | null;
moeLayerCount: number | null;
isDiffusion?: boolean;
} | null>(null);
useEffect(() => {
if (contextFetchKey == null) {
return;
}
let cancelled = false;
void fetchGgufStagedMetadata({
model_path: target.id,
gguf_variant: target.ggufVariant ?? null,
hf_token: hfToken || null,
nativePathToken,
})
.then((dims) => {
if (!cancelled) {
setFetchedStagedDims({ key: contextFetchKey, ...dims });
}
})
.catch(() => {
if (!cancelled) {
setFetchedStagedDims({
key: contextFetchKey,
contextLength: null,
layerCount: null,
moeLayerCount: null,
});
}
const settleWithoutMetadata = () => {
if (!cancelled) {
setFetchedStagedDims({
key: contextFetchKey,
contextLength: null,
layerCount: null,
moeLayerCount: null,
isDiffusion: undefined,
});
}
};
void (async () => {
const preparedToken = await prepareHfTokenForUse(hfToken || null);
if (cancelled) {
return;
}
if (!preparedToken.proceed) {
settleWithoutMetadata();
return;
}
const dims = await fetchGgufStagedMetadata({
model_path: target.id,
gguf_variant: target.ggufVariant ?? null,
hf_token: preparedToken.token,
nativePathToken,
});
if (!cancelled) {
setFetchedStagedDims({ key: contextFetchKey, ...dims });
}
})().catch(() => {
settleWithoutMetadata();
});
return () => {
cancelled = true;
};
@ -744,6 +840,19 @@ export function ModelConfigPage({
]);
const stagedDims =
fetchedStagedDims?.key === contextFetchKey ? fetchedStagedDims : null;
const stagedMetadataPending =
contextFetchKey != null &&
stagedDims == null &&
(config.gpuMemoryMode === "manual" || config.selectedGpuIds != null);
const classifiedIsDiffusion = isDiffusion ? true : stagedDims?.isDiffusion;
const resolvedIsDiffusion = classifiedIsDiffusion === true;
const gpuIndexKind =
pinnableGpuContext(gpuDevices, resolvedIsDiffusion).indexKind ?? null;
useEffect(() => {
setConfig((current) =>
reconcileConfigGpuSelection(current, resolvedIsDiffusion, gpuDevices),
);
}, [gpuDevices, resolvedIsDiffusion]);
const isMtp =
config.speculativeType != null &&
@ -771,9 +880,11 @@ export function ModelConfigPage({
),
maxContext,
);
const setContextLength = (v: number) =>
update({ customContextLength: v });
const baseline = loadedConfig ?? DEFAULT_PER_MODEL_CONFIG;
const setContextLength = (v: number) => update({ customContextLength: v });
const rawBaseline = loadedConfig ?? DEFAULT_PER_MODEL_CONFIG;
const baseline = resolvedIsDiffusion
? withoutUnsupportedDiffusionSettings(rawBaseline, gpuIndexKind)
: rawBaseline;
const atBaseline = perModelConfigsEqual(config, baseline);
// An explicit customContextLength equal to the native ceiling is still an
// override (Reset stays enabled). "At default" means no override at all AND the
@ -802,21 +913,24 @@ export function ModelConfigPage({
// customContextLength stays null. If the user fixes GPU Layers (Manual) and
// remembers, pin that shown context so a later fresh load keeps the fitted
// placement instead of sending native/0 for fixed layers and recreating the OOM.
const loadableConfig = resolvedIsDiffusion
? withoutUnsupportedDiffusionSettings(config, gpuIndexKind)
: config;
const pinFixedLayerContext =
target.isGguf &&
config.gpuMemoryMode === "manual" &&
config.gpuLayers != null &&
config.gpuLayers >= 0 &&
config.customContextLength == null &&
loadableConfig.gpuMemoryMode === "manual" &&
loadableConfig.gpuLayers != null &&
loadableConfig.gpuLayers >= 0 &&
loadableConfig.customContextLength == null &&
activeLoadedContext != null;
// Persisted record: keep config as-is (non-GGUF keeps maxSeqLength null) so
// isDefaultConfig recognises it and clears a remembered override instead of
// pinning the app default.
const runtimeConfig = target.isGguf
? pinFixedLayerContext
? { ...config, customContextLength: activeLoadedContext }
: config
: config;
? { ...loadableConfig, customContextLength: activeLoadedContext }
: loadableConfig
: loadableConfig;
const rememberChanged = remember !== savedRemember;
const persistenceOnly = isActiveModel && atBaseline && rememberChanged;
const primaryActionLabel = persistenceOnly
@ -868,9 +982,12 @@ export function ModelConfigPage({
committedGpuLayers != null ||
committedMoeLayers != null;
const effectiveConfig = hasPending
const committedConfig = hasPending
? { ...config, ...pendingPatch }
: config;
const effectiveConfig = resolvedIsDiffusion
? withoutUnsupportedDiffusionSettings(committedConfig, gpuIndexKind)
: committedConfig;
// pinFixedLayerContext above was computed from the render-time config, before
// the same-click GPU Layers draft was committed. Recompute it from
// effectiveConfig so committing a positive fixed-layer value still pins the
@ -934,7 +1051,7 @@ export function ModelConfigPage({
const effectiveLoadConfig = target.isGguf
? effectiveRuntimeConfig
: { ...effectiveRuntimeConfig, maxSeqLength: effectiveMaxSeqLengthValue };
onRun(effectiveLoadConfig);
onRun(effectiveLoadConfig, classifiedIsDiffusion);
};
return (
@ -1030,6 +1147,8 @@ export function ModelConfigPage({
onEditTemplate={() => setTemplateOpen(true)}
layerCount={stagedDims?.layerCount ?? null}
moeLayerCount={stagedDims?.moeLayerCount ?? null}
isDiffusion={resolvedIsDiffusion}
gpuDevices={gpuDevices}
gpuLayersInputRef={gpuLayersInputRef}
moeLayersInputRef={moeLayersInputRef}
/>
@ -1117,7 +1236,10 @@ export function ModelConfigPage({
type="button"
size="sm"
className="h-8"
disabled={isActiveModel && atBaseline && !rememberChanged}
disabled={
stagedMetadataPending ||
(isActiveModel && atBaseline && !rememberChanged)
}
onClick={handleRun}
>
{primaryActionLabel}

3
studio/frontend/src/features/model-picker/components/model-selector.tsx
View file

@ -535,10 +535,11 @@ function ModelSelectorContent({
key={`${visibleConfigTarget.id}::${visibleConfigTarget.ggufVariant ?? ""}`}
target={visibleConfigTarget}
onBack={() => setConfigTarget(null)}
onRun={(config) =>
onRun={(config, isDiffusion) =>
onSelect(visibleConfigTarget.id, {
...visibleConfigTarget.meta,
config,
isDiffusion,
forceReload: true,
})
}

2
studio/frontend/src/features/model-picker/components/model-selector/types.ts
View file

@ -37,6 +37,8 @@ export interface ModelSelectorChangeMeta {
/** Direct local .gguf file picked without a variant (custom folder / LM
* Studio). Marks it as a GGUF source for the deferred-load staging flow. */
isGguf?: boolean;
/** Staged metadata confirmed the separate DiffusionGemma runner. */
isDiffusion?: boolean;
config?: PerModelConfig;
forceReload?: boolean;
/** Native path token so an active-model reload can reopen a file-picked GGUF. */

3
studio/frontend/src/features/model-picker/components/sidebar-model-config.tsx
View file

@ -11,6 +11,7 @@ interface SidebarModelConfigProps {
modelId: string;
ggufVariant: string | null;
isGguf: boolean;
isDiffusion: boolean;
nativeContextLength: number | null;
loadedContextLength: number | null;
loadedConfig: PerModelConfig;
@ -56,6 +57,7 @@ export function SidebarModelConfig({
modelId,
ggufVariant,
isGguf,
isDiffusion,
nativeContextLength,
loadedContextLength,
loadedConfig,
@ -87,6 +89,7 @@ export function SidebarModelConfig({
loadedConfig={loadedConfig}
loadedContextLength={loadedContextLength}
variant="sidebar"
isDiffusion={isDiffusion}
/>
);
}

5
studio/frontend/src/features/model-picker/hooks/use-active-model-config.ts
View file

@ -29,6 +29,9 @@ export function useActiveModelConfig(): ActiveModelConfigState {
const gpuLayers = useChatRuntimeStore((s) => s.gpuLayers);
const nCpuMoe = useChatRuntimeStore((s) => s.nCpuMoe);
const selectedGpuIds = useChatRuntimeStore((s) => s.selectedGpuIds);
const selectedGpuIndexKind = useChatRuntimeStore(
(s) => s.selectedGpuIndexKind,
);
const isGguf =
activeGgufVariant != null ||
@ -58,6 +61,7 @@ export function useActiveModelConfig(): ActiveModelConfigState {
gpuLayers,
nCpuMoe,
selectedGpuIds,
selectedGpuIndexKind,
};
}, [
checkpoint,
@ -74,6 +78,7 @@ export function useActiveModelConfig(): ActiveModelConfigState {
gpuLayers,
nCpuMoe,
selectedGpuIds,
selectedGpuIndexKind,
]);
return { checkpoint, isGguf, config };

54
studio/frontend/src/features/model-picker/model-config/apply-per-model-config.ts
View file

@ -7,7 +7,7 @@ import {
normalizeSpeculativeType,
readPersistedGpuMemoryMode,
readPersistedSpeculativeType,
reconcilePersistedGpuIds,
reconcilePersistedGpuSelection,
useChatRuntimeStore,
} from "@/features/chat";
import {
@ -20,7 +20,10 @@ function cleanTemplate(value: string | null | undefined): string | null {
return value?.trim() ? value : null;
}
export function applyPerModelConfigToRuntime(config: PerModelConfig): void {
export function applyPerModelConfigToRuntime(
config: PerModelConfig,
options: { isDiffusion?: boolean } = {},
): void {
// Fall back to the standing default when the model has no saved
// maxSeqLength. maxSeqLength is the only per-model field carried on
// params (the rest are reset below), so without this a model with no
@ -32,6 +35,14 @@ export function applyPerModelConfigToRuntime(config: PerModelConfig): void {
if (maxSeqLength !== store.params.maxSeqLength) {
store.setParams({ ...store.params, maxSeqLength });
}
const gpuSelection =
config.selectedGpuIds !== undefined
? reconcilePersistedGpuSelection(
config.selectedGpuIds,
config.selectedGpuIndexKind,
options.isDiffusion,
)
: { ids: null, indexKind: null };
useChatRuntimeStore.setState({
customContextLength: config.customContextLength ?? null,
kvCacheDtype: config.kvCacheDtype ?? null,
@ -40,29 +51,38 @@ export function applyPerModelConfigToRuntime(config: PerModelConfig): void {
readPersistedSpeculativeType(),
specDraftNMax: config.specDraftNMax ?? null,
nParallel: config.nParallel ?? null,
tensorParallel: config.tensorParallel ?? false,
tensorParallel: options.isDiffusion
? false
: (config.tensorParallel ?? false),
chatTemplateOverride: cleanTemplate(config.chatTemplateOverride),
// GPU Memory knobs are per-model (GGUF-only). Absent = defaults; the mode is
// a standing preference so an absent mode falls back to the persisted one.
// The per-GPU split ratio is never remembered, so it always resets. The GPU
// pick is reconciled against the GPUs present now (a saved [1] on a 1-GPU
// host would otherwise be sent and rejected).
gpuMemoryMode: config.gpuMemoryMode ?? readPersistedGpuMemoryMode(),
// A diffusion config is sanitized to gpuMemoryMode "auto" because the mode
// does not apply to it, not because the user chose Auto. Writing that into
// the live standing preference would strand the session on Auto: the load
// itself skips saveGpuMemoryMode for diffusion, so nothing restores it, and
// the next ordinary GGUF loaded without its own config sends this value and
// persists it over the user's Manual. Keep the standing choice instead.
gpuMemoryMode: options.isDiffusion
? readPersistedGpuMemoryMode()
: (config.gpuMemoryMode ?? readPersistedGpuMemoryMode()),
gpuLayers: config.gpuLayers ?? GPU_LAYERS_AUTO,
nCpuMoe: config.nCpuMoe ?? 0,
splitRatio: null,
selectedGpuIds:
config.selectedGpuIds !== undefined
? reconcilePersistedGpuIds(config.selectedGpuIds)
: null,
selectedGpuIds: gpuSelection.ids,
selectedGpuIndexKind: gpuSelection.indexKind,
});
}
export function applyModelLoadConfigToRuntime(
config: PerModelConfig | null | undefined,
options: { isDiffusion?: boolean } = {},
): boolean {
const hasConfig = config != null;
applyPerModelConfigToRuntime(config ?? DEFAULT_PER_MODEL_CONFIG);
applyPerModelConfigToRuntime(config ?? DEFAULT_PER_MODEL_CONFIG, options);
return hasConfig;
}
@ -88,6 +108,7 @@ export function currentRuntimePerModelConfig(
gpuLayers: s.gpuLayers,
nCpuMoe: s.nCpuMoe,
selectedGpuIds: s.selectedGpuIds,
selectedGpuIndexKind: s.selectedGpuIndexKind,
};
}
@ -113,15 +134,22 @@ export function perModelConfigsEqual(
// Serialize the per-model GPU knobs with the same "absent == default"
// coalescing the store applies: mode auto/absent, gpuLayers Auto (< 0) /
// absent, nCpuMoe 0 / absent, and the GPU pick (null / absent = all GPUs).
// absent, nCpuMoe 0 / absent, and null / absent GPU picks as automatic.
export function gpuFieldsSignature(config: PerModelConfig): string {
const gpuSelection =
config.selectedGpuIds == null
? "automatic"
: [
[...config.selectedGpuIds].sort((a, b) => a - b).join(","),
config.selectedGpuIndexKind === undefined
? "physical"
: (config.selectedGpuIndexKind ?? "deferred"),
].join("@");
return [
config.gpuMemoryMode ?? "auto",
config.gpuLayers == null || config.gpuLayers < 0 ? -1 : config.gpuLayers,
config.nCpuMoe ?? 0,
config.selectedGpuIds == null
? "all"
: [...config.selectedGpuIds].sort((a, b) => a - b).join(","),
gpuSelection,
].join("|");
}

18
studio/frontend/src/features/model-picker/model-config/per-model-config.ts
View file

@ -8,6 +8,7 @@ import {
normalizeGgufVariantIdentity,
normalizeModelIdentity,
} from "./model-identity";
import type { GpuIndexKind } from "@/hooks/use-gpu-info";
export interface PerModelConfig {
customContextLength: number | null;
@ -19,13 +20,14 @@ export interface PerModelConfig {
tensorParallel: boolean;
chatTemplateOverride: string | null;
// GPU Memory controls (per-model, GGUF-only), optional so older blobs still
// parse. null selectedGpuIds (all GPUs) is distinct from absent. The --tensor-split
// ratio is deliberately not remembered: it is positionally bound to the exact
// GPU set/order and unvalidated.
// parse. null or absent selectedGpuIds means automatic placement; an array is
// an explicit candidate pool. The --tensor-split ratio is deliberately not
// remembered because it is bound to the exact GPU set and order.
gpuMemoryMode?: "auto" | "manual";
gpuLayers?: number;
nCpuMoe?: number;
selectedGpuIds?: number[] | null;
selectedGpuIndexKind?: GpuIndexKind | null;
}
export const DEFAULT_PER_MODEL_CONFIG: PerModelConfig = {
@ -106,6 +108,7 @@ const STORED_CONFIG_FIELDS = new Set([
"gpuLayers",
"nCpuMoe",
"selectedGpuIds",
"selectedGpuIndexKind",
]);
function normalizeGpuFields(partial: RawConfig): {
@ -113,12 +116,14 @@ function normalizeGpuFields(partial: RawConfig): {
gpuLayers?: number;
nCpuMoe?: number;
selectedGpuIds?: number[] | null;
selectedGpuIndexKind?: GpuIndexKind | null;
} {
const out: {
gpuMemoryMode?: "auto" | "manual";
gpuLayers?: number;
nCpuMoe?: number;
selectedGpuIds?: number[] | null;
selectedGpuIndexKind?: GpuIndexKind | null;
} = {};
// Only "manual" is a real override; persisting "auto" would pin the model and
// stop it following later changes to the global GPU Memory preference.
@ -148,6 +153,13 @@ function normalizeGpuFields(partial: RawConfig): {
) {
out.selectedGpuIds = partial.selectedGpuIds.map((n) => Math.trunc(n));
}
if (
partial.selectedGpuIndexKind === "physical" ||
partial.selectedGpuIndexKind === "vulkan" ||
partial.selectedGpuIndexKind === null
) {
out.selectedGpuIndexKind = partial.selectedGpuIndexKind;
}
return out;
}

121
studio/frontend/src/hooks/gpu-selection.ts Normal file
View file

@ -0,0 +1,121 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export type GpuIndexKind = "physical" | "vulkan";
export interface SystemGpuDevice {
index: number;
indexKind: GpuIndexKind | null;
name: string;
memoryTotalGb: number;
/** Free VRAM at fetch time, or total VRAM when usage is unavailable. */
memoryFreeGb: number;
/** Whether `index` is safe to send as gpu_ids. */
pinnable: boolean;
/** Whether the separate DiffusionGemma runner can use this physical ID. */
diffusionPinnable: boolean;
}
export interface ReconciledGpuSelection {
ids: number[] | null;
indexKind: GpuIndexKind | null;
}
export function sameGpuSelection(
left: ReconciledGpuSelection,
right: ReconciledGpuSelection,
): boolean {
if (left.indexKind !== right.indexKind) return false;
if (left.ids === right.ids) return true;
if (left.ids === null || right.ids === null) return false;
return (
left.ids.length === right.ids.length &&
left.ids.every((id, index) => id === right.ids?.[index])
);
}
export interface PinnableGpuContext {
devices: SystemGpuDevice[] | null;
ids: number[] | null;
indexKind: GpuIndexKind | null | undefined;
}
export function pinnableGpuContext(
devices: SystemGpuDevice[] | null,
forDiffusion = false,
): PinnableGpuContext {
if (devices === null) {
return { devices: null, ids: null, indexKind: undefined };
}
const pinnable = devices.filter((device) =>
forDiffusion ? device.diffusionPinnable : device.pinnable,
);
const indexKind = pinnable[0]?.indexKind ?? null;
if (
pinnable.length <= 1 ||
indexKind === null ||
!pinnable.every((device) => device.indexKind === indexKind)
) {
return { devices: pinnable, ids: [], indexKind: null };
}
return {
devices: pinnable,
ids: pinnable.map((device) => device.index),
indexKind,
};
}
/**
* Selection context when the backend namespace may be known before its device
* membership. A temporarily unavailable Vulkan inventory still proves that
* physical IDs are incompatible, but cannot yet range-check Vulkan ordinals.
*/
export function resolveGpuSelectionContext(
devices: SystemGpuDevice[] | null,
forDiffusion = false,
unavailableIndexKind?: GpuIndexKind,
): PinnableGpuContext {
if (unavailableIndexKind !== undefined) {
if (forDiffusion) {
return { devices: [], ids: [], indexKind: null };
}
return {
devices: [],
ids: null,
indexKind: unavailableIndexKind,
};
}
return pinnableGpuContext(devices, forDiffusion);
}
export function reconcileGpuSelection(
ids: number[] | null,
savedIndexKind: GpuIndexKind | null | undefined,
currentIndexKind: GpuIndexKind | null | undefined,
pinnableIds: number[] | null,
): ReconciledGpuSelection {
if (ids == null) {
return { ids: null, indexKind: null };
}
const expectedIndexKind =
savedIndexKind === undefined ? "physical" : savedIndexKind;
// Namespace knowledge is authoritative even while membership is unavailable.
// This is what prevents a saved CUDA/ROCm ID from becoming Vulkan<i>.
// A null namespace is only safe while discovery is also unresolved.
if (
currentIndexKind !== undefined &&
expectedIndexKind !== currentIndexKind
) {
return { ids: null, indexKind: null };
}
if (currentIndexKind === null) {
return { ids: null, indexKind: null };
}
if (currentIndexKind === undefined || pinnableIds === null) {
return { ids, indexKind: expectedIndexKind };
}
const kept = ids.filter((id) => pinnableIds.includes(id));
return kept.length > 0
? { ids: kept, indexKind: currentIndexKind }
: { ids: null, indexKind: null };
}

21
studio/frontend/src/hooks/system-discovery.ts Normal file
View file

@ -0,0 +1,21 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
interface InferenceGpuDiscovery {
available: boolean;
backend?: string;
}
export function shouldRetrySystemDiscovery(
cacheIsCold: boolean,
inferenceGpu: InferenceGpuDiscovery | undefined,
retrySubscribers: number,
): boolean {
if (retrySubscribers <= 0) {
return false;
}
if (cacheIsCold) {
return true;
}
return inferenceGpu?.backend === "vulkan" && !inferenceGpu.available;
}

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

@ -1,13 +1,32 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { authFetch } from "@/features/auth";
import { useEffect, useState } from "react";
import {
aggregateGpuMemoryTotalGb,
type GpuIndexKind,
type PinnableGpuContext,
type ReconciledGpuSelection,
type SystemGpuDevice,
reconcileGpuSelection,
resolveGpuSelectionContext,
} from "./gpu-selection";
import {
type SystemInfoResponse,
aggregateGpuMemoryTotalGb,
fetchSystemInfo,
getCachedSystemInfo,
subscribeSystemInfo,
} from "./use-system";
export {
pinnableGpuContext,
reconcileGpuSelection,
type GpuIndexKind,
type PinnableGpuContext,
type ReconciledGpuSelection,
type SystemGpuDevice,
} from "./gpu-selection";
export interface GpuInfo {
available: boolean;
budgetKnown: boolean;
@ -16,20 +35,7 @@ export interface GpuInfo {
cpuCore: number;
cpuThread: number;
systemRamAvailableGb: number;
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;
systemRamTotalGb: number;
}
const DEFAULT_GPU: GpuInfo = {
@ -40,31 +46,9 @@ const DEFAULT_GPU: GpuInfo = {
cpuCore: 0,
cpuThread: 0,
systemRamAvailableGb: 0,
systemRamTotalGb: 0
systemRamTotalGb: 0,
};
// 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 fetchSystemOnce(force = false): Promise<SystemInfoResponse | null> {
if (!force && 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}`);
cachedSystem = (await res.json()) as SystemInfoResponse;
return cachedSystem;
} catch {
return null;
} finally {
systemPromise = null;
}
})();
return systemPromise;
}
function toGpuInfo(
data: SystemInfoResponse | null,
source: "gpu" | "inference_gpu" = "gpu",
@ -78,9 +62,7 @@ function toGpuInfo(
systemRamTotalGb: data?.memory?.total_gb ?? 0,
};
const gpuData =
source === "inference_gpu"
? (data?.inference_gpu ?? data?.gpu)
: data?.gpu;
source === "inference_gpu" ? (data?.inference_gpu ?? data?.gpu) : data?.gpu;
const devices = gpuData?.devices ?? [];
if (!gpuData?.available || !devices.length) {
return { ...DEFAULT_GPU, ...base, budgetKnown: data !== null };
@ -106,39 +88,62 @@ function toGpuDevices(data: SystemInfoResponse | null): SystemGpuDevice[] {
// The XPU ban does not apply there, it is about torch-xpu ordinals that no
// applicator speaks; a Vulkan pick does not use them.
const inference = data?.inference_gpu;
if (inference?.backend === "vulkan" && (inference.devices ?? []).length) {
if (inference?.backend === "vulkan") {
// The installed inference backend is confirmed Vulkan, so even an empty
// device list (probe still cold, or transiently failed) must NOT fall
// through to the torch/CUDA inventory below: those physical IDs are
// meaningless to a Vulkan llama-server, and the backend rejects every
// explicit diffusion pin outright while is_vulkan_build is true. Report no
// pinnable/diffusionPinnable devices until the probe succeeds.
if (!(inference.devices ?? []).length) return [];
const picksAccepted = inference.gguf_gpu_ids_supported !== false;
return (inference.devices ?? [])
.filter((d) => typeof d.index === "number")
.map((d) => ({
index: d.index as number,
indexKind: d.index_kind === "vulkan" ? ("vulkan" as const) : null,
name: d.name ?? `GPU ${d.index}`,
memoryTotalGb: d.memory_total_gb ?? 0,
memoryFreeGb: d.vram_free_gb ?? 0,
physicalIndex: picksAccepted && d.index_kind === "vulkan",
pinnable: picksAccepted && d.index_kind === "vulkan",
// The DiffusionGemma runner is torch-side and never speaks ggml
// ordinals, so a Vulkan pick is not usable there.
diffusionPinnable: false,
}));
}
// Otherwise the torch view is the pickable set. Unpinnable configurations
// must hide every pick surface: XPU indices are torch-xpu ordinals no
// applicator speaks, so /load and /validate 400 them, and the backend reports
// gpu.gguf_gpu_ids_supported. Absent support info defaults to pinnable
// (older backend).
const pinnableBackend =
data?.device_backend !== "xpu" &&
data?.gpu?.gguf_gpu_ids_supported !== false;
// must hide every pick surface: the backend reports gguf_gpu_ids_supported,
// and absent support info defaults to pinnable (older backend).
const pinnableBackend = data?.gpu?.gguf_gpu_ids_supported !== false;
// ROCm reuses torch.cuda.* and the same physical-ID path, so the runner takes
// its indices too; only the reported label differs (_backend_label swaps it).
const diffusionBackend =
data?.device_backend === "cuda" || data?.device_backend === "rocm";
return (data?.gpu?.devices ?? [])
.filter((d) => typeof d.index === "number")
.map((d) => ({
index: d.index as number,
indexKind:
d.index_kind === "physical" || d.index_kind === "vulkan"
? d.index_kind
: null,
name: d.name ?? `GPU ${d.index}`,
memoryTotalGb: d.memory_total_gb ?? 0,
memoryFreeGb: d.vram_free_gb ?? 0,
physicalIndex: pinnableBackend && d.index_kind === "physical",
// The XPU ban is about torch-xpu ordinals no applicator speaks, so /load
// and /validate 400 them. A Vulkan ordinal is not one of those, so it
// stays pickable even when this list arrives from an XPU host.
pinnable:
pinnableBackend &&
(d.index_kind === "vulkan" ||
(data?.device_backend !== "xpu" && d.index_kind === "physical")),
diffusionPinnable: diffusionBackend && d.index_kind === "physical",
}));
}
/** Aggregate GPU info from /api/system; shares one module-level fetch across all GPU hooks. */
function useGpuInfoSource(source: "gpu" | "inference_gpu"): GpuInfo {
const cachedSystem = getCachedSystemInfo();
const [gpu, setGpu] = useState<GpuInfo>(
cachedSystem ? toGpuInfo(cachedSystem, source) : DEFAULT_GPU,
);
@ -146,33 +151,29 @@ function useGpuInfoSource(source: "gpu" | "inference_gpu"): GpuInfo {
// 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;
let retryId: number | undefined;
const update = (force = false, retryVulkan = false) => {
fetchSystemOnce(force).then((d) => {
const sync = (data: SystemInfoResponse) => {
if (cancelled) return;
const next = toGpuInfo(data, source);
setGpu((current) =>
JSON.stringify(current) === JSON.stringify(next) ? current : next,
);
};
const update = () => {
fetchSystemInfo().then((d) => {
if (cancelled) return;
if (!d) {
// Once an unavailable Vulkan backend starts polling, a transient API
// failure must preserve the current state and continue the same loop.
if (retryVulkan) {
retryId = window.setTimeout(() => update(true, true), 3000);
}
return;
}
setGpu(toGpuInfo(d, source));
const inferenceGpu = d.inference_gpu;
if (
source === "inference_gpu" &&
inferenceGpu?.backend === "vulkan" &&
!inferenceGpu.available
) {
retryId = window.setTimeout(() => update(true, true), 3000);
}
if (!d) return;
// A cache hit does not publish a new snapshot. Sync it here so a
// consumer mounting while the initial request finishes cannot miss it.
sync(d);
});
};
const unsubscribe = subscribeSystemInfo(sync, {
retryUnavailableVulkan: source === "inference_gpu",
});
update();
return () => {
cancelled = true;
if (retryId !== undefined) window.clearTimeout(retryId);
unsubscribe();
};
}, [source]);
return gpu;
@ -190,6 +191,7 @@ export function useInferenceGpuInfo(): GpuInfo {
/** All backend-visible GPUs (index, name, total VRAM); shares the same fetch. */
export function useGpuDevices(): SystemGpuDevice[] {
const cachedSystem = getCachedSystemInfo();
const [devices, setDevices] = useState<SystemGpuDevice[]>(
cachedSystem ? toGpuDevices(cachedSystem) : [],
);
@ -197,42 +199,87 @@ export function useGpuDevices(): SystemGpuDevice[] {
// 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));
let lastSerialized: string | null = null;
const sync = (data: SystemInfoResponse | null) => {
if (cancelled) return;
const next = toGpuDevices(data);
// Every refresh builds a fresh array, so compare by value or a 3s Vulkan
// retry loop would re-render this hook forever.
const serialized = JSON.stringify(next);
if (serialized === lastSerialized) return;
lastSerialized = serialized;
setDevices(next);
};
const unsubscribe = subscribeSystemInfo(sync, {
retryUnavailableVulkan: true,
});
fetchSystemInfo().then(sync);
return () => {
cancelled = true;
unsubscribe();
};
}, []);
return devices;
}
/**
* 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.
*/
/** Whether device discovery is settled enough to rewrite remembered UI state. */
export function gpuDeviceCacheReady(): boolean {
const cachedSystem = getCachedSystemInfo();
if (cachedSystem === null) {
return false;
}
const inferenceGpu = cachedSystem.inference_gpu;
return !(inferenceGpu?.backend === "vulkan" && !inferenceGpu.available);
}
/** Warm the shared system cache before validating persisted GPU IDs. */
export async function ensureGpuDeviceCache(): Promise<void> {
await fetchSystemOnce();
await fetchSystemInfo();
}
/** Cached pinnable IDs, null before fetch, or [] when pinning is unavailable. */
export function cachedPinnableGpuIndices(
forDiffusion = false,
): number[] | null {
return cachedPinnableGpuContext(forDiffusion).ids;
}
/** Cached index namespace, undefined before fetch and null when unavailable. */
export function cachedPinnableGpuIndexKind(
forDiffusion = false,
): GpuIndexKind | null | undefined {
return cachedPinnableGpuContext(forDiffusion).indexKind;
}
/**
* 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.
* Cached namespace and membership are separate: an unavailable Vulkan probe
* leaves membership unknown while the Vulkan namespace remains authoritative.
*/
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) : [];
export function cachedPinnableGpuContext(
forDiffusion = false,
devices?: SystemGpuDevice[],
): PinnableGpuContext {
const cachedSystem = getCachedSystemInfo();
const unavailableVulkan =
cachedSystem?.inference_gpu?.backend === "vulkan" &&
!cachedSystem.inference_gpu.available;
return resolveGpuSelectionContext(
cachedSystem ? (devices ?? toGpuDevices(cachedSystem)) : null,
forDiffusion,
unavailableVulkan ? "vulkan" : undefined,
);
}
export function reconcileCachedGpuSelection(
ids: number[] | null,
savedIndexKind?: GpuIndexKind | null,
forDiffusion = false,
): ReconciledGpuSelection {
const context = cachedPinnableGpuContext(forDiffusion);
return reconcileGpuSelection(
ids,
savedIndexKind,
context.indexKind,
context.ids,
);
}

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

@ -3,6 +3,7 @@
import { authFetch } from "@/features/auth";
import { useEffect, useState } from "react";
import { shouldRetrySystemDiscovery } from "./system-discovery";
export interface GpuDevice {
index?: number;
@ -20,7 +21,8 @@ export interface GpuDevice {
export interface SystemGpuInfo {
available: boolean;
backend?: string;
/** Whether GGUF loads accept explicit physical GPU IDs. */
/** Whether GGUF loads accept explicit gpu_ids in the device records'
* declared index space. */
gguf_gpu_ids_supported?: boolean;
backend_cuda_visible_devices?: string | null;
parent_visible_gpu_ids?: number[];
@ -74,7 +76,10 @@ export interface SystemInfoResponse {
}
let cachedSystem: SystemInfoResponse | null = null;
let systemFetchPromise: Promise<SystemInfoResponse> | null = null;
let systemFetchPromise: Promise<SystemInfoResponse | null> | null = null;
const systemSubscribers = new Set<(data: SystemInfoResponse) => void>();
let vulkanRetrySubscribers = 0;
let vulkanRetryId: number | null = null;
const DEFAULT_SYSTEM: SystemInfoResponse = {
platform: "Unknown",
@ -88,9 +93,60 @@ const DEFAULT_SYSTEM: SystemInfoResponse = {
ml_packages: {}
};
async function fetchSystemOnce({
export function getCachedSystemInfo(): SystemInfoResponse | null {
return cachedSystem;
}
export function subscribeSystemInfo(
subscriber: (data: SystemInfoResponse) => void,
options: { retryUnavailableVulkan?: boolean } = {},
): () => void {
systemSubscribers.add(subscriber);
if (options.retryUnavailableVulkan) {
vulkanRetrySubscribers += 1;
scheduleVulkanRetry();
}
return () => {
systemSubscribers.delete(subscriber);
if (options.retryUnavailableVulkan) {
vulkanRetrySubscribers = Math.max(0, vulkanRetrySubscribers - 1);
if (vulkanRetrySubscribers === 0 && vulkanRetryId !== null) {
window.clearTimeout(vulkanRetryId);
vulkanRetryId = null;
}
}
};
}
function scheduleVulkanRetry(): void {
if (
!shouldRetrySystemDiscovery(
cachedSystem === null,
cachedSystem?.inference_gpu,
vulkanRetrySubscribers,
)
) {
// A cold subscription schedules before its first request settles. Cancel
// that pending retry as soon as discovery succeeds with a usable inventory
// or a non-Vulkan backend.
if (vulkanRetryId !== null) {
window.clearTimeout(vulkanRetryId);
vulkanRetryId = null;
}
return;
}
if (vulkanRetryId !== null) {
return;
}
vulkanRetryId = window.setTimeout(() => {
vulkanRetryId = null;
void fetchSystemInfo({ force: true });
}, 3000);
}
export async function fetchSystemInfo({
force = false,
}: { force?: boolean } = {}): Promise<SystemInfoResponse> {
}: { force?: boolean } = {}): Promise<SystemInfoResponse | null> {
if (systemFetchPromise) return systemFetchPromise;
if (!force && cachedSystem) return cachedSystem;
@ -101,12 +157,13 @@ async function fetchSystemOnce({
const data = await res.json();
cachedSystem = data as SystemInfoResponse;
systemSubscribers.forEach((subscriber) => subscriber(cachedSystem!));
return cachedSystem;
} catch {
cachedSystem = null;
return DEFAULT_SYSTEM;
return null;
} finally {
systemFetchPromise = null;
scheduleVulkanRetry();
}
})();
@ -131,9 +188,9 @@ export function useSystemInfo({
let timeoutId: number | null = null;
const update = (force: boolean) => {
void fetchSystemOnce({ force })
void fetchSystemInfo({ force })
.then((info) => {
if (!cancelled) setSystemInfo(info);
if (!cancelled && info) setSystemInfo(info);
})
.finally(() => {
if (cancelled || !pollMs) return;

137
studio/frontend/tests/gpu-selection.test.ts Normal file
View file

@ -0,0 +1,137 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import assert from "node:assert/strict";
import test from "node:test";
import {
type GpuIndexKind,
reconcileGpuSelection,
resolveGpuSelectionContext,
sameGpuSelection,
} from "../src/hooks/gpu-selection.ts";
const ids = [0, 1];
function reconcile(
savedIndexKind: GpuIndexKind | null | undefined,
currentIndexKind: GpuIndexKind | null | undefined,
pinnableIds: number[] | null,
) {
return reconcileGpuSelection(
ids,
savedIndexKind,
currentIndexKind,
pinnableIds,
);
}
test("preserves every namespace while the device cache is genuinely cold", () => {
assert.deepEqual(reconcile(undefined, undefined, null), {
ids,
indexKind: "physical",
});
assert.deepEqual(reconcile(null, undefined, null), {
ids,
indexKind: null,
});
assert.deepEqual(reconcile("physical", undefined, null), {
ids,
indexKind: "physical",
});
assert.deepEqual(reconcile("vulkan", undefined, null), {
ids,
indexKind: "vulkan",
});
});
test("known unavailable Vulkan clears picks without a compatible namespace", () => {
const context = resolveGpuSelectionContext([], false, "vulkan");
assert.deepEqual(context, {
devices: [],
ids: null,
indexKind: "vulkan",
});
assert.deepEqual(reconcile(undefined, context.indexKind, context.ids), {
ids: null,
indexKind: null,
});
assert.deepEqual(reconcile("physical", context.indexKind, context.ids), {
ids: null,
indexKind: null,
});
assert.deepEqual(reconcile(null, context.indexKind, context.ids), {
ids: null,
indexKind: null,
});
assert.deepEqual(reconcile("vulkan", context.indexKind, context.ids), {
ids,
indexKind: "vulkan",
});
});
test("known Vulkan suppresses every DiffusionGemma pin while unavailable", () => {
const context = resolveGpuSelectionContext([], true, "vulkan");
assert.deepEqual(context, {
devices: [],
ids: [],
indexKind: null,
});
for (const saved of [undefined, null, "physical", "vulkan"] as const) {
assert.deepEqual(reconcile(saved, context.indexKind, context.ids), {
ids: null,
indexKind: null,
});
}
});
test("recovered inventories filter matching picks and reject mismatched namespaces", () => {
assert.deepEqual(reconcile(null, "vulkan", [1]), {
ids: null,
indexKind: null,
});
assert.deepEqual(reconcile(null, "physical", [0, 1]), {
ids: null,
indexKind: null,
});
assert.deepEqual(reconcile("vulkan", "vulkan", [0]), {
ids: [0],
indexKind: "vulkan",
});
assert.deepEqual(reconcile(undefined, "vulkan", [0, 1]), {
ids: null,
indexKind: null,
});
assert.deepEqual(reconcile("physical", "physical", [1]), {
ids: [1],
indexKind: "physical",
});
assert.deepEqual(reconcile("vulkan", "physical", [0, 1]), {
ids: null,
indexKind: null,
});
});
test("selection equality includes the GPU index namespace", () => {
assert.equal(
sameGpuSelection(
{ ids: [0, 1], indexKind: null },
{ ids: [0, 1], indexKind: "vulkan" },
),
false,
);
assert.equal(
sameGpuSelection(
{ ids: [0, 1], indexKind: "vulkan" },
{ ids: [0, 1], indexKind: "vulkan" },
),
true,
);
assert.equal(
sameGpuSelection(
{ ids: [0], indexKind: "physical" },
{ ids: [1], indexKind: "physical" },
),
false,
);
});

36
studio/frontend/tests/system-discovery.test.ts Normal file
View file

@ -0,0 +1,36 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import assert from "node:assert/strict";
import test from "node:test";
import { shouldRetrySystemDiscovery } from "../src/hooks/system-discovery.ts";
test("retries a cold system cache while an interested subscriber exists", () => {
assert.equal(shouldRetrySystemDiscovery(true, undefined, 1), true);
assert.equal(shouldRetrySystemDiscovery(true, undefined, 0), false);
});
test("retries only an unavailable Vulkan inventory after discovery", () => {
assert.equal(
shouldRetrySystemDiscovery(
false,
{ backend: "vulkan", available: false },
1,
),
true,
);
assert.equal(
shouldRetrySystemDiscovery(
false,
{ backend: "vulkan", available: true },
1,
),
false,
);
assert.equal(
shouldRetrySystemDiscovery(false, { backend: "cuda", available: false }, 1),
false,
);
assert.equal(shouldRetrySystemDiscovery(false, undefined, 1), false);
});

194
studio/install_llama_prebuilt.py
View file

@ -3087,13 +3087,26 @@ def resolve_linux_cuda_choice(
def published_asset_choice_for_kind(
release: PublishedReleaseBundle, install_kind: str
release: PublishedReleaseBundle,
install_kind: str,
*,
host: HostInfo | None = None,
) -> AssetChoice | None:
candidates = sorted(
(artifact for artifact in release.artifacts if artifact.install_kind == install_kind),
key = lambda artifact: (artifact.rank, artifact.asset_name),
)
for artifact in candidates:
if host is not None and install_kind in VULKAN_INSTALL_KINDS:
identity = f"{artifact.bundle_profile or ''} {artifact.asset_name}".lower()
if host.is_arm64:
if not any(token in identity for token in ("arm64", "aarch64")):
continue
elif host.is_x86_64:
if not any(token in identity for token in ("x64", "x86_64", "amd64")):
continue
else:
continue
asset_url = release.assets.get(artifact.asset_name)
if not asset_url:
continue
@ -3104,6 +3117,7 @@ def published_asset_choice_for_kind(
url = asset_url,
source_label = "published",
install_kind = install_kind,
bundle_profile = artifact.bundle_profile,
runtime_line = artifact.runtime_line,
selection_log = list(release.selection_log)
+ [f"published_selection: selected {artifact.asset_name} install_kind={install_kind}"],
@ -3466,6 +3480,17 @@ def resolve_release_asset_choice(
published_choice: AssetChoice | None = None
if host.is_windows and host.is_x86_64:
if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
choices = [
choice
for choice in (
published_asset_choice_for_kind(release, "windows-vulkan", host = host),
published_asset_choice_for_kind(release, "windows-cpu"),
)
if choice is not None
]
if choices:
return apply_approved_hashes(choices, checksums)
# AMD Windows hosts prefer the fork's per-gfx windows-rocm bundle when one
# covers the GPU; otherwise fall through to resolve_asset_choice(). Note
# that on the fork repo the upstream win-hip archive has no approved hash,
@ -5422,9 +5447,9 @@ def resolve_install_attempts(
def _linux_published_attempts(host: HostInfo, bundle: PublishedReleaseBundle) -> list[AssetChoice]:
"""Build the install attempts for a fork Linux host from a manifest-described
bundle: CUDA, per-gfx ROCm, or (non-GPU) CPU. Same selection the upstream
filename path used, just sourced from the manifest instead of reconstructed
from asset names."""
bundle: CUDA, per-gfx ROCm, Vulkan, or CPU. Same selection the upstream filename
path used, just sourced from the manifest instead of reconstructed from asset
names."""
attempts: list[AssetChoice] = []
if host.has_usable_nvidia:
# Prefer the cudart major Unsloth loads at runtime (torch's bundled
@ -5449,6 +5474,10 @@ def _linux_published_attempts(host: HostInfo, bundle: PublishedReleaseBundle) ->
if published_rocm is not None:
attempts.append(published_rocm)
else:
if host.has_intel_gpu and not host.has_physical_nvidia:
vulkan_choice = published_asset_choice_for_kind(bundle, "linux-vulkan", host = host)
if vulkan_choice is not None:
attempts.append(vulkan_choice)
# CPU-only host. A usable-NVIDIA host never reaches here -- if its CUDA
# selection produced nothing we want an empty attempt list so the caller
# source-builds with CUDA, not a CPU-only binary silently installed on a
@ -5799,7 +5828,7 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
["libggml-hip.so*"],
]
if choice.install_kind == "linux-vulkan":
return [
groups = [
["libllama-common.so*"],
["libllama.so*"],
["libggml.so*"],
@ -5812,6 +5841,9 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
["libmtmd.so*"],
["libggml-vulkan.so*"],
]
if choice.source_label == "published":
groups.append(["llama-diffusion-gemma-visual-server"])
return groups
if choice.install_kind in {"windows-cpu", "windows-arm64"}:
return [["llama.dll"]]
if choice.install_kind == "windows-cuda":
@ -5832,7 +5864,10 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
if choice.install_kind in {"windows-hip", "windows-rocm"}:
return [["llama.dll"], ["*hip*.dll"]]
if choice.install_kind == "windows-vulkan":
return [["llama.dll"], ["ggml-vulkan.dll"]]
groups = [["llama.dll"], ["ggml-vulkan.dll"]]
if choice.source_label == "published":
groups.append(["llama-diffusion-gemma-visual-server.exe"])
return groups
return []
@ -5988,6 +6023,10 @@ def validate_prebuilt_choice(
)
log(f"overlaying prebuilt bundle {choice.name} into {install_dir}")
server_path, quantize_path = install_from_archives(choice, host, install_dir, work_dir)
if choice.install_kind in VULKAN_INSTALL_KINDS and not runtime_payload_is_healthy(
install_dir, host, choice
):
raise PrebuiltFallback(f"Vulkan bundle {choice.name} omitted a required runtime component")
preflight_linux_installed_binaries((server_path, quantize_path), install_dir, host)
preflight_macos_installed_binaries((server_path, quantize_path), install_dir, host)
ensure_repo_shape(install_dir)
@ -6188,11 +6227,11 @@ def _normalized_llama_backend(value: str | None) -> str | None:
def llama_backend_from_env() -> str | None:
"""Read an explicit llama.cpp backend preference from the environment.
Only ``UNSLOTH_LLAMA_BACKEND`` is honored. ``UNSLOTH_LLAMA_CPP_BACKEND`` is a separate
setup variable meaning ``auto``/``cpu`` (not a backend name) that setup warns about and
otherwise ignores, so reading it here would force Vulkan behind that warning.
``UNSLOTH_LLAMA_CPP_BACKEND`` is shared with setup.sh/setup.ps1. ``auto`` and
unknown values leave selection automatic. ``cpu`` and ``vulkan`` select those
bundles explicitly, while ``hip``/``rocm`` opt out of automatic Vulkan routing.
"""
return _normalized_llama_backend(os.environ.get("UNSLOTH_LLAMA_BACKEND"))
return _normalized_llama_backend(os.environ.get("UNSLOTH_LLAMA_CPP_BACKEND"))
def resolved_llama_backend(llama_backend: str | None = None) -> str | None:
@ -6202,11 +6241,12 @@ def resolved_llama_backend(llama_backend: str | None = None) -> str | None:
def force_vulkan_requested(llama_backend: str | None = None) -> bool:
"""Whether this run should install the upstream Vulkan llama.cpp prebuilt.
"""Whether this run should install the Vulkan llama.cpp prebuilt.
Triggered by ``UNSLOTH_LLAMA_BACKEND=vulkan``, legacy ``UNSLOTH_FORCE_VULKAN``, or
``--llama-backend vulkan``. Scoped to the llama.cpp backend; the torch/training stack
installs separately and still sees the real GPU.
Triggered by ``UNSLOTH_LLAMA_CPP_BACKEND=vulkan``, legacy ``UNSLOTH_FORCE_VULKAN``,
or ``--llama-backend vulkan``. Scoped to the
llama.cpp backend; the torch/training stack installs separately and still sees
the real GPU.
"""
backend = resolved_llama_backend(llama_backend)
if backend is not None:
@ -6217,6 +6257,7 @@ def force_vulkan_requested(llama_backend: str | None = None) -> bool:
"1",
"true",
"yes",
"on",
)
@ -6350,6 +6391,27 @@ def _vulkan_only_host(host: HostInfo) -> HostInfo:
)
def _vulkan_only_attempts(attempts: Iterable[AssetChoice]) -> list[AssetChoice]:
"""Remove generic CPU fallbacks from an explicit Vulkan install."""
return [
attempt
for attempt in attempts
if attempt.install_kind in ("linux-vulkan", "windows-vulkan")
]
def _vulkan_only_release_plans(plans: Iterable[InstallReleasePlan]) -> list[InstallReleasePlan]:
"""Keep release plans that contain an explicit Vulkan bundle."""
filtered = []
for plan in plans:
attempts = _vulkan_only_attempts(plan.attempts)
if attempts:
filtered.append(dataclasses_replace(plan, attempts = attempts))
if not filtered:
raise PrebuiltFallback("no Vulkan prebuilt bundle attempts were available")
return filtered
def _route_to_vulkan_prebuilt(
host: HostInfo,
published_repo: str,
@ -6358,16 +6420,16 @@ def _route_to_vulkan_prebuilt(
force_cpu: bool,
llama_backend: str | None = None,
) -> tuple[HostInfo, str, str, str | None]:
"""Point a Vulkan-capable host at the upstream ggml-org Vulkan prebuilt.
"""Point a Vulkan-capable host at the selected repository's Vulkan prebuilt.
The unsloth published repo ships only CUDA/ROCm/CPU assets, so Vulkan comes from
UPSTREAM_REPO. Three triggers route here, all suppressed when a CPU flag (--cpu-fallback
or --force-cpu, folded into force_cpu) wins:
* ``UNSLOTH_LLAMA_BACKEND=vulkan`` / ``UNSLOTH_FORCE_VULKAN`` / ``--llama-backend
vulkan`` forces Vulkan over the detected CUDA/ROCm backend;
The default Unsloth release manifest includes Vulkan app bundles, including the
DiffusionGemma visual server. Three triggers route here, all suppressed when a CPU
flag (--cpu-fallback or --force-cpu, folded into force_cpu) wins:
* ``UNSLOTH_LLAMA_CPP_BACKEND=vulkan`` / ``UNSLOTH_FORCE_VULKAN`` /
``--llama-backend vulkan`` forces Vulkan over the detected CUDA/ROCm backend;
* Windows AMD with no HIP-prebuilt gfx arch auto-falls back to Vulkan (#7357);
* an auto-detected Intel GPU with NO physical NVIDIA/ROCm, the purpose of the
has_intel_gpu probe, since the fork manifest ships no Vulkan asset.
has_intel_gpu probe.
Applied by BOTH the install path and the --resolve-prebuilt probe so the "is a prebuilt
available" answer matches what actually gets installed.
@ -6382,50 +6444,61 @@ def _route_to_vulkan_prebuilt(
)
# No PHYSICAL NVIDIA, not merely no usable one: Vulkan ignores CUDA_VISIBLE_DEVICES, so
# auto-routing a host that hides its NVIDIA card would let it grab the reserved GPU.
auto_intel = host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm
auto_intel = (
explicit_backend is None
and host.has_intel_gpu
and not host.has_physical_nvidia
and not host.has_rocm
)
if force_cpu or not (forced or auto_intel or auto_no_hip):
return host, published_repo, published_release_tag, None
if host.is_macos:
if forced:
log(
"UNSLOTH_LLAMA_BACKEND=vulkan is set but ignored on macOS "
"UNSLOTH_LLAMA_CPP_BACKEND=vulkan is set but ignored on macOS "
"(Metal is used; there is no Vulkan prebuilt)"
)
return host, published_repo, published_release_tag, None
if _has_no_vulkan_prebuilt(host):
if forced:
log(
"Vulkan llama.cpp backend requested but ignored on Windows arm64 "
"(upstream ships no Vulkan arm64 prebuilt); keeping the published bundle"
"Vulkan llama.cpp backend requested on Windows arm64, but no "
"compatible Vulkan bundle is published"
)
return host, published_repo, published_release_tag, None
if auto_no_hip:
active = _active_rocm_gfx_target(host) or "unknown"
log(
"Active AMD GPU arch is not supported by the Windows HIP prebuilt "
f"({active}); installing the upstream Vulkan llama.cpp prebuilt instead"
f"({active}); installing the Vulkan llama.cpp prebuilt instead"
)
host = _vulkan_only_host(host)
persist_backend = "vulkan"
elif forced:
log(
"Vulkan llama.cpp backend requested; installing the upstream Vulkan "
"Vulkan llama.cpp backend requested; installing the Vulkan "
"prebuilt instead of the detected GPU backend"
)
host = _vulkan_only_host(host)
persist_backend = "vulkan"
else:
log("Intel GPU detected; installing the upstream Vulkan llama.cpp prebuilt")
log("Intel GPU detected; installing the Vulkan llama.cpp prebuilt")
persist_backend = None
# Swapping the fork for upstream invalidates a fork release pin: the two use
# different tag namespaces (fork b9596-mix-<sha> vs upstream b9596), so a
# pinned fork tag would make the upstream resolver query a nonexistent
# release and fall back to source. Drop it and let the upstream resolver
# pick by the requested llama tag. A pin already on an explicit upstream repo
# (repo unchanged here) is preserved.
if published_repo != UPSTREAM_REPO:
published_release_tag = ""
return host, UPSTREAM_REPO, published_release_tag, persist_backend
# The fork manifest's Vulkan app bundles are x64 only, and the architecture
# filter in published_asset_choice_for_kind rejects them for an ARM64 host, so
# the fork planner returns no Vulkan attempt at all there. Strict Vulkan
# filtering then drops the ARM64 CPU attempt too and the install resolves to
# nothing. Upstream does publish llama-<tag>-bin-ubuntu-vulkan-arm64.tar.gz, so
# keep routing Linux ARM64 there (Windows arm64 exits above via
# _has_no_vulkan_prebuilt, macOS via the Metal branch).
if (published_repo or DEFAULT_PUBLISHED_REPO) == DEFAULT_PUBLISHED_REPO and (
host.is_linux and host.is_arm64
):
# Fork and upstream use different tag namespaces (fork b9596-mix-<sha> vs
# upstream b9596), so carrying a fork pin over would make the upstream
# resolver query a release that does not exist.
return host, UPSTREAM_REPO, "", persist_backend
return host, published_repo, published_release_tag, persist_backend
def diffusion_visual_server_backfill_needed(
@ -6594,6 +6667,9 @@ def install_prebuilt(
) -> None:
# force_cpu drops GPU detection (mechanism, both --cpu-fallback and --force-cpu);
# persist_force_cpu records the deliberate choice so the updater re-asserts it.
if resolved_llama_backend(llama_backend) == "cpu":
force_cpu = True
persist_force_cpu = True
host = detect_host()
host = _apply_host_overrides(
host,
@ -6601,6 +6677,9 @@ def install_prebuilt(
override_rocm_gfx = override_rocm_gfx,
force_cpu = force_cpu,
)
# An explicit Vulkan choice only. The Windows-AMD auto-fallback is a rescue
# from a missing HIP arch, so it must keep the CPU plans it can still use.
strict_vulkan = force_vulkan_requested(llama_backend) and not force_cpu and not host.is_macos
host, published_repo, published_release_tag, persist_llama_backend = _route_to_vulkan_prebuilt(
host,
published_repo,
@ -6628,14 +6707,18 @@ def install_prebuilt(
)
# Single resolver: every fork host selects from the release manifest;
# an explicit ggml-org override selects by asset filename instead. A
# forced-Vulkan host already has published_repo pointed at
# UPSTREAM_REPO above, so the resolver takes the Vulkan asset branch.
# Vulkan host is rewritten above so either resolver takes its Vulkan
# asset branch.
requested_tag, release_plans = resolve_simple_install_release_plans(
llama_tag,
host,
published_repo,
published_release_tag,
)
if strict_vulkan:
# Upstream plans append CPU as a generic fallback. An explicit
# Vulkan selection must fail instead of silently installing it.
release_plans = _vulkan_only_release_plans(release_plans)
if release_plans and existing_install_matches_plan(install_dir, host, release_plans[0]):
current = release_plans[0]
if diffusion_visual_server_backfill_needed(install_dir, host, current.attempts[0]):
@ -6827,10 +6910,12 @@ def parse_args() -> argparse.Namespace:
"--llama-backend",
choices = ("vulkan",),
help = (
"Force the llama.cpp prebuilt backend. vulkan installs the upstream Vulkan "
"bundle and records the choice so Studio updates keep it; ignored on hosts "
"with no Vulkan prebuilt (macOS, Windows arm64). "
"Same effect as UNSLOTH_LLAMA_BACKEND=vulkan / UNSLOTH_FORCE_VULKAN=1."
"Force the llama.cpp prebuilt backend. vulkan installs the Vulkan "
"bundle and records the choice so Studio updates keep it. It is ignored "
"on macOS, which uses Metal, and fails closed on Windows arm64 when no "
"compatible Vulkan bundle is published. "
"Same effect as UNSLOTH_LLAMA_CPP_BACKEND=vulkan / "
"UNSLOTH_FORCE_VULKAN=1."
),
)
resolve_group = parser.add_mutually_exclusive_group()
@ -6919,6 +7004,7 @@ def emit_resolver_output(payload: dict[str, Any], *, output_format: str) -> None
def main() -> int:
args = parse_args()
cpu_backend_requested = resolved_llama_backend(args.llama_backend) == "cpu"
if args.validate_install is not None:
try:
validate_existing_install(
@ -6986,9 +7072,9 @@ def main() -> int:
# Host-aware "is a prebuilt available" probe, no download. Every host now
# plans against the fork (args.published_repo defaults to it); an explicit
# --published-repo overrides. PrebuiltFallback == source build.
# Both flags drop GPU detection; --force-cpu additionally persists (install
# path only). The probe only needs the mechanism, so OR them.
_cpu_mechanism = args.cpu_fallback or args.force_cpu
# CPU fallback and explicit CPU choices drop GPU detection. The probe only
# needs the mechanism; persistence belongs to the install path.
_cpu_mechanism = args.cpu_fallback or args.force_cpu or cpu_backend_requested
host = _apply_host_overrides(
detect_host(),
override_has_rocm = args.has_rocm,
@ -6996,7 +7082,7 @@ def main() -> int:
force_cpu = _cpu_mechanism,
)
# Same Vulkan routing the install path applies, so the probe's answer
# matches what would install (an Intel/forced-Vulkan host -> upstream).
# matches what would install.
host, repo, release_tag, _persist_llama_backend = _route_to_vulkan_prebuilt(
host,
args.published_repo,
@ -7008,6 +7094,12 @@ def main() -> int:
_requested, plans = resolve_simple_install_release_plans(
args.resolve_prebuilt, host, repo, release_tag
)
if (
force_vulkan_requested(args.llama_backend)
and not _cpu_mechanism
and not host.is_macos
):
plans = _vulkan_only_release_plans(plans)
choice = plans[0].attempts[0] if plans and plans[0].attempts else None
if choice is None:
payload = {"prebuilt_available": False, "repo": repo}
@ -7040,10 +7132,10 @@ def main() -> int:
published_release_tag = args.published_release_tag or "",
override_has_rocm = args.has_rocm,
override_rocm_gfx = args.rocm_gfx,
# Both drop GPU detection; only --force-cpu (deliberate) is recorded so the
# updater re-asserts it. --cpu-fallback stays transient and heals to GPU.
force_cpu = args.cpu_fallback or args.force_cpu,
persist_force_cpu = args.force_cpu,
# Explicit CPU choices persist so the updater re-asserts them.
# --cpu-fallback stays transient and heals to GPU.
force_cpu = args.cpu_fallback or args.force_cpu or cpu_backend_requested,
persist_force_cpu = args.force_cpu or cpu_backend_requested,
llama_backend = args.llama_backend,
instruction_cleanup_root = install_arg.absolute(),
)

83
studio/setup.ps1
View file

@ -33,9 +33,9 @@ $PackageDir = Split-Path -Parent $ScriptDir
# errors. Only use "master" temporarily when the latest release is missing
# support for a new model architecture.
#
# UNSLOTH_LLAMA_CPP_BACKEND : "auto" (default) or "cpu". When "cpu", forces
# the CPU-only prebuilt bundle on GPU hosts. Fixes Intel iGPU Vulkan
# crashes (#7213).
# UNSLOTH_LLAMA_CPP_BACKEND : "auto" (default), "cpu", "vulkan", "hip", or
# "rocm". "cpu" forces the CPU-only prebuilt. "vulkan" selects Vulkan even
# when CUDA or ROCm is detected. "hip"/"rocm" opts out of automatic Vulkan.
$DefaultLlamaPrForce = ""
$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp"
$DefaultLlamaTag = "latest"
@ -3712,6 +3712,18 @@ $ResolvedSourceUrl = $LlamaSource
$ResolvedSourceRef = $RequestedLlamaTag
$ResolvedSourceRefKind = "tag"
$ResolvedLlamaTag = $RequestedLlamaTag
$sourceLlamaBackend = "$($env:UNSLOTH_LLAMA_CPP_BACKEND)".Trim().ToLowerInvariant()
$sourceLegacyForceVulkan = "$($env:UNSLOTH_FORCE_VULKAN)".Trim().ToLowerInvariant()
$explicitVulkanSourceBuild = (
-not $IsMacOS -and
(
$sourceLlamaBackend -eq "vulkan" -or
(
$sourceLlamaBackend -notin @("cpu", "vulkan", "hip", "rocm") -and
$sourceLegacyForceVulkan -in @("1", "true", "yes", "on")
)
)
)
if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
$NeedLlamaSourceBuild = $true
@ -3869,6 +3881,11 @@ if ($LocalLlamaCppSrc) {
if ($LocalLlamaCppLinked) {
# local directory linked above; skip prebuilt install
} elseif ($explicitVulkanSourceBuild -and $NeedLlamaSourceBuild) {
Write-Host ""
step "llama.cpp" "Vulkan was explicitly requested, but this installation requires a source build" "Red"
substep "Vulkan source builds are not supported by this installer; use the prebuilt Vulkan bundle or unset the Vulkan override" "Yellow"
exit 1
} elseif ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
Write-Host ""
substep "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt llama.cpp install" "Yellow"
@ -3936,14 +3953,43 @@ if ($LocalLlamaCppLinked) {
if ($env:UNSLOTH_LLAMA_RELEASE_TAG) {
$prebuiltArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG)
}
# UNSLOTH_LLAMA_CPP_BACKEND=cpu (case-insensitive, whitespace-trimmed) forces the
# CPU-only prebuilt via --force-cpu (persisted so updates keep it). Fixes Intel
# iGPU Vulkan crash (#7213).
$llamaBackend = "$($env:UNSLOTH_LLAMA_CPP_BACKEND)".Trim().ToLowerInvariant()
# The backend override is case-insensitive and whitespace-trimmed. cpu
# maps to the persisted --force-cpu choice. vulkan is consumed directly
# by install_llama_prebuilt.py and does not change the torch backend.
$llamaBackend = $sourceLlamaBackend
$legacyForceVulkan = $sourceLegacyForceVulkan
$windowsArm64 = (
$env:OS -eq "Windows_NT" -and
(
"$($env:PROCESSOR_ARCHITECTURE)".ToUpperInvariant() -eq "ARM64" -or
"$($env:PROCESSOR_ARCHITEW6432)".ToUpperInvariant() -eq "ARM64"
)
)
$explicitVulkanBackend = $false
if ($llamaBackend -eq "cpu") {
$prebuiltArgs += "--force-cpu"
} elseif ($llamaBackend -and $llamaBackend -ne "auto") {
Write-Host "[WARN] Ignoring UNSLOTH_LLAMA_CPP_BACKEND='$($env:UNSLOTH_LLAMA_CPP_BACKEND)' (expected 'auto' or 'cpu')" -ForegroundColor Yellow
} elseif ($llamaBackend -eq "vulkan") {
if ($IsMacOS) {
Write-Host "[WARN] Vulkan has no effect on macOS; the universal build uses Metal" -ForegroundColor Yellow
} elseif ($windowsArm64) {
throw "Vulkan was requested, but no Windows ARM64 Vulkan bundle is published. Unset UNSLOTH_LLAMA_CPP_BACKEND or compile llama.cpp from source."
} else {
$prebuiltArgs += @("--llama-backend", "vulkan")
$explicitVulkanBackend = $true
Write-Host " llama.cpp Vulkan selected for GGUF inference; the PyTorch training backend is unchanged" -ForegroundColor Cyan
}
} elseif ($llamaBackend -and $llamaBackend -notin @("auto", "hip", "rocm")) {
Write-Host "[WARN] Ignoring UNSLOTH_LLAMA_CPP_BACKEND='$llamaBackend' (expected 'auto', 'cpu', 'vulkan', 'hip', or 'rocm')" -ForegroundColor Yellow
}
if (
-not $IsMacOS -and
$llamaBackend -notin @("cpu", "vulkan", "hip", "rocm") -and
$legacyForceVulkan -in @("1", "true", "yes", "on")
) {
if ($windowsArm64) {
throw "Vulkan was requested, but no Windows ARM64 Vulkan bundle is published. Unset UNSLOTH_FORCE_VULKAN or compile llama.cpp from source."
}
$explicitVulkanBackend = $true
}
$prevEAPPrebuilt = $ErrorActionPreference
$ErrorActionPreference = "Continue"
@ -4006,14 +4052,27 @@ if ($LocalLlamaCppLinked) {
if (Test-Path -LiteralPath $_cand) { $PreservedLlamaServerFound = $true; break }
}
if (-not $PreservedLlamaServerFound) { $script:LlamaCppDegraded = $true }
# A preserved CUDA/ROCm/CPU server does not satisfy an explicit Vulkan
# request, and it leaves LlamaCppDegraded false, so without this the
# run reports success on the backend the user asked to replace.
if ($explicitVulkanBackend) {
step "llama.cpp" "Vulkan was explicitly requested, so the installer will not keep the existing backend" "Red"
exit 1
}
} else {
step "llama.cpp" "prebuilt install failed (continuing)" "Yellow"
step "llama.cpp" "prebuilt install failed" "Yellow"
Write-LlamaFailureLog -Output $prebuiltOutput
if (Test-Path -LiteralPath $LlamaCppDir) {
substep "Prebuilt update failed; existing install was restored or cleaned before source build fallback" "Yellow"
}
substep "Prebuilt llama.cpp path unavailable or failed validation -- falling back to source build" "Yellow"
$NeedLlamaSourceBuild = $true
if ($explicitVulkanBackend) {
step "llama.cpp" "Vulkan was explicitly requested, so the installer will not substitute a CUDA, ROCm, or CPU source build" "Red"
substep "Check the download error above or try a different UNSLOTH_LLAMA_RELEASE_TAG" "Yellow"
exit 1
} else {
substep "Prebuilt llama.cpp path unavailable or failed validation -- falling back to source build" "Yellow"
$NeedLlamaSourceBuild = $true
}
}
}

77
studio/setup.sh
View file

@ -37,9 +37,11 @@ fi
# Only use "master" temporarily when the latest release
# is missing support for a new model architecture.
#
# UNSLOTH_LLAMA_CPP_BACKEND : "auto" (default) or "cpu". When "cpu", forces
# the CPU-only prebuilt bundle on GPU hosts.
# Fixes Intel iGPU Vulkan crashes (#7213).
# UNSLOTH_LLAMA_CPP_BACKEND : "auto" (default), "cpu", "vulkan", "hip", or
# "rocm". "cpu" forces the CPU-only prebuilt. "vulkan"
# selects Vulkan even when CUDA or ROCm is detected.
# "hip"/"rocm" keeps the detected HIP backend and opts
# out of automatic Vulkan fallback.
# ──────────────────────────────────────────────────────────────────────────
_DEFAULT_LLAMA_PR_FORCE=""
_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp"
@ -1272,6 +1274,20 @@ _LLAMA_FORCE_COMPILE="${UNSLOTH_LLAMA_FORCE_COMPILE:-0}"
_REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-${_DEFAULT_LLAMA_TAG}}"
_HOST_SYSTEM="$(uname -s 2>/dev/null || true)"
_HOST_MACHINE="$(uname -m 2>/dev/null || true)"
_source_backend_choice="$(printf '%s' "${UNSLOTH_LLAMA_CPP_BACKEND:-auto}" | awk '{$1=$1; print tolower($0)}')"
_source_legacy_force_vulkan="$(printf '%s' "${UNSLOTH_FORCE_VULKAN:-}" | awk '{$1=$1; print tolower($0)}')"
_explicit_vulkan_source_build=false
if [ "$_HOST_SYSTEM" != "Darwin" ]; then
case "$_source_backend_choice" in
vulkan) _explicit_vulkan_source_build=true ;;
cpu|hip|rocm) ;;
*)
case "$_source_legacy_force_vulkan" in
1|true|yes|on) _explicit_vulkan_source_build=true ;;
esac
;;
esac
fi
# Pick the release repo install_llama_prebuilt.py plans against. Every host this
# installer supports now pulls its llama.cpp prebuilt from the unslothai fork: it
@ -1407,6 +1423,10 @@ fi
if [ "$_LOCAL_LLAMA_CPP_LINKED" = true ]; then
: # local directory linked above; skip prebuilt install
elif [ "$_explicit_vulkan_source_build" = true ] && [ "$_NEED_LLAMA_SOURCE_BUILD" = true ]; then
step "llama.cpp" "Vulkan was explicitly requested, but this installation requires a source build" "$C_ERR"
substep "Vulkan source builds are not supported by this installer; use the prebuilt Vulkan bundle or unset the Vulkan override"
exit 1
elif [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then
step "llama.cpp" "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt" "$C_WARN"
_NEED_LLAMA_SOURCE_BUILD=true
@ -1447,11 +1467,10 @@ else
# through to the CPU prebuilt instead of breaking the install.
_PREBUILT_CMD+=(--has-rocm)
fi
# UNSLOTH_LLAMA_CPP_BACKEND=cpu (case-insensitive, trimmed) forces the CPU-only
# prebuilt via --force-cpu, bypassing Vulkan/CUDA/ROCm. Fixes Intel iGPU crash (#7213).
# No effect on macOS: the universal bundle already runs on CPU (Metal is a runtime
# -ngl choice), so warn instead of writing a misleading forced-CPU marker.
_llama_backend="$(printf '%s' "${UNSLOTH_LLAMA_CPP_BACKEND:-auto}" | awk '{$1=$1; print tolower($0)}')"
# The normalized override affects llama.cpp only, not the training backend.
_llama_backend="$_source_backend_choice"
_legacy_force_vulkan="$_source_legacy_force_vulkan"
_explicit_vulkan_backend=false
case "$_llama_backend" in
cpu)
if [ "$_HOST_SYSTEM" = "Darwin" ]; then
@ -1460,9 +1479,28 @@ else
_PREBUILT_CMD+=(--force-cpu)
fi
;;
""|auto) ;;
*) step "llama.cpp" "Ignoring UNSLOTH_LLAMA_CPP_BACKEND='$UNSLOTH_LLAMA_CPP_BACKEND' (expected 'auto' or 'cpu')" "$C_WARN" >&2 ;;
vulkan)
if [ "$_HOST_SYSTEM" = "Darwin" ]; then
step "llama.cpp" "Vulkan has no effect on macOS; the universal build uses Metal" "$C_WARN" >&2
else
_PREBUILT_CMD+=(--llama-backend vulkan)
_explicit_vulkan_backend=true
step "llama.cpp" "Vulkan selected for GGUF inference; the PyTorch training backend is unchanged" "$C_OK"
fi
;;
""|auto|hip|rocm) ;;
*) step "llama.cpp" "Ignoring UNSLOTH_LLAMA_CPP_BACKEND='$_llama_backend' (expected 'auto', 'cpu', 'vulkan', 'hip', or 'rocm')" "$C_WARN" >&2 ;;
esac
if [ "$_HOST_SYSTEM" != "Darwin" ]; then
case "$_llama_backend" in
cpu|vulkan|hip|rocm) ;;
*)
case "$_legacy_force_vulkan" in
1|true|yes|on) _explicit_vulkan_backend=true ;;
esac
;;
esac
fi
_PREBUILT_LOG="$(mktemp)"
set +e
if _is_verbose; then
@ -1502,15 +1540,28 @@ else
substep "free up disk or move UNSLOTH_STUDIO_HOME/TMPDIR to a larger volume, then re-run"
_LLAMA_CPP_NO_SPACE=true
_has_local_llama_server "$LLAMA_CPP_DIR" || _LLAMA_CPP_DEGRADED=true
# A preserved CUDA/ROCm/CPU server does not satisfy an explicit Vulkan
# request, and it leaves _LLAMA_CPP_DEGRADED false, so without this the
# run reports success on the backend the user asked to replace.
if [ "$_explicit_vulkan_backend" = true ]; then
step "llama.cpp" "Vulkan was explicitly requested, so the installer will not keep the existing backend" "$C_ERR"
exit 1
fi
else
step "llama.cpp" "prebuilt install failed (continuing)" "$C_WARN"
step "llama.cpp" "prebuilt install failed" "$C_WARN"
print_llama_error_log "$_PREBUILT_LOG"
rm -f "$_PREBUILT_LOG"
if [ -d "$LLAMA_CPP_DIR" ]; then
substep "prebuilt update failed; existing install restored"
fi
substep "falling back to source build"
_NEED_LLAMA_SOURCE_BUILD=true
if [ "$_explicit_vulkan_backend" = true ]; then
step "llama.cpp" "Vulkan was explicitly requested, so the installer will not substitute a ROCm or CPU source build" "$C_ERR"
substep "check the download error above or try a different UNSLOTH_LLAMA_RELEASE_TAG"
exit 1
else
substep "falling back to source build"
_NEED_LLAMA_SOURCE_BUILD=true
fi
fi
fi

20
tests/studio/load_freeze/test_load_orchestrator.py
View file

@ -3,6 +3,7 @@
from __future__ import annotations
import asyncio
import importlib.util
import os
import re
import socket
@ -49,7 +50,24 @@ import logging as _logging # noqa: E402
_loggers_stub = types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: _logging.getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
sys.modules.setdefault("structlog", types.ModuleType("structlog"))
# structlog is a hard studio.txt requirement, but it is only imported lazily, so a
# bare setdefault here used to park an empty placeholder BEFORE anything imported
# the real package -- and it then shadowed it for the rest of the session. Every
# later file importing a studio module that calls structlog.get_logger at module
# scope (routes.inference -> core.inference.external_provider, utils.mlx_repair)
# blew up with AttributeError, but only when this file was collected first, so the
# same test passed alone and failed under `pytest tests/studio`. Only stub when the
# package is genuinely missing, and give the stub the attribute those callers use.
# Guard on sys.modules FIRST: another test module may have parked its own bare
# stub, and find_spec() raises ValueError on a module whose __spec__ is None.
# Anything already there (real or stub) is left alone; only a genuinely absent
# package gets stubbed.
if "structlog" not in sys.modules and importlib.util.find_spec("structlog") is None:
_structlog_stub = types.ModuleType("structlog")
_structlog_stub.get_logger = lambda *args, **kwargs: _logging.getLogger(
args[0] if args else "structlog"
)
sys.modules["structlog"] = _structlog_stub
import httpx # noqa: E402

191
tests/studio/test_autoload_hf_token_preflight.py Normal file
View file

@ -0,0 +1,191 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Background auto-load must prepare the stored HF token before its GGUF
metadata preflight.
The Hub rejects an invalid Authorization header with 401 even for a PUBLIC
repo. ``fetchGgufStagedMetadata`` posts to the same /api/inference/validate
endpoint ``validateModel`` uses, and ``parseJsonOrThrow`` turns a non-OK
response into a throw. In ``loadAutoLoadCandidate`` that preflight runs BEFORE
``validateModel``, and every call site of ``loadAutoLoadCandidate`` is wrapped
in ``catch { hadNonTrustFailure = true; continue; }``. So a stale saved token
made auto-load skip a cached model that would have loaded anonymously, without
ever reaching validateModel's "continue anonymously / replace token" recovery.
The real classification block is sliced verbatim out of chat-adapter.ts and run
under node, so this asserts on the token value that actually reaches the
request rather than on the presence of a symbol.
"""
import json
import os
import shutil
import subprocess
import tempfile
import textwrap
from pathlib import Path
import pytest
WORKDIR = Path(__file__).resolve().parents[2]
def _source_path(relative_path: str) -> Path:
direct = WORKDIR / relative_path
if direct.exists():
return direct
return WORKDIR / "unsloth_repo" / relative_path
ADAPTER = _source_path("studio/frontend/src/features/chat/api/chat-adapter.ts")
TEMP = WORKDIR / "temp" / "autoload_hf_token_preflight"
def _require_node():
if shutil.which("node") is None:
pytest.skip("node not available")
if not ADAPTER.exists():
pytest.skip("studio chat sources not present")
result = subprocess.run(
["node", "--experimental-strip-types", "--version"],
capture_output = True,
text = True,
timeout = 5,
)
if result.returncode != 0:
pytest.skip("node --experimental-strip-types not available")
def _classification_slice() -> str:
"""The verbatim `isDiffusion` classification block from loadAutoLoadCandidate.
Anchored on the declaration and on the `effectiveGpuIds` statement that
consumes it, so the slice tracks either the prepared-token form or the
older raw-token ternary.
"""
src = ADAPTER.read_text(encoding = "utf-8")
anchor = src.index("async function loadAutoLoadCandidate(")
starts = [
pos
for pos in (
src.find("let isDiffusion", anchor),
src.find("const isDiffusion", anchor),
)
if pos != -1
]
assert starts, "could not locate the isDiffusion classification block"
start = min(starts)
end = src.index("const effectiveGpuIds", start)
return src[start:end].rstrip()
def _run(script: str, harness: str):
_require_node()
TEMP.mkdir(parents = True, exist_ok = True)
workdir = Path(tempfile.mkdtemp(prefix = "run", dir = TEMP))
(workdir / "harness.ts").write_text(harness, encoding = "utf-8")
(workdir / "run.mts").write_text(script, encoding = "utf-8")
env = dict(os.environ, NODE_NO_WARNINGS = "1")
result = subprocess.run(
["node", "--experimental-strip-types", "--no-warnings", "run.mts"],
cwd = str(workdir),
capture_output = True,
text = True,
timeout = 30,
env = env,
)
assert result.returncode == 0, f"stderr: {result.stderr}\nstdout: {result.stdout}"
last = [line for line in result.stdout.strip().splitlines() if line.strip()][-1]
return json.loads(last)
_HARNESS_TEMPLATE = """\
// Real classification block, sliced verbatim from chat-adapter.ts.
export async function classify(ctx: any) {{
const {{
candidate,
config,
modelPath,
hfToken,
prepareHfTokenForUse,
fetchGgufStagedMetadata,
}} = ctx;
{slice}
return isDiffusion;
}}
"""
_STALE_TOKEN = "hf_staleTokenFromAnEarlierSession"
def _harness() -> str:
return _HARNESS_TEMPLATE.format(slice = textwrap.indent(_classification_slice(), " "))
_SCRIPT = textwrap.dedent(
"""
import { classify } from "./harness.ts";
const sent: Array<string | null> = [];
// Mirrors prepareHfTokenForUse: an invalid stored token, with the user's
// one-shot "continue anonymously" choice, resolves to a null token.
const prepareHfTokenForUse = async (token: string | null) => {
if (!token) return { proceed: true, token: null };
return { proceed: true, token: null };
};
// Mirrors the Hub via /api/inference/validate + parseJsonOrThrow: any
// non-null Authorization value here is the stale token, and the Hub 401s
// on it even though the repo is public.
const fetchGgufStagedMetadata = async (payload: any) => {
sent.push(payload.hf_token ?? null);
if (payload.hf_token != null) {
throw new Error("401 Unauthorized: Invalid credentials in Authorization header");
}
return { isDiffusion: true };
};
let threw: string | null = null;
let isDiffusion: boolean | null = null;
try {
isDiffusion = await classify({
candidate: { kind: "gguf", ggufVariant: "Q4_K_M" },
config: { selectedGpuIds: [0] },
modelPath: "unsloth/some-public-gguf",
hfToken: %s,
prepareHfTokenForUse,
fetchGgufStagedMetadata,
});
} catch (e) {
threw = String((e as Error).message);
}
console.log(JSON.stringify({ sent, threw, isDiffusion }));
"""
)
def test_autoload_preflight_sends_the_prepared_token_not_the_stale_one():
out = _run(_SCRIPT % json.dumps(_STALE_TOKEN), _harness())
assert out["threw"] is None, (
"a stale saved token aborted the auto-load metadata preflight; the "
f"candidate would be skipped: {out['threw']}"
)
assert out["sent"] == [None], (
"the GGUF metadata preflight must send the prepared token, not the raw "
f"stored one; it sent {out['sent']!r}"
)
assert out["isDiffusion"] is True
def test_autoload_preflight_is_skipped_without_a_remembered_gpu_pick():
"""No remembered GPU selection means no preflight and no token use at all."""
script = _SCRIPT % json.dumps(_STALE_TOKEN)
script = script.replace("selectedGpuIds: [0]", "selectedGpuIds: null")
out = _run(script, _harness())
assert out["sent"] == []
assert out["threw"] is None
assert out["isDiffusion"] is False

274
tests/studio/test_model_picker_contracts.py
View file

@ -188,7 +188,13 @@ def test_active_model_config_round_trips_gpu_fields():
a sidebar/hub-gear reload cannot silently reset manual GPU settings, and
"Remember settings" cannot persist a GPU-less config over a saved one."""
src = _read("features/model-picker/hooks/use-active-model-config.ts")
for field in ("gpuMemoryMode", "gpuLayers", "nCpuMoe", "selectedGpuIds"):
for field in (
"gpuMemoryMode",
"gpuLayers",
"nCpuMoe",
"selectedGpuIds",
"selectedGpuIndexKind",
):
assert field in src, field
assert "if (!isGguf)" in src and "return base" in src
for rel in (
@ -202,6 +208,24 @@ def test_active_model_config_round_trips_gpu_fields():
assert "export function gpuFieldsSignature" in shared
def test_deferred_gpu_pick_keeps_its_index_namespace():
"""A remembered pick restored before GPU discovery must keep its namespace
until load-time reconciliation, or Vulkan IDs can be reused as physical IDs."""
store = _read("features/chat/stores/chat-runtime-store.ts")
assert "selectedGpuIndexKind: GpuIndexKind | null;" in store
apply = _read("features/model-picker/model-config/apply-per-model-config.ts")
assert "selectedGpuIndexKind: s.selectedGpuIndexKind" in apply
assert "config.selectedGpuIndexKind === undefined" in apply
assert 'config.selectedGpuIndexKind ?? "physical"' not in apply
runtime = _read("features/chat/hooks/use-chat-model-runtime.ts")
assert "stateBeforeUnload.selectedGpuIndexKind," in runtime
compare = _read("features/chat/shared-composer.tsx")
assert "store.selectedGpuIndexKind," in compare
def test_gpu_picker_round_trips_requested_pool_not_fitted_subset():
"""A GGUF fit may narrow [0, 1] to [0], but load/status hydration must keep
[0, 1] as the editable pool so a later reload can grow back onto GPU 1."""
@ -209,10 +233,20 @@ def test_gpu_picker_round_trips_requested_pool_not_fitted_subset():
assert types.count("requested_gpu_ids?: number[] | null") >= 2
store = _read("features/chat/stores/chat-runtime-store.ts")
assert "resp.requested_gpu_ids ?? resp.gpu_ids ?? null" in store
assert 'hasOwnProperty.call(resp, "requested_gpu_ids")' in store
assert "const reportedGpuIds = requestedGpuIdsFromResponse(resp)" in store
assert "loadedGpuIndexKind: GpuIndexKind | null;" in store
assert "loadedGpuIndexKind: gpuIds == null ? null : (gpuIndexKind ?? null)" in store
# A cold discovery cache reports gpuIndexKind === undefined (deferred, not
# rejected); only a warm cache's definitive null should drop the pin.
assert "reportedGpuIds != null && gpuIndexKind !== null" in store
status = _read("features/chat/lib/apply-inference-status-to-store.ts")
assert "status.requested_gpu_ids ?? status.gpu_ids ?? null" in status
assert "const incomingGpuFields = loadedGpuMemoryFields(status)" in status
assert "const incomingGpuIds = incomingGpuFields.loadedGpuIds" in status
assert "const incomingGpuIndexKind = incomingGpuFields.loadedGpuIndexKind" in status
assert status.count("prevState.loadedGpuIndexKind") >= 2
assert status.count("sameGpuSelection(") >= 2
def test_compare_load_uses_each_models_gpu_config():
@ -221,7 +255,9 @@ def test_compare_load_uses_each_models_gpu_config():
assert "ownConfig.gpuLayers ?? compareLoadKnobs.gpuLayers" in src
assert "ownConfig.nCpuMoe ?? compareLoadKnobs.nCpuMoe" in src
assert "if (ownConfig.selectedGpuIds != null)" in src
assert "reconcilePersistedGpuIds(ownConfig.selectedGpuIds)" in src
assert "ownConfig.selectedGpuIndexKind," in src
assert "compareLoadKnobs.selectedGpuIndexKind," in src
assert src.count("resolvedIsDiffusion === true") >= 2
for field in (
"gpu_memory_mode: effectiveGpuMemoryMode",
"gpu_layers: effectiveGpuLayers",
@ -230,6 +266,57 @@ def test_compare_load_uses_each_models_gpu_config():
):
assert field in src
page = _read("features/chat/chat-page.tsx")
assert page.count("isDiffusion: meta.isDiffusion") >= 2
assert "isDiffusion: globalIsDiffusion" in page
def test_diffusion_load_paths_disable_tensor_parallel():
"""Every frontend path that classifies DiffusionGemma must send tensor
parallelism as false so the ignored setting cannot force repeat reloads."""
compare = _read("features/chat/shared-composer.tsx")
compact_compare = " ".join(compare.split())
assert "const effectiveTensorParallel = resolvedIsDiffusion ? false :" in compact_compare
assert compare.count("tensor_parallel: effectiveTensorParallel") == 2
runtime = " ".join(_read("features/chat/hooks/use-chat-model-runtime.ts").split())
assert "const loadTensorParallel = targetIsDiffusion ? false :" in runtime
apply = " ".join(_read("features/model-picker/model-config/apply-per-model-config.ts").split())
assert "tensorParallel: options.isDiffusion ? false :" in apply
adapter = _read("features/chat/api/chat-adapter.ts")
autoload = adapter.split("async function loadAutoLoadCandidate", 1)[1]
autoload = autoload.split("async function autoLoadSmallestModel", 1)[0]
compact_autoload = " ".join(autoload.split())
assert "config.selectedGpuIds != null || config.tensorParallel === true" in compact_autoload
assert "const effectiveTensorParallel = isDiffusion ? false :" in compact_autoload
assert autoload.count("tensor_parallel: effectiveTensorParallel") == 2
def test_diffusion_load_keeps_the_standing_gpu_memory_mode():
"""A diffusion config is sanitized to gpuMemoryMode "auto" because the mode does
not apply to it, not because the user picked Auto. Applying that sanitized value
to the runtime store would strand the session on Auto: persistGpuMemoryModeOnLoad
deliberately skips diffusion responses, so nothing writes the standing preference
back, and the next ordinary GGUF loaded without its own config sends the stale
"auto" and persists it over the user's Manual."""
apply = " ".join(_read("features/model-picker/model-config/apply-per-model-config.ts").split())
assert (
"gpuMemoryMode: options.isDiffusion ? readPersistedGpuMemoryMode() : "
"(config.gpuMemoryMode ?? readPersistedGpuMemoryMode())," in apply
)
# The unconditional form is what leaked the sanitized mode into the store.
assert "gpuMemoryMode: config.gpuMemoryMode ?? readPersistedGpuMemoryMode()," not in apply
# The other half of the contract: the diffusion load itself must not persist.
store = " ".join(_read("features/chat/stores/chat-runtime-store.ts").split())
assert "if (resp.is_gguf && !resp.is_diffusion) saveGpuMemoryMode(mode);" in store
# And the sanitizer that produces the "auto" this guards against still runs.
page = " ".join(_read("features/model-picker/components/model-config-page.tsx").split())
assert "withoutUnsupportedDiffusionSettings(committedConfig, gpuIndexKind)" in page
def test_active_native_gguf_metadata_uses_path_token():
src = _read("features/model-picker/components/model-config-page.tsx")
@ -393,7 +480,7 @@ def test_fixed_layer_gguf_pins_displayed_context():
instead of sending native/0 and recreating the OOM."""
src = _read("features/model-picker/components/model-config-page.tsx")
assert "const pinFixedLayerContext =" in src
assert 'config.gpuMemoryMode === "manual"' in src
assert 'loadableConfig.gpuMemoryMode === "manual"' in src
assert "customContextLength: activeLoadedContext" in src
@ -469,7 +556,7 @@ def test_reset_persists_null_max_length_and_substitutes_only_for_load():
assert "const effectiveLoadConfig" in src
# The persisted record is saved from effectiveRuntimeConfig; the load request
# carries effectiveLoadConfig (with any committed context input).
assert "onRun(effectiveLoadConfig)" in src
assert "onRun(effectiveLoadConfig, classifiedIsDiffusion)" in src
assert "savePerModelConfig(" in src
@ -581,7 +668,7 @@ def test_compare_pane_non_gguf_falls_back_to_app_default():
assert (
"const effectiveMaxSeqLength = ownConfig.customContextLength ?? "
"normalizeMaxSeqLength(ownConfig.maxSeqLength) ?? "
"(isGgufLoad ? 0 : DEFAULT_MAX_SEQ_LENGTH);" in src
"(targetIsGguf ? 0 : DEFAULT_MAX_SEQ_LENGTH);" in src
)
# The buggy fallback to the active model's shared runtime value must not return.
assert "(isGgufLoad ? 0 : maxSeqLength)" not in src
@ -597,6 +684,140 @@ def test_default_gpu_mode_clears_manual_knobs():
assert "gpuLayers: undefined," in src
assert "nCpuMoe: undefined," in src
assert "selectedGpuIds: undefined," in src
assert "selectedGpuIndexKind: undefined," in src
def test_requested_gpu_pick_survives_fit_narrowing_and_namespace_changes():
store = _read("features/chat/stores/chat-runtime-store.ts")
assert "requestedGpuIdsFromResponse(resp)" in store
gpu_selection = _read("hooks/gpu-selection.ts")
assert 'savedIndexKind === undefined ? "physical" : savedIndexKind' in gpu_selection
assert "currentIndexKind !== undefined" in gpu_selection
assert "expectedIndexKind !== currentIndexKind" in gpu_selection
assert "A null namespace is only safe while discovery is also unresolved" in gpu_selection
assert (
"Namespace knowledge is authoritative even while membership is unavailable" in gpu_selection
)
gpu_info = _read("hooks/use-gpu-info.ts")
assert "cachedPinnableGpuContext" in gpu_info
assert 'unavailableVulkan ? "vulkan" : undefined' in gpu_info
assert "cachedPinnableGpuIndexKind" in gpu_info
config = _read("features/model-picker/model-config/per-model-config.ts")
assert '"selectedGpuIndexKind"' in config
def test_staged_gpu_pick_reconciles_with_async_namespace_discovery():
page = _read("features/model-picker/components/model-config-page.tsx")
assert "reconcileGpuSelection(" in page
assert "setConfig((current)" in page
assert "reconcileConfigGpuSelection(" in page
assert "gpuIndexKind" in page
assert "pinnableDevices" in page
assert "reconciled.ids === null ? undefined : reconciled.indexKind" in " ".join(page.split())
assert "selectsAll ? null : next" not in page
gpu_info = _read("hooks/use-gpu-info.ts")
assert 'inferenceGpu?.backend === "vulkan"' in gpu_info
assert "!inferenceGpu.available" in gpu_info
def test_compare_classifies_gguf_before_reconciling_gpu_ids():
compare = _read("features/chat/shared-composer.tsx")
metadata = compare.index("fetchGgufStagedMetadata(")
reconcile = compare.index("reconcilePersistedGpuIds(", metadata)
validate = compare.index("const validation = await validateModel(", reconcile)
load = compare.index("const resp = await loadModel(", validate)
assert metadata < reconcile < validate < load
assert compare.count("resolvedIsDiffusion") >= 3
assert "prepareHfTokenForUse(" in compare
page = _read("features/model-picker/components/model-config-page.tsx")
assert "isDiffusion?: boolean" in page
assert "onRun(effectiveLoadConfig, classifiedIsDiffusion)" in page
def test_model_config_prepares_hf_token_before_gguf_metadata_preflight():
"""Settings classification must use the same stale-token recovery as load."""
page = _read("features/model-picker/components/model-config-page.tsx")
assert 'import { prepareHfTokenForUse } from "@/features/hf-auth";' in page
effect = page.split("// Fetch GGUF header dims", 1)[1]
effect = effect.split("const stagedDims =", 1)[0]
prepare = effect.index("prepareHfTokenForUse(hfToken || null)")
metadata = effect.index("fetchGgufStagedMetadata({", prepare)
assert prepare < metadata
assert "if (!preparedToken.proceed)" in effect
cancel = effect.index("if (!preparedToken.proceed)")
assert "settleWithoutMetadata();" in effect[cancel:metadata]
assert effect.count("settleWithoutMetadata();") == 2
assert "hf_token: preparedToken.token" in effect
assert "hf_token: hfToken" not in effect
def test_chat_load_prepares_hf_token_before_gguf_metadata_preflight():
"""The single-model load path classifies a GGUF via fetchGgufStagedMetadata
before validateModel/loadModel run. The Hub rejects an invalid Authorization
header with 401 even for a PUBLIC repo, so that preflight must prepare the
token like every other caller; otherwise a stale saved token aborts the whole
load instead of offering the "continue anonymously / replace token" recovery.
"""
runtime = _read("features/chat/hooks/use-chat-model-runtime.ts")
prepare = runtime.index("prepareHfTokenForUse(")
metadata = runtime.index("fetchGgufStagedMetadata({", prepare)
assert prepare < metadata
# The raw store token must not be handed to the preflight.
assert "hf_token: preparedToken.token" in runtime
assert (
"hf_token: useChatRuntimeStore.getState().hfToken" not in runtime
), "GGUF metadata preflight must not send the unprepared stored token"
assert 'throw new Error("Model load cancelled.")' in runtime
def test_chat_autoload_prepares_hf_token_before_gguf_metadata_preflight():
"""Background autoload performs the same GGUF classification probe as the
interactive paths, so a stale saved token must take the same recovery path.
"""
adapter = _read("features/chat/api/chat-adapter.ts")
autoload = adapter.split("async function loadAutoLoadCandidate", 1)[1]
autoload = autoload.split("async function autoLoadSmallestModel", 1)[0]
prepare = autoload.index("prepareHfTokenForUse(")
metadata = autoload.index("fetchGgufStagedMetadata({", prepare)
assert prepare < metadata
assert "hf_token: preparedToken.token" in autoload
assert 'throw new Error("Model load cancelled.")' in autoload
def test_cpu_only_llama_build_hides_gpu_picker():
src = (WORKDIR / "studio" / "backend" / "main.py").read_text(encoding = "utf-8")
assert "and not LlamaCppBackend._backend_lacks_gpu_lib()" in src
def test_vulkan_inference_devices_do_not_replace_global_gpu_info():
backend = (WORKDIR / "studio" / "backend" / "main.py").read_text(encoding = "utf-8")
assert '"gguf_gpu_ids_supported": gpu_ids_supported' in backend
assert '"inference_gpu": inference_gpu_info' in backend
frontend = _read("hooks/use-gpu-info.ts")
assert 'inference?.backend === "vulkan"' in frontend
assert "data?.gpu?.devices ?? []" in frontend
def test_diffusion_picker_hides_and_clears_unsupported_memory_modes():
api = _read("features/chat/api/chat-api.ts")
assert "isDiffusion: res.is_diffusion ?? false" in api
page = _read("features/model-picker/components/model-config-page.tsx")
assert 'className={isDiffusion ? "hidden" : ROW_CLASS}' in page
assert "withoutUnsupportedDiffusionSettings(config, gpuIndexKind)" in page
assert "reconcileConfigGpuSelection(" in page
assert "resolvedIsDiffusion" in page
assert "stagedMetadataPending ||" in page
assert "config.selectedGpuIds != null" in page
for field in (
'gpuMemoryMode: "auto"',
"gpuLayers: undefined",
"nCpuMoe: undefined",
"tensorParallel: false",
"selectedGpuIds: undefined",
"selectedGpuIndexKind: undefined",
):
assert field in page
def test_legacy_migration_is_idempotent_and_non_destructive():
@ -862,8 +1083,10 @@ def test_failed_switch_rollback_restores_the_slot_intent_not_the_resolved_count(
"selection.previousConfig ? (selection.previousConfig.nParallel ?? null) "
": useChatRuntimeStore.getState().nParallel;" in runtime
)
# Matched on the call prefix, not the whole call: the staged apply also
# carries the resolved diffusion flag. Only the ordering is the contract.
assert runtime.index("const previousNParallel") < runtime.index(
"applyPerModelConfigToRuntime(pendingLoadConfig);"
"applyPerModelConfigToRuntime(pendingLoadConfig,"
), "a config staged on the selection must not replace it either"
picker = " ".join(_read("features/chat/chat-page.tsx").split())
assert (
@ -886,14 +1109,31 @@ def test_vulkan_inference_devices_are_the_pickable_set():
applicator speaks, and a Vulkan pick does not use them.
"""
src = " ".join(_read("hooks/use-gpu-info.ts").split())
# The Vulkan inventory is consulted first, and only when it has devices.
assert (
"const inference = data?.inference_gpu; "
'if (inference?.backend === "vulkan" && (inference.devices ?? []).length) {' in src
)
# The Vulkan inventory is authoritative even while its probe is temporarily
# empty. Falling through would expose physical CUDA/ROCm IDs in an ordinal
# picker and make DiffusionGemma offer a selection the route rejects.
assert "const inference = data?.inference_gpu; " 'if (inference?.backend === "vulkan") {' in src
# A confirmed-Vulkan backend with no enumerated devices yet must return no
# devices, not fall through to the torch/CUDA inventory below.
assert "if (!(inference.devices ?? []).length) return [];" in src
# Pinnable on the ggml ordinal space, gated on the backend's own support flag.
assert "const picksAccepted = inference.gguf_gpu_ids_supported !== false;" in src
assert 'physicalIndex: picksAccepted && d.index_kind === "vulkan",' in src
# The torch fallback keeps its physical-only gate and the XPU ban.
assert 'data?.device_backend !== "xpu" &&' in src
assert 'physicalIndex: pinnableBackend && d.index_kind === "physical",' in src
assert 'pinnable: picksAccepted && d.index_kind === "vulkan",' in src
# A Vulkan ordinal is not a CUDA ID, so the torch-side diffusion runner
# cannot take the pick even when llama-server can.
assert "diffusionPinnable: false," in src
# The torch fallback keeps the XPU ban for torch ordinals only: a Vulkan
# ordinal stays pickable even when this list arrives from an XPU host.
assert (
"pinnable: pinnableBackend && "
'(d.index_kind === "vulkan" || '
'(data?.device_backend !== "xpu" && d.index_kind === "physical")),' in src
)
# Only a physical index is ever handed to the diffusion runner, and ROCm
# counts: it reuses torch.cuda.* and the same physical-ID path, so excluding
# it would hide the picker on every multi-GPU ROCm host.
assert (
"const diffusionBackend = "
'data?.device_backend === "cuda" || data?.device_backend === "rocm";' in src
)
assert 'diffusionPinnable: diffusionBackend && d.index_kind === "physical",' in src

8
tests/studio_setup_ps1/Studio.Setup.Vs2026.Tests.ps1
View file

@ -27,10 +27,16 @@ BeforeAll {
if (-not $script:SetupPs1) { throw "Could not locate studio/setup.ps1 (set SETUP_PS1_PATH)." }
Write-Host "setup.ps1 under test: $script:SetupPs1"
# Test-VCRedistInstalled consults Get-HostMachineArch before trusting the
# System32 DLL, and dot-sourcing functions one at a time does not carry that
# callee in. The registry hit returns early, so only the clean-box cases reach
# the call and fail there with "Get-HostMachineArch is not recognized"; the
# same list in the VC++ round-trip job gained the helper in #7597.
foreach ($fn in @('Resolve-VsGeneratorFromLabel', 'Find-VsBuildTools', 'Get-VcBuildCustomizationsDir',
'Test-CmakeSupportsGenerator', 'Get-CmakeVersion', 'Test-CmakeListsGenerator',
'Test-CmakeCanDriveGenerator', 'Get-FallbackVsGenerator',
'Ensure-BuildToolsForLlamaSourceBuild', 'Test-VCRedistInstalled')) {
'Ensure-BuildToolsForLlamaSourceBuild', 'Get-HostMachineArch',
'Test-VCRedistInstalled')) {
$src = Get-FunctionSource -Path $script:SetupPs1 -Name $fn
if (-not $src) { throw "Function '$fn' not found in $script:SetupPs1 - cannot test the real code." }
. ([scriptblock]::Create($src))