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

Merge remote-tracking branch 'origin/main' into pr-6933-studio-launcher-install-id

Browse source
This commit is contained in:
Lyxot 2026-07-07 23:16:16 +08:00
commit e287449bc5
No known key found for this signature in database
32 changed files with 2194 additions and 178 deletions
Download patch file Download diff file Expand all files Collapse all files

10
install.ps1
View file

@ -2181,7 +2181,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -2195,7 +2195,7 @@ exit 0
}
}
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -2261,7 +2261,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
@ -2273,7 +2273,7 @@ exit 0
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" }
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -2301,7 +2301,7 @@ exit 0
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)

172
install.sh
View file

@ -1529,6 +1529,81 @@ elif [ "$OS" = "macos" ]; then
fi
tauri_diag_marker "$_TAURI_INITIAL_GPU_BRANCH" "none"
# AMD GPU name from the Windows host via WMI, or empty. Discrete cards aren't in
# /proc/cpuinfo, so ask Windows. Cached ("-" = negative), self-contained, bounded
# to 10s. Defined here so the reroute below can use it before _run_bounded exists.
_WSL_AMD_GPU_NAME_CACHE=""
_wsl_amd_gpu_name() {
if [ -n "$_WSL_AMD_GPU_NAME_CACHE" ]; then
[ "$_WSL_AMD_GPU_NAME_CACHE" = "-" ] && return 1
printf '%s' "$_WSL_AMD_GPU_NAME_CACHE"; return 0
fi
command -v powershell.exe >/dev/null 2>&1 || { _WSL_AMD_GPU_NAME_CACHE="-"; return 1; }
_wag_ps="(Get-CimInstance Win32_VideoController | Where-Object { \$_.Name -match 'AMD|Radeon' } | Select-Object -First 1).Name"
if command -v timeout >/dev/null 2>&1; then
_wag_n="$(timeout 10 powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')"
else
_wag_n="$(powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')"
fi
if [ -n "$_wag_n" ]; then _WSL_AMD_GPU_NAME_CACHE="$_wag_n"; printf '%s' "$_wag_n"; return 0; fi
_WSL_AMD_GPU_NAME_CACHE="-"; return 1
}
# ── Bounded command runner ──
# Runs a command under a 10s timeout when the `timeout` binary is available,
# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during
# driver init or after a reset) from hanging the installer: a timed-out probe
# exits nonzero and is treated exactly like a failed probe. No-op semantics on
# hosts without `timeout` (e.g. macOS) or when the probe is healthy.
_run_bounded() {
if command -v timeout >/dev/null 2>&1; then
timeout 10 "$@"
else
"$@"
fi
}
# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every
# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to
# the AMD card). Unset means all devices visible. nvidia-smi ignores this env
# var, so the probes below cannot see the distinction on their own.
_cvd_hides_nvidia() {
[ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1
_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]')
[ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ]
}
# ── NVIDIA usable-GPU helper ──
# Returns 0 (true) if an NVIDIA GPU is present and usable.
# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs,
# which the NVIDIA driver populates on Linux regardless of nvidia-smi state
# -- handles PATH gaps, subprocess timeouts, and driver init races that
# could otherwise cause nvidia-smi to fail and silence NVIDIA detection.
# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches
# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs.
_has_usable_nvidia_gpu() {
if _cvd_hides_nvidia; then
return 1
fi
_nvsmi=""
if command -v nvidia-smi >/dev/null 2>&1; then
_nvsmi="nvidia-smi"
elif [ -x "/usr/bin/nvidia-smi" ]; then
_nvsmi="/usr/bin/nvidia-smi"
fi
if [ -n "$_nvsmi" ]; then
if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
return 0
fi
fi
# Fallback: NVIDIA driver exposes one subdir per GPU under this path.
if [ -d /proc/driver/nvidia/gpus ] && \
[ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
return 0
fi
return 1
}
# Strix Halo ROCm-on-WSL only targets Ubuntu 24.04. On a newer distro (e.g. 26.04)
# with a 24.04 distro present, re-run the install there and stop; else fall through
# to CPU + the `wsl --install` hint below (never auto-create a distro). Runs before
@ -1539,7 +1614,15 @@ _maybe_reroute_strixhalo_to_2404() {
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
[ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0
[ -e /dev/dxg ] || return 0
grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0
# A usable NVIDIA GPU (common on hybrid AMD+NVIDIA hosts) means the CUDA path works on
# this distro, so don't reroute for AMD. _has_usable_nvidia_gpu (moved above) honors
# CUDA_VISIBLE_DEVICES=""/-1 and the /proc/driver/nvidia fallback for PATH/timeout gaps.
if _has_usable_nvidia_gpu; then return 0; fi
# Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes.
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
# Already ROCm-on-WSL? leave a working GPU alone, whatever the version.
if [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ]; then
return 0
@ -2005,61 +2088,6 @@ _has_amd_rocm_gpu() {
return 1
}
# ── Bounded command runner ──
# Runs a command under a 10s timeout when the `timeout` binary is available,
# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during
# driver init or after a reset) from hanging the installer: a timed-out probe
# exits nonzero and is treated exactly like a failed probe. No-op semantics on
# hosts without `timeout` (e.g. macOS) or when the probe is healthy.
_run_bounded() {
if command -v timeout >/dev/null 2>&1; then
timeout 10 "$@"
else
"$@"
fi
}
# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every
# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to
# the AMD card). Unset means all devices visible. nvidia-smi ignores this env
# var, so the probes below cannot see the distinction on their own.
_cvd_hides_nvidia() {
[ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1
_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]')
[ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ]
}
# ── NVIDIA usable-GPU helper ──
# Returns 0 (true) if an NVIDIA GPU is present and usable.
# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs,
# which the NVIDIA driver populates on Linux regardless of nvidia-smi state
# -- handles PATH gaps, subprocess timeouts, and driver init races that
# could otherwise cause nvidia-smi to fail and silence NVIDIA detection.
# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches
# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs.
_has_usable_nvidia_gpu() {
if _cvd_hides_nvidia; then
return 1
fi
_nvsmi=""
if command -v nvidia-smi >/dev/null 2>&1; then
_nvsmi="nvidia-smi"
elif [ -x "/usr/bin/nvidia-smi" ]; then
_nvsmi="/usr/bin/nvidia-smi"
fi
if [ -n "$_nvsmi" ]; then
if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
return 0
fi
fi
# Fallback: NVIDIA driver exposes one subdir per GPU under this path.
if [ -d /proc/driver/nvidia/gpus ] && \
[ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
return 0
fi
return 1
}
# ── Detect GPU and choose PyTorch index URL ──
# Mirrors Get-TorchIndexUrl in install.ps1.
# On CPU-only machines this returns the cpu index, avoiding the solver
@ -2373,19 +2401,19 @@ _persist_rocm_wsl_dropin() {
fi
}
# _wsl_amd_gpu_name is defined earlier so both the reroute and this bootstrap can use it.
_maybe_bootstrap_rocm_wsl() {
[ "${OS:-}" = "wsl" ] || return 0
[ "${SKIP_TORCH:-false}" = "false" ] || return 0
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
# Leave any already-usable GPU completely alone (NVIDIA, or working ROCm).
if _has_usable_nvidia_gpu; then return 0; fi
# "Usable ROCm" here = rocminfo enumerates the gfx1151 agent. Don't use the
# generic _has_amd_rocm_gpu: its broad gfx match accepts "gfx11-generic" and
# would skip this bootstrap while the real GPU is still unusable. awk consumes
# all input, so rocminfo isn't SIGPIPE'd like `grep -q` would under pipefail.
# Usable ROCm = rocminfo enumerates a real GPU agent: gfx[1-9] (excludes gfx000,
# the CPU agent) and not the "gfx11-generic" fallback. awk consumes all input so
# rocminfo isn't SIGPIPE'd like `grep -q` under pipefail.
_ensure_rocm_probe_env
if command -v rocminfo >/dev/null 2>&1 && \
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/ && !/generic/{found=1} END{exit !found}'; then
# rocminfo may work only via the transient env _ensure_rocm_probe_env
# just set, which dies with the installer. Persist the drop-in so login
# shells (Studio, llama.cpp) inherit it -- else a reinstall over an
@ -2395,9 +2423,12 @@ _maybe_bootstrap_rocm_wsl() {
fi
# WSL GPU passthrough device must exist (present on any WSL2 GPU host).
[ -e /dev/dxg ] || return 0
# Only Strix Halo (gfx1151): rocminfo can't tell us the arch yet, so match
# the CPU model string WSL exposes (e.g. "AMD Ryzen AI Max+ ... Radeon 8060S").
grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0
# Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also
# ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo.
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
command -v bash >/dev/null 2>&1 || return 0
# Fast path: already configured (librocdxg present) but launched from a
@ -2415,7 +2446,8 @@ _maybe_bootstrap_rocm_wsl() {
fi
echo ""
substep "Detected AMD Strix Halo (Radeon 8000S) in WSL with no ROCm runtime yet." "$C_WARN"
_rw_gpu="$(_wsl_amd_gpu_name 2>/dev/null || true)"; [ -n "$_rw_gpu" ] || _rw_gpu="an AMD GPU"
substep "Detected ${_rw_gpu} in WSL with no ROCm runtime yet." "$C_WARN"
substep "Setting up ROCm-on-WSL (ROCm 7.2 + librocdxg) automatically to enable this GPU."
substep "One-time, uses sudo and a large download. (skip: re-run with UNSLOTH_SKIP_ROCM_WSL_SETUP=1)"
@ -2720,7 +2752,7 @@ if [ "$_MIGRATED" = true ]; then
# to prevent transitive torch resolution.
run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
"unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1"
# Resolve pydantic WITH deps so pip pins pydantic-core to the
# matching version (no-torch-runtime.txt below is --no-deps).
# All transitive deps are torch-free.
@ -2735,7 +2767,7 @@ if [ "$_MIGRATED" = true ]; then
# overrides file, so UV_OVERRIDE is unset and this positional is the only cover.
run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" ${_MLX_LM_EXCLUDE_ARG:-}
"unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" ${_MLX_LM_EXCLUDE_ARG:-}
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@ -2939,7 +2971,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--upgrade-package unsloth --upgrade-package unsloth-zoo \
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
"unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1"
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd_retry "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
@ -2957,7 +2989,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
--upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1"
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -2989,7 +3021,7 @@ else
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."

6
pyproject.toml
View file

@ -73,7 +73,7 @@ triton = [
]
huggingfacenotorch = [
"unsloth_zoo>=2026.6.7",
"unsloth_zoo>=2026.7.1",
"wheel>=0.42.0",
"packaging",
"numpy",
@ -94,7 +94,7 @@ huggingfacenotorch = [
]
huggingface = [
"unsloth[huggingfacenotorch]",
"unsloth_zoo>=2026.6.7",
"unsloth_zoo>=2026.7.1",
"torchvision",
"unsloth[triton]",
]
@ -579,7 +579,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
"unsloth_zoo>=2026.6.7",
"unsloth_zoo>=2026.7.1",
"packaging",
"tyro",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",

57
scripts/install_rocm_wsl_strixhalo.sh
View file

@ -3,13 +3,14 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
#
# ──────────────────────────────────────────────────────────────────────────────
# Enable ROCm-on-WSL for AMD Strix Halo (Radeon 8060S / gfx1151)
# Enable ROCm-on-WSL for AMD GPUs (Strix Halo/Point APUs AND discrete Radeon RX
# 7000/9000). Verified on gfx1151 (Radeon 8060S) and gfx1200 (Radeon RX 9060 XT).
# ──────────────────────────────────────────────────────────────────────────────
# install.sh already routes gfx1151 to the right ROCm wheels once a ROCm runtime
# is present; what it does NOT do is install AMD's ROCm userspace + the WSL DXG
# bridge. This helper automates that Linux-side prerequisite on Ubuntu 24.04
# WSL2 and is invoked by install.sh when it sees a Strix Halo APU in WSL (via
# /dev/dxg) but no ROCm runtime yet. Fully idempotent (re-run just re-verifies).
# install.sh routes the detected arch to the right ROCm wheels once a runtime exists;
# what it does NOT do is install AMD's ROCm userspace + the WSL DXG bridge (librocdxg).
# This helper does that Linux-side prerequisite on Ubuntu 24.04 WSL2, invoked by
# install.sh when it sees an AMD GPU via /dev/dxg but no ROCm yet. Arch-agnostic: the
# arch is auto-detected from rocminfo (override UNSLOTH_WSL_GFX=gfx1200). Idempotent.
#
# Manual, admin-gated Windows prerequisite: an AMD Adrenalin driver with
# production ROCDXG/WSL support (26.2.2+). install.ps1 offers to update it. Once
@ -34,10 +35,12 @@ set -euo pipefail
# ── Tunables (override via env) ──────────────────────────────────────────────
ROCM_VER="${UNSLOTH_WSL_ROCM_VER:-7.2.1}" # ROCm release to install
GFX="gfx1151"
# GPU arch: empty = auto-detect from rocminfo after install (override UNSLOTH_WSL_GFX=gfx1200).
# The ROCm + librocdxg setup is arch-agnostic; only verify + the smoke test need the arch.
GFX="${UNSLOTH_WSL_GFX:-}"
LIBROCDXG_REF="${UNSLOTH_LIBROCDXG_REF:-develop}" # ROCm/librocdxg git ref to build
# AMD's gfx1151 wheel index (same one install.sh uses); only for the smoke test.
TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${GFX}/"
# AMD's wheel index for the (optional) smoke test; resolved after arch detection.
TORCH_INDEX=""
# Optional torch smoke test (throwaway venv). OFF by default: install.sh installs
# torch itself into the real venv right after, so a duplicate download is wasteful.
SMOKE_TEST="${UNSLOTH_WSL_SMOKE_TEST:-0}"
@ -220,12 +223,12 @@ $SUDO ldconfig
say "Persisting ROCm-on-WSL environment"
_envfile="/etc/profile.d/unsloth-rocm-wsl.sh"
$SUDO tee "$_envfile" >/dev/null <<EOF
# >>> Unsloth ROCm-on-WSL (gfx1151) >>>
# >>> Unsloth ROCm-on-WSL >>>
export HSA_ENABLE_DXG_DETECTION=1
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
export PATH="${ROCM_DIR}/bin:\${PATH}"
export LD_LIBRARY_PATH="${ROCM_DIR}/lib:\${LD_LIBRARY_PATH:-}"
# <<< Unsloth ROCm-on-WSL (gfx1151) <<<
# <<< Unsloth ROCm-on-WSL <<<
EOF
# also drop into ~/.bashrc for interactive shells
if [ -n "${HOME:-}" ] && ! grep -q "Unsloth ROCm-on-WSL" "${HOME}/.bashrc" 2>/dev/null; then
@ -237,32 +240,50 @@ export PATH="${ROCM_DIR}/bin:${PATH}"
export LD_LIBRARY_PATH="${ROCM_DIR}/lib:${LD_LIBRARY_PATH:-}"
# ── Step 5: verify the runtime enumerates the GPU ────────────────────────────
say "Verifying rocminfo sees ${GFX}"
say "Verifying rocminfo enumerates the GPU over DXG"
# Capture rocminfo into a var BEFORE grepping: piping into `grep -q` SIGPIPEs
# rocminfo on first match, which under `set -o pipefail` turns a successful match
# into a pipeline failure. Match the gfx1151 ISA "Name:" agent exactly (not a
# broad gfx1[0-9]) so a generic fallback ISA or unrelated RDNA GPU can't pass.
# into a pipeline failure.
_rocminfo_out="$(rocminfo 2>/dev/null || true)"
if ! printf '%s\n' "$_rocminfo_out" | grep -qE "Name:[[:space:]]*${GFX}([^0-9]|$)"; then
# GPU agents advertise an ISA "Name: gfxNNNN". Match gfx[1-9] (excludes gfx000, the CPU
# agent), drop the "gfx*-generic" fallback ISA, and take the first real GPU arch.
_detected_gfx="$(printf '%s\n' "$_rocminfo_out" | grep -E 'Name:[[:space:]]*gfx[1-9]' | grep -v 'generic' | grep -oE 'gfx[1-9][0-9a-z]*' | head -1 || true)"
if [ -z "$_detected_gfx" ]; then
printf '%s\n' "$_rocminfo_out" | head -25 >&2 || true
die "rocminfo did not enumerate a ${GFX} GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run."
die "rocminfo did not enumerate any GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run."
fi
# Honour a caller-pinned arch (sanity-check via a consuming grep, not grep -q: under
# pipefail -q would SIGPIPE printf on large output and misreport the arch); else adopt.
if [ -n "$GFX" ] && ! printf '%s\n' "$_rocminfo_out" | grep -E "Name:[[:space:]]*${GFX}([^0-9]|$)" >/dev/null; then
die "rocminfo enumerated '${_detected_gfx}' but not the requested UNSLOTH_WSL_GFX='${GFX}'."
fi
GFX="${GFX:-$_detected_gfx}"
# Display-only summary: best-effort (|| true) so head's early pipe-close under
# `set -o pipefail` can't fail the bootstrap after verification already passed.
printf '%s\n' "$_rocminfo_out" | grep -E 'Marketing Name|Device Type|Compute Unit' | grep -iE "Radeon|GPU|Compute" | head -3 || true
note "ROCm-on-WSL runtime is live for ${GFX}."
# ── Step 6 (optional): torch smoke test from the gfx1151 index ───────────────
# ── Step 6 (optional): torch smoke test from AMD's per-arch wheel index ───────
if [ "$SMOKE_TEST" = "1" ]; then
say "Smoke-testing PyTorch on ${GFX} (throwaway venv)"
# Map the detected arch to AMD's repo.amd.com wheel family index.
case "$GFX" in
gfx1200|gfx1201) _fam="gfx120X-all" ;;
gfx1100|gfx1101|gfx1102|gfx1103) _fam="gfx110X-all" ;;
*) _fam="$GFX" ;; # gfx1150/gfx1151/gfx90a: own index
esac
TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${_fam}/"
_venv="${HOME}/.unsloth/rocm-smoketest"
rm -rf "$_venv"; python3 -m venv "$_venv"
"$_venv/bin/pip" install --quiet --upgrade pip
# gfx1151 index is primary (torch + triton); PyPI only an extra for pure-py
# AMD arch index is primary (torch + triton); PyPI only an extra for pure-py
# deps. The constraint keeps pip on the ROCm wheel, not a newer PyPI CUDA torch.
"$_venv/bin/pip" install --index-url "$TORCH_INDEX" \
--extra-index-url https://pypi.org/simple "$TORCH_CONSTRAINT" || \
die "torch install from ${TORCH_INDEX} failed."
# WSL: torch's bundled ROCr must load the DXG bridge -- drop librocdxg into torch/lib.
_tlib="$("$_venv/bin/python" -c 'import torch,os;print(os.path.join(os.path.dirname(torch.__file__),"lib"))' 2>/dev/null || true)"
[ -d "$_tlib" ] && cp -f "${ROCM_DIR}"/lib/librocdxg.so* "$_tlib"/ 2>/dev/null || true
"$_venv/bin/python" - <<'PY'
import torch
ok = torch.cuda.is_available()

9
studio/backend/assets/configs/inference_defaults.json
View file

@ -235,6 +235,13 @@
"min_p": 0.1,
"repetition_penalty": 1.0
},
"deepseek-v4": {
"temperature": 1.0,
"top_p": 1.0,
"top_k": -1,
"min_p": 0.0,
"repetition_penalty": 1.0
},
"deepseek-r1": {
"temperature": 0.6,
"top_p": 0.95,
@ -394,7 +401,7 @@
"phi-4", "phi-3",
"mistral-nemo", "mistral-small", "mistral-large", "magistral", "ministral",
"devstral", "pixtral",
"deepseek-r1", "deepseek-v3", "deepseek-ocr",
"deepseek-v4", "deepseek-r1", "deepseek-v3", "deepseek-ocr",
"glm-5", "glm-4",
"nemotron",
"minimax-m2.7", "minimax-m2.5", "minimax",

2
studio/backend/core/inference/defaults.py
View file

@ -8,6 +8,7 @@ import utils.hardware.hardware as hw
DEFAULT_MODELS_GGUF = [
"unsloth/Qwen3.6-27B-MTP-GGUF",
"unsloth/Qwen3.6-35B-A3B-MTP-GGUF",
"unsloth/DeepSeek-V4-Flash-GGUF",
"unsloth/gemma-4-E2B-it-GGUF",
"unsloth/gemma-4-E4B-it-GGUF",
"unsloth/gemma-4-31B-it-GGUF",
@ -27,6 +28,7 @@ DEFAULT_MODELS_GGUF = [
DEFAULT_MODELS_STANDARD = [
"unsloth/Qwen3.6-27B-MTP-GGUF",
"unsloth/Qwen3.6-35B-A3B-MTP-GGUF",
"unsloth/DeepSeek-V4-Flash-GGUF",
"unsloth/gemma-4-E2B-it-GGUF",
"unsloth/gemma-4-E4B-it-GGUF",
"unsloth/gemma-4-31B-it-GGUF",

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

@ -686,6 +686,16 @@ def detect_reasoning_flags(
else []
)
if effort_levels:
# DeepSeek-V4's encoder accepts reasoning_effort {'high', 'max'} but its
# template only branches on 'max', so the literal scan misses 'high'. Add it
# (matched on whole repo-name segments, so 'deepseek-v40' won't false-match)
# to expose the full none/high/max ladder instead of none/max.
segments = re.split(r"[-_.]", (model_identifier or "").lower().split("/")[-1])
is_dsv4 = "deepseek4" in segments or any(
a == "deepseek" and b == "v4" for a, b in zip(segments, segments[1:])
)
if is_dsv4 and "high" not in effort_levels:
effort_levels = sorted(set(effort_levels) | {"high"}, key = _REASONING_EFFORT_SCALE.index)
# GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort
# level among a discrete set (e.g. 'high' | 'max'). Distinct from
# gpt-oss (reasoning_effort only, no on/off gate) and Qwen
@ -1741,9 +1751,13 @@ class LlamaCppBackend:
# 'low' effort the way gpt-oss does (those models genuinely
# cannot disable).
thinking_off = enable_thinking is False or reasoning_effort == "none"
if enable_thinking is not None or reasoning_effort == "none":
# A named effort level implies thinking on, so emit enable_thinking
# even if the caller sent only reasoning_effort (else the template
# defaults it off and the requested level never renders).
effort_on = reasoning_effort in self._reasoning_effort_levels
if enable_thinking is not None or reasoning_effort == "none" or effort_on:
kwargs["enable_thinking"] = not thinking_off
if not thinking_off and reasoning_effort in self._reasoning_effort_levels:
if not thinking_off and effort_on:
kwargs["reasoning_effort"] = reasoning_effort
elif self._reasoning_style == "reasoning_effort":
if reasoning_effort in ("none", "low", "medium", "high"):
@ -3129,6 +3143,12 @@ class LlamaCppBackend:
_CTX_COMPUTE_BYTES_PER_EMBD = 2.25 # quantized KV, regular attention (dequant scratch)
_CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25 # quantized KV, MLA (compressed attn: measured 0.94x)
_CTX_COMPUTE_F16_MASK_SAFETY = 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok)
# DeepSeek-V4 (deepseek4): its lightning indexer + sparse attention reserve a large
# context-scaling compute buffer the rates above miss (present even with an f16
# cache). Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k -> ~65.5 GiB at 1M. Without
# it auto-fit commits the full 1M train context, OOMs the reserve, and spills to CPU.
_DSV4_CTX_COMPUTE_FLAT_BYTES = 2 * 1024**3 # ctx-independent indexer scratch
_DSV4_CTX_COMPUTE_BYTES_PER_TOK = 72000 # per token at ub=512 (~72 GiB at 1M)
def _estimate_compute_buffer_bytes(
self,
@ -3178,6 +3198,14 @@ class LlamaCppBackend:
if n_embd <= 0 or n_ctx <= 0:
return 0
ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH))
if getattr(self, "_architecture", None) == "deepseek4":
# DSV4 indexer/CSA buffer (see constants): flat + linear, ub-scaled. Fires
# for any KV type -- the indexer scratch is present even with an f16 cache.
ub_scale = ub / self._DEFAULT_N_UBATCH
return int(
self._DSV4_CTX_COMPUTE_FLAT_BYTES
+ self._DSV4_CTX_COMPUTE_BYTES_PER_TOK * n_ctx * ub_scale
)
if _kv_bytes_per_elem(cache_type_kv) < 2.0:
# Quantized cache: the dequant scratch dominates and scales with n_embd.
# MLA (compressed KV) needs far less of it: measured 0.94 x n_embd on

117
studio/backend/core/rag/embeddings.py
View file

@ -63,6 +63,87 @@ def _install_torchao_stub_once() -> None:
install_torchao_windows_rocm_stub()
class UnsafeEmbeddingModelError(RuntimeError):
"""Raised when the embedding model repo is flagged unsafe. A distinct type so the
llama-server fallback paths re-raise it instead of masking a security block as a
routine ST failure."""
def _ambient_hf_token() -> str | None:
"""The HF token the loader itself would use (HF_TOKEN env or the cached login), so
the scan can reach a gated/private repo instead of failing open. None if unavailable."""
try:
from huggingface_hub import get_token
return get_token()
except Exception:
return None
def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]:
"""The module directories a SentenceTransformer load reads weights from, taken from
the repo's ``modules.json`` (each module's non-empty ``path``, e.g. ``0_Transformer``).
ST deserializes ``pytorch_model.bin`` from these dirs, so they are load roots for the
security scan: a flagged pickle directly under one must block. Returns () on any
failure (no modules.json, offline, malformed) so the guard never bricks the embedder.
"""
try:
import json
from utils.paths import is_local_path
if is_local_path(name):
from pathlib import Path
from utils.paths import normalize_path
path = Path(normalize_path(name)).expanduser() / "modules.json"
if not path.is_file():
return ()
data = json.loads(path.read_text())
else:
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError
try:
local = hf_hub_download(name, "modules.json", token = token or None)
except EntryNotFoundError:
return ()
data = json.loads(open(local).read())
subdirs = []
for module in data or ():
sub = str((module or {}).get("path", "")).strip().strip("/")
if sub:
subdirs.append(sub)
return tuple(dict.fromkeys(subdirs))
except Exception:
return ()
def _guard_model_security(name: str) -> None:
"""Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside
SentenceTransformer regardless of trust_remote_code. Defense in depth behind the
/settings gate (a name can also arrive via env/default); local paths and unreachable
scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error.
"""
try:
from utils.security import evaluate_file_security, security_load_subdirs
token = _ambient_hf_token()
# Union the audio-model load roots with the ST module dirs so a flagged pickle
# directly under a Transformer module dir (0_Transformer/) blocks instead of
# passing as an unreferenced nested shard.
load_subdirs = tuple(
dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token)))
)
blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked
except Exception:
return
if blocked:
raise UnsafeEmbeddingModelError(
f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security "
"scan; refusing to load. Set a different RAG embedding model."
)
def _get(model_name: str | None = None):
"""Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16
for a ~1.5x speedup at negligible accuracy loss."""
@ -75,6 +156,7 @@ def _get(model_name: str | None = None):
device = _device()
logger.info("loading embedding model %s on %s", name, device)
_guard_model_security(name)
_model = SentenceTransformer(
name, device = device, model_kwargs = {"torch_dtype": "float16"}
)
@ -159,6 +241,8 @@ class _SentenceTransformersBackend:
):
try:
return _st_encode(texts, model_name = model_name, normalize = normalize)
except UnsafeEmbeddingModelError:
raise # a security block must hard-fail, not fall back to llama-server
except Exception as st_err: # noqa: BLE001 - runtime ST/CUDA encode failure
# ST loaded but this encode blew up; swap the process to the llama-server
# embedder (so later encodes stay in one space) and retry.
@ -222,6 +306,8 @@ def _build_st_backend_or_fallback():
try:
backend.warm(model_name = None)
return backend
except UnsafeEmbeddingModelError:
raise # a security block must hard-fail, not fall back to llama-server
except Exception as st_err: # noqa: BLE001 - any ST/torch import or load failure
fallback = _try_make_llama_backend()
if fallback is None:
@ -290,6 +376,37 @@ def _reset_backend() -> None:
_backend_key = None
def active_backend_is_llama() -> bool:
"""True when this process actually embeds via the llama-server (GGUF) backend.
Reflects the ACTUAL built backend once one exists: an ``auto`` install that
resolves to sentence-transformers but then falls back to llama-server at
runtime (``_build_st_backend_or_fallback`` on a torch/CUDA load failure, or
``_switch_to_llama_fallback`` on an encode failure) loads only inert GGUF, so
callers gating on the ST pickle must see llama here. Before any backend is
built, defers to the resolver (``auto`` -> ``_resolve_auto()``, else the raw
key) exactly as a fresh process would. Never raises: a backend probe must not
block saving a model."""
try:
with _backend_lock:
backend = _backend
if backend is not None:
# A backend exists: report what it ACTUALLY is. A concrete
# sentence-transformers backend must return False even if the
# resolver would now pick llama, so its pickle stays gated. If the
# llama import fails we cannot be llama, so fall to the safe False.
try:
from .embed_llama_server import LlamaServerBackend
except Exception: # noqa: BLE001 - llama plumbing import must never block
return False
return isinstance(backend, LlamaServerBackend)
raw = (config.EMBED_BACKEND or "auto").strip().lower()
key = _resolve_auto() if raw in _AUTO_ALIASES else raw
return key in _LLAMA_ALIASES
except Exception: # noqa: BLE001 - a backend probe must never block saving
return False
def warm(model_name: str | None = None) -> None:
"""Eagerly load the embedder so the first real request isn't slow."""
_get_backend().warm(model_name = model_name)

78
studio/backend/routes/settings.py
View file

@ -260,17 +260,29 @@ def _embedding_model_response() -> EmbeddingModelResponse:
)
def _llama_backend_active() -> bool:
"""True when this install embeds via the llama-server (GGUF) backend."""
from core.rag import config as rag_config
from core.rag import embeddings
def _ambient_hf_token() -> Optional[str]:
"""The HF token the loader would use (HF_TOKEN env or the cached login), so a gated
repo is scanned rather than failing open. None if unavailable."""
try:
raw = (rag_config.EMBED_BACKEND or "auto").strip().lower()
key = embeddings._resolve_auto() if raw in embeddings._AUTO_ALIASES else raw
from huggingface_hub import get_token
return get_token()
except Exception:
return None
def _llama_backend_active() -> bool:
"""True when this install actually embeds via the llama-server (GGUF) backend.
Delegates to the embeddings module so a runtime fallback from
sentence-transformers to llama-server (after a torch/CUDA load or encode
failure) is honored: in that state the process loads only inert GGUF, so the
ST pickle gate below must not hard-block a repo whose GGUF companion is clean.
Before any backend is built this still reflects the resolver."""
from core.rag import embeddings
try:
return embeddings.active_backend_is_llama()
except Exception: # noqa: BLE001 - backend probe must never block saving
return False
return key in embeddings._LLAMA_ALIASES
def _resolves_as_local_gguf(model: str) -> bool:
@ -357,6 +369,8 @@ def update_embedding_model(
"""Set the RAG embedding model. Unless ``force`` is set, the repo is verified
to be an embedding model via HF metadata; an unverifiable model (wrong type,
typo, gated repo, or no network) returns 409 so the UI can offer "save anyway".
A repo flagged unsafe by HF's security scan returns 403 instead: a hard block
that ``force`` cannot bypass, so the UI must not offer "save anyway".
Documents indexed under the previous model must be re-uploaded."""
from utils.models import is_embedding_model
@ -370,15 +384,51 @@ def update_embedding_model(
event = "settings.update_embedding_model_failed",
log = logger,
) from exc
hf_token = (payload.hf_token or "").strip() or None
# The env/default model needs no verification; saving it is a no-op override.
# A local GGUF on the llama-server backend is accepted as-is: it is exactly
# what the backend loads, and HF metadata cannot verify a local path.
if (
model != default_embedding_model()
and not payload.force
and not (_llama_backend_active() and _resolves_as_local_gguf(model))
):
hf_token = (payload.hf_token or "").strip() or None
is_local_gguf = _llama_backend_active() and _resolves_as_local_gguf(model)
# The pickle gate only matters for the sentence-transformers backend, which is what
# deserializes pickles. On the llama-server backend the embedder loads GGUF files
# (inert) from effective_gguf_repo(), so scanning the ST repo's pickle here would
# wrongly reject a custom repo whose GGUF companion is clean; the GGUF availability
# checks below cover that path instead.
scan_st_pickle = (
model != default_embedding_model() and not is_local_gguf and not _llama_backend_active()
)
if scan_st_pickle:
# Malware/pickle gate before we persist a repo the embedder later loads with
# SentenceTransformer. Runs even under force (force only skips the is-embedding
# type check for offline/local repos HF cannot verify); local paths and
# unreachable scans fail open inside evaluate_file_security.
from utils.security import evaluate_file_security, security_load_subdirs
from core.rag.embeddings import _st_module_subdirs
# Fall back to the loader's own token so a gated/private repo is actually scanned
# (a token-less scan fails open for exactly the repo that would still load).
scan_token = hf_token or _ambient_hf_token()
# Include the ST module dirs (0_Transformer/) so a flagged pickle directly under
# one blocks instead of passing as an unreferenced nested shard.
load_subdirs = tuple(
dict.fromkeys(
(
*security_load_subdirs(model, scan_token),
*_st_module_subdirs(model, scan_token),
)
)
)
if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked:
# 403, not 409: the client routes every 409 into the forceable "save anyway"
# flow, but this block is a hard, non-forceable security refusal.
raise HTTPException(
status_code = 403,
detail = (
f"{model!r} is flagged as unsafe by Hugging Face's security scan and "
"cannot be used as the embedding model."
),
)
if model != default_embedding_model() and not payload.force and not is_local_gguf:
from core.rag import config as rag_config
# A GGUF-named repo on the llama-server backend is loaded from its .gguf

61
studio/backend/tests/test_compute_buffer.py
View file

@ -65,12 +65,14 @@ def _backend(
vocab = 248320,
embd = 5120,
mla = None,
arch = None,
):
"""Backend with just the dims the compute-buffer estimate reads."""
b = LlamaCppBackend.__new__(LlamaCppBackend)
b._vocab_size = vocab
b._embedding_length = embd
b._key_length_mla = mla # non-None -> MLA (compressed attention)
b._architecture = arch # GGUF general.architecture (e.g. 'deepseek4')
return b
@ -290,3 +292,62 @@ class TestContextBufferMLA:
b = _backend(embd = 6144, mla = 256)
est = b._compute_buffer_ctx_bytes(754688, cache_type_kv = "q8_0") / MIB
assert est <= 4141 * 1.7
class TestContextBufferDSV4:
"""DeepSeek-V4 (deepseek4) reserves a large lightning-indexer / sparse-attention
compute buffer the KQ-mask and MLA rates miss (present even with an f16 cache).
Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k ctx, ~65.5 GiB at 1M. The auto-fit
must see this so it does not commit the full 1M train context and OOM (spilling
to CPU at ~4 tok/s)."""
_MEASURED_1M_GIB = 65.5 # 70353790464 B compute-graph reserve that OOM'd at 1M ctx
GIB = 1024**3
def test_covers_measured_1m_buffer(self):
b = _backend(embd = 4096, arch = "deepseek4")
gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB
assert gib >= self._MEASURED_1M_GIB, f"under-reserved {gib:.1f} < {self._MEASURED_1M_GIB}"
def test_not_wildly_over_at_1m(self):
# Within ~1.3x of measured so the fit still grants a large (~256k) context.
b = _backend(embd = 4096, arch = "deepseek4")
gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB
assert gib <= self._MEASURED_1M_GIB * 1.3
def test_fires_for_f16_cache(self):
# The bug: an f16 (default) cache took the tiny mask-only path. DSV4 must
# reserve GiB, not the ~MiB a non-DSV4 model reserves at the same ctx.
dsv4 = _backend(embd = 4096, arch = "deepseek4")._compute_buffer_ctx_bytes(
262144, cache_type_kv = "f16"
)
other = _backend(embd = 4096, arch = "qwen3")._compute_buffer_ctx_bytes(
262144, cache_type_kv = "f16"
)
assert dsv4 > 40 * other
def test_cache_type_independent(self):
# Indexer scratch is present for an f16 and a quantized cache alike.
b = _backend(embd = 4096, arch = "deepseek4")
assert b._compute_buffer_ctx_bytes(
262144, cache_type_kv = "f16"
) == b._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0")
def test_flat_floor_at_small_ctx(self):
# ~2 GiB indexer scratch present even at tiny ctx (covers the measured 16k ~2 GiB).
b = _backend(embd = 4096, arch = "deepseek4")
assert b._compute_buffer_ctx_bytes(16384, cache_type_kv = "f16") / self.GIB >= 2.0
def test_scales_with_context_and_ubatch(self):
b = _backend(embd = 4096, arch = "deepseek4")
assert b._compute_buffer_ctx_bytes(131072) > b._compute_buffer_ctx_bytes(65536)
assert b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) > b._compute_buffer_ctx_bytes(
131072, n_ubatch = 256
)
def test_non_dsv4_unchanged(self):
# Regression guard: a non-deepseek4 model keeps the mask-only f16 rate.
b = _backend(embd = 4096, arch = "llama")
per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "f16") / 100000
expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY
assert per_tok == pytest.approx(expected, rel = 1e-6)

