diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3393f2b4be..7fa117980b 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -31,8 +31,12 @@ from core.inference.llama_server_args import ( extra_args_disable_mmproj, parse_cache_override, parse_ctx_override, + parse_split_mode_override, resolve_cache_type_kv, resolve_requested_ctx, + resolve_tensor_parallel, + strip_shadowing_flags, + strip_split_mode_only, ) from core.tool_healing import ( _TC_END_TAG_RE, @@ -705,6 +709,8 @@ class LlamaCppBackend: self._supports_preserve_thinking: bool = False self._supports_tools: bool = False self._cache_type_kv: Optional[str] = None + # Whether --split-mode tensor was applied on the active load. + self._tensor_parallel: bool = False self._reasoning_default: bool = True self._speculative_type: Optional[str] = None # Canonical UI-facing mode the user requested @@ -1003,6 +1009,11 @@ class LlamaCppBackend: def cache_type_kv(self) -> Optional[str]: return self._cache_type_kv + @property + def tensor_parallel(self) -> bool: + """Whether --split-mode tensor is active on the loaded server.""" + return self._tensor_parallel + @property def speculative_type(self) -> Optional[str]: return self._speculative_type @@ -1606,6 +1617,23 @@ class LlamaCppBackend: # buffers; 0.90 dropped 91-94% fits to CPU offload (#5106). _GPU_PIN_VRAM_FRACTION = 0.95 + # Per-GPU compute-graph buffer to reserve in tensor mode (MiB). This is the + # logits buffer (n_batch x vocab) + activation scratch that llama.cpp sizes + # via graph_reserve -- it is roughly EQUAL on every device (not proportional + # to the tensor split) and independent of context. Measured ~2.3 GB + # (gemma-3-27B) to ~3.8 GB (gemma-4-31B) on a 256k-vocab model; we reserve a + # conservative headroom above that. It is (a) subtracted from each GPU's free + # VRAM before computing --tensor-split, so the roomier GPU absorbs more + # weight and the smallest GPU keeps room for KV, and (b) reserved per device + # when capping context. The auto-fallback to layer split covers any + # underestimate. NOTE: scales with the model's vocab / batch size; tune if a + # large-vocab model OOMs at load. + _TENSOR_PARALLEL_BUFFER_RESERVE_MIB = 5120 + + # KV cache types llama.cpp accepts in tensor mode. A quantized KV cache + # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. + _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + @staticmethod def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]: """Return DLL dirs from pip-installed CUDA wheels under @@ -1908,6 +1936,7 @@ class LlamaCppBackend: ctx_checkpoints: int = 0, kv_on_gpu: bool = True, mtp_engaged: bool = False, + budget_frac: Optional[float] = None, ) -> int: """Return the largest context length that fits in GPU VRAM. @@ -1941,8 +1970,11 @@ class LlamaCppBackend: ctx_checkpoints = ctx_checkpoints, ) - # MTP engaged: carve the drafter's reserve out of the fit budget. - budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if mtp_engaged else 0.0) + # MTP engaged: carve the drafter's reserve out of the fit budget. Callers + # can override outright (tensor-parallel mode passes a fatter margin), so + # only compute a default when none was supplied. + if budget_frac is None: + budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if mtp_engaged else 0.0) budget_bytes = available_mib * 1024 * 1024 * budget_frac model_footprint = model_size_bytes @@ -2751,6 +2783,16 @@ class LlamaCppBackend: """ lowered = (output or "").lower() + # Tensor parallelism (--split-mode tensor) is arch-gated in llama.cpp; + # unsupported architectures abort the load with this marker. Point the + # user at the toggle instead of a generic invalid-GGUF/OOM message. + if "split_mode_tensor not implemented" in lowered: + return ( + "Tensor parallelism is not supported for this model's " + "architecture. Turn off Tensor Parallelism in the model " + "settings and reload." + ) + # Detect Ollama source up front so the arch branch can keep the # Ollama hint instead of the generic "unsupported arch" message. gguf = gguf_path or "" @@ -2807,6 +2849,97 @@ class LlamaCppBackend: "Check that the GGUF file is valid and you have enough memory." ) + def _plan_tensor_parallel( + self, + gpus: list[tuple[int, int]], + model_size: int, + target_ctx: int, + cache_type_kv: Optional[str] = None, + n_parallel: int = 1, + mtp_engaged: bool = False, + max_target_ctx: Optional[int] = None, + ) -> tuple[int, int, list[int], Optional[list[int]]]: + """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed. + + ``gpus`` is a list of ``(gpu_index, free_mib)``; ``model_size`` is the + weight size in bytes; ``target_ctx`` is the context to fit (the explicit + request, or the model's native length for auto). ``max_target_ctx`` is + the native/hardware ceiling used only for the UI bound (defaults to + ``target_ctx``). Returns + ``(effective_ctx, max_available_ctx, gpu_indices, tensor_split)``. + + Policy (assumes >= 2 GPUs; the caller drops the toggle below that): + - Cap context to the KV that fits the pooled VRAM after the weights and + one per-device compute-graph buffer (``_TENSOR_PARALLEL_BUFFER_RESERVE_MIB``). + llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only + cap, honored even for an explicit ``-c``. It is more accurate than the + 0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused. + - ``tensor_split`` is None (llama.cpp's even default, safe for every arch + incl. Gemma 3n which GGML_ASSERTs on a weighted split) when an even + share fits the smallest GPU; otherwise it is weighted by + ``(free - buffer)`` so the roomier GPU absorbs more weight and the + smallest GPU keeps room for KV. + """ + # Drop GPUs that can't hold the per-device compute-graph buffer; they'd + # OOM in tensor mode. load_model already filters before calling, so this + # is defense-in-depth that also keeps the pure function self-contained + # (and unit-testable without a GPU). + reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + usable_gpus = [g for g in gpus if g[1] >= reserve_mib] + gpu_indices = sorted(idx for idx, _ in usable_gpus) + if len(gpu_indices) < 2: + # Tensor parallelism is meaningless on <2 GPUs (the caller drops the + # toggle before this); be defensive and never emit a split here. + return ( + target_ctx if target_ctx > 0 else 4096, + target_ctx if target_ctx > 0 else 4096, + gpu_indices, + None, + ) + free_by_idx = {idx: free for idx, free in usable_gpus} + pool_mib = sum(free_by_idx.values()) + kv_budget_b = (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size + if mtp_engaged: + # MTP keeps a draft model + its own KV cache on GPU. + kv_budget_b -= 2 * 1024**3 + + def _fit_ctx(ctx: int) -> int: + # Largest context whose KV fits the pooled budget. Floors small, but + # never raises an explicit ctx above what was asked. + if self._can_estimate_kv() and ctx > 0: + ctx_floor = min(2048, ctx) + if kv_budget_b <= 0: + # Weights + buffers exceed the pool -> floor; the load then + # falls back to layer split. + return ctx_floor + kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel) + if kv_at <= kv_budget_b: + return ctx + return max(ctx_floor, int(ctx * kv_budget_b / kv_at)) + # KV size unknown -> can't prove a safe cap; floor. + return min(4096, ctx) if ctx > 0 else 4096 + + # max_available_ctx is the hardware ceiling for the UI bound, sized from + # the native context independent of an explicit small -c (which only + # caps effective_ctx). + max_ctx_target = max_target_ctx if (max_target_ctx and max_target_ctx > 0) else target_ctx + max_available_ctx = _fit_ctx(max_ctx_target) + effective_ctx = min(_fit_ctx(target_ctx), max_available_ctx) + + min_free_mib = min(free_by_idx.values()) + kv_bytes = ( + self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = n_parallel) + if (self._can_estimate_kv() and effective_ctx > 0) + else 0 + ) + even_share_mib = (model_size + kv_bytes) / len(gpu_indices) / (1024 * 1024) + tensor_split: Optional[list[int]] = None + if even_share_mib > (min_free_mib - reserve_mib): + adj = [max(0, int(free_by_idx[i] - reserve_mib)) for i in gpu_indices] + if sum(adj) > 0: + tensor_split = adj + return effective_ctx, max_available_ctx, gpu_indices, tensor_split + @staticmethod def _is_projector_incompatibility(output: str) -> bool: """True when llama-server aborted because it cannot load the model's @@ -2929,6 +3062,7 @@ class LlamaCppBackend: cache_type_kv: Optional[str] = None, speculative_type: Optional[str] = None, spec_draft_n_max: Optional[int] = None, + tensor_parallel: bool = False, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # caller compat, unused n_parallel: int = 1, @@ -2960,6 +3094,7 @@ class LlamaCppBackend: cache_type_kv = cache_type_kv, speculative_type = speculative_type, spec_draft_n_max = spec_draft_n_max, + tensor_parallel = tensor_parallel, chat_template_override = chat_template_override, extra_args = extra_args, is_vision = is_vision, @@ -3099,10 +3234,43 @@ class LlamaCppBackend: requested_ctx = resolve_requested_ctx(extra_args, n_ctx) cache_override = parse_cache_override(extra_args) cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv) + # A user --split-mode in extras last-wins-overrides the + # toggle, so reconcile it back into tensor_parallel state. + split_mode_override = parse_split_mode_override(extra_args) + tensor_parallel = resolve_tensor_parallel(extra_args, tensor_parallel) + # Tensor mode aborts on a quantized KV cache, so drop it for the + # tensor attempt (and strip any inherited/explicit --cache-type + # that would re-impose it when appended last). The layer-split + # fallback re-runs with tensor_parallel False and keeps the type. + if ( + tensor_parallel + and cache_type_kv + and cache_type_kv.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES + ): + logger.info( + "Tensor parallelism requires a non-quantized KV cache; " + "ignoring cache type %s for the tensor attempt.", + cache_type_kv, + ) + cache_type_kv = None + if extra_args: + extra_args = strip_shadowing_flags( + extra_args, + strip_context = False, + strip_cache = True, + strip_spec = False, + strip_template = False, + strip_split_mode = False, + ) if ctx_override is not None and ctx_override > 0: logger.info(f"User --ctx-size {ctx_override} honored; skipping auto-reduce") if cache_override is not None: logger.info(f"User --cache-type-k/-v {cache_override} honored for KV estimate") + if split_mode_override is not None: + logger.info( + f"User --split-mode {split_mode_override} honored; " + "reconciled into tensor_parallel state" + ) effective_ctx = requested_ctx if requested_ctx > 0 else (self._context_length or 0) max_available_ctx = self._context_length or effective_ctx gpus: list[tuple[int, int]] = [] @@ -3164,6 +3332,8 @@ class LlamaCppBackend: # Auto n_ctx=0 (native): prefer fewer GPUs with reduced # context, since multi-GPU is slower. gpu_indices, use_fit = None, True + # Per-GPU weight proportions for tensor mode (None = even). + tp_tensor_split: Optional[list[int]] = None explicit_ctx = requested_ctx > 0 # MTP draft model lives outside the main estimates; carve # its reserve out of every fit budget and pin threshold so @@ -3171,10 +3341,67 @@ class LlamaCppBackend: _mtp_reserve = _MTP_VRAM_RESERVE_FRAC if _mtp_will_engage else 0.0 _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _mtp_reserve - if gpus and self._can_estimate_kv() and effective_ctx > 0: - # Largest hardware-aware cap from the native context - # across all usable GPU subsets (for UI bounds), - # independent of the requested context. + # Tensor mode allocates a compute-graph buffer on every + # participating GPU, so a GPU with less free VRAM than that + # reserve can't host it and would OOM at load. Drop those + # from the tensor-parallel set up front (gpu_indices below + # becomes the CUDA_VISIBLE_DEVICES mask, so they're excluded + # from llama-server entirely, not just given zero weight). + tp_gpus = gpus + if tensor_parallel: + reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + tp_gpus = [g for g in gpus if g[1] >= reserve_mib] + + if tensor_parallel and len(tp_gpus) < 2: + # Tensor parallelism needs >= 2 usable GPUs. On a single + # GPU --split-mode tensor is a no-op; with 0 GPUs (CPU-only + # or probe failed) it must not reach llama-server; and a + # GPU below the buffer reserve can't participate. Drop the + # flag and fall through to normal layer/CPU allocation. + logger.info( + "Tensor parallelism requested but only %d of %d GPU(s) " + "have enough free VRAM for the compute buffer; " + "ignoring (needs >= 2).", + len(tp_gpus), + len(gpus), + ) + tensor_parallel = False + # A user --split-mode tensor in extras is appended after + # Studio's flags, so it would still reach llama-server and + # fail here; strip it so the downgrade actually applies. + extra_args = strip_split_mode_only(extra_args) + + if tensor_parallel and tp_gpus: + # Tensor-parallel allocation: use all usable GPUs, weight + # the split by (free - buffer), and cap context to the + # pooled VRAM after weights + per-device compute-graph + # buffers. See _plan_tensor_parallel for the policy. + target_ctx = ( + effective_ctx + if explicit_ctx + else (self._context_length or effective_ctx) + ) + ( + effective_ctx, + max_available_ctx, + gpu_indices, + tp_tensor_split, + ) = self._plan_tensor_parallel( + tp_gpus, + model_size, + target_ctx, + cache_type_kv = cache_type_kv, + n_parallel = n_parallel, + mtp_engaged = _mtp_will_engage, + # Report the UI ceiling from native ctx, not the + # explicit small request. + max_target_ctx = self._context_length or target_ctx, + ) + use_fit = False + elif gpus and self._can_estimate_kv() and effective_ctx > 0: + # Compute the largest hardware-aware cap from the model's + # native context across all usable GPU subsets (for UI + # bounds), independent of the currently requested context. native_ctx_for_cap = self._context_length or effective_ctx if native_ctx_for_cap > 0: ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) @@ -3297,6 +3524,7 @@ class LlamaCppBackend: except Exception as e: logger.warning(f"GPU selection failed ({e}), using --fit on") gpu_indices, use_fit = None, True + tp_tensor_split = None effective_ctx = requested_ctx # fall back to original launch_mmproj_path = None @@ -3396,6 +3624,27 @@ class LlamaCppBackend: else: self._cache_type_kv = None + # Tensor parallelism: split the model across GPUs by tensor + # rather than by layer. Multi-GPU only -- a no-op on a single + # GPU. Default (layer split) is left implicit by omitting the + # flag. See llama.cpp --split-mode. + if tensor_parallel: + cmd.extend(["--split-mode", "tensor"]) + if tp_tensor_split and len(tp_tensor_split) > 1: + cmd.extend( + [ + "--tensor-split", + ",".join(str(int(x)) for x in tp_tensor_split), + ] + ) + self._tensor_parallel = True + logger.info( + "Tensor parallelism: --split-mode tensor, --tensor-split %s", + tp_tensor_split, + ) + else: + self._tensor_parallel = False + # Speculative decoding. See _build_speculative_flags for the # mode resolution, benchmarks, and llama.cpp references. launch_mtp_draft_path = self._resolve_launch_mtp_path( @@ -4172,6 +4421,7 @@ class LlamaCppBackend: is_vision: bool, gguf_path: Optional[str] = None, spec_draft_n_max: Optional[int] = None, + tensor_parallel: bool = False, mtp_draft_path: Optional[str] = None, ) -> bool: """True iff the live server already satisfies these load kwargs. @@ -4208,6 +4458,12 @@ class LlamaCppBackend: if _norm(self._cache_type_kv) != _norm(cache_type_kv): return False + # Reconcile a user --split-mode in extras (load_model does the same), so + # an extras-driven tensor load isn't seen as a mismatch that needlessly + # kills/reloads a healthy server. + if self._tensor_parallel != resolve_tensor_parallel(extra_args, tensor_parallel): + return False + # Compare on the canonical requested mode. With --spec-type in # extra_args the backend stores None; mirror that here. if _extra_args_set_spec_type(extra_args): @@ -4279,6 +4535,12 @@ class LlamaCppBackend: return None return saw_gpu_buffer + def load_cancelled(self) -> bool: + """True if a load was cancelled (e.g. via unload/_cancel_event) and not + yet consumed by the next load_model. Lets the tensor->layer fallback + avoid restarting a load the user just cancelled.""" + return self._cancel_event.is_set() + def unload_model(self) -> bool: """Terminate the subprocess and cancel any in-flight download.""" self._cancel_event.set() @@ -4311,6 +4573,7 @@ class LlamaCppBackend: self._supports_preserve_thinking = False self._supports_tools = False self._cache_type_kv = None + self._tensor_parallel = False self._speculative_type = None self._requested_spec_mode = None self._spec_draft_n_max = None diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 00f8c66d5c..69a86fa3ba 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -109,6 +109,7 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]: out.append(token) parse_ctx_override(out) parse_cache_override(out) + parse_split_mode_override(out) return out @@ -157,8 +158,20 @@ _TEMPLATE_FLAGS: frozenset[str] = frozenset( "--no-jinja", } ) +# Multi-GPU split mode shadows the Tensor Parallelism toggle +# (--split-mode tensor). Pass-through stays allowed so users keep the +# row/none/layer modes the toggle doesn't expose, but it's stripped on +# inherit and reconciled into the round-tripped tensor_parallel state. +# --tensor-split is coupled to the split mode and is stripped with it: Studio +# owns the tensor-mode split ratios, so an inherited/stale --tensor-split must +# not last-wins-override Studio's computed asymmetric split. +_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 -_SHADOWING_FLAGS: frozenset[str] = _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS +_SHADOWING_FLAGS: frozenset[str] = ( + _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS +) # Shadowing flags that take no value -- strip the flag only, not the next token. _BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"}) @@ -213,11 +226,12 @@ def resolve_requested_ctx(args: Optional[Iterable[str]], fallback_n_ctx: int) -> return override if override is not None else fallback_n_ctx -def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: - """Return the last-wins cache type if extras pass cache flags. +def _last_flag_value(args: Optional[Iterable[str]], flags: frozenset[str]) -> Optional[str]: + """Return the last-wins string value among ``flags`` in extras, or None. - Recognises -ctk (key) and -ctv (value); treats both as one setting, - since Studio's KV estimate has a single cache_type_kv knob. + Handles both ``--flag=value`` and ``--flag value`` forms and raises if a + matched flag has no (or an empty) value. Shared by the single-knob + last-wins parsers (cache type, split mode). """ if not args: return None @@ -228,7 +242,7 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: while i < n: tok = tokens[i] flag = _flag_name(tok) - if flag is None or flag not in _CACHE_FLAGS: + if flag is None or flag not in flags: i += 1 continue @@ -249,6 +263,17 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: return override +def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: + """Return the last-wins cache type if extras pass cache flags. + + Mirrors parse_ctx_override but for cache type. Recognises both -ctk + (key) and -ctv (value). When both flags appear, returns the last-wins + value, treating key and value cache flags as the same setting because + Studio's KV estimate has a single cache_type_kv knob. + """ + return _last_flag_value(args, _CACHE_FLAGS) + + def resolve_cache_type_kv( args: Optional[Iterable[str]], fallback_cache_type_kv: Optional[str] ) -> Optional[str]: @@ -260,6 +285,30 @@ def resolve_cache_type_kv( return override if override is not None else fallback_cache_type_kv +def parse_split_mode_override(args: Optional[Iterable[str]]) -> Optional[str]: + """Return the last-wins ``--split-mode`` / ``-sm`` value from extras. + + Mirrors parse_cache_override for the multi-GPU split mode. Returns the + raw mode string (e.g. ``tensor`` / ``row`` / ``none`` / ``layer``), or + None when extras don't set it. + """ + return _last_flag_value(args, _SPLIT_MODE_FLAGS) + + +def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_parallel: bool) -> bool: + """Return the tensor-parallel state load_model should treat as requested. + + A user-supplied ``--split-mode`` in extras last-wins-overrides the + toggle, so reconcile it back into the boolean: any explicit split mode + means tensor-parallel is on iff that mode is ``tensor``. Falls back to + the toggle value when extras don't set it. + """ + override = parse_split_mode_override(args) + if override is None: + return fallback_tensor_parallel + return override.strip().lower() == "tensor" + + _MMPROJ_DISABLE_FLAGS: frozenset[str] = frozenset({"--no-mmproj", "--no-mmproj-auto"}) _MMPROJ_ENABLE_FLAGS: frozenset[str] = frozenset({"--mmproj-auto"}) @@ -289,12 +338,15 @@ def strip_shadowing_flags( strip_cache: bool = True, strip_spec: bool = True, strip_template: bool = True, + strip_split_mode: bool = True, ) -> list[str]: """Strip flags that shadow first-class Studio settings. Used when inheriting a previous load's ``llama_extra_args`` so an - inherited `-c 4096` can't override the current `max_seq_length` (same for - cache / spec / template). Each ``strip_*`` toggle controls one group. + inherited `-c 4096` can't override the current `max_seq_length` + (same for cache / spec / template / split-mode). Each ``strip_*`` + toggle controls one group; the route only strips groups whose + first-class field the caller actually supplied. """ shadowing: set[str] = set() if strip_context: @@ -305,6 +357,8 @@ def strip_shadowing_flags( shadowing |= _SPEC_FLAGS if strip_template: shadowing |= _TEMPLATE_FLAGS + if strip_split_mode: + shadowing |= _SPLIT_SHADOWING_FLAGS tokens = [str(a) for a in (args or [])] out: list[str] = [] @@ -325,3 +379,20 @@ def strip_shadowing_flags( else: i += 1 return out + + +def strip_split_mode_only(args: Optional[Iterable[str]]) -> Optional[list[str]]: + """Remove the split-mode group (``--split-mode`` / ``-sm`` and the coupled + ``--tensor-split`` / ``-ts``) from ``args``, keeping every other shadow flag. + Preserves a None/empty input so the inherit-vs-explicit-empty distinction + survives. Used where tensor mode is being forced off (downgrade / fallback).""" + if not args: + return args + return strip_shadowing_flags( + args, + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = True, + ) diff --git a/studio/backend/core/inference/tensor_fallback.py b/studio/backend/core/inference/tensor_fallback.py new file mode 100644 index 0000000000..73687165b8 --- /dev/null +++ b/studio/backend/core/inference/tensor_fallback.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tensor-parallel -> layer-split auto-fallback for GGUF loads. + +Kept in its own module (no FastAPI / httpx deps) so the orchestration can be +unit-tested with a fake loader, without a GPU or a running llama-server. +""" + +from __future__ import annotations + +import logging +from typing import Awaitable, Callable, Optional + +from core.inference.llama_server_args import ( + resolve_tensor_parallel, + strip_split_mode_only, +) + +logger = logging.getLogger(__name__) + + +async def load_with_tensor_fallback( + attempt_load: Callable[[bool, Optional[list[str]]], Awaitable[bool]], + *, + requested_tensor: bool, + extra_args: Optional[list[str]], + label: str = "", + cancelled: Optional[Callable[[], bool]] = None, +) -> bool: + """Run a GGUF load with the tensor-parallel -> layer-split auto-fallback. + + ``attempt_load(tensor_parallel, extra_args)`` performs one load and returns + True on success; it *raises* on a hard crash (llama-server aborts on some + archs / older builds), which is treated the same as a False return. + + Tensor mode can be requested by the toggle or by a ``--split-mode tensor`` + in ``extra_args`` (an allowed shadow flag), so the retry is keyed on whether + tensor mode is actually engaged, and it strips ``--split-mode`` from the + extras so the layer retry can't relaunch the same failing tensor load. A + non-tensor load keeps its original contract and propagates exceptions. + + ``cancelled()`` distinguishes a real tensor-start failure from a user + cancellation: ``attempt_load`` also returns False when the load was + cancelled, so without this the helper would restart a load the user just + cancelled. + """ + tensor_requested = resolve_tensor_parallel(extra_args, requested_tensor) + try: + success = await attempt_load(requested_tensor, extra_args) + except Exception as exc: + if not tensor_requested: + raise + logger.warning("Tensor-parallel load raised for '%s': %s", label, exc) + success = False + + if success or not tensor_requested: + return success + + # The first attempt returned False because the user cancelled, not because + # tensor mode is unsupported -- do not relaunch the cancelled load. + if cancelled is not None and cancelled(): + return success + + logger.warning( + "Tensor-parallel load failed for '%s'; retrying with layer split " + "(this model may not support tensor parallelism)", + label, + ) + return await attempt_load(False, strip_split_mode_only(extra_args)) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 039a5d75dd..5a648bc999 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -87,6 +87,15 @@ class LoadRequest(BaseModel): "'mtp' or 'mtp+ngram'." ), ) + tensor_parallel: bool = Field( + False, + description = ( + "Split the model across GPUs by tensor (--split-mode tensor) " + "instead of by layer for GGUF models. Only affects multi-GPU " + "setups, where it can make generation significantly faster. " + "No effect on a single GPU. Ignored for non-GGUF models." + ), + ) llama_extra_args: Optional[List[str]] = Field( None, description = ( @@ -224,6 +233,10 @@ class LoadResponse(BaseModel): "None when the platform default is in effect." ), ) + tensor_parallel: bool = Field( + False, + description = "Whether tensor-parallel split (--split-mode tensor) is active.", + ) class UnloadResponse(BaseModel): @@ -339,6 +352,10 @@ class InferenceStatusResponse(BaseModel): "None when the platform default is in effect." ), ) + tensor_parallel: bool = Field( + False, + description = "Whether tensor-parallel split (--split-mode tensor) is active.", + ) llama_cpp_supports_mtp: bool = Field( True, description = ( diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9ff30cd6db..7b17e450d1 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -596,9 +596,11 @@ try: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, ) + from core.inference.tensor_fallback import load_with_tensor_fallback from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import ( @@ -627,9 +629,11 @@ except ImportError: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, ) + from core.inference.tensor_fallback import load_with_tensor_fallback from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import ( @@ -1159,6 +1163,21 @@ def _normalise_settings_str(value: Optional[str]) -> Optional[str]: return value +def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[str]]) -> bool: + """Whether an inherited --split-mode should be stripped on reload. + + The binary Tensor Parallelism toggle can't carry --split-mode's row/none/ + layer modes, so only strip when the toggle overrides it: tensor being turned + on, or the inherited mode is tensor (toggle turning it off). Non-tensor modes + survive. Shared by the inheritance strip and the already-loaded stale check + so they agree on what reload would do. + """ + fields_set = getattr(request, "model_fields_set", set()) + return "tensor_parallel" in fields_set and ( + request.tensor_parallel or resolve_tensor_parallel(backend_extra, False) + ) + + def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaCppBackend) -> bool: """True iff every runtime setting on the request matches the loaded server. Caller has already checked model+variant+is_loaded. See #5401.""" @@ -1170,6 +1189,24 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC llama_backend.cache_type_kv ): return False + # Reconcile a user --split-mode in extras into the effective tensor state. + # When the request omits llama_extra_args ("inherit"), compare using the + # stored extras stripped the way the reload strips them, so an extras-driven + # tensor load isn't seen as a mismatch that needlessly reloads the server. + backend_extra = list(llama_backend.extra_args) if llama_backend.extra_args else [] + effective_extra = ( + request.llama_extra_args + if request.llama_extra_args is not None + else strip_shadowing_flags( + backend_extra, + strip_split_mode = _should_strip_split_mode(request, backend_extra), + ) + ) + if ( + resolve_tensor_parallel(effective_extra, request.tensor_parallel) + != llama_backend.tensor_parallel + ): + return False # Spec decoding works on vision models too (MTP is mmproj-compatible, # llama.cpp #22673; the old ``not is_vision`` gate is gone), so compare # the real requested mode -- coercing vision to ``off`` here used to @@ -1188,10 +1225,19 @@ def _request_matches_loaded_settings(request: LoadRequest, llama_backend: LlamaC # llama_extra_args=None means "inherit"; only an explicit differing list # forces a reload. On the inherit path, refuse to match if stored extras # contain any shadow flag, so the reload path strips them rather than - # leaving a stale override in effect. - backend_extra = list(llama_backend.extra_args) if llama_backend.extra_args else [] + # leaving a stale override in effect. (backend_extra computed above.) if request.llama_extra_args is None: - if backend_extra and strip_shadowing_flags(backend_extra) != backend_extra: + # Mirror the reload's conditional split-mode strip, so a preserved + # non-tensor mode (row/none/layer) isn't seen as stale and doesn't + # trigger a needless reload of a healthy server. + if ( + backend_extra + and strip_shadowing_flags( + backend_extra, + strip_split_mode = _should_strip_split_mode(request, backend_extra), + ) + != backend_extra + ): return False else: if list(request.llama_extra_args) != backend_extra: @@ -1360,6 +1406,7 @@ async def load_model( chat_template = llama_backend.chat_template, speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, + tensor_parallel = llama_backend.tensor_parallel, ) else: if ( @@ -1490,6 +1537,9 @@ async def load_model( "speculative_type" in fields_set or "spec_draft_n_max" in fields_set ), strip_template = "chat_template_override" in fields_set, + strip_split_mode = _should_strip_split_mode( + request, llama_backend.extra_args + ), ) try: extra_llama_args = validate_extra_args(stripped) @@ -1514,22 +1564,25 @@ async def load_model( # during the (potentially long) GGUF download + llama-server start. _n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1) + # Load kwargs common to HF and local modes; the two differ only by + # the model-source args (hf_repo/-token vs gguf_path/mmproj). + _common_load_kwargs = dict( + model_identifier = config.identifier, + is_vision = config.is_vision, + n_ctx = request.max_seq_length, + chat_template_override = request.chat_template_override, + cache_type_kv = request.cache_type_kv, + speculative_type = request.speculative_type, + spec_draft_n_max = request.spec_draft_n_max, + n_parallel = _n_parallel, + extra_args = extra_llama_args, + ) if config.gguf_hf_repo: # HF mode: download via huggingface_hub then start llama-server - success = await asyncio.to_thread( - llama_backend.load_model, + _source_load_kwargs = dict( hf_repo = config.gguf_hf_repo, hf_variant = config.gguf_variant, hf_token = request.hf_token, - model_identifier = config.identifier, - is_vision = config.is_vision, - n_ctx = request.max_seq_length, - chat_template_override = request.chat_template_override, - cache_type_kv = request.cache_type_kv, - speculative_type = request.speculative_type, - spec_draft_n_max = request.spec_draft_n_max, - n_parallel = _n_parallel, - extra_args = extra_llama_args, ) else: # Local mode: llama-server loads via -m @@ -1548,8 +1601,7 @@ async def load_model( except HTTPException as exc: logger.warning("Dropping MTP drafter for native load: %s", exc.detail) config.gguf_mtp_file = None - success = await asyncio.to_thread( - llama_backend.load_model, + _source_load_kwargs = dict( gguf_path = config.gguf_file, mmproj_path = config.gguf_mmproj_file, mtp_draft_path = config.gguf_mtp_file, @@ -1557,17 +1609,36 @@ async def load_model( # the same string the inheritance check at the top of /load # uses (#5401 followup). hf_variant = config.gguf_variant, - model_identifier = config.identifier, - is_vision = config.is_vision, - n_ctx = request.max_seq_length, - chat_template_override = request.chat_template_override, - cache_type_kv = request.cache_type_kv, - speculative_type = request.speculative_type, - spec_draft_n_max = request.spec_draft_n_max, - n_parallel = _n_parallel, - extra_args = extra_llama_args, ) + # Run a single load attempt with the given tensor flag + extras. + async def _attempt_gguf_load( + tensor_parallel: bool, attempt_extra_args: Optional[list[str]] + ) -> bool: + attempt_kwargs = { + **_common_load_kwargs, + "extra_args": attempt_extra_args, + } + return await asyncio.to_thread( + llama_backend.load_model, + **_source_load_kwargs, + **attempt_kwargs, + tensor_parallel = tensor_parallel, + ) + + # Tensor parallelism is arch-gated in llama.cpp and crashes some loads + # outright (e.g. Gemma 3n aborts with a GGML_ASSERT). The helper auto- + # falls back to layer split so the checkbox never blocks a model from + # loading; the response reports the backend's actual tensor_parallel + # state so the UI toggle reflects the fallback. + success = await load_with_tensor_fallback( + _attempt_gguf_load, + requested_tensor = request.tensor_parallel, + extra_args = extra_llama_args, + label = config.identifier, + cancelled = llama_backend.load_cancelled, + ) + if not success: raise HTTPException( status_code = 500, @@ -1612,6 +1683,7 @@ async def load_model( chat_template = llama_backend.chat_template, speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, + tensor_parallel = llama_backend.tensor_parallel, ) # ── Standard path: load via Unsloth/transformers ────────── @@ -2100,6 +2172,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)): chat_template_override = llama_backend.chat_template_override, speculative_type = llama_backend.requested_spec_mode, spec_draft_n_max = llama_backend.spec_draft_n_max, + tensor_parallel = llama_backend.tensor_parallel, llama_cpp_supports_mtp = _supports_mtp, spec_fallback_reason = llama_backend.spec_fallback_reason, llama_cpp_prebuilt_stale = _stale, diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index 6ae9d21e47..f3a3ea1ec4 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -25,8 +25,11 @@ _spec.loader.exec_module(_lsa) is_managed_flag = _lsa.is_managed_flag parse_cache_override = _lsa.parse_cache_override parse_ctx_override = _lsa.parse_ctx_override +parse_split_mode_override = _lsa.parse_split_mode_override resolve_cache_type_kv = _lsa.resolve_cache_type_kv +resolve_tensor_parallel = _lsa.resolve_tensor_parallel strip_shadowing_flags = _lsa.strip_shadowing_flags +strip_split_mode_only = _lsa.strip_split_mode_only extra_args_disable_mmproj = _lsa.extra_args_disable_mmproj validate_extra_args = _lsa.validate_extra_args @@ -510,6 +513,112 @@ def test_strip_shadowing_flags_defaults_strip_everything(): assert out == [] +# ── --split-mode (Tensor Parallelism toggle) ───────────────────────── +# Soft-shadowed exactly like --cache-type-*: pass-through allowed (keeps +# the row/none/layer modes the boolean toggle doesn't expose), stripped +# on inherit, and reconciled back into the round-tripped tensor_parallel +# state. + + +@pytest.mark.parametrize( + "args", + [ + ["--split-mode", "tensor"], + ["--split-mode", "row"], + ["--split-mode", "none"], + ["--split-mode", "layer"], + ["-sm", "tensor"], + ["--split-mode=row"], + ["-sm=tensor"], + ], +) +def test_split_mode_passes_through(args): + # Not denylisted -- a user keeps row/none/layer via extras. + assert validate_extra_args(args) == args + + +def test_split_mode_is_not_managed(): + assert is_managed_flag("--split-mode") is False + assert is_managed_flag("-sm") is False + + +@pytest.mark.parametrize( + "args,expected", + [ + (None, None), + ([], None), + (["--top-k", "20"], None), + (["--split-mode", "tensor"], "tensor"), + (["--split-mode", "row"], "row"), + (["-sm", "none"], "none"), + (["--split-mode=layer"], "layer"), + (["-sm=tensor"], "tensor"), + # last-wins when supplied twice + (["-sm", "row", "--split-mode", "tensor"], "tensor"), + ], +) +def test_parse_split_mode_override(args, expected): + assert parse_split_mode_override(args) == expected + + +@pytest.mark.parametrize( + "args", + [ + ["--split-mode"], + ["-sm"], + ["--split-mode", "-c", "4096"], # next token is a flag, not a value + ], +) +def test_parse_split_mode_override_rejects_malformed_values(args): + with pytest.raises(ValueError, match = "split-mode|'-sm'"): + parse_split_mode_override(args) + + +def test_validate_extra_args_rejects_malformed_split_mode(): + # Validation catches a value-less --split-mode at the boundary, + # mirroring the early --ctx-size / --cache-type checks. + with pytest.raises(ValueError, match = "split-mode"): + validate_extra_args(["--split-mode"]) + + +@pytest.mark.parametrize( + "args,fallback,expected", + [ + # No override -> fall back to the toggle value, both directions. + (["--top-k", "20"], True, True), + (["--top-k", "20"], False, False), + (None, True, True), + ([], False, False), + # Explicit override wins: tensor -> on, anything else -> off, + # regardless of the toggle fallback. + (["--split-mode", "tensor"], False, True), + (["-sm", "tensor"], False, True), + (["--split-mode", "row"], True, False), + (["--split-mode", "none"], True, False), + (["--split-mode", "layer"], True, False), + (["--split-mode=tensor"], False, True), + # Case-insensitive on the mode string. + (["--split-mode", "TENSOR"], False, True), + # last-wins across multiple --split-mode flags. + (["-sm", "tensor", "--split-mode", "row"], True, False), + ], +) +def test_resolve_tensor_parallel(args, fallback, expected): + assert resolve_tensor_parallel(args, fallback) is expected + + +def test_strip_shadowing_flags_drops_split_mode_when_requested(): + out = strip_shadowing_flags( + ["--split-mode", "row", "--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = True, + ) + assert out == ["--top-k", "20"] + + def test_extra_args_disable_mmproj_detects_flag(): assert extra_args_disable_mmproj(["--no-mmproj"]) is True assert extra_args_disable_mmproj(["--threads", "12", "--no-mmproj"]) is True @@ -540,6 +649,86 @@ def test_strip_shadowing_flags_drops_model_draft_with_spec(): assert out == ["--top-k", "20"] +def test_strip_shadowing_flags_keeps_split_mode_when_not_requested(): + # No tensor_parallel field supplied on the Apply -> an inherited + # --split-mode survives (mirrors the chat-template keep behavior). + out = strip_shadowing_flags( + ["--split-mode", "row", "--top-k", "20"], + strip_context = True, + strip_cache = True, + strip_spec = True, + strip_template = True, + strip_split_mode = False, + ) + assert out == ["--split-mode", "row", "--top-k", "20"] + + +def test_strip_shadowing_flags_drops_split_mode_short_alias_and_equals(): + assert strip_shadowing_flags(["-sm", "tensor", "--top-k", "20"], strip_split_mode = True) == [ + "--top-k", + "20", + ] + assert strip_shadowing_flags(["--split-mode=row", "--seed", "-1"], strip_split_mode = True) == [ + "--seed", + "-1", + ] + + +def test_strip_shadowing_flags_defaults_strip_split_mode_too(): + # The route's already-loaded comparator (no kwargs) must see a stored + # --split-mode as a shadowing flag so it forces a reload. + assert strip_shadowing_flags(["--split-mode", "tensor"]) == [] + + +@pytest.mark.parametrize( + "args", + [ + ["--split-mode", "tensor", "-c", "4096"], + ["-sm", "tensor", "-c", "4096"], + ["--split-mode=tensor", "-c", "4096"], + ["-sm=tensor", "-c", "4096"], + ], +) +def test_strip_split_mode_only_keeps_other_shadow_flags(args): + # Every --split-mode form (long/short, space/=) is dropped; -c survives. + assert strip_split_mode_only(args) == ["-c", "4096"] + + +def test_strip_split_mode_only_preserves_none_and_empty(): + # None means "inherit"; [] means "explicit empty" -- both must round-trip. + assert strip_split_mode_only(None) is None + assert strip_split_mode_only([]) == [] + + +def test_strip_shadowing_flags_drops_tensor_split_with_split_mode(): + # --tensor-split is coupled to the split mode: stripped together so a stale + # ratio can't override Studio's computed tensor split. Other flags survive. + out = strip_shadowing_flags( + ["--split-mode", "row", "--tensor-split", "1,1", "--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + strip_split_mode = True, + ) + assert out == ["--top-k", "20"] + + +def test_strip_shadowing_flags_keeps_tensor_split_when_not_requested(): + # strip_split_mode=False keeps the whole split group (mode + ratios). + assert strip_shadowing_flags( + ["--tensor-split", "1,1", "--top-k", "20"], strip_split_mode = False + ) == ["--tensor-split", "1,1", "--top-k", "20"] + + +def test_strip_split_mode_only_drops_tensor_split_too(): + # Downgrade / layer fallback must drop the coupled --tensor-split (all forms). + assert strip_split_mode_only( + ["--split-mode", "tensor", "--tensor-split", "1,1", "-c", "4096"] + ) == ["-c", "4096"] + assert strip_split_mode_only(["-sm=tensor", "-ts=3,1"]) == [] + + def test_strip_shadowing_flags_keeps_model_draft_without_spec(): out = strip_shadowing_flags( ["--model-draft", "/custom/mtp.gguf"], diff --git a/studio/backend/tests/test_studio_api.py b/studio/backend/tests/test_studio_api.py index 27f695b744..928b636e3e 100644 --- a/studio/backend/tests/test_studio_api.py +++ b/studio/backend/tests/test_studio_api.py @@ -149,6 +149,7 @@ def test_help_output(): "--host", "--frontend", "--silent", + "--tensor-parallel", ]: assert flag in out, f"Missing flag {flag!r} in --help output" print(" PASS --help shows all flags") diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py new file mode 100644 index 0000000000..30bfb91a08 --- /dev/null +++ b/studio/backend/tests/test_tensor_parallel.py @@ -0,0 +1,578 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Backend contract for the Tensor Parallelism toggle. + +The toggle threads a single ``tensor_parallel`` bool from the chat UI +through the load request to a ``--split-mode tensor`` llama-server flag, +and round-trips it back via the load/status responses so the switch +reflects what is actually running. These tests pin: + + * the pydantic request/response/status contract (snake_case key, + default False), + * the backend ``tensor_parallel`` property and its reset on unload, + * the ``_already_in_target_state`` reload-detection branch, and + * that ``--split-mode tensor`` is emitted only behind the toggle. +""" + +from __future__ import annotations + +import asyncio +import inspect +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Same external-dep stubs as the other llama_cpp unit tests so importing +# the backend doesn't drag in structlog / httpx / loggers. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +_structlog_stub = _types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +sys.modules.setdefault("structlog", _structlog_stub) + +_httpx_stub = _types.ModuleType("httpx") +for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) +_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) +_httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__exit__": lambda s, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference import llama_cpp as llama_cpp_module +from core.inference.llama_cpp import LlamaCppBackend +from core.inference.llama_server_args import resolve_tensor_parallel +from core.inference.tensor_fallback import load_with_tensor_fallback +from models.inference import ( + InferenceStatusResponse, + LoadRequest, + LoadResponse, +) + + +# ── Pydantic contract (snake_case key, default False) ──────────────── + + +def test_load_request_defaults_tensor_parallel_false(): + req = LoadRequest(model_path = "owner/repo") + assert req.tensor_parallel is False + + +def test_load_request_accepts_tensor_parallel(): + req = LoadRequest(model_path = "owner/repo", tensor_parallel = True) + assert req.tensor_parallel is True + + +def test_load_request_round_trips_json_key(): + # The frontend sends the snake_case key verbatim. + req = LoadRequest.model_validate({"model_path": "owner/repo", "tensor_parallel": True}) + assert req.tensor_parallel is True + assert req.model_dump()["tensor_parallel"] is True + + +@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) +def test_response_models_emit_tensor_parallel(model_cls): + # Default False, and the key is always present in the JSON body. + if model_cls is LoadResponse: + default = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + ) + on = model_cls( + status = "loaded", + model = "owner/repo", + display_name = "repo", + inference = {}, + tensor_parallel = True, + ) + else: + default = model_cls() + on = model_cls(tensor_parallel = True) + assert default.model_dump()["tensor_parallel"] is False + assert on.model_dump()["tensor_parallel"] is True + + +# ── Backend property + reset ───────────────────────────────────────── + + +class _FakeProcess: + """Stand-in for subprocess.Popen so _kill_process is a no-op.""" + + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + def poll(self): + return 0 + + +def test_tensor_parallel_property_defaults_false(): + assert LlamaCppBackend().tensor_parallel is False + + +def test_tensor_parallel_property_reflects_field(): + backend = LlamaCppBackend() + backend._tensor_parallel = True + assert backend.tensor_parallel is True + + +def test_unload_resets_tensor_parallel(): + backend = LlamaCppBackend() + backend._process = _FakeProcess() + backend._tensor_parallel = True + backend.unload_model() + assert backend.tensor_parallel is False + + +# ── _already_in_target_state reload-detection branch ───────────────── + + +def _loaded_backend(tensor_parallel: bool) -> LlamaCppBackend: + backend = LlamaCppBackend() + backend._process = _FakeProcess() # is_loaded only checks "is not None" + backend._healthy = True + backend._model_identifier = "owner/repo" + backend._hf_variant = "Q4_K_M" + backend._requested_n_ctx = 8192 + backend._cache_type_kv = None + backend._requested_spec_mode = "auto" + backend._chat_template_override = None + backend._is_vision = False + backend._extra_args = None + backend._gguf_path = None + backend._tensor_parallel = tensor_parallel + return backend + + +def _target_state(backend: LlamaCppBackend, tensor_parallel: bool) -> bool: + return backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = None, + is_vision = False, + tensor_parallel = tensor_parallel, + ) + + +@pytest.mark.parametrize("flag", [True, False]) +def test_already_in_target_state_matches_same_tensor_parallel(flag): + assert _target_state(_loaded_backend(flag), flag) is True + + +@pytest.mark.parametrize( + "loaded,requested", + [(False, True), (True, False)], +) +def test_already_in_target_state_reloads_on_tensor_parallel_change(loaded, requested): + # Flipping the toggle either direction must force a reload so the + # command is rebuilt with/without --split-mode tensor. + assert _target_state(_loaded_backend(loaded), requested) is False + + +def test_already_in_target_state_reconciles_split_mode_extras(): + # Tensor engaged via --split-mode in extras (boolean omitted/default False) + # must match a server already running tensor mode -- no spurious reload. + backend = _loaded_backend(tensor_parallel = True) + backend._extra_args = ["--split-mode", "tensor"] + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = "auto", + chat_template_override = None, + extra_args = ["--split-mode", "tensor"], + is_vision = False, + tensor_parallel = False, + ) + is True + ) + + +# ── --split-mode tensor is emitted only behind the toggle ──────────── + + +def _load_model_source() -> str: + return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model) + + +def test_split_mode_tensor_is_gated_on_the_toggle(): + src = _load_model_source() + assert ( + 'cmd.extend(["--split-mode", "tensor"])' in src + ), "the tensor-parallel flag emission must be present in load_model" + # The emission lives behind `if tensor_parallel:` -- it must never be + # part of the unconditional base cmd list. + base_start = src.find("cmd = [") + base_end = src.find("\n ]", base_start) + base_block = src[base_start:base_end] if base_end > base_start else "" + assert ( + "--split-mode" not in base_block + ), "--split-mode must be conditional, not in the base cmd list" + gate = src.find("if tensor_parallel:") + emit = src.find('cmd.extend(["--split-mode", "tensor"])') + assert 0 <= gate < emit, "emission must sit under `if tensor_parallel:`" + + +def test_proportional_tensor_split_is_emitted_in_tensor_mode(): + # Asymmetric GPUs (e.g. 48 GB + 24 GB) OOM the smaller card under the + # even default; the allocator weights --tensor-split by free VRAM. Pin + # that the flag is emitted from inside the tensor-parallel block. + src = _load_model_source() + assert '"--tensor-split"' in src + gate = src.find("if tensor_parallel:") + ts = src.find('"--tensor-split"') + nxt_else = src.find("self._tensor_parallel = False") + assert 0 <= gate < ts < nxt_else, "--tensor-split must be emitted under `if tensor_parallel:`" + + +# ── tensor-mode allocation: conservative VRAM budget ───────────────── + + +def _kv_seeded_backend() -> LlamaCppBackend: + # Minimal GGUF metadata so _can_estimate_kv() is True (legacy KV path). + backend = LlamaCppBackend() + backend._n_layers = 32 + backend._embedding_length = 4096 + backend._n_heads = 32 + backend._n_kv_heads = 8 + backend._context_length = 131072 + return backend + + +def test_fit_context_budget_frac_override_is_tighter(): + backend = _kv_seeded_backend() + model_size = 8 * 1024**3 + pool_mib = 24 * 1024 # tight enough that KV capping bites + + fit_default = backend._fit_context_to_vram(131072, pool_mib, model_size, "f16") + fit_tp = backend._fit_context_to_vram(131072, pool_mib, model_size, "f16", budget_frac = 0.80) + assert fit_tp < 131072, "expected the context to be capped at this VRAM tier" + assert fit_tp <= fit_default, "a tighter budget must not allow MORE context" + # Omitting the override must reproduce the default budget exactly. + assert backend._fit_context_to_vram(131072, pool_mib, model_size, "f16") == fit_default + + +# ── unsupported-arch load failure -> clean message ─────────────────── + + +def test_split_mode_tensor_arch_failure_message(): + msg = LlamaCppBackend._classify_llama_start_failure( + "llama_model_create: LLAMA_SPLIT_MODE_TENSOR not implemented for " + "architecture 'deepseek2'", + None, + "unsloth/DeepSeek-V3-GGUF", + ) + assert "Tensor parallelism is not supported" in msg + + +def test_unrelated_arch_failure_not_hijacked_by_tensor_message(): + msg = LlamaCppBackend._classify_llama_start_failure( + "unknown model architecture: 'flux'", "/models/flux.gguf", None + ) + assert "Tensor parallelism" not in msg + + +# ── _plan_tensor_parallel: the allocation math (pure, no model/GPU) ─── +# Seeded full-attention KV (~128 KiB/token) via _kv_seeded_backend, so the +# context cap + split are deterministic. Asserts relationships rather than +# magic numbers so the KV estimate can evolve without breaking these. + +_GB = 1024**3 +_ASYM = [(0, 48000), (1, 24000)] # asymmetric pool, 72000 MiB +_SYM = [(0, 24000), (1, 24000)] # symmetric pool + + +def _plan( + model_gb, + target = 131072, + gpus = _ASYM, + mtp = False, +): + b = _kv_seeded_backend() + return b, b._plan_tensor_parallel(gpus, int(model_gb * _GB), target, mtp_engaged = mtp) + + +def _kv_budget_b(model_gb, gpus = _ASYM): + reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + return (sum(f for _, f in gpus) - len(gpus) * reserve) * 1024 * 1024 - int(model_gb * _GB) + + +def test_tp_plan_weighted_split_on_asymmetric_big_model(): + b, (ec, mac, gi, ts) = _plan(50) + reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + assert gi == [0, 1] + # split weighted by (free - buffer), not raw free + assert ts == [48000 - reserve, 24000 - reserve] + assert ec < 131072 # capped below native + + +def test_tp_plan_even_split_when_model_fits(): + # A small model whose even share fits the smallest GPU -> llama.cpp's even + # default (None), which is safe for archs that crash on a weighted split. + _, (ec, mac, gi, ts) = _plan(4) + assert ts is None + + +def test_tp_plan_symmetric_gpus_use_even_split(): + _, (ec, mac, gi, ts) = _plan(8, gpus = _SYM) + assert ts is None + + +def test_tp_plan_context_fits_pool_budget_no_oom(): + b, (ec, mac, gi, ts) = _plan(50) + # the chosen context's KV must fit the pooled budget (weights + buffers) + assert b._estimate_kv_cache_bytes(ec) <= _kv_budget_b(50) + + +def test_tp_plan_uses_available_vram_not_wasteful(): + # when the cap engages, the chosen context nearly fills the budget + b, (ec, mac, gi, ts) = _plan(50) + assert b._estimate_kv_cache_bytes(ec) >= 0.9 * _kv_budget_b(50) + + +def test_tp_plan_weights_exceed_pool_floors_context(): + # 70 GB > pool minus per-GPU reserves -> floor (triggers layer fallback) + _, (ec, mac, gi, ts) = _plan(70) + assert ec == 2048 + + +def test_tp_plan_floor_never_exceeds_explicit_small_context(): + # An explicit context below the 2048 floor must not be raised: a caller + # asking for 1024 should not have KV sized for 2048 (avoidable OOM). + _, (ec, mac, gi, ts) = _plan(70, target = 1024) # weights exceed pool -> floor path + assert ec == 1024 + _, (ec2, *_rest) = _plan(50, target = 1024) # cap path with a tiny budget + assert ec2 <= 1024 + + +def test_tp_plan_explicit_context_honored_when_it_fits(): + _, (ec, mac, gi, ts) = _plan(50, target = 8192) + assert ec == 8192 + + +def test_tp_plan_explicit_context_capped_when_too_large(): + _, (ec, mac, gi, ts) = _plan(50, target = 131072) + assert 2048 <= ec < 131072 + + +def test_tp_plan_max_available_ctx_reports_native_not_explicit_ctx(): + # An explicit small ctx caps effective_ctx but the UI ceiling + # (max_available_ctx) must reflect the native/hardware cap, not the request. + b = _kv_seeded_backend() + ec, mac, _gi, _ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 8192, max_target_ctx = 131072) + _, native_mac, *_ = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + assert ec == 8192 # explicit request honored for the load + assert mac == native_mac > ec # ceiling reflects the hardware cap + + +def test_tp_plan_mtp_reserves_extra_and_shrinks_context(): + _, (ec_no, *_rest) = _plan(50) + _, (ec_mtp, *_rest) = _plan(50, mtp = True) + assert ec_mtp < ec_no + + +def test_tp_plan_no_kv_metadata_floors_context(): + b = LlamaCppBackend() # no KV metadata -> can't size safely + ec, mac, gi, ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + assert ec <= 4096 + + +def test_tp_plan_single_gpu_never_splits(): + # The toggle is a no-op without >= 2 GPUs (most dev/CI machines). Even if + # the planner is reached, it must not emit a tensor split. + b = _kv_seeded_backend() + ec, mac, gi, ts = b._plan_tensor_parallel([(0, 24000)], int(8 * _GB), 8192) + assert ts is None + assert gi == [0] + + +def test_tp_plan_zero_gpus_never_splits(): + b = _kv_seeded_backend() + ec, mac, gi, ts = b._plan_tensor_parallel([], int(8 * _GB), 8192) + assert ts is None + assert gi == [] + + +def test_tp_plan_drops_gpu_below_buffer_reserve(): + # A GPU with less free VRAM than the per-device compute-buffer reserve + # can't host tensor mode; it's excluded, which here leaves <2 usable -> no + # split (and gpu_indices reflects only the usable device). + b = _kv_seeded_backend() + reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + ec, mac, gi, ts = b._plan_tensor_parallel([(0, 48000), (1, reserve - 1)], int(8 * _GB), 8192) + assert gi == [0] + assert ts is None + + +# ── route auto-fallback survives a *raised* tensor-load crash ───────── +# A tensor-incompatible model makes load_model RAISE (Gemma 3n aborts) rather +# than return False. The /load fallback helper must catch that and retry with +# layer split -- stripping any --split-mode from the extras so the retry can't +# relaunch tensor -- while a non-tensor load propagates its exception. These +# exercise the real helper with a fake loader (no GPU, no llama-server). + + +class _RecordingLoader: + """Fake ``attempt_load``: crashes whenever tensor mode is effectively + engaged (via the bool or a ``--split-mode`` in extras), like a real + tensor-incompatible model; succeeds on layer split.""" + + def __init__(self): + self.calls: list[tuple] = [] + + async def __call__(self, tensor_parallel, extra_args): + self.calls.append((tensor_parallel, list(extra_args) if extra_args else extra_args)) + if resolve_tensor_parallel(extra_args, tensor_parallel): + raise RuntimeError("llama-server failed to start") + return True + + +def test_tensor_fallback_retries_layer_on_crash(): + loader = _RecordingLoader() + ok = asyncio.run( + load_with_tensor_fallback(loader, requested_tensor = True, extra_args = None, label = "m") + ) + assert ok is True + # tensor first (crashes), then layer split. + assert [c[0] for c in loader.calls] == [True, False] + + +def test_tensor_fallback_no_retry_on_success(): + calls: list[bool] = [] + + async def _ok(tensor_parallel, extra_args): + calls.append(tensor_parallel) + return True + + ok = asyncio.run( + load_with_tensor_fallback(_ok, requested_tensor = True, extra_args = None, label = "m") + ) + assert ok is True + assert calls == [True] # no fallback when the tensor load succeeds + + +def test_tensor_fallback_retries_when_tensor_returns_false(): + # load_model can signal failure by *returning False* (not only by raising); + # that must trigger the layer-split retry just like a crash does. + calls: list[bool] = [] + + async def _false_on_tensor(tensor_parallel, extra_args): + calls.append(tensor_parallel) + return not resolve_tensor_parallel(extra_args, tensor_parallel) + + ok = asyncio.run( + load_with_tensor_fallback( + _false_on_tensor, requested_tensor = True, extra_args = None, label = "m" + ) + ) + assert ok is True + assert calls == [True, False] + + +def test_tensor_fallback_returns_false_when_both_attempts_fail(): + # Tensor fails and the layer retry also fails -> the helper returns False so + # the route raises its own HTTP 500 (it does not crash mid-flight). + calls: list[bool] = [] + + async def _always_false(tensor_parallel, extra_args): + calls.append(tensor_parallel) + return False + + ok = asyncio.run( + load_with_tensor_fallback(_always_false, requested_tensor = True, extra_args = None, label = "m") + ) + assert ok is False + assert calls == [True, False] # tried tensor, then layer split + + +def test_tensor_fallback_skips_layer_retry_when_cancelled(): + # load_model returns False on a user cancellation too. When cancelled() is + # True, the helper must NOT relaunch the load the user just cancelled. + calls: list[bool] = [] + + async def _false_on_tensor(tensor_parallel, extra_args): + calls.append(tensor_parallel) + return False + + ok = asyncio.run( + load_with_tensor_fallback( + _false_on_tensor, + requested_tensor = True, + extra_args = None, + label = "m", + cancelled = lambda: True, + ) + ) + assert ok is False + assert calls == [True] # no layer-split retry after cancellation + + +@pytest.mark.parametrize( + "extras", + [ + ["--split-mode", "tensor", "-c", "4096"], + ["-sm", "tensor", "-c", "4096"], + ["--split-mode=tensor", "-c", "4096"], + ["-sm=tensor", "-c", "4096"], + ], +) +def test_tensor_fallback_strips_split_mode_from_extras_on_retry(extras): + # Tensor engaged via extras (boolean False); the retry must drop every + # --split-mode form (long/short, space/=) but keep the user's other flags, + # else resolve_tensor_parallel re-enables tensor and relaunches the crash. + loader = _RecordingLoader() + ok = asyncio.run( + load_with_tensor_fallback(loader, requested_tensor = False, extra_args = extras, label = "m") + ) + assert ok is True + assert len(loader.calls) == 2 + assert loader.calls[1][1] == ["-c", "4096"] # split-mode stripped, -c kept + + +def test_tensor_fallback_propagates_non_tensor_crash(): + async def _always_raise(tensor_parallel, extra_args): + raise RuntimeError("bad model") + + with pytest.raises(RuntimeError, match = "bad model"): + asyncio.run( + load_with_tensor_fallback( + _always_raise, requested_tensor = False, extra_args = None, label = "m" + ) + ) diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index ce9f1c6544..5903bef107 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1244,6 +1244,8 @@ async function autoLoadSmallestModel(): Promise<{ ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false), kvCacheDtype: loadResp.cache_type_kv ?? null, loadedKvCacheDtype: loadResp.cache_type_kv ?? null, + tensorParallel: loadResp.tensor_parallel ?? false, + loadedTensorParallel: loadResp.tensor_parallel ?? false, defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, loadedChatTemplateOverride: null, @@ -1405,6 +1407,8 @@ async function autoLoadSmallestModel(): Promise<{ ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false), kvCacheDtype: loadResp.cache_type_kv ?? null, loadedKvCacheDtype: loadResp.cache_type_kv ?? null, + tensorParallel: loadResp.tensor_parallel ?? false, + loadedTensorParallel: loadResp.tensor_parallel ?? false, defaultChatTemplate: loadResp.chat_template ?? null, chatTemplateOverride: null, loadedIsMultimodal: isMultimodalResponse(loadResp), diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 59ac2ae6b4..6eb59382b1 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -531,6 +531,11 @@ export function ChatSettingsPanel({ const kvCacheDtype = useChatRuntimeStore((s) => s.kvCacheDtype); const setKvCacheDtype = useChatRuntimeStore((s) => s.setKvCacheDtype); const loadedKvCacheDtype = useChatRuntimeStore((s) => s.loadedKvCacheDtype); + const tensorParallel = useChatRuntimeStore((s) => s.tensorParallel); + const setTensorParallel = useChatRuntimeStore((s) => s.setTensorParallel); + const loadedTensorParallel = useChatRuntimeStore( + (s) => s.loadedTensorParallel, + ); const customContextLength = useChatRuntimeStore((s) => s.customContextLength); const setCustomContextLength = useChatRuntimeStore( (s) => s.setCustomContextLength, @@ -551,7 +556,9 @@ export function ChatSettingsPanel({ const ctxDirty = customContextLength !== null; const specDirty = speculativeType !== loadedSpeculativeType; const specDraftDirty = specDraftNMax !== loadedSpecDraftNMax; - const modelSettingsDirty = kvDirty || ctxDirty || specDirty || specDraftDirty; + const tpDirty = tensorParallel !== (loadedTensorParallel ?? false); + const modelSettingsDirty = + kvDirty || ctxDirty || specDirty || specDraftDirty || tpDirty; const loadedChatTemplateOverride = useChatRuntimeStore( (s) => s.loadedChatTemplateOverride, ); @@ -1027,6 +1034,24 @@ export function ChatSettingsPanel({ /> )} +
+
+ + Tensor Parallelism + + + No effect on a single GPU. On multi-GPU setups, improves + tokens/sec during generation when using dense models. MoE + models don't benefit and can be much slower. + +
+ +
)} {!isGguf && params.checkpoint && ( @@ -1081,6 +1106,7 @@ export function ChatSettingsPanel({ setKvCacheDtype(loadedKvCacheDtype); setSpeculativeType(loadedSpeculativeType); setSpecDraftNMax(loadedSpecDraftNMax); + setTensorParallel(loadedTensorParallel ?? false); setChatTemplateOverride(loadedChatTemplateOverride); }} className="h-7 px-3 text-[12px] font-medium tracking-nav text-muted-foreground" diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 23ae1d80aa..70cba21afc 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -515,6 +515,7 @@ export function useChatModelRuntime() { ggufContextLength, speculativeType, specDraftNMax, + tensorParallel, activePresetSource, activeGgufVariant, } = useChatRuntimeStore.getState(); @@ -543,6 +544,7 @@ export function useChatModelRuntime() { cache_type_kv: kvCacheDtype, speculative_type: speculativeType, spec_draft_n_max: specDraftNMax, + tensor_parallel: tensorParallel, }); // If cancelled while loading, don't update UI to show @@ -570,6 +572,7 @@ export function useChatModelRuntime() { } } const loadedKv = loadResponse.cache_type_kv ?? null; + const loadedTp = loadResponse.tensor_parallel ?? false; const loadedSpec = normalizeSpeculativeType( loadResponse.speculative_type, ); @@ -626,6 +629,8 @@ export function useChatModelRuntime() { : resolveToolsEnabledOnLoad(supportsTools)), kvCacheDtype: loadedKv, loadedKvCacheDtype: loadedKv, + tensorParallel: loadedTp, + loadedTensorParallel: loadedTp, speculativeType: loadedSpec, loadedSpeculativeType: loadedSpec, specDraftNMax: loadResponse.spec_draft_n_max ?? null, @@ -700,6 +705,9 @@ export function useChatModelRuntime() { gguf_variant: previousVariant, trust_remote_code: previousModelRequiresTrustRemoteCode || trustRemoteCode, + // Restore the previous model in the split mode it was running, + // not the default layer split. + tensor_parallel: stateBeforeUnload.loadedTensorParallel ?? false, }); useChatRuntimeStore.setState({ activeNativePathToken: previousActiveNativePathToken ?? null, diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts index c02d304dcb..332bde6a3c 100644 --- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -170,6 +170,11 @@ export function applyActiveModelStatusToStore( kvCacheDtype: status.cache_type_kv, loadedKvCacheDtype: status.cache_type_kv, }), + ...(status.tensor_parallel !== undefined && + prevState.loadedTensorParallel === null && { + tensorParallel: status.tensor_parallel, + loadedTensorParallel: status.tensor_parallel, + }), ...(status.chat_template_override !== undefined && prevState.loadedChatTemplateOverride === null && prevState.chatTemplateOverride === null && { diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 74ab4a1b9e..5ebc02ab51 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -942,6 +942,8 @@ export function SharedComposer({ gguf_variant: sel.ggufVariant ?? null, trust_remote_code: trustRemoteCode, chat_template_override: effectiveChatTemplateOverride, + // Honor the Tensor Parallelism toggle on compare loads too. + tensor_parallel: currentStore.tensorParallel, }); const store = useChatRuntimeStore.getState(); store.setCheckpoint( @@ -957,6 +959,8 @@ export function SharedComposer({ reasoningStyle: resp.reasoning_style ?? "enable_thinking", supportsPreserveThinking: resp.supports_preserve_thinking ?? false, supportsTools: resp.supports_tools ?? false, + tensorParallel: resp.tensor_parallel ?? false, + loadedTensorParallel: resp.tensor_parallel ?? false, loadedIsMultimodal: isMultimodalResponse(resp), }); // Sync the models[] entry with the load response so attach/send gates diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 2f0e3027a8..2d23777378 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -448,6 +448,10 @@ type ChatRuntimeStore = { /** User --spec-draft-n-max override (null = platform default). */ specDraftNMax: number | null; loadedSpecDraftNMax: number | null; + /** Tensor-parallel split (--split-mode tensor) toggle, GGUF multi-GPU only. */ + tensorParallel: boolean; + /** Backend-reported tensor-parallel state; null until first hydrated. */ + loadedTensorParallel: boolean | null; loadedIsMultimodal: boolean; customContextLength: number | null; defaultChatTemplate: string | null; @@ -531,6 +535,7 @@ type ChatRuntimeStore = { setKvCacheDtype: (dtype: string | null) => void; setSpeculativeType: (type: string | null) => void; setSpecDraftNMax: (value: number | null) => void; + setTensorParallel: (value: boolean) => void; setCustomContextLength: (v: number | null) => void; setChatTemplateOverride: (template: string | null) => void; setPendingAudio: (base64: string, name: string) => void; @@ -809,6 +814,8 @@ export const useChatRuntimeStore = create((set, get) => ({ specFallbackReason: null, specDraftNMax: null, loadedSpecDraftNMax: null, + tensorParallel: false, + loadedTensorParallel: null, loadedIsMultimodal: false, customContextLength: null, defaultChatTemplate: null, @@ -1022,6 +1029,8 @@ export const useChatRuntimeStore = create((set, get) => ({ specFallbackReason: null, specDraftNMax: null, loadedSpecDraftNMax: null, + tensorParallel: false, + loadedTensorParallel: null, loadedIsMultimodal: false, customContextLength: null, defaultChatTemplate: null, @@ -1210,6 +1219,7 @@ export const useChatRuntimeStore = create((set, get) => ({ setKvCacheDtype: (kvCacheDtype) => set({ kvCacheDtype }), setSpeculativeType: (speculativeType) => set({ speculativeType }), setSpecDraftNMax: (specDraftNMax) => set({ specDraftNMax }), + setTensorParallel: (tensorParallel) => set({ tensorParallel }), setCustomContextLength: (customContextLength) => set({ customContextLength }), setChatTemplateOverride: (chatTemplateOverride) => set({ chatTemplateOverride }), diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 96caa38bd4..b3a7b16b5f 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -56,6 +56,11 @@ export interface LoadModelRequest { * when speculative_type resolves to "mtp" or "mtp+ngram". */ spec_draft_n_max?: number | null; + /** + * Split the model across GPUs by tensor (--split-mode tensor) instead + * of by layer for GGUF models. Multi-GPU only; no effect on a single GPU. + */ + tensor_parallel?: boolean | null; } export interface ValidateModelResponse { @@ -134,6 +139,8 @@ export interface LoadModelResponse { /** Canonical UI-facing mode the load request resolved to. See LoadModelRequest. */ speculative_type?: string | null; spec_draft_n_max?: number | null; + /** Whether tensor-parallel split (--split-mode tensor) is active. */ + tensor_parallel?: boolean; } export interface UnloadModelRequest { @@ -174,6 +181,8 @@ export interface InferenceStatusResponse { /** Canonical UI-facing mode currently active. See LoadModelRequest. */ speculative_type?: string | null; spec_draft_n_max?: number | null; + /** Whether tensor-parallel split (--split-mode tensor) is active. */ + tensor_parallel?: boolean; /** * Why MTP was disabled on the loaded model despite being requested. * "binary_no_mtp" / "binary_outdated" -> updating llama.cpp would re-enable diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts index f11fed6e13..a067855fd5 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts @@ -264,6 +264,8 @@ async function loadLocalModelSelection( cache_type_kv: null, // biome-ignore lint/style/useNamingConvention: api schema speculative_type: null, + // biome-ignore lint/style/useNamingConvention: api schema + tensor_parallel: false, }); toast.success(`Loaded ${modelLabel}`, { id: toastId, duration: 2000 }); return null; diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index ffdf36bedf..4608457a96 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -577,6 +577,7 @@ def _load_model_via_http( gguf_variant: Optional[str], max_seq_length: int, load_in_4bit: bool, + tensor_parallel: bool = False, llama_extra_args: Optional[List[str]] = None, timeout: int = 600, ) -> dict: @@ -592,6 +593,8 @@ def _load_model_via_http( } if gguf_variant: payload["gguf_variant"] = gguf_variant + if tensor_parallel: + payload["tensor_parallel"] = True if llama_extra_args: payload["llama_extra_args"] = list(llama_extra_args) @@ -940,6 +943,15 @@ def run( "--cloudflare/--no-cloudflare", help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 (default on).", ), + tensor_parallel: bool = typer.Option( + False, + "--tensor-parallel/--no-tensor-parallel", + help = ( + "Split a GGUF across GPUs by tensor (--split-mode tensor) instead of " + "by layer. Multi-GPU only (no effect on one GPU); dense models gain " + "decode speed, MoE usually don't." + ), + ), ): """Start Studio, load a model, print an API key -- one-liner server. @@ -956,6 +968,7 @@ def run( unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --top-k 20 --seed 42 --parallel 8 unsloth studio run --model some-model --chat-template-file /path/to/tpl.jinja + unsloth studio run --model unsloth/Qwen3-27B-GGUF --gguf-variant Q8_0 --tensor-parallel """ extra_llama_args: List[str] = list(ctx.args) if ctx.args else [] @@ -1061,6 +1074,7 @@ def run( args.extend(["--parallel", str(parallel)]) # Forward the explicit polarity (same rationale as --load-in-4bit above). args.append("--cloudflare" if cloudflare else "--no-cloudflare") + args.append("--tensor-parallel" if tensor_parallel else "--no-tensor-parallel") # llama-server pass-through extras → child ctx.args → load payload. if extra_llama_args: args.extend(extra_llama_args) @@ -1126,6 +1140,7 @@ def run( gguf_variant = gguf_variant, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit, + tensor_parallel = tensor_parallel, llama_extra_args = extra_llama_args, ) except RuntimeError as exc: