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

Studio: Add Tensor-Parallel llama.cpp support (#6040)

Browse source 
* Studio: Add Tensor-Parallel llama.cpp support

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: harden Tensor-Parallel fallback and GPU selection

* Studio: reconcile split-mode extras and harden tensor-split planning

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: reconcile split-mode extras in backend duplicate-load guard

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: preserve inherited non-tensor split modes on reload

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: honor cancellation in tensor fallback, preserve tensor mode on rollback, and don't raise an explicit small context

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: reconcile split-mode in reload check and strip it on tensor downgrade

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Strip --tensor-split alongside --split-mode so inherited ratios don't override the tensor planner

An inherited or stale --tensor-split in llama_extra_args was appended after
Studio's computed --tensor-split and won last in llama.cpp, re-introducing the
asymmetric-GPU OOM tensor mode is meant to prevent. Group -ts/--tensor-split
into the split-mode shadow set so it is stripped on inherit and on the layer
fallback; parse_split_mode_override still keys on the mode value only.

* Drop quantized KV for the tensor attempt and report native max context

Tensor mode aborts on a quantized KV cache, so a user with q8_0/q4_1 etc. who
enabled Tensor Parallelism silently fell back to layer split. Clear the cache
type (and strip inherited/explicit --cache-type) for the tensor attempt only;
the layer fallback re-runs with tensor off and keeps the user's choice.

Also report max_available_ctx from the native context, not an explicit small
-c, so the context slider no longer warns too early in tensor mode.

* Reconcile inherited split-mode extras in the already-loaded check

When a same-model load omitted llama_extra_args, the tensor comparison resolved
the raw (None) request and treated an inherited --split-mode tensor server as a
mismatch, forcing a needless reload. Compare using the stored extras stripped
the same way the reload strips them.

* Pass tensor_parallel through compare-mode loads

The generalized compare path loaded each GGUF without tensor_parallel, so
compare ran layer split even with the toggle on and left the settings sheet
stale. Send the toggle and hydrate the loaded state from the response, matching
the main chat and recipe load paths.

* Add --tensor-parallel flag to unsloth studio run

The headless one-liner could only reach tensor mode by passing --split-mode
tensor as a raw llama.cpp extra. Add a first-class --tensor-parallel/
--no-tensor-parallel option that sets the tensor_parallel field on the
/api/inference/load payload, forwarded through the studio-venv re-exec like the
other polarity flags. Matches the web UI toggle and the API field.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
This commit is contained in:
oobabooga 2026-06-12 08:00:52 -03:00 • committed by GitHub
commit 72e67ae5a6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 1386 additions and 41 deletions
Download patch file Download diff file Expand all files Collapse all files

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

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

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

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

70
studio/backend/core/inference/tensor_fallback.py Normal file
View file

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

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

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

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

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

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

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

1
studio/backend/tests/test_studio_api.py
View file

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

578
studio/backend/tests/test_tensor_parallel.py Normal file
View file

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

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

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

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

@ -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({
/>
</div>
)}
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Tensor Parallelism
</span>
<InfoHint>
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.
</InfoHint>
</div>
<Switch
className="panel-switch shrink-0"
checked={tensorParallel}
onCheckedChange={setTensorParallel}
data-test-id="tensor-parallel-switch"
/>
</div>
</>
)}
{!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"

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

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

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

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

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

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

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

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

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

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

2
studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts
View file

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

15
unsloth_cli/commands/studio.py
View file

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