181
studio/backend/tests/test_deepseek_v4_thinking_effort.py Normal file
View file

@ -0,0 +1,181 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""DeepSeek-V4-Flash reasoning toggle: None / High / Max.
The GGUF template gates thinking with ``enable_thinking`` and only branches
``reasoning_effort`` on ``'max'`` (an escalation layered over plain thinking).
Detection used to return the single level ``['max']``, so the UI collapsed to
None / Max and the plain-thinking tier was unreachable. Detection now surfaces
``'high'`` as that plain tier, giving None / High / Max. These tests pin the
classifier, the GLM-style parity case, and the full request-kwargs -> rendered
prompt path for each state (the model itself is too large to load here).
"""
from __future__ import annotations
import sys
from pathlib import Path
from types import SimpleNamespace
import pytest
_backend_root = Path(__file__).resolve().parent.parent
if str(_backend_root) not in sys.path:
sys.path.insert(0, str(_backend_root))
# Faithful slice of the DeepSeek-V4-Flash GGUF template: the enable_thinking
# gate, the sole ``reasoning_effort == 'max'`` escalation, and the plain-think
# fallback. Any non-'max' effort renders as ordinary thinking.
DEEPSEEK_V4_TEMPLATE = """
{%- if not thinking is defined -%}
{%- if enable_thinking is defined -%}
{%- set thinking = enable_thinking -%}
{%- else -%}
{%- set thinking = false -%}
{%- endif -%}
{%- endif -%}
{%- if not reasoning_effort is defined -%}
{%- set reasoning_effort = none -%}
{%- endif -%}
{{- bos_token -}}
{%- if thinking and reasoning_effort == 'max' -%}
{{- 'Reasoning Effort: Absolute maximum with no shortcuts permitted.\\n\\n' -}}
{%- endif -%}
{%- for message in messages -%}
{{- '<|User|>' + (message['content'] or '') -}}
{%- endfor -%}
{%- if add_generation_prompt -%}
{{- '<|Assistant|>' -}}
{%- if thinking -%}{{- '<think>' -}}{%- else -%}{{- '</think>' -}}{%- endif -%}
{%- endif -%}
"""
# GLM-5.2-style: branches on two effort literals, so 'high' already exists as
# the sub-'max' tier and detection must leave the pair untouched.
GLM_STYLE_TEMPLATE = """
{%- if enable_thinking -%}
{%- if reasoning_effort == 'high' -%}{{- 'H' -}}
{%- elif reasoning_effort == 'max' -%}{{- 'M' -}}
{%- endif -%}
{%- endif -%}
"""
# A ['max']-only template under a non-deepseek id: the synthetic 'high' is scoped
# to deepseek-v4, so this must stay ['max'] (no phantom 'high').
NON_DEEPSEEK_MAX_ONLY_TEMPLATE = DEEPSEEK_V4_TEMPLATE
# A template whose sole effort literal is a sub-'max' level: the guard targets
# only the ['max']-alone case, so a lone 'high' stays a singleton.
HIGH_ONLY_TEMPLATE = """
{%- if enable_thinking and reasoning_effort == 'high' -%}{{- 'H' -}}{%- endif -%}
"""
def _render(template: str, **kwargs) -> str:
jinja2 = pytest.importorskip("jinja2")
env = jinja2.Environment()
tmpl = env.from_string(template)
return tmpl.render(bos_token = "<BOS>", add_generation_prompt = True, **kwargs)
# -- Classifier -------------------------------------------------------
def test_deepseek_v4_surfaces_high_as_plain_tier():
"""Sole 'max' escalation expands to ['high', 'max'] so None/High/Max show."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash")
assert flags["supports_reasoning"] is True
assert flags["reasoning_style"] == "enable_thinking_effort"
assert flags["reasoning_effort_levels"] == ["high", "max"]
def test_glm_style_two_level_template_unchanged():
"""A template that already names a sub-'max' tier is left as-is."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(GLM_STYLE_TEMPLATE, "unsloth/GLM-5.2")
assert flags["reasoning_style"] == "enable_thinking_effort"
assert flags["reasoning_effort_levels"] == ["high", "max"]
def test_synthetic_high_scoped_to_deepseek_v4():
"""The same ['max']-only template under a non-deepseek id keeps ['max']."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(NON_DEEPSEEK_MAX_ONLY_TEMPLATE, "vendor/OtherHybrid-GGUF")
assert flags["reasoning_effort_levels"] == ["max"]
def test_guard_does_not_fire_for_sub_max_singleton():
"""The expansion targets only ['max']; a lone 'high' stays a singleton."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(HIGH_ONLY_TEMPLATE, "custom/high-only")
assert flags["reasoning_effort_levels"] == ["high"]
# -- Request kwargs -> rendered prompt, for each state ----------------
def _kwargs_for(flags: dict, enable_thinking, reasoning_effort):
"""Drive the real backend method with a shim carrying the detected flags."""
from core.inference.llama_cpp import LlamaCppBackend
shim = SimpleNamespace(
_supports_reasoning = flags["supports_reasoning"],
_reasoning_always_on = flags["reasoning_always_on"],
_reasoning_style = flags["reasoning_style"],
_reasoning_effort_levels = flags["reasoning_effort_levels"],
_supports_preserve_thinking = flags["supports_preserve_thinking"],
)
build = LlamaCppBackend._request_reasoning_kwargs.__get__(shim)
return build(enable_thinking, reasoning_effort, None) or {}
def _flags():
from core.inference.llama_cpp import detect_reasoning_flags
return detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash")
def test_none_state_renders_non_thinking():
"""UI 'None' -> enable_thinking=false -> closed </think>, no preamble."""
kwargs = _kwargs_for(_flags(), enable_thinking = False, reasoning_effort = None)
assert kwargs == {"enable_thinking": False}
out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs)
assert out.endswith("</think>")
assert "Absolute maximum" not in out
def test_high_state_renders_plain_thinking():
"""UI 'High' -> et=true, effort=high -> open <think>, no max preamble."""
kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "high")
assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"}
out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs)
assert out.endswith("<think>")
assert "Absolute maximum" not in out
def test_max_state_injects_max_preamble():
"""UI 'Max' -> et=true, effort=max -> open <think> plus the max preamble."""
kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "max")
assert kwargs == {"enable_thinking": True, "reasoning_effort": "max"}
out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs)
assert out.endswith("<think>")
assert "Absolute maximum" in out
def test_high_effort_alone_enables_thinking():
"""API caller sending only reasoning_effort='high' (no enable_thinking) still
gets thinking on, so the newly exposed High mode renders correctly."""
kwargs = _kwargs_for(_flags(), enable_thinking = None, reasoning_effort = "high")
assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"}
out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs)
assert out.endswith("<think>")
assert "Absolute maximum" not in out

365
studio/backend/tests/test_embedding_model_security_gate.py Normal file
View file

@ -0,0 +1,365 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""The RAG embedding model must pass the malware/pickle gate before it is persisted or
loaded. A flagged repo (or any repo saved with force) previously reached
SentenceTransformer unscanned, bypassing the normal model-load protections."""
from pathlib import Path
import sys
import types as _types
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
import routes.settings as settings
class _Decision:
def __init__(self, blocked):
self.blocked = blocked
def _security_stub(blocked):
mod = _types.ModuleType("utils.security")
mod.evaluate_file_security = lambda *a, **k: _Decision(blocked)
mod.security_load_subdirs = lambda *a, **k: ()
return mod
@pytest.fixture
def client(monkeypatch):
# The settings scan unions in the ST module dirs read from modules.json; keep it
# offline and deterministic for the endpoint tests that use this fixture.
import core.rag.embeddings as embeddings
monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ())
saved: dict = {}
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
monkeypatch.setattr(settings, "_llama_backend_active", lambda: False)
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
app = FastAPI()
app.include_router(settings.router)
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
return TestClient(app, raise_server_exceptions = False), saved
def test_flagged_repo_is_blocked_even_with_force(client, monkeypatch):
c, saved = client
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
r = c.put(
"/embedding-model", json = {"embedding_model": "attacker/malicious-embed", "force": True}
)
# 403, not the forceable 409, so the client does not offer "save anyway".
assert r.status_code == 403
assert "model" not in saved # force must not persist a flagged repo
def test_flagged_repo_is_blocked_without_force(client, monkeypatch):
c, saved = client
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
r = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"})
assert r.status_code == 403
assert "model" not in saved
def test_hard_block_uses_non_forceable_status(client, monkeypatch):
# The forceable verification path uses 409; the hard security block must be distinct
# (403) so the frontend never routes it into the "save anyway" force flow.
c, _saved = client
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
blocked = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"})
assert blocked.status_code == 403
# A verification failure (not-an-embedding-model) stays forceable at 409.
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
monkeypatch.setattr(settings, "is_embedding_model", lambda *a, **k: False, raising = False)
import utils.models as _models
monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False)
unverified = c.put("/embedding-model", json = {"embedding_model": "acme/not-an-embedder"})
assert unverified.status_code == 409
def test_llama_backend_skips_the_st_pickle_scan(monkeypatch):
# On the llama-server backend the embedder loads GGUF (inert), not the ST repo's
# pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here.
saved: dict = {}
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
monkeypatch.setattr(settings, "_llama_backend_active", lambda: True)
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
# force skips the GGUF availability checks; the ST pickle gate is what we assert is skipped.
called = {"scanned": False}
mod = _types.ModuleType("utils.security")
def _fail(*a, **k):
called["scanned"] = True
return _Decision(True)
mod.evaluate_file_security = _fail
mod.security_load_subdirs = lambda *a, **k: ()
monkeypatch.setitem(sys.modules, "utils.security", mod)
app = FastAPI()
app.include_router(settings.router)
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
c = TestClient(app, raise_server_exceptions = False)
r = c.put(
"/embedding-model",
json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True},
)
assert r.status_code == 200
assert called["scanned"] is False # the ST pickle scan never ran on the llama path
assert saved.get("model") == "attacker/flagged-st-clean-gguf"
def test_runtime_llama_fallback_skips_the_st_pickle_scan(monkeypatch):
# auto resolves to sentence-transformers (GPU present) but the embedder fell back to
# llama-server at runtime (torch/CUDA load or encode failure), so the process now loads
# only inert GGUF. The real _llama_backend_active() must reflect that cached fallback,
# so a flagged ST repo with a clean GGUF companion must not be hard-blocked here.
import core.rag.embeddings as embeddings
from core.rag.embed_llama_server import LlamaServerBackend
# Simulate the runtime fallback: the process-wide backend is a LlamaServerBackend even
# though the auto resolver would still say sentence-transformers.
monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend())
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers")
monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ())
saved: dict = {}
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
# Deliberately do NOT monkeypatch settings._llama_backend_active: this test exercises the
# real delegation to embeddings.active_backend_is_llama() so the cached fallback is honored.
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
called = {"scanned": False}
mod = _types.ModuleType("utils.security")
def _fail(*a, **k):
called["scanned"] = True
return _Decision(True)
mod.evaluate_file_security = _fail
mod.security_load_subdirs = lambda *a, **k: ()
monkeypatch.setitem(sys.modules, "utils.security", mod)
app = FastAPI()
app.include_router(settings.router)
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
c = TestClient(app, raise_server_exceptions = False)
r = c.put(
"/embedding-model",
json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True},
)
assert r.status_code == 200
assert called["scanned"] is False # the ST pickle scan never ran on the llama fallback
assert saved.get("model") == "attacker/flagged-st-clean-gguf"
def test_active_backend_is_llama_reflects_cache_and_resolver(monkeypatch):
# active_backend_is_llama() reports the ACTUAL built backend when one exists, and defers
# to the resolver (fresh-process behavior) when none has been built yet.
import core.rag.embeddings as embeddings
import core.rag.config as rag_config
from core.rag.embed_llama_server import LlamaServerBackend
# A cached llama backend wins even when auto would resolve to sentence-transformers.
monkeypatch.setattr(rag_config, "EMBED_BACKEND", "auto")
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers")
monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend())
assert embeddings.active_backend_is_llama() is True
# A cached ST backend reports False even when the resolver now picks llama, so its
# pickle stays gated (the cached backend, not the resolver, is what actually embeds).
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server")
monkeypatch.setattr(embeddings, "_backend", embeddings._SentenceTransformersBackend())
assert embeddings.active_backend_is_llama() is False
# No cached backend -> the resolver decides, unchanged from before.
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers")
monkeypatch.setattr(embeddings, "_backend", None)
assert embeddings.active_backend_is_llama() is False # auto -> sentence-transformers
monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server")
assert embeddings.active_backend_is_llama() is True # auto -> llama-server
# An explicit (non-auto) key is honored verbatim without a cached backend.
monkeypatch.setattr(rag_config, "EMBED_BACKEND", "llama-server")
assert embeddings.active_backend_is_llama() is True
def test_settings_scan_scopes_module_subdirs(monkeypatch):
# The settings scan must pass the ST module dirs (0_Transformer/) as load roots so a
# pickle directly under one blocks; assert those subdirs reach evaluate_file_security.
saved: dict = {}
monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed")
monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v)
monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v))
monkeypatch.setattr(settings, "_llama_backend_active", lambda: False)
monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False)
monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", ""))
monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model"))
import core.rag.embeddings as embeddings
monkeypatch.setattr(
embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",)
)
seen = {}
def _capture(*a, **k):
seen["subdirs"] = tuple(k.get("load_subdirs") or ())
return _Decision(False)
mod = _types.ModuleType("utils.security")
mod.security_load_subdirs = lambda *a, **k: ()
mod.evaluate_file_security = _capture
monkeypatch.setitem(sys.modules, "utils.security", mod)
app = FastAPI()
app.include_router(settings.router)
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
c = TestClient(app, raise_server_exceptions = False)
r = c.put(
"/embedding-model", json = {"embedding_model": "acme/embed-with-module-dir", "force": True}
)
assert r.status_code == 200
assert "0_Transformer" in seen["subdirs"]
def test_clean_repo_saves_under_force(client, monkeypatch):
c, saved = client
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
r = c.put("/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True})
assert r.status_code == 200
assert saved.get("model") == "acme/clean-embed"
def test_load_sink_refuses_flagged_model(monkeypatch):
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True))
import core.rag.embeddings as embeddings
with pytest.raises(embeddings.UnsafeEmbeddingModelError):
embeddings._guard_model_security("attacker/malicious-embed")
def test_load_sink_allows_clean_model(monkeypatch):
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
import core.rag.embeddings as embeddings
embeddings._guard_model_security("acme/clean-embed") # no raise
def test_sink_threads_ambient_token_into_scan(monkeypatch):
# A gated repo set via env/default has no request token; the guard must feed the
# loader's own token to the scan, or it fails open for the repo that still loads.
seen = {}
mod = _types.ModuleType("utils.security")
mod.security_load_subdirs = (
lambda name, token = None: seen.setdefault("subdirs_token", token) or ()
)
mod.evaluate_file_security = lambda *a, **k: seen.setdefault(
"scan_token", k.get("hf_token")
) or _Decision(False)
monkeypatch.setitem(sys.modules, "utils.security", mod)
import core.rag.embeddings as embeddings
monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: "hf_ambient")
embeddings._guard_model_security("acme/gated-embed")
assert seen["scan_token"] == "hf_ambient"
assert seen["subdirs_token"] == "hf_ambient"
def test_sink_scopes_st_module_subdirs_into_scan(monkeypatch):
# A flagged pickle directly under a Transformer module dir (0_Transformer/) must
# reach the scan as a load root; assert the guard unions the module dirs into
# load_subdirs so evaluate_file_security treats such a pickle as root-level.
seen = {}
def _capture(*a, **k):
seen["subdirs"] = tuple(k.get("load_subdirs") or ())
return _Decision(False)
mod = _types.ModuleType("utils.security")
mod.security_load_subdirs = lambda name, token = None: ()
mod.evaluate_file_security = _capture
monkeypatch.setitem(sys.modules, "utils.security", mod)
import core.rag.embeddings as embeddings
monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: None)
monkeypatch.setattr(
embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",)
)
embeddings._guard_model_security("acme/embed-with-module-dir")
assert "0_Transformer" in seen["subdirs"]
def test_st_module_subdirs_reads_local_modules_json(tmp_path, monkeypatch):
# The helper must parse each module's non-empty "path" from a local repo's
# modules.json and drop the root-level ("") Transformer entry.
import json
import core.rag.embeddings as embeddings
(tmp_path / "modules.json").write_text(
json.dumps(
[
{"idx": 0, "name": "0", "path": "0_Transformer", "type": "..."},
{"idx": 1, "name": "1", "path": "1_Pooling", "type": "..."},
{"idx": 2, "name": "2", "path": "", "type": "..."},
]
)
)
subdirs = embeddings._st_module_subdirs(str(tmp_path), None)
assert subdirs == ("0_Transformer", "1_Pooling")
def test_st_module_subdirs_swallows_errors(monkeypatch):
# Any failure (no modules.json, offline, malformed) returns () so the guard never
# bricks the embedder.
import huggingface_hub
import core.rag.embeddings as embeddings
def _boom(*a, **k):
raise RuntimeError("offline")
monkeypatch.setattr(huggingface_hub, "hf_hub_download", _boom)
assert embeddings._st_module_subdirs("acme/no-such-repo-xyz", None) == ()
def test_security_block_is_not_swallowed_by_llama_fallback(monkeypatch):
# The ST encode fallback must re-raise a security block, not swap to llama-server.
import core.rag.embeddings as embeddings
def _boom(*a, **k):
raise embeddings.UnsafeEmbeddingModelError("flagged")
monkeypatch.setattr(embeddings, "_st_encode", _boom)
monkeypatch.setattr(
embeddings,
"_switch_to_llama_fallback",
lambda err: pytest.fail("security block must not fall back to llama-server"),
)
with pytest.raises(embeddings.UnsafeEmbeddingModelError):
embeddings._SentenceTransformersBackend().encode(["hi"])

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

@ -1900,7 +1900,7 @@ class TestApiMonitorProviderAndCompletionStreams:
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
timeout = 5.0,
)
assert isinstance(response, _SameTaskStreamingResponse)
@ -2044,7 +2044,7 @@ class TestApiMonitorProviderAndCompletionStreams:
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
timeout = 5.0,
)
assert isinstance(response, _SameTaskStreamingResponse)
gate.set()
@ -2107,7 +2107,7 @@ class TestApiMonitorProviderAndCompletionStreams:
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
timeout = 5.0,
)
assert isinstance(response, _SameTaskStreamingResponse)
@ -2190,7 +2190,7 @@ class TestApiMonitorProviderAndCompletionStreams:
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
timeout = 5.0,
)
assert isinstance(response, _SameTaskStreamingResponse)
@ -2252,7 +2252,7 @@ class TestApiMonitorProviderAndCompletionStreams:
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
timeout = 5.0,
)
assert isinstance(response, _SameTaskStreamingResponse)
assert cancel_id in inf_mod._CANCEL_REGISTRY
@ -2323,13 +2323,13 @@ class TestApiMonitorProviderAndCompletionStreams:
monitor_id = monitor_id,
)
)
await asyncio.wait_for(entered.wait(), timeout = 0.2)
await asyncio.wait_for(entered.wait(), timeout = 5.0)
assert cancel_id in inf_mod._CANCEL_REGISTRY
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
await asyncio.wait_for(cancelled.wait(), timeout = 0.2)
await asyncio.wait_for(cancelled.wait(), timeout = 5.0)
assert cancel_id not in inf_mod._CANCEL_REGISTRY
asyncio.run(_run())
@ -2389,13 +2389,13 @@ class TestApiMonitorProviderAndCompletionStreams:
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
timeout = 5.0,
)
assert isinstance(response, _SameTaskStreamingResponse)
assert cancel_id in inf_mod._CANCEL_REGISTRY
gate.set()
await asyncio.wait_for(returned.wait(), timeout = 0.2)
await asyncio.wait_for(returned.wait(), timeout = 5.0)
await asyncio.sleep(0)
await response._unstarted_cleanup()
assert upstream_response.is_closed

38
studio/backend/tests/test_safetensors_capability_advertise.py
View file

@ -48,6 +48,21 @@ reasoning_effort: {{ reasoning_effort }}
"""
# DeepSeek-V4-Flash: an enable_thinking on/off gate PLUS a reasoning_effort
# 'max' preamble. The shipped template only *branches* on 'max' ('high' renders
# identically to thinking-on-without-the-preamble), so the literal scan alone
# would surface only ['max']; the classifier adds 'high' for deepseek-v4 to
# expose the encoder's full none/high/max ladder.
DEEPSEEK_V4_TEMPLATE = (
"{%- if not thinking is defined %}"
"{%- if enable_thinking is defined %}{%- set thinking = enable_thinking %}"
"{%- else %}{%- set thinking = false %}{%- endif %}{%- endif %}\n"
"{%- if thinking and reasoning_effort == 'max' %}"
"{{- 'Reasoning Effort: Absolute maximum' }}{%- endif %}\n"
"{%- for message in messages %}{{- message.content }}{%- endfor %}"
)
PLAIN_TEMPLATE = """
{%- for message in messages %}
{{- message.role + ': ' + message.content + '\\n' }}
@ -90,6 +105,29 @@ def test_detect_reasoning_flags_none_template_returns_all_false():
assert flags["reasoning_style"] == "enable_thinking"
def test_detect_reasoning_flags_deepseek_v4_exposes_none_high_max():
"""DeepSeek-V4-Flash: enable_thinking gate + reasoning_effort 'max' preamble.
Classified as the hybrid style with the full none/high/max ladder even
though the template only branches on 'max'."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash-GGUF")
assert flags["supports_reasoning"] is True
assert flags["reasoning_style"] == "enable_thinking_effort"
assert flags["reasoning_effort_levels"] == ["high", "max"]
assert flags["reasoning_always_on"] is False
def test_detect_reasoning_flags_non_deepseek_v4_effort_only_max_not_injected():
"""The 'high' injection is scoped to deepseek-v4: a different model whose
template only branches on 'max' keeps ['max'] (no phantom 'high')."""
from core.inference.llama_cpp import detect_reasoning_flags
flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "vendor/OtherHybrid-GGUF")
assert flags["reasoning_style"] == "enable_thinking_effort"
assert flags["reasoning_effort_levels"] == ["max"]
def test_detect_safetensors_features_passes_template_through_to_classifier():
"""Route wrapper forwards a real template to the inner classifier."""
from routes.inference import _detect_safetensors_features

13
studio/backend/tests/test_security_gate_consistency.py
View file

@ -99,3 +99,16 @@ def test_malware_and_consent_gates_cover_the_lora_base():
if runs_gate and not resolves_base:
offenders.append(f"{rel} runs a load gate but never resolves the LoRA base")
assert not offenders, "\n".join(offenders)
def test_rag_embedding_path_runs_the_malware_gate():
"""The RAG embedding model is set through /settings and later loaded by
SentenceTransformer, which deserializes pickles; both sites must run the malware gate
or a flagged repo loads unscanned (bypassing the normal model-load protections)."""
offenders = []
for rel in ("routes/settings.py", "core/rag/embeddings.py"):
if "evaluate_file_security(" not in (_BACKEND / rel).read_text():
offenders.append(
f"{rel} loads/persists an embedding model without evaluate_file_security"
)
assert not offenders, "\n".join(offenders)

483
studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx Normal file
View file

@ -0,0 +1,483 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"use client";
import {
Sheet,
SheetContent,
SheetDescription,
SheetHeader,
SheetTitle,
} from "@/components/ui/sheet";
import {
customProviderDisplayName,
parseExternalModelId,
useChatPreferencesStore,
useChatRuntimeStore,
useExternalProvidersStore,
} from "@/features/chat";
import { cn } from "@/lib/utils";
import { FileDatabaseIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useMessage, useMessageTiming } from "@assistant-ui/react";
import type { FC, ReactNode } from "react";
type ResponseDetailsMetadata = {
modelId?: string;
modelLabel?: string;
responseModelId?: string;
providerId?: string;
providerName?: string;
providerType?: string;
startedAt?: number;
finishedAt?: number;
durationMs?: number;
sessionId?: string | null;
cancelId?: string;
toolCalls?: string[];
tools?: Record<string, boolean | undefined>;
};
type ContextUsageMetadata = {
promptTokens?: number;
completionTokens?: number;
totalTokens?: number;
cachedTokens?: number;
cacheWriteTokens?: number;
modelId?: string;
};
type MessageCustomMetadata = {
responseDetails?: ResponseDetailsMetadata;
contextUsage?: ContextUsageMetadata;
serverTimings?: Record<string, unknown>;
reasoningDuration?: number;
};
function asNumber(value: unknown): number | undefined {
return typeof value === "number" && Number.isFinite(value)
? value
: undefined;
}
function formatNumber(value: number | undefined): string | null {
return value == null ? null : value.toLocaleString();
}
function formatMs(value: number | undefined): string | null {
if (value == null) return null;
if (value < 1000) return `${Math.round(value)}ms`;
return `${(value / 1000).toFixed(2)}s`;
}
function formatRate(value: number | undefined): string | null {
if (value == null) return null;
return `${value.toFixed(1)} tok/s`;
}
function formatDate(value: Date | number | string | undefined): string | null {
if (value == null) return null;
const date = value instanceof Date ? value : new Date(value);
if (Number.isNaN(date.getTime())) return null;
return new Intl.DateTimeFormat(undefined, {
dateStyle: "medium",
timeStyle: "medium",
}).format(date);
}
const TOOL_CATEGORY_LABELS: Record<string, string> = {
search: "Search",
fetch: "Fetch",
code: "Code",
images: "Images",
mcp: "MCP",
docs: "Docs",
artifacts: "Canvas",
};
const TOOL_CALL_LABELS: Record<string, string> = {
web_search: "Search",
web_fetch: "Fetch",
code_execution: "Code",
python: "Python",
terminal: "Terminal",
image_generation: "Images",
search_knowledge_base: "Docs",
render_html: "Canvas",
};
function uniqueValues(values: string[]): string[] {
return Array.from(new Set(values));
}
function toolCategoryFromCall(toolName: string): string | null {
const normalized = toolName.toLowerCase();
if (normalized === "web_search") return "search";
if (normalized === "web_fetch") return "fetch";
if (
normalized === "code_execution" ||
normalized === "python" ||
normalized === "terminal"
) {
return "code";
}
if (normalized === "image_generation") return "images";
if (normalized === "search_knowledge_base") return "docs";
if (normalized === "render_html") return "artifacts";
if (normalized.startsWith("mcp__")) return "mcp";
return null;
}
function formatToolCallName(toolName: string): string {
const normalized = toolName.toLowerCase();
if (TOOL_CALL_LABELS[normalized]) return TOOL_CALL_LABELS[normalized];
if (normalized.startsWith("mcp__")) return `MCP: ${toolName.slice(5)}`;
return toolName
.replace(/[_-]+/g, " ")
.replace(/\b\w/g, (letter) => letter.toUpperCase());
}
function toolCallsFromContent(content: unknown): string[] {
if (!Array.isArray(content)) return [];
return uniqueValues(
content
.map((part) =>
part && typeof part === "object" && "type" in part
? (part as { type?: unknown; toolName?: unknown })
: null,
)
.filter(
(part): part is { type: "tool-call"; toolName: string } =>
part?.type === "tool-call" &&
typeof part.toolName === "string" &&
part.toolName.length > 0,
)
.map((part) => part.toolName),
);
}
function enabledTools(
tools: Record<string, boolean | undefined> | undefined,
toolCalls: string[],
): string | null {
if (!tools && toolCalls.length === 0) return null;
const activeKeys = new Set<string>();
for (const key of Object.keys(TOOL_CATEGORY_LABELS)) {
if (tools?.[key] === true) activeKeys.add(key);
}
for (const toolName of toolCalls) {
const key = toolCategoryFromCall(toolName);
if (key) activeKeys.add(key);
}
const active = Object.keys(TOOL_CATEGORY_LABELS)
.filter((key) => activeKeys.has(key))
.map((key) => TOOL_CATEGORY_LABELS[key]);
return active.length > 0 ? active.join(", ") : "None";
}
function calledTools(toolCalls: string[]): string | null {
if (toolCalls.length === 0) return null;
return uniqueValues(toolCalls.map(formatToolCallName)).join(", ");
}
function DetailSection({
title,
children,
}: {
title: string;
children: ReactNode;
}) {
return (
<section className="rounded-md bg-muted/45 p-3">
<h3 className="mb-2 font-heading text-foreground text-sm">{title}</h3>
<div className="grid gap-2">{children}</div>
</section>
);
}
function DetailRow({
label,
value,
mono = false,
}: {
label: string;
value: ReactNode | null | undefined;
mono?: boolean;
}) {
if (value == null || value === "") return null;
return (
<div className="grid grid-cols-[8.5rem_minmax(0,1fr)] items-start gap-3 text-[13px]">
<span className="text-muted-foreground">{label}</span>
<span
className={cn(
"min-w-0 break-words text-right text-foreground",
mono && "font-mono tabular-nums",
)}
>
{value}
</span>
</div>
);
}
function useResponseModelDisplay() {
const message = useMessage();
const models = useChatRuntimeStore((s) => s.models);
const providers = useExternalProvidersStore((s) => s.providers);
const custom = (
message.metadata as Record<string, unknown> | undefined
)?.custom as MessageCustomMetadata | undefined;
const responseDetails = custom?.responseDetails;
const usage = custom?.contextUsage;
const serverTimings = custom?.serverTimings;
const recordedModelId =
responseDetails?.responseModelId ??
responseDetails?.modelId ??
usage?.modelId;
const parsedExternal = parseExternalModelId(recordedModelId);
const provider = parsedExternal
? providers.find((candidate) => candidate.id === parsedExternal.providerId)
: null;
const modelSummary = models.find(
(candidate) => candidate.id === recordedModelId,
);
const modelLabel =
responseDetails?.modelLabel ??
responseDetails?.responseModelId ??
parsedExternal?.modelId ??
modelSummary?.name ??
recordedModelId ??
"Not recorded";
const providerLabel =
responseDetails?.providerName ??
provider?.name ??
(responseDetails?.providerType
? customProviderDisplayName(responseDetails.providerType)
: parsedExternal
? customProviderDisplayName(provider?.providerType)
: recordedModelId
? "Local model"
: null);
return {
message,
custom,
responseDetails,
usage,
serverTimings,
modelLabel,
providerLabel,
};
}
export const MessageResponseModelBadge: FC<{ className?: string }> = ({
className,
}) => {
const showResponseModel = useChatPreferencesStore(
(state) => state.showResponseModel,
);
const { modelLabel, providerLabel } = useResponseModelDisplay();
if (!showResponseModel || modelLabel === "Not recorded") {
return null;
}
return (
<span
className={cn(
"aui-response-model-badge inline-flex min-h-5 max-w-full items-center text-muted-foreground/80 text-xs font-medium leading-5 opacity-0 transition-opacity duration-150 group-hover/assistant-message:opacity-100 group-focus-within/assistant-message:opacity-100",
className,
)}
title={providerLabel ? `${modelLabel} - ${providerLabel}` : modelLabel}
>
<span className="min-w-0 truncate align-middle">{modelLabel}</span>
</span>
);
};
export const MessageResponseDetailsSheet: FC<{
open: boolean;
onOpenChange: (open: boolean) => void;
}> = ({ open, onOpenChange }) => {
const timing = useMessageTiming();
const {
message,
responseDetails,
usage,
serverTimings,
modelLabel,
providerLabel,
} = useResponseModelDisplay();
const promptTokens =
usage?.promptTokens ?? asNumber(serverTimings?.prompt_n);
const completionTokens =
usage?.completionTokens ??
timing?.tokenCount ??
asNumber(serverTimings?.predicted_n);
const totalTokens =
usage?.totalTokens ??
(promptTokens != null && completionTokens != null
? promptTokens + completionTokens
: undefined);
const totalTime =
responseDetails?.durationMs ?? timing?.totalStreamTime ?? undefined;
const summaryLabel =
modelLabel === "Not recorded" ? "Model not recorded" : `Used ${modelLabel}`;
const messageToolCalls = toolCallsFromContent(message.content);
const toolCalls =
responseDetails?.toolCalls && responseDetails.toolCalls.length > 0
? responseDetails.toolCalls
: messageToolCalls;
return (
<Sheet open={open} onOpenChange={onOpenChange}>
<SheetContent
side="right"
className="w-[min(28rem,100vw)] p-0 sm:max-w-[28rem]"
>
<SheetHeader className="border-b p-4">
<SheetTitle className="flex items-center gap-2 pr-10 font-heading text-base">
<HugeiconsIcon
icon={FileDatabaseIcon}
strokeWidth={1.75}
className="size-icon text-chat-icon-fg"
/>
Response details
</SheetTitle>
<SheetDescription className="sr-only">
Timing, model, token, and tool details for this response.
</SheetDescription>
</SheetHeader>
<div className="flex min-h-0 flex-1 flex-col gap-4 overflow-y-auto p-4">
<div className="min-w-0 rounded-md border border-border/70 bg-card p-3">
<p className="min-w-0 break-words font-heading text-foreground text-sm">
{summaryLabel}
</p>
{providerLabel ? (
<p className="mt-1 min-w-0 break-words text-muted-foreground text-xs">
{providerLabel}
</p>
) : null}
</div>
<DetailSection title="Response">
<DetailRow label="Model" value={modelLabel} />
<DetailRow
label="Requested"
value={
responseDetails?.modelId &&
responseDetails.modelId !== responseDetails.responseModelId
? responseDetails.modelId
: null
}
/>
<DetailRow label="Provider" value={providerLabel} />
<DetailRow label="Message ID" value={message.id} mono={true} />
<DetailRow label="Created" value={formatDate(message.createdAt)} />
<DetailRow
label="Started"
value={formatDate(responseDetails?.startedAt)}
/>
<DetailRow
label="Finished"
value={formatDate(responseDetails?.finishedAt)}
/>
</DetailSection>
<DetailSection title="Tokens">
<DetailRow label="Prompt" value={formatNumber(promptTokens)} mono />
<DetailRow
label="Output"
value={formatNumber(completionTokens)}
mono
/>
<DetailRow label="Total" value={formatNumber(totalTokens)} mono />
<DetailRow
label="Cache hits"
value={formatNumber(
usage?.cachedTokens ?? asNumber(serverTimings?.cache_n),
)}
mono
/>
<DetailRow
label="Cache writes"
value={formatNumber(usage?.cacheWriteTokens)}
mono
/>
</DetailSection>
<DetailSection title="Timing">
<DetailRow label="Total" value={formatMs(totalTime)} mono />
<DetailRow
label="First token"
value={formatMs(timing?.firstTokenTime)}
mono
/>
<DetailRow
label="Prompt eval"
value={formatMs(asNumber(serverTimings?.prompt_ms))}
mono
/>
<DetailRow
label="Generation"
value={formatMs(asNumber(serverTimings?.predicted_ms))}
mono
/>
<DetailRow
label="Speed"
value={formatRate(
asNumber(serverTimings?.predicted_per_second) ??
timing?.tokensPerSecond,
)}
mono
/>
<DetailRow
label="Chunks"
value={formatNumber(timing?.totalChunks)}
mono
/>
<DetailRow
label="Tool calls"
value={formatNumber(timing?.toolCallCount)}
mono
/>
</DetailSection>
<DetailSection title="Tools">
<DetailRow
label="Enabled"
value={enabledTools(responseDetails?.tools, toolCalls)}
/>
<DetailRow label="Called" value={calledTools(toolCalls)} />
<DetailRow
label="Confirmation"
value={
responseDetails?.tools?.confirmToolCalls === true
? "On"
: responseDetails?.tools?.confirmToolCalls === false
? "Off"
: null
}
/>
<DetailRow
label="Bypass"
value={
responseDetails?.tools?.bypassPermissions === true
? "On"
: responseDetails?.tools?.bypassPermissions === false
? "Off"
: null
}
/>
<DetailRow label="Session" value={responseDetails?.sessionId} mono />
<DetailRow label="Run ID" value={responseDetails?.cancelId} mono />
</DetailSection>
</div>
</SheetContent>
</Sheet>
);
};

10
studio/frontend/src/components/assistant-ui/reasoning.tsx
View file

@ -6,6 +6,7 @@
/* eslint-disable react-refresh/only-export-components */
import { MarkdownText } from "@/components/assistant-ui/markdown-text";
import { MessageResponseModelBadge } from "@/components/assistant-ui/message-response-details-sheet";
import {
Collapsible,
CollapsibleContent,
@ -390,14 +391,17 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
onOpenChange={handleOpenChange}
variant={variant}
>
<div className="flex items-center gap-2">
<div className="flex min-w-0 items-center gap-2">
<ReasoningTrigger
className="min-w-0 flex-1"
className="min-w-0 flex-none"
active={isReasoningStreaming}
// Prefer server timing when available.
duration={persistedDuration || duration}
/>
<div className="flex w-16 shrink-0 justify-end">
<span className="hidden min-w-0 max-w-[12rem] group-hover/assistant-message:inline-flex group-focus-within/assistant-message:inline-flex sm:max-w-[16rem]">
<MessageResponseModelBadge className="min-w-0" />
</span>
<div className="ml-auto flex w-16 shrink-0 justify-end">
{isOpen && !isReasoningStreaming && (
<ReasoningCopyButton startIndex={startIndex} endIndex={endIndex} />
)}

118
studio/frontend/src/components/assistant-ui/thread.tsx
View file

@ -12,6 +12,10 @@ import {
import { downloadImagePart } from "@/components/assistant-ui/image";
import { MarkdownText } from "@/components/assistant-ui/markdown-text";
import { MessageHtmlArtifacts } from "@/components/assistant-ui/message-html-artifacts";
import {
MessageResponseDetailsSheet,
MessageResponseModelBadge,
} from "@/components/assistant-ui/message-response-details-sheet";
import { MessageTiming } from "@/components/assistant-ui/message-timing";
import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning";
import { RagSourcesGroup } from "@/components/assistant-ui/rag-sources";
@ -3564,6 +3568,9 @@ const AssistantMessage: FC = () => {
const aui = useAui();
const messageId = useAuiState(({ message }) => message.id);
const messageContent = useAuiState(({ message }) => message.content);
const hasReasoningParts = useAuiState(({ message }) =>
message.parts.some((part) => part.type === "reasoning"),
);
const incognito = useChatRuntimeStore((s) => s.incognito);
// Use global store for editing state to ensure a single source of truth
@ -3620,7 +3627,7 @@ const AssistantMessage: FC = () => {
return (
<MessagePrimitive.Root
className="aui-assistant-message-root relative mx-auto min-w-0 w-full max-w-(--thread-content-max-width) pt-0.5 pb-4 text-[15.5px] [font-weight:410] tracking-[0.01em] dark:tracking-[0.02em]"
className="group/assistant-message aui-assistant-message-root relative mx-auto min-w-0 w-full max-w-(--thread-content-max-width) pt-0.5 pb-4 text-[15.5px] [font-weight:410] tracking-[0.01em] dark:tracking-[0.02em]"
data-role="assistant"
>
<div className="aui-assistant-message-content wrap-break-word min-w-0 text-[#0d0d0d] dark:text-foreground leading-relaxed">
@ -3649,6 +3656,11 @@ const AssistantMessage: FC = () => {
</div>
) : (
<>
{!hasReasoningParts ? (
<div className="pointer-events-none relative h-0 min-w-0">
<MessageResponseModelBadge className="absolute -top-6 left-0 max-w-[min(22rem,100%)]" />
</div>
) : null}
<GeneratingIndicator />
<CancelledIndicator />
<DiffusionCanvas />
@ -3893,58 +3905,76 @@ const EditAssistantMessageButton: FC = () => {
const AssistantActionBar: FC = () => {
const { forkMessage, forkDisabled } = useForkMessageAction();
const [detailsOpen, setDetailsOpen] = useState(false);
return (
<ActionBarPrimitive.Root
hideWhenRunning={true}
className="aui-assistant-action-bar-root col-start-3 row-start-2 flex items-center gap-1 text-chat-icon-fg [&_button:not([data-slot=message-timing-trigger])]:size-8 [&_button]:!rounded-full [&_button:hover]:bg-chat-icon-bg-hover [&_button:hover]:text-chat-icon-fg-hover"
>
<CopyButton />
<EditAssistantMessageButton />
<ActionBarPrimitive.Reload asChild={true}>
<TooltipIconButton tooltip="Refresh">
<RefreshCwIcon strokeWidth={1.75} className="size-icon" />
</TooltipIconButton>
</ActionBarPrimitive.Reload>
<ForkCountBadge />
<DeleteMessageButton />
<ActionBarMorePrimitive.Root>
<ActionBarMorePrimitive.Trigger asChild={true}>
<TooltipIconButton
tooltip="More"
className="data-[state=open]:bg-accent"
>
<MoreHorizontalIcon strokeWidth={1.75} className="size-icon" />
<>
<ActionBarPrimitive.Root
hideWhenRunning={true}
className="aui-assistant-action-bar-root col-start-3 row-start-2 flex items-center gap-1 text-chat-icon-fg [&_button:not([data-slot=message-timing-trigger])]:size-8 [&_button]:!rounded-full [&_button:hover]:bg-chat-icon-bg-hover [&_button:hover]:text-chat-icon-fg-hover"
>
<CopyButton />
<EditAssistantMessageButton />
<ActionBarPrimitive.Reload asChild={true}>
<TooltipIconButton tooltip="Refresh">
<RefreshCwIcon strokeWidth={1.75} className="size-icon" />
</TooltipIconButton>
</ActionBarMorePrimitive.Trigger>
<ActionBarMorePrimitive.Content
side="bottom"
align="start"
onCloseAutoFocus={(e) => e.preventDefault()}
className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none"
>
<ActionBarMorePrimitive.Item
disabled={forkDisabled}
onSelect={() => void forkMessage()}
className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50"
</ActionBarPrimitive.Reload>
<ForkCountBadge />
<DeleteMessageButton />
<ActionBarMorePrimitive.Root>
<ActionBarMorePrimitive.Trigger asChild={true}>
<TooltipIconButton
tooltip="More"
className="data-[state=open]:bg-accent"
>
<MoreHorizontalIcon strokeWidth={1.75} className="size-icon" />
</TooltipIconButton>
</ActionBarMorePrimitive.Trigger>
<ActionBarMorePrimitive.Content
side="bottom"
align="start"
onCloseAutoFocus={(e) => e.preventDefault()}
className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none"
>
<GitBranchIcon strokeWidth={1.75} className="size-icon" />
Fork in new chat
</ActionBarMorePrimitive.Item>
<ActionBarPrimitive.ExportMarkdown asChild={true}>
<ActionBarMorePrimitive.Item className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground">
<ActionBarMorePrimitive.Item
disabled={forkDisabled}
onSelect={() => void forkMessage()}
className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50"
>
<GitBranchIcon strokeWidth={1.75} className="size-icon" />
Fork in new chat
</ActionBarMorePrimitive.Item>
<ActionBarPrimitive.ExportMarkdown asChild={true}>
<ActionBarMorePrimitive.Item className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground">
<HugeiconsIcon
icon={Download01Icon}
strokeWidth={1.75}
className="size-icon"
/>
Export as Markdown
</ActionBarMorePrimitive.Item>
</ActionBarPrimitive.ExportMarkdown>
<ActionBarMorePrimitive.Item
onSelect={() => setDetailsOpen(true)}
className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground"
>
<HugeiconsIcon
icon={Download01Icon}
icon={FileDatabaseIcon}
strokeWidth={1.75}
className="size-icon"
/>
Export as Markdown
See response details
</ActionBarMorePrimitive.Item>
</ActionBarPrimitive.ExportMarkdown>
</ActionBarMorePrimitive.Content>
</ActionBarMorePrimitive.Root>
<MessageTiming side="top" className="h-8 px-2" />
</ActionBarPrimitive.Root>
</ActionBarMorePrimitive.Content>
</ActionBarMorePrimitive.Root>
<MessageTiming side="top" className="h-8 px-2" />
</ActionBarPrimitive.Root>
<MessageResponseDetailsSheet
open={detailsOpen}
onOpenChange={setDetailsOpen}
/>
</>
);
};

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

@ -140,6 +140,32 @@ interface ServerTimings {
diffusion_steps_per_second?: number;
}
interface ResponseDetailsMetadata {
modelId: string;
modelLabel: string;
responseModelId: string;
providerId?: string;
providerName: string;
providerType: string;
startedAt: number;
finishedAt: number;
durationMs: number;
sessionId?: string;
cancelId: string;
toolCalls: string[];
tools: {
search: boolean;
fetch: boolean;
code: boolean;
images: boolean;
mcp: boolean;
docs: boolean;
artifacts: boolean;
confirmToolCalls: boolean;
bypassPermissions: boolean;
};
}
type RunMessages = Parameters<ChatModelAdapter["run"]>[0]["messages"];
type RunMessage = RunMessages[number];
@ -1769,6 +1795,9 @@ export function createOpenAIStreamAdapter(
(provider) => provider.id === externalSelection.providerId,
)
: null;
const selectedModelSummary = runtime.models.find(
(model) => model.id === params.checkpoint,
);
const externalApiKey = externalProvider
? getExternalProviderApiKey(externalProvider.id).trim()
: "";
@ -2151,6 +2180,7 @@ export function createOpenAIStreamAdapter(
let waitingFirstChunk = true;
let firstTokenSettled = false;
const streamStartTime = Date.now();
let responseModelId = externalSelection?.modelId ?? params.checkpoint;
let firstTokenTime: number | undefined;
let totalChunks = 0;
let resolveFirstToken: (() => void) | null = null;
@ -2372,6 +2402,59 @@ export function createOpenAIStreamAdapter(
const externalBackendProviderType = toExternalBackendProviderType(
externalProvider?.providerType,
);
const buildResponseDetails = (
finishedAt: number,
): ResponseDetailsMetadata => ({
modelId: params.checkpoint,
modelLabel:
(isExternalRequest || responseModelId !== params.checkpoint
? responseModelId
: selectedModelSummary?.name || responseModelId) ||
params.checkpoint ||
"Unknown model",
responseModelId:
responseModelId ||
externalSelection?.modelId ||
params.checkpoint,
...(externalProvider?.id ? { providerId: externalProvider.id } : {}),
providerName:
externalProvider?.name ??
(isExternalRequest ? "External provider" : "Local model"),
providerType: externalProvider?.providerType ?? "local",
startedAt: streamStartTime,
finishedAt,
durationMs: finishedAt - streamStartTime,
...(sandboxSessionId ? { sessionId: sandboxSessionId } : {}),
cancelId,
toolCalls: Array.from(
new Set(
toolCallParts
.map((part) => part.toolName)
.filter(
(toolName): toolName is string =>
typeof toolName === "string" && toolName.length > 0,
),
),
),
tools: {
search:
webSearchEnabledForThisTurn ||
(!isExternalRequest && supportsTools && toolsEnabled),
fetch: webFetchEnabledForThisTurn,
code:
codeExecEnabledForThisTurn ||
(!isExternalRequest && supportsTools && codeToolsEnabled),
images: imageGenerationEnabledForThisTurn,
mcp: !isExternalRequest && supportsTools && mcpEnabledForChat,
docs:
!isExternalRequest &&
supportsTools &&
(ragEnabled || projectRagEnabled),
artifacts: renderHtmlToolEnabledForThisTurn,
confirmToolCalls,
bypassPermissions,
},
});
const externalCapabilities = getProviderCapabilities(
externalProvider?.providerType,
);
@ -2768,6 +2851,11 @@ export function createOpenAIStreamAdapter(
const stream = streamChatCompletions(requestPayload, abortSignal);
for await (const chunk of stream) {
const chunkModel = (chunk as { model?: unknown }).model;
if (typeof chunkModel === "string" && chunkModel.length > 0) {
responseModelId = chunkModel;
}
// Handle tool status events
const toolStatusText = (
chunk as unknown as { _toolStatus?: string }
@ -3435,11 +3523,12 @@ export function createOpenAIStreamAdapter(
});
}
const finishedAt = Date.now();
const finalTiming = buildTiming(
streamStartTime,
totalChunks,
serverPromptEvalTime ?? firstTokenTime,
Date.now() - streamStartTime,
finishedAt - streamStartTime,
finalTokenCount,
toolCallParts.length,
finalTokPerSec,
@ -3475,6 +3564,7 @@ export function createOpenAIStreamAdapter(
modelId: params.checkpoint,
}
: undefined,
responseDetails: buildResponseDetails(finishedAt),
timing: finalTiming,
},
},

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

@ -26,7 +26,12 @@ export {
type PlusMenuItemId,
} from "./stores/plus-menu-prefs-store";
export { useChatModelRuntime } from "./hooks/use-chat-model-runtime";
export { isExternalModelId } from "./external-providers";
export {
customProviderDisplayName,
isExternalModelId,
parseExternalModelId,
} from "./external-providers";
export { useExternalProvidersStore } from "./stores/external-providers-store";
export { ChatSearchDialog } from "./components/chat-search-dialog";
export { setTrainingCompareHandoff } from "./lib/training-compare-handoff";
export type { ProjectRecord } from "./types";

7
studio/frontend/src/features/chat/stores/chat-preferences-store.ts
View file

@ -7,11 +7,14 @@ import { persist } from "zustand/middleware";
// Client-side chat UI prefs kept in localStorage, not the chat DB.
// confirmDeleteChats: when off, deleting a chat skips the confirm dialog.
// showModelDisclaimer: when off, hide the "LLMs can make mistakes" footer note.
// showResponseModel: when on, assistant responses show the producing model.
export interface ChatPreferencesState {
confirmDeleteChats: boolean;
setConfirmDeleteChats: (value: boolean) => void;
showModelDisclaimer: boolean;
setShowModelDisclaimer: (value: boolean) => void;
showResponseModel: boolean;
setShowResponseModel: (value: boolean) => void;
}
export const useChatPreferencesStore = create<ChatPreferencesState>()(
@ -23,6 +26,9 @@ export const useChatPreferencesStore = create<ChatPreferencesState>()(
showModelDisclaimer: true,
setShowModelDisclaimer: (showModelDisclaimer) =>
set({ showModelDisclaimer }),
showResponseModel: false,
setShowResponseModel: (showResponseModel) =>
set({ showResponseModel }),
}),
{
name: "unsloth_chat_preferences",
@ -32,6 +38,7 @@ export const useChatPreferencesStore = create<ChatPreferencesState>()(
...current,
confirmDeleteChats: saved?.confirmDeleteChats ?? true,
showModelDisclaimer: saved?.showModelDisclaimer ?? true,
showResponseModel: saved?.showResponseModel ?? false,
};
},
},

9
studio/frontend/src/features/settings/api/embedding-model.ts
View file

@ -23,6 +23,10 @@ type ApiEmbeddingModelSettings = {
* (wrong type, gated repo, or offline). Retry with force to save anyway. */
export class EmbeddingModelVerificationError extends Error {}
/** 403 from the backend: the repo is flagged unsafe by Hugging Face's security scan.
* A hard block; force cannot bypass it, so it must not enter the "save anyway" flow. */
export class EmbeddingModelBlockedError extends Error {}
function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings {
return {
embeddingModel: settings.embedding_model,
@ -56,6 +60,11 @@ export async function updateEmbeddingModelSettings(
force: options?.force ?? false,
}),
});
if (res.status === 403) {
throw new EmbeddingModelBlockedError(
await readFastApiError(res, "This model is blocked by a security scan"),
);
}
if (res.status === 409) {
throw new EmbeddingModelVerificationError(
await readFastApiError(res, "Could not verify the embedding model"),

15
studio/frontend/src/features/settings/tabs/chat-tab.tsx
View file

@ -213,6 +213,12 @@ export function ChatTab() {
const setShowModelDisclaimer = useChatPreferencesStore(
(state) => state.setShowModelDisclaimer,
);
const showResponseModel = useChatPreferencesStore(
(state) => state.showResponseModel,
);
const setShowResponseModel = useChatPreferencesStore(
(state) => state.setShowResponseModel,
);
useEffect(() => {
void countAllChats().then(setCount);
@ -412,6 +418,15 @@ export function ChatTab() {
onCheckedChange={setShowModelDisclaimer}
/>
</SettingsRow>
<SettingsRow
label="Show response model"
description="Show model metadata in assistant responses."
>
<Switch
checked={showResponseModel}
onCheckedChange={setShowResponseModel}
/>
</SettingsRow>
</SettingsSection>
<SettingsSection title={t("settings.chat.artifacts.title")}>

6
studio/frontend/src/features/settings/tabs/general-tab.tsx
View file

@ -42,6 +42,7 @@ import {
updatePreviewSharing,
} from "../api/preview-sharing";
import {
EmbeddingModelBlockedError,
type EmbeddingModelSettings,
EmbeddingModelVerificationError,
loadEmbeddingModelSettings,
@ -410,7 +411,10 @@ export function GeneralTab() {
description: t("settings.general.rag.reindexWarning"),
});
} catch (error) {
if (error instanceof EmbeddingModelVerificationError) {
// A hard security block cannot be forced; keep the "save anyway" action hidden.
if (error instanceof EmbeddingModelBlockedError) {
setEmbeddingModelNeedsForce(false);
} else if (error instanceof EmbeddingModelVerificationError) {
setEmbeddingModelNeedsForce(true);
}
setEmbeddingModelError(

84
tests/_zoo_rocm_spoof.py Normal file
View file

@ -0,0 +1,84 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team.
"""ROCm/RDNA spoof: present torch as an AMD Radeon (RDNA 2/3/4) card on a
GPU-less host, so hip paths (device_type -> "hip", llama.cpp ROCm bundle) are
testable in CPU-only CI with no AMD hardware. The ROCm sibling of
_zoo_aggressive_cuda_spoof.py: it reuses that spoof's torch.cuda no-op machinery
and overlays the AMD identity (torch.version.hip, gcnArchName, Radeon name).
Apply BEFORE importing unsloth/unsloth_zoo, since DEVICE_TYPE is cached there.
"""
from __future__ import annotations
import importlib.util
import os
import sys
# gfx -> (marketing name, (capability major, minor), torch.version.hip). hip is
# the ROCm build torch was made against (RDNA2/3 ship 6.x; gfx1102/115x/RDNA4 7.2).
_PROFILES: dict[str, tuple[str, tuple[int, int], str]] = {
"gfx1030": ("AMD Radeon RX 6900 XT", (10, 3), "6.4.43483"), # RDNA2
"gfx1031": ("AMD Radeon RX 6700 XT", (10, 3), "6.4.43483"),
"gfx1032": ("AMD Radeon RX 6600", (10, 3), "6.4.43483"),
"gfx1034": ("AMD Radeon RX 6400", (10, 3), "6.4.43483"),
"gfx1100": ("AMD Radeon RX 7900 XTX", (11, 0), "6.4.43483"), # RDNA3
"gfx1101": ("AMD Radeon RX 7800 XT", (11, 0), "6.4.43483"),
"gfx1102": ("AMD Radeon RX 7600", (11, 0), "7.2.1"),
"gfx1150": ("AMD Radeon 890M", (11, 5), "7.2.1"), # RDNA3.5 APU
"gfx1151": ("AMD Radeon 8060S", (11, 5), "7.2.1"),
"gfx1200": ("AMD Radeon RX 9060 XT", (12, 0), "7.2.1"), # RDNA4
"gfx1201": ("AMD Radeon RX 9070 XT", (12, 0), "7.2.1"),
}
def _cuda_spoof():
"""Load the sibling CUDA spoof by path (robust to sys.path), so we reuse its
torch.cuda machinery instead of duplicating it."""
if "_zoo_aggressive_cuda_spoof" in sys.modules:
return sys.modules["_zoo_aggressive_cuda_spoof"]
path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_zoo_aggressive_cuda_spoof.py")
spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", path)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
sys.modules["_zoo_aggressive_cuda_spoof"] = mod
return mod
def apply(gfx: str = "gfx1100", device_count: int = 1) -> None:
"""Present torch as `gfx`. Re-callable to switch arch (identity is overlaid;
the underlying no-op machinery is applied once)."""
import torch
if gfx not in _PROFILES:
raise KeyError(f"Unknown gfx {gfx!r}; known: {', '.join(_PROFILES)}")
name, cap, hip = _PROFILES[gfx]
_cuda_spoof().apply() # is_available/device_count/streams/rng/amp/...
# Overlay the AMD identity on top of the (NVIDIA-shaped) CUDA spoof.
torch.version.hip = hip
torch.version.cuda = None
torch.cuda.device_count = lambda: device_count
torch.cuda.get_device_name = lambda *a, **k: name
torch.cuda.get_device_capability = lambda *a, **k: cap
torch.cuda.get_arch_list = lambda: [gfx]
class _Props:
pass
_p = _Props()
_p.name = name
_p.gcnArchName = f"{gfx}:sramecc-:xnack-" # ROCm advertises feature flags
_p.major, _p.minor = cap
_p.total_memory = 16 * 1024**3
_p.multi_processor_count = 40
_p.warp_size = 32 # RDNA wavefront (CDNA is 64)
_p.is_integrated = gfx in ("gfx1150", "gfx1151")
_p.is_multi_gpu_board = False
torch.cuda.get_device_properties = lambda *a, **k: _p
if __name__ == "__main__":
apply()
import torch
print("ROCm spoof applied:", torch.version.hip, torch.cuda.get_device_properties(0).gcnArchName)

84
tests/studio/install/test_rocm_rdna_routing.py Normal file
View file

@ -0,0 +1,84 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team.
"""RDNA 2/3/4 routing, validated on CPU-only CI with no AMD hardware.
tests/_zoo_rocm_spoof.py presents torch as each Radeon gfx arch, then we assert
unsloth_zoo routes it: device_type -> "hip", llama.cpp target -> ("rocm", gfx),
and the per-family ROCm bundle suffix. The torch-facing checks run in a
subprocess so the spoof never leaks into sibling tests and DEVICE_TYPE (cached
at import) resolves from a clean process.
"""
from __future__ import annotations
import json
import subprocess
import sys
from pathlib import Path
import pytest
pytest.importorskip("torch")
pytest.importorskip("unsloth_zoo")
_TESTS_DIR = Path(__file__).resolve().parents[2] # tests/
# gfx -> (expected llama.cpp target, expected ROCm bundle family).
_ARCHES = {
"gfx1030": (("rocm", "gfx1030"), "gfx103X"), # RDNA2
"gfx1031": (("rocm", "gfx1031"), "gfx103X"),
"gfx1032": (("rocm", "gfx1032"), "gfx103X"),
"gfx1034": (("rocm", "gfx1034"), "gfx103X"),
"gfx1100": (("rocm", "gfx1100"), "gfx110X"), # RDNA3
"gfx1101": (("rocm", "gfx1101"), "gfx110X"),
"gfx1102": (("rocm", "gfx1102"), "gfx110X"),
"gfx1150": (("rocm", "gfx1150"), "gfx1150"), # RDNA3.5 APU (self-family)
"gfx1151": (("rocm", "gfx1151"), "gfx1151"),
"gfx1200": (("rocm", "gfx1200"), "gfx120X"), # RDNA4
"gfx1201": (("rocm", "gfx1201"), "gfx120X"),
}
# Child: spoof each arch, then record device_type once (fresh import) and the
# live llama.cpp target per arch. Emits one JSON line the parent parses.
_CHILD = """
import json, sys
sys.path.insert(0, {tests!r})
import _zoo_rocm_spoof as spoof
arches = {arches!r}
spoof.apply(arches[0])
from unsloth_zoo.device_type import get_device_type, is_hip
device_type = [get_device_type(), is_hip()]
from unsloth_zoo import llama_cpp as lc
targets = {{}}
for gfx in arches:
spoof.apply(gfx)
targets[gfx] = list(lc._detect_gpu_target())
print("RESULT " + json.dumps({{"device_type": device_type, "targets": targets}}))
"""
@pytest.fixture(scope = "module")
def routed():
code = _CHILD.format(tests = str(_TESTS_DIR), arches = list(_ARCHES))
proc = subprocess.run([sys.executable, "-c", code], capture_output = True, text = True)
line = next((l for l in proc.stdout.splitlines() if l.startswith("RESULT ")), None)
assert line, f"child produced no result.\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}"
return json.loads(line[len("RESULT ") :])
@pytest.mark.parametrize("gfx", list(_ARCHES))
def test_detect_gpu_target(routed, gfx):
# RDNA card is routed to its ROCm gfx target (drives the llama.cpp bundle).
assert tuple(routed["targets"][gfx]) == _ARCHES[gfx][0]
def test_device_type_is_hip(routed):
# An RDNA card must resolve the compute device_type to "hip".
assert routed["device_type"] == ["hip", True]
@pytest.mark.parametrize("gfx", list(_ARCHES))
def test_rocm_gfx_family(gfx):
# Pure mapping (no torch): each gfx picks the right per-family ROCm bundle.
from unsloth_zoo import llama_cpp as lc
assert lc._rocm_gfx_family(gfx) == _ARCHES[gfx][1]

48
tests/studio/install/test_rocm_support.py
View file

@ -3426,8 +3426,9 @@ class TestInstallShDropinPersistence:
def test_gate5_early_return_persists_dropin(self):
"""The rocminfo-already-works early return must call the persist helper before returning."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
# The persist call must precede `return 0` at the rocminfo gfx1151 gate.
gate = source.find("Name:[[:space:]]*gfx1151")
# The persist call must precede `return 0` at the rocminfo GPU-agent gate
# (uniquely identified by the `!/generic/` clause the other probes lack).
gate = source.find("Name:[[:space:]]*gfx[1-9]/ && !/generic/")
assert gate != -1
window = source[gate : gate + 900]
assert "_persist_rocm_wsl_dropin" in window
@ -3441,6 +3442,49 @@ class TestInstallShDropinPersistence:
assert "profile.d/unsloth-rocm-wsl.sh" in body
_STRIXHALO_WSL_PATH = PACKAGE_ROOT / "scripts" / "install_rocm_wsl_strixhalo.sh"
class TestWslRerouteNvidiaGuard:
"""_maybe_reroute_strixhalo_to_2404 must skip the AMD reroute on hybrid AMD+NVIDIA hosts by
reusing _has_usable_nvidia_gpu (CUDA_VISIBLE_DEVICES-aware + /proc/driver/nvidia fallback),
which must be defined before the reroute's call site so it is actually available."""
def test_reroute_calls_nvidia_helper_before_amd_signal(self):
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
start = source.find("_maybe_reroute_strixhalo_to_2404()")
assert start != -1
body = source[start : start + 1200]
nv = body.find("_has_usable_nvidia_gpu")
wmi = body.find("_wsl_amd_gpu_name")
assert nv != -1, "reroute must consult _has_usable_nvidia_gpu before deciding to reroute"
assert wmi != -1
# The NVIDIA guard must precede the AMD/WMI signal and return early.
assert nv < wmi
assert body.find("return 0", nv) < wmi
def test_nvidia_helper_and_deps_defined_before_reroute_callsite(self):
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
call = source.find("\n_maybe_reroute_strixhalo_to_2404 || true")
assert call != -1
for fn in ("_run_bounded() {", "_cvd_hides_nvidia() {", "_has_usable_nvidia_gpu() {"):
idx = source.find(fn)
assert idx != -1 and idx < call, f"{fn} must be defined before the reroute call"
class TestStrixhaloGfxOverridePipefail:
"""The UNSLOTH_WSL_GFX override check must use a consuming grep, not grep -q: under
`set -o pipefail` an early -q exit SIGPIPEs printf and misreports the arch on large output."""
def test_gfx_override_uses_consuming_grep(self):
source = _STRIXHALO_WSL_PATH.read_text(encoding = "utf-8")
idx = source.find('grep -E "Name:[[:space:]]*${GFX}')
assert idx != -1, "GFX override must use a consuming grep -E (not grep -q)"
line = source[idx : source.find("\n", idx)]
assert ">/dev/null" in line
assert 'grep -qE "Name:[[:space:]]*${GFX}' not in source
class TestLlamaCppRuntimeWslOrdering:
"""The serve-time launcher mirrors binary_env: system HIP before the bundle dir on WSL."""

93
tests/studio/test_chat_response_details_ui_contract.py Normal file
View file

@ -0,0 +1,93 @@
"""Static contract for the chat response-details action and metadata."""
from __future__ import annotations
import re
from pathlib import Path
REPO = Path(__file__).resolve().parents[2]
THREAD_TSX = REPO / "studio/frontend/src/components/assistant-ui/thread.tsx"
DETAILS_TSX = (
REPO / "studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx"
)
REASONING_TSX = REPO / "studio/frontend/src/components/assistant-ui/reasoning.tsx"
ADAPTER_TS = REPO / "studio/frontend/src/features/chat/api/chat-adapter.ts"
CHAT_PREFS_TS = REPO / "studio/frontend/src/features/chat/stores/chat-preferences-store.ts"
CHAT_TAB_TSX = REPO / "studio/frontend/src/features/settings/tabs/chat-tab.tsx"
def test_assistant_more_menu_exposes_response_details_action():
src = THREAD_TSX.read_text()
assert "MessageResponseDetailsSheet" in src
assert "See response details" in src
assert "setDetailsOpen(true)" in src
def test_response_details_sheet_uses_unsloth_sheet_and_key_sections():
src = DETAILS_TSX.read_text()
assert "SheetContent" in src
assert "Response details" in src
assert "MessageResponseModelBadge" in src
assert "showResponseModel" in src
assert "ChipIcon" not in src
assert "s.params.checkpoint" not in src
assert "Not recorded" in src
assert "min-w-0 break-words font-heading" in src
assert "toolCallsFromContent(message.content)" in src
assert 'label="Called"' in src
for section in ["Response", "Tokens", "Timing", "Tools"]:
assert f'title="{section}"' in src
for field in ["Model", "Provider", "Total", "Cache hits", "Enabled", "Called"]:
assert f'label="{field}"' in src
def test_response_model_chip_is_user_configurable_and_rendered_in_metadata_rows():
prefs_src = CHAT_PREFS_TS.read_text()
chat_tab_src = CHAT_TAB_TSX.read_text()
thread_src = THREAD_TSX.read_text()
reasoning_src = REASONING_TSX.read_text()
assert "showResponseModel: boolean" in prefs_src
assert "showResponseModel: false" in prefs_src
assert "showResponseModel: saved?.showResponseModel ?? false" in prefs_src
assert "Show response model" in chat_tab_src
assert "setShowResponseModel" in chat_tab_src
assert "aui-response-model-badge inline-flex min-h-5" in DETAILS_TSX.read_text()
assert "leading-5" in DETAILS_TSX.read_text()
assert "group-hover/assistant-message:opacity-100" in DETAILS_TSX.read_text()
assert "MessageResponseModelBadge" in thread_src
assert "hasReasoningParts" in thread_src
assert "group/assistant-message aui-assistant-message-root" in thread_src
assert "pointer-events-none relative h-0" in thread_src
assert "MessageResponseModelBadge" in reasoning_src
assert 'className="min-w-0 flex-none"' in reasoning_src
assert "hidden min-w-0 max-w-[12rem]" in reasoning_src
assert "group-hover/assistant-message:inline-flex" in reasoning_src
def test_response_details_metadata_is_persisted_without_backend_schema_change():
src = ADAPTER_TS.read_text()
assert "interface ResponseDetailsMetadata" in src
assert "buildResponseDetails" in src
assert "responseDetails: buildResponseDetails(finishedAt)" in src
assert "toolCalls: Array.from(" in src
assert "!isExternalRequest && supportsTools && toolsEnabled" in src
assert "!isExternalRequest && supportsTools && codeToolsEnabled" in src
assert re.search(r"selectedModelSummary\?\.name\s*\|\|\s*responseModelId", src)
assert "providerName" in src
assert "cancelId" in src
metadata_block = src[
src.find("interface ResponseDetailsMetadata") : src.find("type RunMessages")
]
builder_block = src[
src.find("const buildResponseDetails") : src.find("const externalCapabilities")
]
for forbidden in [
"encrypted_api_key",
"externalApiKey",
"apiKey",
"providerKey",
"secret",
]:
assert forbidden not in metadata_block
assert forbidden not in builder_block

138
tests/utils/test_rope_scaling_drift.py
View file

@ -257,6 +257,63 @@ def test_recompute_helper_scales_on_cpu():
), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled."
def test_extended_rotary_reads_config_factor():
# LlamaExtendedRotaryEmbedding must honor the config factor, not hardcode 8
# (Llama-3.2 uses 32); otherwise the subclass path re-drops scaling (#2405).
from types import SimpleNamespace
from unsloth.models.llama import LlamaExtendedRotaryEmbedding
rot = object.__new__(LlamaExtendedRotaryEmbedding)
rot.base = ROPE_THETA
rot.dim = HEAD_DIM
rot._unsloth_rope_config = SimpleNamespace(
rope_scaling = {
"rope_type": "llama3",
"factor": 32.0,
"low_freq_factor": 1.0,
"high_freq_factor": 4.0,
"original_max_position_embeddings": 8192,
}
)
vanilla = _vanilla_inv_freq()
scaled = rot._apply_inv_freq_scaling(vanilla).reshape(-1)
ratio = float(vanilla[-1]) / float(scaled[-1])
assert abs(ratio - 32.0) < 1e-3, (
f"LlamaExtendedRotaryEmbedding ignored config factor 32 (ratio {ratio}); the "
"low-frequency band must be divided by the config factor (issue #2405)."
)
def test_extended_rotary_reads_rope_parameters_v5():
# transformers v5 stores scaling under rope_parameters (rope_scaling is a
# back-compat shim that may be removed); the factor must still be read.
from types import SimpleNamespace
from unsloth.models.llama import LlamaExtendedRotaryEmbedding
rot = object.__new__(LlamaExtendedRotaryEmbedding)
rot.base = ROPE_THETA
rot.dim = HEAD_DIM
rot._unsloth_rope_config = SimpleNamespace(
rope_scaling = None,
rope_parameters = {
"rope_type": "llama3",
"factor": 32.0,
"low_freq_factor": 1.0,
"high_freq_factor": 4.0,
"original_max_position_embeddings": 8192,
},
)
vanilla = _vanilla_inv_freq()
scaled = rot._apply_inv_freq_scaling(vanilla).reshape(-1)
ratio = float(vanilla[-1]) / float(scaled[-1])
assert abs(ratio - 32.0) < 1e-3, (
f"Extended rotary ignored rope_parameters factor 32 (ratio {ratio}); v5 "
"keeps the factor under rope_parameters, not rope_scaling."
)
def _cos_at_position(rot, position):
"""cos row at one position, built like _set_cos_sin_cache but CPU-only."""
inv_freq = rot.inv_freq.float().cpu()
@ -324,6 +381,87 @@ def test_extended_cache_keeps_scaling_after_growth():
)
def _blank_nonpersistent_buffers(module):
"""Mimic transformers v5 meta-load: overwrite non-persistent buffers with garbage."""
for name, buf in list(module.named_buffers()):
leaf = module
*parents, attr = name.split(".")
for part in parents:
leaf = getattr(leaf, part)
if attr in getattr(leaf, "_non_persistent_buffers_set", set()):
setattr(leaf, attr, torch.rand_like(buf))
def _build_llama3_rotary():
from unsloth.models import llama as llama_mod
config = _make_config(LLAMA3_ROPE_SCALING)
return llama_mod.LlamaRotaryEmbedding(config = config), config
def _build_longrope_rotary():
from types import SimpleNamespace
from unsloth.models import llama as llama_mod
short_factor, long_factor = [1.05] * 48, [1.3] * 48
rot = llama_mod.LongRopeRotaryEmbedding(
dim = 96,
max_position_embeddings = 131072,
original_max_position_embeddings = 4096,
base = ROPE_THETA,
short_factor = short_factor,
long_factor = long_factor,
)
config = SimpleNamespace(
rope_scaling = {
"rope_type": "longrope",
"short_factor": short_factor,
"long_factor": long_factor,
"original_max_position_embeddings": 4096,
}
)
return rot, config
@requires_cuda
@pytest.mark.parametrize(
"build", [_build_llama3_rotary, _build_longrope_rotary], ids = ["llama3", "longrope"]
)
def test_v5_blank_repair_roundtrip(build):
# Build scaled -> blank non-persistent buffers (what transformers v5 does on
# load) -> run the repair -> every buffer must return to its scaled value.
# Family-agnostic: encodes no scaling math, so it guards any rotary that
# keeps scaling in a buffer (issue #2405 / PR #6907).
from unsloth.models import loader
# The repair only runs on transformers v5 (it is what blanks the buffers);
# on v4 _fix_rope_inv_freq is a no-op, so the round-trip cannot restore.
if not loader._NEEDS_ROPE_FIX:
pytest.skip("transformers < 5 does not blank rope buffers; repair is a no-op")
rot, config = build()
snapshot = {name: buf.detach().clone() for name, buf in rot.named_buffers()}
assert snapshot, "rotary registers no buffers; nothing to guard"
_blank_nonpersistent_buffers(rot)
assert any(
not torch.equal(rot.get_buffer(name), snapshot[name]) for name in snapshot
), "blanking changed no buffer; the round-trip would be vacuous"
wrapper = torch.nn.Module()
wrapper.add_module("rotary_emb", rot)
wrapper.config = config
loader._fix_rope_inv_freq(wrapper)
for name in snapshot:
assert torch.allclose(
rot.get_buffer(name).cpu(), snapshot[name].cpu(), rtol = 1e-4, atol = 1e-6
), (
f"{name} was not restored to its scaled value by loader._fix_rope_inv_freq "
"after the transformers v5 buffer blank (issue #2405 / PR #6907)."
)
def test_object_style_rope_scaling_does_not_crash():
# Object-style rope_scaling must be normalized, not .get()'d directly.
from dataclasses import dataclass

3
unsloth/models/_utils.py
View file

@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
__version__ = "2026.6.9"
__version__ = "2026.7.1"
__all__ = [
"SUPPORTS_BFLOAT16",
@ -2834,6 +2834,7 @@ def patch_llama_rope_scaling(
dim = self.head_dim,
max_position_embeddings=self.max_position_embeddings,
base=self.rope_theta,
config=self.config,
)
elif scaling_type == "longrope":
self.rotary_emb = {longrope_rope_function}(

17
unsloth/models/llama.py
View file

@ -1930,11 +1930,18 @@ class LlamaExtendedRotaryEmbedding(LlamaRotaryEmbedding):
# From https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/api/model.py#L41
def _apply_inv_freq_scaling(self, freqs: torch.Tensor):
# Values obtained from grid search
scale_factor = 8
low_freq_factor = 1
high_freq_factor = 4
old_context_len = 8192 # original llama3 length
# llama3 factors from config; Llama-3.1 defaults when built without one
# (legacy codegen path). Hardcoding 8 is wrong for e.g. Llama-3.2 (32).
# v5 renames rope_scaling -> rope_parameters; read either so the factor
# survives even if the rope_scaling back-compat shim is dropped.
config = getattr(self, "_unsloth_rope_config", None)
rope_scaling = _rope_scaling_as_dict(
getattr(config, "rope_scaling", None) or getattr(config, "rope_parameters", None) or {}
)
scale_factor = rope_scaling.get("factor", 8)
low_freq_factor = rope_scaling.get("low_freq_factor", 1)
high_freq_factor = rope_scaling.get("high_freq_factor", 4)
old_context_len = rope_scaling.get("original_max_position_embeddings", 8192)
low_freq_wavelen = old_context_len / low_freq_factor
high_freq_wavelen = old_context_len / high_freq_factor

3
unsloth/models/rl.py
View file

@ -1370,6 +1370,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"):
# [TODO] See https://fengyao.notion.site/off-policy-rl
# https://github.com/huggingface/trl/pull/3867 (August 7th)
"vllm_importance_sampling_correction": False,
# TRL >= 1.7.0 enables the MoE router aux loss by default (0.001); the optimized
# GRPO forward does not compute it, so default off. Opt in via router_aux_loss_coef > 0.
"router_aux_loss_coef": 0.0,
}
for k, v in replacements.items():
x = f"{k}( = [^,\n]{{1,}})?,\n"