Merge branch 'main' into pip

install.ps1

@@ -749,7 +749,6 @@ shell.Run cmd, 0, False
     } else {
         step "gpu" "none (chat-only / GGUF)" "Yellow"
         substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow"
-        substep "https://www.nvidia.com/Download/index.aspx" "Yellow"
     }
 
     # ── Choose the correct PyTorch index URL based on driver CUDA version ──
@@ -777,10 +776,10 @@ shell.Run cmd, 0, False
     # ── Print CPU-only hint when no GPU detected ──
     if (-not $SkipTorch -and $TorchIndexUrl -like "*/cpu") {
         Write-Host ""
-        Write-Host "  NOTE: No NVIDIA GPU detected." -ForegroundColor Yellow
-        Write-Host "  Installing CPU-only PyTorch. If you only need GGUF chat/inference,"
-        Write-Host "  re-run with --no-torch for a faster, lighter install:"
-        Write-Host "    .\install.ps1 --no-torch"
+        substep "No NVIDIA GPU detected." "Yellow"
+        substep "Installing CPU-only PyTorch. If you only need GGUF chat/inference," "Yellow"
+        substep "re-run with --no-torch for a faster, lighter install:" "Yellow"
+        substep ".\install.ps1 --no-torch" "Yellow"
         Write-Host ""
     }
 
@@ -820,7 +819,7 @@ shell.Run cmd, 0, False
         if ($SkipTorch) {
             # No-torch: install unsloth + unsloth-zoo with --no-deps, then
             # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
-            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.18" unsloth-zoo }
             if ($baseInstallExit -eq 0) {
                 $NoTorchReq = Find-NoTorchRuntimeFile
                 if ($NoTorchReq) {
@@ -828,7 +827,7 @@ shell.Run cmd, 0, False
                 }
             }
         } else {
-            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.18" unsloth-zoo }
         }
         if ($baseInstallExit -ne 0) {
             Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@@ -858,7 +857,7 @@ shell.Run cmd, 0, False
         if ($SkipTorch) {
             # No-torch: install unsloth + unsloth-zoo with --no-deps, then
             # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
-            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.3.16" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.3.18" unsloth-zoo }
             if ($baseInstallExit -eq 0) {
                 $NoTorchReq = Find-NoTorchRuntimeFile
                 if ($NoTorchReq) {
@@ -866,7 +865,7 @@ shell.Run cmd, 0, False
                 }
             }
         } elseif ($StudioLocalInstall) {
-            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.18" unsloth-zoo }
         } else {
             $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName" }
         }
@@ -887,7 +886,7 @@ shell.Run cmd, 0, False
         # Fallback: GPU detection failed to produce a URL -- let uv resolve torch
         substep "installing unsloth (this may take a few minutes)..."
         if ($StudioLocalInstall) {
-            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto }
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.18" --torch-backend=auto }
             if ($baseInstallExit -ne 0) {
                 Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
                 return

install.sh

@@ -949,6 +949,17 @@ if [ -x "$VENV_DIR/bin/python" ]; then
     substep "${VENV_DIR}"
 fi
 
+# Default torch constraint -- tightened for Python 3.13+ on arm64 macOS
+# (torch <2.6 has no cp313 macOS arm64 wheels)
+TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
+if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
+    _PY_MINOR=$("$VENV_DIR/bin/python" -c \
+        "import sys; print(sys.version_info.minor)" 2>/dev/null || echo "0")
+    if [ "$_PY_MINOR" -ge 13 ] 2>/dev/null; then
+        TORCH_CONSTRAINT="torch>=2.6,<2.11.0"
+    fi
+fi
+
 # ── Resolve repo root (for --local installs) ──
 _REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)"
 
@@ -1029,7 +1040,7 @@ if [ "$_MIGRATED" = true ]; then
         # to prevent transitive torch resolution.
         run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
             --reinstall-package unsloth --reinstall-package unsloth-zoo \
-            "unsloth>=2026.3.16" unsloth-zoo
+            "unsloth>=2026.3.18" unsloth-zoo
         _NO_TORCH_RT="$(_find_no_torch_runtime)"
         if [ -n "$_NO_TORCH_RT" ]; then
             run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
@@ -1037,7 +1048,7 @@ if [ "$_MIGRATED" = true ]; then
     else
         run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
             --reinstall-package unsloth --reinstall-package unsloth-zoo \
-            "unsloth>=2026.3.16" unsloth-zoo
+            "unsloth>=2026.3.18" unsloth-zoo
     fi
     if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
         substep "overlaying local repo (editable)..."
@@ -1049,7 +1060,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
         substep "skipping PyTorch (--no-torch or Intel Mac x86_64)." "$C_WARN"
     else
         substep "installing PyTorch ($TORCH_INDEX_URL)..."
-        run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \
+        run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
             --index-url "$TORCH_INDEX_URL"
     fi
     # Fresh: Step 2 - install unsloth, preserving pre-installed torch
@@ -1059,7 +1070,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
         # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
         run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
             --upgrade-package unsloth --upgrade-package unsloth-zoo \
-            "unsloth>=2026.3.16" unsloth-zoo
+            "unsloth>=2026.3.18" unsloth-zoo
         _NO_TORCH_RT="$(_find_no_torch_runtime)"
         if [ -n "$_NO_TORCH_RT" ]; then
             run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
@@ -1070,7 +1081,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
         fi
     elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
         run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \
-            --upgrade-package unsloth "unsloth>=2026.3.16" unsloth-zoo
+            --upgrade-package unsloth "unsloth>=2026.3.18" unsloth-zoo
         substep "overlaying local repo (editable)..."
         run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
     else
@@ -1081,7 +1092,7 @@ else
     # Fallback: GPU detection failed to produce a URL -- let uv resolve torch
     substep "installing unsloth (this may take a few minutes)..."
     if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
-        run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.16" --torch-backend=auto
+        run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.3.18" --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
     else

install_gemma4_mlx.sh

@@ -0,0 +1,215 @@
+#!/bin/bash
+set -e
+
+# ============================================================
+# Gemma 4 MLX — One-command setup + inference
+#
+# Usage:
+#   bash install_gemma4_mlx.sh [--venv-dir DIR]
+#
+# This script:
+#   1. Creates a Python virtual environment
+#   2. Installs uv, mlx, mlx-lm, transformers
+#   3. Downloads gemma4.py and gemma4_text.py from unsloth repo
+#   4. Installs them into mlx-lm's models directory
+# ============================================================
+
+# ── Output style (inspired by unsloth/install.sh) ─────────────
+RULE=""
+_rule_i=0
+while [ "$_rule_i" -lt 52 ]; do
+    RULE="${RULE}─"
+    _rule_i=$((_rule_i + 1))
+done
+
+if [ -n "${NO_COLOR:-}" ]; then
+    C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST=
+elif [ -t 1 ] || [ -n "${FORCE_COLOR:-}" ]; then
+    _ESC="$(printf '\033')"
+    C_TITLE="${_ESC}[38;5;117m"
+    C_DIM="${_ESC}[38;5;245m"
+    C_OK="${_ESC}[38;5;108m"
+    C_WARN="${_ESC}[38;5;136m"
+    C_ERR="${_ESC}[91m"
+    C_RST="${_ESC}[0m"
+else
+    C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST=
+fi
+
+step()    { printf "  ${C_DIM}%-18.18s${C_RST}${3:-$C_OK}%s${C_RST}\n" "$1" "$2"; }
+substep() { printf "  ${C_DIM}%-18s${2:-$C_DIM}%s${C_RST}\n" "" "$1"; }
+fail()    { step "error" "$1" "$C_ERR"; exit 1; }
+
+# ── Parse flags ───────────────────────────────────────────────
+VENV_DIR=""
+_next_is_venv=false
+
+for arg in "$@"; do
+    if [ "$_next_is_venv" = true ]; then
+        VENV_DIR="$arg"
+        _next_is_venv=false
+        continue
+    fi
+    case "$arg" in
+        --venv-dir)  _next_is_venv=true ;;
+    esac
+done
+
+# Default venv location
+if [ -z "$VENV_DIR" ]; then
+    VENV_DIR="$HOME/.unsloth/unsloth_gemma4_mlx"
+fi
+
+# ── Banner ────────────────────────────────────────────────────
+echo ""
+printf "  ${C_TITLE}%s${C_RST}\n" "💎 Gemma 4 MLX Installer"
+printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
+echo ""
+
+# ── Platform check ────────────────────────────────────────────
+if [ "$(uname)" != "Darwin" ]; then
+    fail "MLX requires macOS with Apple Silicon. Detected: $(uname)"
+fi
+
+_ARCH=$(uname -m)
+if [ "$_ARCH" != "arm64" ]; then
+    step "warning" "Apple Silicon recommended (detected: $_ARCH)" "$C_WARN"
+fi
+
+step "platform" "macOS ($_ARCH)"
+
+# ── Detect Python ─────────────────────────────────────────────
+PYTHON=""
+for _candidate in python3.12 python3.11 python3.13 python3; do
+    if command -v "$_candidate" >/dev/null 2>&1; then
+        PYTHON="$_candidate"
+        break
+    fi
+done
+
+if [ -z "$PYTHON" ]; then
+    fail "Python 3 not found. Install via: brew install python@3.12"
+fi
+
+_PY_VERSION=$("$PYTHON" -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}')")
+step "python" "$PYTHON ($_PY_VERSION)"
+
+# ── Create virtual environment ────────────────────────────────
+if [ -x "$VENV_DIR/bin/python" ]; then
+    step "venv" "using existing environment"
+    substep "$VENV_DIR"
+else
+    step "venv" "creating virtual environment"
+    substep "$VENV_DIR"
+    mkdir -p "$(dirname "$VENV_DIR")"
+    "$PYTHON" -m venv "$VENV_DIR"
+fi
+
+# ── Install uv ───────────────────────────────────────────────
+if ! command -v uv >/dev/null 2>&1; then
+    step "uv" "installing uv package manager..."
+    _uv_tmp=$(mktemp)
+    curl -LsSf "https://astral.sh/uv/install.sh" -o "$_uv_tmp"
+    sh "$_uv_tmp" </dev/null >/dev/null 2>&1
+    rm -f "$_uv_tmp"
+    if [ -f "$HOME/.local/bin/env" ]; then
+        . "$HOME/.local/bin/env"
+    fi
+    export PATH="$HOME/.local/bin:$PATH"
+    substep "done"
+else
+    step "uv" "found $(uv --version 2>/dev/null || echo 'uv')"
+fi
+
+_VENV_PY="$VENV_DIR/bin/python"
+
+# ── Repo config ──────────────────────────────────────────────
+BRANCH="fix/ui-fix"
+REPO_URL="https://raw.githubusercontent.com/unslothai/unsloth/refs/heads/${BRANCH}"
+
+# ── Install dependencies ──────────────────────────────────────
+step "install" "installing mlx, mlx-lm..."
+uv pip install --python "$_VENV_PY" -q mlx mlx-lm 2>/dev/null
+substep "done"
+
+TRANSFORMERS_WHL="transformers-5.5.0-py3-none-any.whl"
+TRANSFORMERS_GH="git+https://github.com/huggingface/transformers.git@v5.5-release"
+
+step "install" "installing transformers>=5.5.0..."
+if uv pip install --python "$_VENV_PY" -q "$TRANSFORMERS_GH" 2>/dev/null; then
+    substep "installed from huggingface/transformers v5.5-release"
+elif uv pip install --python "$_VENV_PY" -q "transformers>=5.5.0" 2>/dev/null; then
+    substep "installed from PyPI"
+else
+    substep "not on PyPI, trying unsloth branch..."
+    _whl_tmp=$(mktemp -d)/"${TRANSFORMERS_WHL}"
+    if curl -fsSL "${REPO_URL}/${TRANSFORMERS_WHL}" -o "$_whl_tmp" 2>/dev/null && \
+       uv pip install --python "$_VENV_PY" -q "$_whl_tmp" 2>/dev/null; then
+        substep "installed from branch ${BRANCH}"
+    elif [ -f "./${TRANSFORMERS_WHL}" ]; then
+        substep "using local ./${TRANSFORMERS_WHL}"
+        uv pip install --python "$_VENV_PY" -q "./${TRANSFORMERS_WHL}"
+    else
+        rm -f "$_whl_tmp" 2>/dev/null
+        step "install" "skipping transformers — could not find >=5.5.0" "$C_WARN"
+        substep "tried: huggingface/transformers v5.5-release, PyPI, branch ${BRANCH}, local ./${TRANSFORMERS_WHL}"
+    fi
+    rm -f "$_whl_tmp" 2>/dev/null
+fi
+
+# ── Find mlx-lm models directory ─────────────────────────────
+MLX_MODELS=$("$_VENV_PY" -c "import mlx_lm; print(mlx_lm.__path__[0])")/models
+step "models dir" "$MLX_MODELS"
+
+# ── Download and install Gemma 4 model files ──────────────────
+
+step "download" "installing Gemma 4 model files..."
+
+_install_model_file() {
+    _fname="$1"
+    if curl -fsSL "${REPO_URL}/unsloth/models/${_fname}" -o "${MLX_MODELS}/${_fname}" 2>/dev/null; then
+        substep "downloaded ${_fname} from branch ${BRANCH}"
+    elif [ -f "./${_fname}" ]; then
+        substep "using local ./${_fname}"
+        cp "./${_fname}" "${MLX_MODELS}/${_fname}"
+    else
+        fail "Could not install ${_fname}. Tried:
+                    1) ${REPO_URL}/unsloth/models/${_fname}
+                    2) Local file ./${_fname}
+
+                    To fix, download the file manually and place it in the current directory,
+                    then re-run this script."
+    fi
+}
+
+_install_model_file "gemma4.py"
+_install_model_file "gemma4_text.py"
+
+# Verify files were installed correctly
+if "$_VENV_PY" -c "from mlx_lm.models.gemma4_text import ProportionalRoPE" 2>/dev/null; then
+    substep "model files verified"
+else
+    fail "Model files installed but verification failed (ProportionalRoPE import error).
+                    Try manually from: https://github.com/unslothai/unsloth/tree/feature/${BRANCH}"
+fi
+
+# ── Done ──────────────────────────────────────────────────────
+echo ""
+printf "  ${C_TITLE}%s${C_RST}\n" "Gemma 4 MLX installed!"
+printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
+echo ""
+step "available models" "unsloth/gemma-4-E2B-it-UD-MLX-4bit (/BF16)"
+substep "unsloth/gemma-4-E4B-it-UD-MLX-4bit (/BF16)"
+echo ""
+step "venv activate" "source ${VENV_DIR}/bin/activate"
+echo ""
+step "quick start" "python -m mlx_lm chat --model unsloth/gemma-4-E2B-it-UD-MLX-4bit --max-tokens 200"
+echo ""
+step "python API" "from mlx_lm import load, generate"
+substep "model, tokenizer = load('unsloth/gemma-4-E2B-it-UD-MLX-4bit')"
+substep "messages = [{'role': 'user', 'content': 'Hello!'}]"
+substep "prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)"
+substep "print(generate(model, tokenizer, prompt=prompt, max_tokens=200))"
+echo ""
+printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
+echo ""

studio/backend/core/inference/_html_to_md.py

@@ -17,18 +17,26 @@ from html.parser import HTMLParser
 
 __all__ = ["html_to_markdown"]
 
-_SKIP_TAGS = frozenset({"script", "style", "head", "noscript", "svg", "math"})
+_SKIP_TAGS = frozenset(
+    {
+        "script",
+        "style",
+        "head",
+        "noscript",
+        "svg",
+        "math",
+        "nav",
+        "footer",
+    }
+)
 _BLOCK_TAGS = frozenset(
     {
         "p",
         "div",
         "section",
         "article",
-        "header",
-        "footer",
         "main",
         "aside",
-        "nav",
         "figure",
         "figcaption",
         "details",

studio/backend/core/inference/defaults.py

@@ -6,6 +6,14 @@
 import utils.hardware.hardware as hw
 
 DEFAULT_MODELS_GGUF = [
+    "unsloth/gemma-4-E2B-it-GGUF",
+    "unsloth/gemma-4-E4B-it-GGUF",
+    "unsloth/gemma-4-31B-it-GGUF",
+    "unsloth/gemma-4-26B-A4B-it-GGUF",
+    "unsloth/Qwen3.5-4B-GGUF",
+    "unsloth/Qwen3.5-9B-GGUF",
+    "unsloth/Qwen3.5-35B-A3B-GGUF",
+    "unsloth/Qwen3.5-0.8B-GGUF",
     "unsloth/Llama-3.2-1B-Instruct-GGUF",
     "unsloth/Llama-3.2-3B-Instruct-GGUF",
     "unsloth/Llama-3.1-8B-Instruct-GGUF",
@@ -15,6 +23,18 @@ DEFAULT_MODELS_GGUF = [
 ]
 
 DEFAULT_MODELS_STANDARD = [
+    "unsloth/gemma-4-E2B-it-GGUF",
+    "unsloth/gemma-4-E4B-it-GGUF",
+    "unsloth/gemma-4-31B-it-GGUF",
+    "unsloth/gemma-4-26B-A4B-it-GGUF",
+    "unsloth/Qwen3.5-4B-GGUF",
+    "unsloth/Qwen3.5-9B-GGUF",
+    "unsloth/Qwen3.5-35B-A3B-GGUF",
+    "unsloth/Qwen3.5-0.8B-GGUF",
+    "unsloth/gemma-4-E2B-it",
+    "unsloth/gemma-4-E4B-it",
+    "unsloth/gemma-4-31B-it",
+    "unsloth/gemma-4-26B-A4B-it",
     "unsloth/Qwen3-4B-Instruct-2507",
     "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
     "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit",

studio/backend/core/inference/llama_cpp.py

@@ -27,6 +27,52 @@ import httpx
 
 logger = get_logger(__name__)
 
+# ── Pre-compiled patterns for plan-without-action re-prompt ──
+# Forward-looking intent signals that indicate the model is
+# describing what it *will* do rather than giving a final answer.
+_INTENT_SIGNAL = re.compile(
+    r"(?i)("
+    # Direct intent: "I'll ...", "I will ...", "Let me ...", "I am going to ..."
+    # Handles both straight and curly apostrophes.
+    # Excludes "I can", "I should", "I want to", "let's" which
+    # appear frequently in direct answers / explanations.
+    r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b"
+    r"|"
+    # Step/plan framing: "First ...", "Step 1:", "Here's my plan"
+    r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))"
+    r"|"
+    # "Now I" / "Next I" patterns
+    r"\b(?:now i|next i)\b"
+    r")"
+)
+_MAX_REPROMPTS = 3
+_REPROMPT_MAX_CHARS = 2000
+
+# ── Pre-compiled patterns for GGUF shard detection ───────────
+_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$")
+_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
+
+# Model size extraction (shared with routes/inference.py)
+from utils.models import extract_model_size_b as _extract_model_size_b
+
+# ── Pre-compiled patterns for tool XML stripping ─────────────
+_TOOL_CLOSED_PATS = [
+    re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
+    re.compile(r"<function=\w+>.*?</function>", re.DOTALL),
+]
+_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
+    re.compile(r"<tool_call>.*$", re.DOTALL),
+    re.compile(r"<function=\w+>.*$", re.DOTALL),
+]
+
+# ── Pre-compiled patterns for tool-call XML parsing ──────────
+_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
+_TC_FUNC_START_RE = re.compile(r"<function=(\w+)>\s*")
+_TC_END_TAG_RE = re.compile(r"</tool_call>")
+_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
+_TC_PARAM_START_RE = re.compile(r"<parameter=(\w+)>\s*")
+_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
+
 
 class LlamaCppBackend:
     """
@@ -61,6 +107,15 @@ class LlamaCppBackend:
         self._n_kv_heads: Optional[int] = None
         self._n_heads: Optional[int] = None
         self._embedding_length: Optional[int] = None
+        # Architecture-aware KV fields (8 new fields for 5-path estimation)
+        self._kv_key_length: Optional[int] = None
+        self._kv_value_length: Optional[int] = None
+        self._sliding_window: Optional[int] = None
+        self._full_attention_interval: Optional[int] = None
+        self._kv_lora_rank: Optional[int] = None
+        self._key_length_mla: Optional[int] = None
+        self._ssm_inner_size: Optional[int] = None
+        self._ssm_state_size: Optional[int] = None
         self._lock = threading.Lock()
         self._stdout_lines: list[str] = []
         self._stdout_thread: Optional[threading.Thread] = None
@@ -107,6 +162,11 @@ class LlamaCppBackend:
         """Return the maximum context currently available on this hardware."""
         return self._max_context_length or self._context_length
 
+    @property
+    def native_context_length(self) -> Optional[int]:
+        """Return the model's native context length from GGUF metadata."""
+        return self._context_length
+
     @property
     def chat_template(self) -> Optional[str]:
         return self._chat_template
@@ -228,14 +288,11 @@ class LlamaCppBackend:
     @staticmethod
     def _get_gguf_size_bytes(model_path: str) -> int:
         """Get total GGUF size in bytes, including split shards."""
-        import re
-
         main = Path(model_path)
         total = main.stat().st_size
 
         # Check for split shards (e.g., model-00001-of-00003.gguf)
-        shard_pat = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$")
-        m = shard_pat.match(main.name)
+        m = _SHARD_FULL_RE.match(main.name)
         if m:
             prefix, _, num_total = m.group(1), m.group(2), m.group(3)
             sibling_pat = re.compile(
@@ -306,11 +363,11 @@ class LlamaCppBackend:
         """Pick GPU(s) for a model based on estimated VRAM and free memory.
 
         ``model_size_bytes`` should include both model weights and estimated
-        KV cache.  The 70% threshold provides headroom for compute buffers,
+        KV cache.  The 90% threshold provides headroom for compute buffers,
         CUDA context, and other runtime overhead.
 
         Returns (gpu_indices, use_fit):
-          - ([1], False)       model fits on 1 GPU at 70% of free
+          - ([1], False)       model fits on 1 GPU at 90% of free
           - ([1, 2], False)    model needs 2 GPUs
           - (None, True)       model too large, let --fit handle it
         """
@@ -322,8 +379,8 @@ class LlamaCppBackend:
         # Sort GPUs by free memory descending
         ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
 
-        # Try fitting on 1 GPU (70% of free memory threshold)
-        if ranked[0][1] * 0.70 >= model_size_mib:
+        # Try fitting on 1 GPU (90% of free memory threshold)
+        if ranked[0][1] * 0.90 >= model_size_mib:
             return [ranked[0][0]], False
 
         # Try fitting on N GPUs (accumulate free memory from most-free)
@@ -331,7 +388,7 @@ class LlamaCppBackend:
         selected = []
         for idx, free_mib in ranked:
             selected.append(idx)
-            cumulative += free_mib * 0.70
+            cumulative += free_mib * 0.90
             if cumulative >= model_size_mib:
                 return sorted(selected), False
 
@@ -347,10 +404,17 @@ class LlamaCppBackend:
 
     def _can_estimate_kv(self) -> bool:
         """True if we have enough GGUF metadata to estimate KV cache size."""
-        return (
-            self._n_layers is not None
-            and self._embedding_length is not None
-            and (self._n_kv_heads is not None or self._n_heads is not None)
+        if self._n_layers is None:
+            return False
+        # MLA: kv_lora_rank is sufficient (K-only cache)
+        if self._kv_lora_rank is not None:
+            return True
+        # New-style: need both explicit key AND value dimensions
+        if self._kv_key_length is not None and self._kv_value_length is not None:
+            return True
+        # Legacy: need embedding_length + head count
+        return self._embedding_length is not None and (
+            self._n_kv_heads is not None or self._n_heads is not None
         )
 
     def _estimate_kv_cache_bytes(
@@ -358,14 +422,20 @@ class LlamaCppBackend:
     ) -> int:
         """Estimate KV cache VRAM for a given context length.
 
+        Uses 5-path architecture-aware estimation:
+          1. MLA      -- compressed KV latent + RoPE, K-only (no separate V)
+          2. Hybrid   -- only attention layers need KV (Mamba layers don't)
+          3. SWA      -- sliding-window layers cache min(ctx, window) tokens
+          4. GQA      -- standard full KV with explicit key/value dimensions
+          5. Legacy   -- fallback using embed // n_heads
+
         Returns 0 if metadata is insufficient for estimation.
         """
         if not self._can_estimate_kv() or n_ctx <= 0:
             return 0
 
         n_layers = self._n_layers  # type: ignore[assignment]
-        n_kv_heads = self._n_kv_heads or self._n_heads  # type: ignore[assignment]
-        head_dim = self._embedding_length // self._n_heads if self._n_heads else 128  # type: ignore[operator]
+        n_kv = self._n_kv_heads or self._n_heads or 1  # type: ignore[assignment]
 
         # Bytes per element depends on KV cache quantization
         bpe = {
@@ -380,8 +450,60 @@ class LlamaCppBackend:
             "iq4_nl": 0.5625,
         }.get(cache_type_kv or "f16", 2.0)
 
-        # K + V caches: 2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe
-        return int(2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe)
+        # Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5)
+        # MLA stores one compressed KV latent per token/layer (shared across heads).
+        # V is reconstructed from the latent on the fly -- no separate V cache.
+        # key_length = kv_lora_rank + rope_dim (the full compressed representation).
+        # MLA GGUFs set head_count_kv=1; default to 1 if absent to avoid
+        # falling back to n_heads (e.g., 128 for DeepSeek-V3) which would 128x.
+        if self._kv_lora_rank is not None:
+            n_kv_mla = self._n_kv_heads or 1
+            rope_dim = self._key_length_mla or 64
+            key_len = self._kv_key_length or (self._kv_lora_rank + rope_dim)
+            return int(n_layers * n_ctx * n_kv_mla * key_len * bpe)
+
+        key_len = self._kv_key_length
+        val_len = self._kv_value_length
+
+        # Path 2: Hybrid Mamba/Attention (Qwen3.5-27B, Qwen3.5-35B-A3B)
+        # Only 1 in N layers is attention; the rest are Mamba (no KV cache).
+        if (
+            self._ssm_inner_size is not None
+            and self._full_attention_interval is not None
+        ):
+            fai = self._full_attention_interval
+            n_attn = -(-n_layers // fai) if fai > 0 else n_layers  # ceiling division
+            if key_len is not None and val_len is not None:
+                return int(n_attn * n_ctx * n_kv * (key_len + val_len) * bpe)
+            head_dim = self._embedding_length // self._n_heads if self._n_heads else 128  # type: ignore[operator]
+            return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe)
+
+        # Path 3: Sliding Window (Gemma-3, gpt-oss)
+        # SWA layers only cache min(ctx, window) tokens; global layers cache full ctx.
+        # Most SWA architectures use few global layers (e.g., Gemma-3 uses 1 in 6).
+        # Without an explicit field, we conservatively assume 1/4 of layers are global
+        # which is still far more accurate than the legacy formula (which ignores SWA).
+        if (
+            self._sliding_window is not None
+            and self._sliding_window > 0
+            and key_len is not None
+            and val_len is not None
+        ):
+            swa = self._sliding_window
+            n_global = max(1, n_layers // 4)
+            n_swa = n_layers - n_global
+            kv_per_token = n_kv * (key_len + val_len) * bpe
+            return int(
+                n_global * n_ctx * kv_per_token + n_swa * min(n_ctx, swa) * kv_per_token
+            )
+
+        # Path 4: Standard GQA with explicit key/value dimensions
+        if key_len is not None and val_len is not None:
+            return int(n_layers * n_ctx * n_kv * (key_len + val_len) * bpe)
+
+        # Path 5: Legacy fallback (old GGUFs without explicit dimensions)
+        head_dim = self._embedding_length // self._n_heads if self._n_heads else 128  # type: ignore[operator]
+        return int(2 * n_kv * head_dim * n_layers * n_ctx * bpe)
 
     def _fit_context_to_vram(
         self,
@@ -393,8 +515,8 @@ class LlamaCppBackend:
     ) -> int:
         """Return the largest context length that fits in GPU VRAM.
 
-        Uses 70% of available VRAM as the budget (matching _select_gpus
-        threshold -- 30% reserved for compute buffers, CUDA context,
+        Uses 90% of available VRAM as the budget (matching _select_gpus
+        threshold -- 10% reserved for compute buffers, CUDA context,
         scratch space, flash-attn workspace, etc.).
         If the model weights alone don't fit, returns min_ctx unchanged.
         """
@@ -406,7 +528,7 @@ class LlamaCppBackend:
             )
             return requested_ctx
 
-        budget_bytes = available_mib * 1024 * 1024 * 0.70
+        budget_bytes = available_mib * 1024 * 1024 * 0.90
         model_footprint = model_size_bytes
 
         # Check if requested context already fits
@@ -460,8 +582,6 @@ class LlamaCppBackend:
 
         Returns (first_shard_filename, total_size_bytes) or None if nothing fits.
         """
-        import re
-
         try:
             from huggingface_hub import get_paths_info, list_repo_files
 
@@ -477,10 +597,9 @@ class LlamaCppBackend:
             size_map = {p.path: (p.size or 0) for p in path_infos}
 
             # Group files by variant: shards share a prefix before -NNNNN-of-NNNNN
-            shard_pat = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
             variants: dict[str, list[str]] = {}
             for f in gguf_files:
-                m = shard_pat.match(f)
+                m = _SHARD_RE.match(f)
                 key = m.group(1) if m else f
                 variants.setdefault(key, []).append(f)
 
@@ -585,6 +704,14 @@ class LlamaCppBackend:
         self._n_kv_heads = None
         self._n_heads = None
         self._embedding_length = None
+        self._kv_key_length = None
+        self._kv_value_length = None
+        self._sliding_window = None
+        self._full_attention_interval = None
+        self._kv_lora_rank = None
+        self._key_length_mla = None
+        self._ssm_inner_size = None
+        self._ssm_state_size = None
 
         try:
             WANTED = {"general.architecture", "tokenizer.chat_template"}
@@ -619,6 +746,15 @@ class LlamaCppBackend:
                                     f"{arch}.attention.head_count_kv": "n_kv_heads",
                                     f"{arch}.attention.head_count": "n_heads",
                                     f"{arch}.embedding_length": "embedding_length",
+                                    # Architecture-aware KV cache fields
+                                    f"{arch}.attention.key_length": "kv_key_length",
+                                    f"{arch}.attention.value_length": "kv_value_length",
+                                    f"{arch}.attention.sliding_window": "sliding_window",
+                                    f"{arch}.full_attention_interval": "full_attention_interval",
+                                    f"{arch}.attention.kv_lora_rank": "kv_lora_rank",
+                                    f"{arch}.attention.key_length_mla": "key_length_mla",
+                                    f"{arch}.ssm.inner_size": "ssm_inner_size",
+                                    f"{arch}.ssm.state_size": "ssm_state_size",
                                 }
                             elif key == "tokenizer.chat_template":
                                 self._chat_template = val_s
@@ -674,7 +810,9 @@ class LlamaCppBackend:
                 # Detect tool calling support from chat template
                 tool_markers = [
                     "{%- if tools %}",
+                    "{%- if tools -%}",
                     "{% if tools %}",
+                    "{% if tools -%}",
                     '"role" == "tool"',
                     "'role' == 'tool'",
                     'message.role == "tool"',
@@ -714,7 +852,6 @@ class LlamaCppBackend:
         gguf_extra_shards: list[str] = []
         if hf_variant:
             try:
-                import re
                 from huggingface_hub import list_repo_files
 
                 files = list_repo_files(hf_repo, token = hf_token)
@@ -729,11 +866,10 @@ class LlamaCppBackend:
                 )
                 if gguf_files:
                     gguf_filename = gguf_files[0]
-                    shard_pat = re.compile(r"^(.*)-\d{5}-of-(\d{5})\.gguf$")
-                    m = shard_pat.match(gguf_filename)
+                    m = _SHARD_FULL_RE.match(gguf_filename)
                     if m:
                         prefix = m.group(1)
-                        total = m.group(2)
+                        total = m.group(3)
                         sibling_pat = re.compile(
                             r"^"
                             + re.escape(prefix)
@@ -790,10 +926,7 @@ class LlamaCppBackend:
                             f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)"
                         )
                         gguf_filename = fallback_file
-                        import re as _re
-
-                        _shard_pat = _re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
-                        _m = _shard_pat.match(gguf_filename)
+                        _m = _SHARD_RE.match(gguf_filename)
                         _prefix = _m.group(1) if _m else None
                         if _prefix:
                             gguf_extra_shards = sorted(
@@ -1042,7 +1175,7 @@ class LlamaCppBackend:
                             )
                             kv = self._estimate_kv_cache_bytes(capped, cache_type_kv)
                             total_mib = (model_size + kv) / (1024 * 1024)
-                            if total_mib <= pool_mib * 0.70:
+                            if total_mib <= pool_mib * 0.90:
                                 best_cap = max(best_cap, capped)
                         if best_cap > 0:
                             max_available_ctx = best_cap
@@ -1071,7 +1204,7 @@ class LlamaCppBackend:
                                     capped, cache_type_kv
                                 )
                                 total_mib = (model_size + kv) / (1024 * 1024)
-                                if total_mib <= pool_mib * 0.70:
+                                if total_mib <= pool_mib * 0.90:
                                     effective_ctx = capped
                                     gpu_indices = sorted(idx for idx, _ in subset)
                                     use_fit = False
@@ -1090,7 +1223,7 @@ class LlamaCppBackend:
                             )
                             kv = self._estimate_kv_cache_bytes(capped, cache_type_kv)
                             total_mib = (model_size + kv) / (1024 * 1024)
-                            if total_mib <= pool_mib * 0.70:
+                            if total_mib <= pool_mib * 0.90:
                                 effective_ctx = capped
                                 gpu_indices = sorted(idx for idx, _ in subset)
                                 use_fit = False
@@ -1196,17 +1329,12 @@ class LlamaCppBackend:
             # Qwen3.5 models below 9B (0.8B, 2B, 4B) disable thinking by default.
             # Only 9B and larger enable thinking.
             if self._supports_reasoning:
-                import re
-
                 thinking_default = True
                 mid = (model_identifier or "").lower()
                 if "qwen3.5" in mid:
-                    # Extract size like "0.8b", "4b", "35b" etc.
-                    size_match = re.search(r"(\d+\.?\d*)\s*b", mid)
-                    if size_match:
-                        size_val = float(size_match.group(1))
-                        if size_val < 9:
-                            thinking_default = False
+                    size_val = _extract_model_size_b(mid)
+                    if size_val is not None and size_val < 9:
+                        thinking_default = False
                 self._reasoning_default = thinking_default
                 cmd.extend(
                     [
@@ -1422,6 +1550,14 @@ class LlamaCppBackend:
             self._n_kv_heads = None
             self._n_heads = None
             self._embedding_length = None
+            self._kv_key_length = None
+            self._kv_value_length = None
+            self._sliding_window = None
+            self._full_attention_interval = None
+            self._kv_lora_rank = None
+            self._key_length_mla = None
+            self._ssm_inner_size = None
+            self._ssm_state_size = None
             # Clean up temp chat template file
             if hasattr(self, "_chat_template_file") and self._chat_template_file:
                 try:
@@ -1671,13 +1807,11 @@ class LlamaCppBackend:
         Closing tags (</tool_call>, </function>, </parameter>) are all optional
         since models frequently omit them.
         """
-        import re
-
         tool_calls = []
 
         # Pattern 1: JSON inside <tool_call> tags.
         # Use balanced-brace extraction that skips braces inside JSON strings.
-        for m in re.finditer(r"<tool_call>\s*\{", content):
+        for m in _TC_JSON_START_RE.finditer(content):
             brace_start = m.end() - 1  # position of the opening {
             depth, i = 0, brace_start
             in_string = False
@@ -1727,7 +1861,7 @@ class LlamaCppBackend:
             # boundaries.  We avoid using </function> as a boundary because
             # code parameter values can contain that literal string.
             # After extracting, we trim a trailing </function> if present.
-            func_starts = list(re.finditer(r"<function=(\w+)>\s*", content))
+            func_starts = list(_TC_FUNC_START_RE.finditer(content))
             for idx, fm in enumerate(func_starts):
                 func_name = fm.group(1)
                 body_start = fm.end()
@@ -1737,7 +1871,7 @@ class LlamaCppBackend:
                     if idx + 1 < len(func_starts)
                     else len(content)
                 )
-                end_tag = re.search(r"</tool_call>", content[body_start:])
+                end_tag = _TC_END_TAG_RE.search(content[body_start:])
                 if end_tag:
                     body_end = body_start + end_tag.start()
                 else:
@@ -1745,20 +1879,20 @@ class LlamaCppBackend:
                 body_end = min(body_end, next_func)
                 body = content[body_start:body_end]
                 # Trim trailing </function> if present (it's the real closing tag)
-                body = re.sub(r"\s*</function>\s*$", "", body)
+                body = _TC_FUNC_CLOSE_RE.sub("", body)
 
                 # Step 2: Extract parameters from body.
                 # For single-parameter functions (the common case: code, command,
                 # query), use body end as the only boundary to avoid false matches
                 # on </parameter> inside code strings.
                 arguments = {}
-                param_starts = list(re.finditer(r"<parameter=(\w+)>\s*", body))
+                param_starts = list(_TC_PARAM_START_RE.finditer(body))
                 if len(param_starts) == 1:
                     # Single parameter: value is everything from after the tag
                     # to end of body, trimming any trailing </parameter>.
                     pm = param_starts[0]
                     val = body[pm.end() :]
-                    val = re.sub(r"\s*</parameter>\s*$", "", val)
+                    val = _TC_PARAM_CLOSE_RE.sub("", val)
                     arguments[pm.group(1)] = val.strip()
                 else:
                     for pidx, pm in enumerate(param_starts):
@@ -1772,7 +1906,7 @@ class LlamaCppBackend:
                         )
                         val = body[val_start:next_param]
                         # Trim trailing </parameter> if present
-                        val = re.sub(r"\s*</parameter>\s*$", "", val)
+                        val = _TC_PARAM_CLOSE_RE.sub("", val)
                         arguments[param_name] = val.strip()
 
                 tc = {
@@ -2145,22 +2279,10 @@ class LlamaCppBackend:
         _accumulated_predicted_ms = 0.0
         _accumulated_predicted_n = 0
 
-        # ── Shared patterns for stripping tool XML from streamed content ──
-        import re as _re_tool
-
-        _TOOL_CLOSED_PATTERNS = [
-            _re_tool.compile(r"<tool_call>.*?</tool_call>", _re_tool.DOTALL),
-            _re_tool.compile(r"<function=\w+>.*?</function>", _re_tool.DOTALL),
-        ]
-        _TOOL_ALL_PATTERNS = _TOOL_CLOSED_PATTERNS + [
-            _re_tool.compile(r"<tool_call>.*$", _re_tool.DOTALL),
-            _re_tool.compile(r"<function=\w+>.*$", _re_tool.DOTALL),
-        ]
-
         def _strip_tool_markup(text: str, *, final: bool = False) -> str:
             if not auto_heal_tool_calls:
                 return text
-            patterns = _TOOL_ALL_PATTERNS if final else _TOOL_CLOSED_PATTERNS
+            patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
             for pat in patterns:
                 text = pat.sub("", text)
             return text.strip() if final else text
@@ -2180,7 +2302,19 @@ class LlamaCppBackend:
         # identical call succeeded).
         _tool_call_history: list[tuple[str, bool]] = []  # (key, failed)
 
-        for iteration in range(max_tool_iterations):
+        # ── Re-prompt on plan-without-action ─────────────────
+        # When the model describes what it intends to do (forward-looking
+        # language) without actually calling a tool, re-prompt once.
+        # Only triggers on responses that signal intent/planning -- a
+        # direct answer like "4" or "Hello!" will not match.
+        # Pattern is compiled once at module level (_INTENT_SIGNAL).
+        _reprompt_count = 0
+
+        # Reserve extra iterations for re-prompts so they don't
+        # consume the caller's tool-call budget.  Only add the
+        # extra slot when tool iterations are actually allowed.
+        _extra = _MAX_REPROMPTS if max_tool_iterations > 0 else 0
+        for iteration in range(max_tool_iterations + _extra):
             if cancel_event is not None and cancel_event.is_set():
                 return
 
@@ -2491,6 +2625,57 @@ class LlamaCppBackend:
                             content_accum,
                         )
                     if not _safety_tc:
+                        # ── Re-prompt on plan-without-action ──
+                        # If the model described what it intends to do
+                        # (forward-looking language) without calling any
+                        # tool, nudge it to act.  Only fires once per
+                        # request and only on short responses that
+                        # contain intent signals -- a direct answer
+                        # like "4" or "Hello!" won't trigger this.
+                        # Use content if available, otherwise fall back
+                        # to reasoning text (reasoning-only stalls).
+                        _stripped = content_accum.strip()
+                        if not _stripped:
+                            _stripped = reasoning_accum.strip()
+                        if (
+                            tools
+                            and _reprompt_count < _MAX_REPROMPTS
+                            and 0 < len(_stripped) < _REPROMPT_MAX_CHARS
+                            and _INTENT_SIGNAL.search(_stripped)
+                        ):
+                            _reprompt_count += 1
+                            logger.info(
+                                f"Re-prompt {_reprompt_count}/{_MAX_REPROMPTS}: "
+                                f"model responded without calling tools "
+                                f"({len(_stripped)} chars)"
+                            )
+                            conversation.append(
+                                {
+                                    "role": "assistant",
+                                    "content": _stripped,
+                                }
+                            )
+                            conversation.append(
+                                {
+                                    "role": "user",
+                                    "content": (
+                                        "STOP. Do NOT write code or explain. "
+                                        "You MUST call a tool NOW. "
+                                        "Call web_search or python immediately."
+                                    ),
+                                }
+                            )
+                            # Accumulate tokens and timing from this iteration
+                            _fu_r = _iter_usage or {}
+                            _accumulated_completion_tokens += _fu_r.get(
+                                "completion_tokens", 0
+                            )
+                            _it_r = _iter_timings or {}
+                            _accumulated_predicted_ms += _it_r.get("predicted_ms", 0)
+                            _accumulated_predicted_n += _it_r.get("predicted_n", 0)
+                            yield {"type": "status", "text": ""}
+                            continue
+
                         # Content was already streamed.  Yield metadata.
                         yield {"type": "status", "text": ""}
                         _fu = _iter_usage or {}
@@ -2723,10 +2908,15 @@ class LlamaCppBackend:
                         _error_prefixes
                     )
                     _tool_call_history.append((_tc_key, _is_error))
+                    # Strip image sentinel before feeding result to the LLM
+                    # (the full result with sentinel is still yielded via
+                    # tool_end so the frontend can extract image paths).
                     _result_content = result
+                    if "\n__IMAGES__:" in _result_content:
+                        _result_content = _result_content.rsplit("\n__IMAGES__:", 1)[0]
                     if _is_error:
                         _result_content = (
-                            result + "\n\nThe tool call encountered an issue. "
+                            _result_content + "\n\nThe tool call encountered an issue. "
                             "Please try a different approach or rephrase your request."
                         )
 
@@ -2740,6 +2930,8 @@ class LlamaCppBackend:
                         tool_msg["tool_call_id"] = tool_call_id
                     conversation.append(tool_msg)
 
+                # Clear tool status badge before next generation iteration
+                yield {"type": "status", "text": ""}
                 # Continue the loop to let model respond with context
                 continue
 

studio/backend/core/inference/orchestrator.py

@@ -109,12 +109,13 @@ class InferenceOrchestrator:
         self._top_models_ready.wait(timeout = 5)
         top_gguf = self._top_gguf_cache or []
         top_hub = self._top_hub_cache or []
-        # GGUFs first, then hub models, then static fallbacks.
+        # Curated static defaults first (editorial picks like new models),
+        # then HF download-ranked models to backfill.
         # Send extras so the frontend still has 4 per category
         # after removing already-downloaded models.
         result: list[str] = []
         seen: set[str] = set()
-        for m in top_gguf + top_hub + self._static_models:
+        for m in self._static_models + top_gguf + top_hub:
             if m not in seen:
                 result.append(m)
                 seen.add(m)

studio/backend/core/inference/tools.py

@@ -8,20 +8,28 @@ Supports web search (DuckDuckGo), Python code execution, and terminal commands.
 """
 
 import ast
+import http.client
 import os
 
 os.environ["UNSLOTH_IS_PRESENT"] = "1"
 
+import random
+import ssl
 import subprocess
 import sys
 import tempfile
 import threading
+import urllib.request
 
 from loggers import get_logger
 
 logger = get_logger(__name__)
 
 _EXEC_TIMEOUT = 300  # 5 minutes
+
+# Strict raster-image allowlist for sandbox file serving.
+# No .svg (XSS risk via embedded scripts), no .html, no .pdf.
+_IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"})
 _MAX_OUTPUT_CHARS = 8000  # truncate long output
 _BASH_BLOCKED_WORDS = {"rm", "sudo", "dd", "chmod", "mkfs", "shutdown", "reboot"}
 
@@ -154,8 +162,74 @@ def execute_tool(
     return f"Unknown tool: {name}"
 
 
-_MAX_PAGE_CHARS = 16000  # limit fetched page text
-_MAX_FETCH_BYTES = _MAX_PAGE_CHARS * 4 + 1  # cap raw download size
+_MAX_PAGE_CHARS = 16000  # limit fetched page text (after HTML-to-MD conversion)
+# Raw download cap.  Must be larger than _MAX_PAGE_CHARS because SSR pages
+# embed large <head> sections (CSS, JS, SVGs) that are stripped during
+# HTML-to-Markdown conversion.  512 KB is enough to reach article content
+# on GitBook / Next.js / Docusaurus pages whose <head> alone can be 200 KB.
+_MAX_FETCH_BYTES = 512 * 1024
+
+_USER_AGENTS = (
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
+    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
+    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:133.0) Gecko/20100101 Firefox/133.0",
+    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:133.0) Gecko/20100101 Firefox/133.0",
+    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/18.2 Safari/605.1.15",
+)
+
+_tls_ctx = ssl.create_default_context()
+
+
+class _NoRedirect(urllib.request.HTTPRedirectHandler):
+    def redirect_request(self, req, fp, code, msg, headers, newurl):
+        return None
+
+
+class _PinnedHTTPSConnection(http.client.HTTPSConnection):
+    """HTTPS connection that connects to a pinned IP but uses a different
+    hostname for SNI and certificate verification.
+
+    The SSRF IP-pinning rewrites URLs to raw IPs.  A normal HTTPSConnection
+    would then send no SNI and verify the cert against the IP, both of which
+    fail.  This subclass splits the two concerns: TCP connects to the pinned
+    IP (``host`` parameter) while TLS uses ``sni_hostname`` for the
+    ClientHello and cert check.
+    """
+
+    def __init__(self, host: str, *, sni_hostname: str, **kwargs):
+        super().__init__(host, **kwargs)
+        self._sni_hostname = sni_hostname
+
+    def connect(self):
+        # TCP connect to the pinned IP stored in self.host (+ tunnel if
+        # a proxy is configured via set_tunnel, though we do not use one).
+        http.client.HTTPConnection.connect(self)
+        # TLS handshake with the real hostname for SNI + cert verification.
+        self.sock = self._context.wrap_socket(
+            self.sock,
+            server_hostname = self._sni_hostname,
+        )
+
+
+class _SNIHTTPSHandler(urllib.request.HTTPSHandler):
+    """HTTPS handler that sends the correct SNI hostname during TLS handshake.
+
+    The SSRF IP-pinning rewrites URLs to raw IPs, which breaks SNI and cert
+    verification.  This handler returns a ``_PinnedHTTPSConnection`` that
+    connects to the pinned IP but verifies TLS against the original hostname.
+    """
+
+    def __init__(self, hostname: str):
+        super().__init__(context = _tls_ctx)
+        self._sni_hostname = hostname
+
+    def https_open(self, req):
+        return self.do_open(self._sni_connection, req)
+
+    def _sni_connection(self, host, **kwargs):
+        kwargs["context"] = _tls_ctx
+        return _PinnedHTTPSConnection(host, sni_hostname = self._sni_hostname, **kwargs)
 
 
 def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str]:
@@ -215,33 +289,32 @@ def _fetch_page_text(
         return reason
 
     try:
-        import urllib.request
         from urllib.error import HTTPError as _HTTPError
         from urllib.parse import urljoin, urlunparse
 
-        # Disable auto-redirect so we can validate each hop for SSRF.
-        # urllib raises HTTPError for 3xx when the handler returns None,
-        # so we catch that and extract the Location header manually.
-        class _NoRedirect(urllib.request.HTTPRedirectHandler):
-            def redirect_request(self, req, fp, code, msg, headers, newurl):
-                return None
-
-        opener = urllib.request.build_opener(_NoRedirect)
-        max_bytes = max_chars * 4 + 1
+        max_bytes = _MAX_FETCH_BYTES
         current_url = url
         current_host = parsed.hostname
+        ua = random.choice(_USER_AGENTS)
 
         for _hop in range(5):
             # Pin to the validated IP to prevent DNS rebinding.
             # Rewrite the URL to use the IP and set the Host header.
             cp = urlparse(current_url)
-            ip_netloc = f"{pinned_ip}:{cp.port}" if cp.port else pinned_ip
+            # Bracket IPv6 addresses so the netloc is valid in a URL.
+            ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip
+            ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str
             pinned_url = urlunparse(cp._replace(netloc = ip_netloc))
 
+            opener = urllib.request.build_opener(
+                _NoRedirect,
+                _SNIHTTPSHandler(current_host),
+            )
+
             req = urllib.request.Request(
                 pinned_url,
                 headers = {
-                    "User-Agent": "UnslothStudio/1.0",
+                    "User-Agent": ua,
                     "Host": current_host,
                 },
             )
@@ -522,6 +595,12 @@ def _check_code_safety(code: str) -> str | None:
     """
     safe, info = _check_signal_escape_patterns(code)
     if not safe:
+        # SyntaxError from ast.parse -- let these through so the subprocess
+        # produces a normal Python traceback instead of a misleading
+        # "unsafe code detected" message.
+        if info.get("error"):
+            return None
+
         reasons = [
             item.get("description", "") for item in info.get("signal_tampering", [])
         ]
@@ -565,6 +644,17 @@ def _python_exec(
 
     tmp_path = None
     workdir = _get_workdir(session_id)
+    # Snapshot image mtimes so we detect both new and overwritten files.
+    _before: dict[str, int] = {}
+    if os.path.isdir(workdir):
+        for _name in os.listdir(workdir):
+            if os.path.splitext(_name)[1].lower() in _IMAGE_EXTS:
+                _p = os.path.join(workdir, _name)
+                if os.path.isfile(_p):
+                    try:
+                        _before[_name] = os.stat(_p).st_mtime_ns
+                    except OSError:
+                        pass
     try:
         fd, tmp_path = tempfile.mkstemp(
             suffix = ".py", prefix = "studio_exec_", dir = workdir
@@ -600,7 +690,29 @@ def _python_exec(
         result = output or ""
         if proc.returncode != 0:
             result = f"Exit code {proc.returncode}:\n{result}"
-        return _truncate(result) if result.strip() else "(no output)"
+        result = _truncate(result) if result.strip() else "(no output)"
+
+        # Detect new or overwritten image files and append sentinel for frontend
+        if session_id and os.path.isdir(workdir):
+            new_images = []
+            for _name in os.listdir(workdir):
+                if os.path.splitext(_name)[1].lower() not in _IMAGE_EXTS:
+                    continue
+                _p = os.path.join(workdir, _name)
+                if not os.path.isfile(_p):
+                    continue
+                try:
+                    _mtime = os.stat(_p).st_mtime_ns
+                except OSError:
+                    continue
+                if _name not in _before or _mtime != _before[_name]:
+                    new_images.append(_name)
+            if new_images:
+                import json as _json
+
+                result += f"\n__IMAGES__:{_json.dumps(sorted(new_images))}"
+
+        return result
 
     except Exception as e:
         return f"Execution error: {e}"

studio/backend/core/training/worker.py

@@ -390,11 +390,16 @@ def run_training_process(
     # Some newer architectures (e.g. NemotronH) have config parsing bugs in
     # transformers that require trust_remote_code=True as a workaround.
     # Only auto-enable for unsloth/* prefixed models (trusted source).
+    # Exclude Gemma 4 since it is a native transformers 5.5 model and
+    # trust_remote_code=True would bypass the compiler (disabling fused CE).
     from utils.transformers_version import needs_transformers_5
 
+    _lowered = model_name.lower()
+    _is_native_t5 = any(x in _lowered for x in ("gemma-4", "gemma4"))
     if (
         needs_transformers_5(model_name)
-        and model_name.lower().startswith("unsloth/")
+        and _lowered.startswith("unsloth/")
+        and not _is_native_t5
         and not config.get("trust_remote_code", False)
     ):
         config["trust_remote_code"] = True

studio/backend/main.py

@@ -121,13 +121,13 @@ async def lifespan(app: FastAPI):
     if storage.ensure_default_admin():
         bootstrap_pw = storage.get_bootstrap_password()
         app.state.bootstrap_password = bootstrap_pw
+
+        bootstrap_path = storage.DB_PATH.parent / ".bootstrap_password"
         print("\n" + "=" * 60)
         print("DEFAULT ADMIN ACCOUNT CREATED")
-        print(
-            "Sign in with the seeded credentials and change the password immediately:\n"
-        )
         print(f"    username: {storage.DEFAULT_ADMIN_USERNAME}")
-        print(f"    password: {bootstrap_pw}\n")
+        print(f"    password saved to: {bootstrap_path}")
+        print("    Open the Studio UI to sign in and change it.")
         print("=" * 60 + "\n")
     else:
         app.state.bootstrap_password = storage.get_bootstrap_password()

studio/backend/models/inference.py

@@ -139,6 +139,10 @@ class LoadResponse(BaseModel):
     max_context_length: Optional[int] = Field(
         None, description = "Maximum context length currently available on this hardware"
     )
+    native_context_length: Optional[int] = Field(
+        None,
+        description = "Model's native context length from GGUF metadata (not capped by VRAM)",
+    )
     supports_reasoning: bool = Field(
         False,
         description = "Whether model supports thinking/reasoning mode (enable_thinking)",
@@ -217,6 +221,10 @@ class InferenceStatusResponse(BaseModel):
         None,
         description = "Maximum context length currently available for the active model",
     )
+    native_context_length: Optional[int] = Field(
+        None,
+        description = "Model's native context length from GGUF metadata (not capped by VRAM)",
+    )
 
 
 # =====================================================================

studio/backend/requirements/no-torch-runtime.txt

@@ -28,6 +28,21 @@ peft>=0.18.0,!=0.11.0
 huggingface_hub>=0.34.0
 hf_transfer
 diffusers
+
+# Transitive deps required because this file is installed with --no-deps.
+# Without these, `from transformers import AutoConfig` fails at import time.
+regex
+typing_extensions
+filelock
+httpx
+httpcore
+certifi
+idna
+anyio
+sniffio
+h11
+
+tokenizers
 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.3.0
 trl>=0.18.2,!=0.19.0,<=0.24.0
 sentence-transformers

studio/backend/routes/inference.py

@@ -5,11 +5,12 @@
 Inference API routes for model loading and text generation.
 """
 
+import os
 import sys
 import time
 import uuid
 from pathlib import Path
-from fastapi import APIRouter, Depends, HTTPException, Request
+from fastapi import APIRouter, Depends, HTTPException, Request, status
 from fastapi.responses import StreamingResponse, JSONResponse
 from typing import Optional
 import json
@@ -21,6 +22,9 @@ import threading
 
 import re as _re
 
+# Model size extraction (shared with core/inference/llama_cpp.py)
+from utils.models import extract_model_size_b as _extract_model_size_b
+
 
 def _friendly_error(exc: Exception) -> str:
     """Extract a user-friendly message from known llama-server errors."""
@@ -90,6 +94,14 @@ from datetime import date as _date
 
 router = APIRouter()
 
+# Appended to tool-use nudge to discourage plan-without-action
+_TOOL_ACTION_NUDGE = (
+    " IMPORTANT: Always call tools directly -- never write code yourself."
+    " Never describe what you plan to do -- just call the tool immediately."
+    " For any code request, call the python tool. For any factual question, call web_search."
+    " Do NOT output code blocks -- use the python tool instead."
+)
+
 # Regex for stripping leaked tool-call XML from assistant messages/stream
 _TOOL_XML_RE = _re.compile(
     r"<tool_call>.*?</tool_call>|<function=\w+>.*?</function>",
@@ -163,6 +175,7 @@ async def load_model(
                     inference = inference_config,
                     context_length = llama_backend.context_length,
                     max_context_length = llama_backend.max_context_length,
+                    native_context_length = llama_backend.native_context_length,
                     supports_reasoning = llama_backend.supports_reasoning,
                     reasoning_always_on = llama_backend.reasoning_always_on,
                     chat_template = llama_backend.chat_template,
@@ -298,6 +311,7 @@ async def load_model(
                 inference = inference_config,
                 context_length = llama_backend.context_length,
                 max_context_length = llama_backend.max_context_length,
+                native_context_length = llama_backend.native_context_length,
                 supports_reasoning = llama_backend.supports_reasoning,
                 reasoning_always_on = llama_backend.reasoning_always_on,
                 supports_tools = llama_backend.supports_tools,
@@ -637,6 +651,7 @@ async def get_status(
                 supports_tools = llama_backend.supports_tools,
                 context_length = llama_backend.context_length,
                 max_context_length = llama_backend.max_context_length,
+                native_context_length = llama_backend.native_context_length,
             )
 
         # Otherwise, report Unsloth backend status
@@ -1092,12 +1107,20 @@ async def openai_chat_completions(
 
             _date_line = f"The current date is {_date.today().isoformat()}."
 
-            _web_tips = (
-                "When you search and find a relevant URL in the results, "
-                "fetch its full content by calling web_search with the url parameter. "
-                "Do not repeat the same search query. If a search returns "
-                "no useful results, try rephrasing or fetching a result URL directly."
-            )
+            # Small models (<9B) struggle with multi-step search plans,
+            # so simplify the web tips to avoid plan-then-stall behavior.
+            _model_size_b = _extract_model_size_b(model_name)
+            _is_small_model = _model_size_b is not None and _model_size_b < 9
+
+            if _is_small_model:
+                _web_tips = "Do not repeat the same search query."
+            else:
+                _web_tips = (
+                    "When you search and find a relevant URL in the results, "
+                    "fetch its full content by calling web_search with the url parameter. "
+                    "Do not repeat the same search query. If a search returns "
+                    "no useful results, try rephrasing or fetching a result URL directly."
+                )
             _code_tips = (
                 "Use code execution for math, calculations, data processing, "
                 "or to parse and analyze information from tool results."
@@ -1129,6 +1152,7 @@ async def openai_chat_completions(
                 _nudge = ""
 
             if _nudge:
+                _nudge += _TOOL_ACTION_NUDGE
                 # Append nudge to system prompt (preserve user's prompt)
                 if system_prompt:
                     system_prompt = system_prompt.rstrip() + "\n\n" + _nudge
@@ -1205,7 +1229,14 @@ async def openai_chat_completions(
                             break
 
                         if event["type"] == "status":
+                            # Empty status marks an iteration boundary
+                            # in the GGUF tool loop (e.g. after a
+                            # re-prompt).  Reset the cumulative cursor
+                            # so the next assistant turn streams cleanly.
+                            if not event["text"]:
+                                prev_text = ""
                             # Emit tool status as a custom SSE event
+                            # (including empty ones to clear UI badges)
                             status_data = json.dumps(
                                 {
                                     "type": "tool_status",
@@ -1652,6 +1683,94 @@ async def openai_chat_completions(
             raise HTTPException(status_code = 500, detail = str(e))
 
 
+# =====================================================================
+# Sandbox file serving  (/sandbox/{session_id}/{filename})
+# =====================================================================
+
+_SANDBOX_MEDIA_TYPES = {
+    ".png": "image/png",
+    ".jpg": "image/jpeg",
+    ".jpeg": "image/jpeg",
+    ".gif": "image/gif",
+    ".webp": "image/webp",
+    ".bmp": "image/bmp",
+}
+
+
+@router.get("/sandbox/{session_id}/{filename}")
+async def serve_sandbox_file(
+    session_id: str,
+    filename: str,
+    request: Request,
+    token: Optional[str] = None,
+):
+    """
+    Serve image files created by Python tool execution.
+
+    Accepts auth via Authorization header OR ?token= query param
+    (needed because <img src> cannot send custom headers).
+    """
+    from fastapi.responses import FileResponse
+
+    # ── Authentication (header or query param) ──────────────────
+    auth_header = request.headers.get("authorization")
+    if auth_header and auth_header.lower().startswith("bearer "):
+        jwt_token = auth_header[7:]
+    elif token:
+        jwt_token = token
+    else:
+        raise HTTPException(
+            status_code = status.HTTP_401_UNAUTHORIZED,
+            detail = "Missing authentication token",
+        )
+    from fastapi.security import HTTPAuthorizationCredentials
+
+    creds = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = jwt_token)
+    await get_current_subject(creds)
+
+    # ── Filename sanitization ───────────────────────────────────
+    safe_filename = os.path.basename(filename)
+    if not safe_filename or safe_filename in (".", ".."):
+        raise HTTPException(status_code = 404, detail = "Not found")
+
+    # ── Extension allowlist ─────────────────────────────────────
+    ext = os.path.splitext(safe_filename)[1].lower()
+    media_type = _SANDBOX_MEDIA_TYPES.get(ext)
+    if not media_type:
+        raise HTTPException(
+            status_code = status.HTTP_403_FORBIDDEN,
+            detail = "File type not allowed",
+        )
+
+    # ── Path containment check ──────────────────────────────────
+    home = os.path.expanduser("~")
+    sandbox_root = os.path.realpath(os.path.join(home, "studio_sandbox"))
+    safe_session = os.path.basename(session_id.replace("..", ""))
+    if not safe_session:
+        raise HTTPException(status_code = 404, detail = "Not found")
+
+    file_path = os.path.realpath(
+        os.path.join(sandbox_root, safe_session, safe_filename)
+    )
+    if not file_path.startswith(sandbox_root + os.sep):
+        raise HTTPException(
+            status_code = status.HTTP_403_FORBIDDEN,
+            detail = "Access denied",
+        )
+
+    if not os.path.isfile(file_path):
+        raise HTTPException(status_code = 404, detail = "Not found")
+
+    return FileResponse(
+        path = file_path,
+        media_type = media_type,
+        headers = {
+            "Cache-Control": "private, no-store",
+            "X-Content-Type-Options": "nosniff",
+        },
+    )
+
+
 # =====================================================================
 # OpenAI-Compatible Models Listing  (/models → /v1/models)
 # =====================================================================

studio/backend/tests/test_kv_cache_estimation.py

@@ -0,0 +1,929 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+"""Tests for 5-path architecture-aware KV cache VRAM estimation.
+
+Covers the GGUF metadata parser, _can_estimate_kv gate, all 5 estimation
+paths (MLA, Hybrid Mamba, Sliding Window, Standard GQA, Legacy), KV cache
+quantization, edge cases, and lifecycle (init/unload/reparse).
+
+Requires no GPU, network, or external libraries beyond pytest.
+Cross-platform: Linux, macOS, Windows, WSL.
+"""
+
+import io
+import struct
+import sys
+import types as _types
+from pathlib import Path
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# Stub heavy / unavailable external dependencies before importing the
+# module under test.  Same pattern as test_native_context_length.py.
+# ---------------------------------------------------------------------------
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+# loggers
+_loggers_stub = _types.ModuleType("loggers")
+_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
+sys.modules.setdefault("loggers", _loggers_stub)
+
+# structlog
+_structlog_stub = _types.ModuleType("structlog")
+sys.modules.setdefault("structlog", _structlog_stub)
+
+# httpx
+_httpx_stub = _types.ModuleType("httpx")
+for _exc_name in (
+    "ConnectError",
+    "TimeoutException",
+    "ReadTimeout",
+    "ReadError",
+    "RemoteProtocolError",
+    "CloseError",
+):
+    setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
+
+
+class _FakeTimeout:
+    def __init__(self, *a, **kw):
+        pass
+
+
+_httpx_stub.Timeout = _FakeTimeout
+_httpx_stub.Client = type(
+    "Client",
+    (),
+    {
+        "__init__": lambda self, **kw: None,
+        "__enter__": lambda self: self,
+        "__exit__": lambda self, *a: None,
+    },
+)
+sys.modules.setdefault("httpx", _httpx_stub)
+
+from core.inference.llama_cpp import LlamaCppBackend
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes:
+    """Build a minimal GGUF v3 binary blob with the given KV metadata.
+
+    Only supports UINT32 (type 4), UINT64 (type 10), and STRING (type 8)
+    values, which is all the metadata parser reads.
+    """
+    buf = io.BytesIO()
+    # Header: magic, version, tensor_count, kv_count
+    buf.write(struct.pack("<I", 0x46554747))  # GGUF magic
+    buf.write(struct.pack("<I", 3))  # version 3
+    buf.write(struct.pack("<Q", 0))  # tensor_count
+    buf.write(struct.pack("<Q", len(kv_pairs)))
+
+    for key, val in kv_pairs.items():
+        key_bytes = key.encode("utf-8")
+        buf.write(struct.pack("<Q", len(key_bytes)))
+        buf.write(key_bytes)
+        if isinstance(val, str):
+            buf.write(struct.pack("<I", 8))  # STRING
+            val_bytes = val.encode("utf-8")
+            buf.write(struct.pack("<Q", len(val_bytes)))
+            buf.write(val_bytes)
+        elif isinstance(val, int):
+            if val <= 0xFFFFFFFF:
+                buf.write(struct.pack("<I", 4))  # UINT32
+                buf.write(struct.pack("<I", val))
+            else:
+                buf.write(struct.pack("<I", 10))  # UINT64
+                buf.write(struct.pack("<Q", val))
+        else:
+            raise TypeError(f"Unsupported value type: {type(val)}")
+    return buf.getvalue()
+
+
+def _backend_from_gguf(arch: str, fields: dict) -> LlamaCppBackend:
+    """Create a LlamaCppBackend with parsed GGUF metadata from given fields."""
+    kv = {"general.architecture": arch}
+    for k, v in fields.items():
+        kv[f"{arch}.{k}"] = v
+    import tempfile, os
+
+    data = _make_gguf_bytes(arch, kv)
+    fd, path = tempfile.mkstemp(suffix = ".gguf")
+    try:
+        os.write(fd, data)
+        os.close(fd)
+        b = LlamaCppBackend()
+        b._read_gguf_metadata(path)
+        return b
+    finally:
+        os.unlink(path)
+
+
+# ---------------------------------------------------------------------------
+# A. GGUF Parser Tests
+# ---------------------------------------------------------------------------
+
+
+class TestGGUFParserNewFields:
+    """Verify that the 8 new architecture-aware fields are correctly parsed."""
+
+    @pytest.mark.parametrize(
+        "field,gguf_key,value",
+        [
+            ("_kv_key_length", "attention.key_length", 128),
+            ("_kv_value_length", "attention.value_length", 128),
+            ("_sliding_window", "attention.sliding_window", 1024),
+            ("_full_attention_interval", "full_attention_interval", 4),
+            ("_kv_lora_rank", "attention.kv_lora_rank", 512),
+            ("_key_length_mla", "attention.key_length_mla", 256),
+            ("_ssm_inner_size", "ssm.inner_size", 6144),
+            ("_ssm_state_size", "ssm.state_size", 128),
+        ],
+    )
+    def test_field_parsed(self, field, gguf_key, value):
+        b = _backend_from_gguf("testarch", {gguf_key: value})
+        assert getattr(b, field) == value
+
+    def test_missing_fields_are_none(self):
+        b = _backend_from_gguf("testarch", {"block_count": 10})
+        for attr in [
+            "_kv_key_length",
+            "_kv_value_length",
+            "_sliding_window",
+            "_full_attention_interval",
+            "_kv_lora_rank",
+            "_key_length_mla",
+            "_ssm_inner_size",
+            "_ssm_state_size",
+        ]:
+            assert getattr(b, attr) is None
+
+    def test_all_13_fields_parsed_together(self):
+        fields = {
+            "context_length": 131072,
+            "block_count": 62,
+            "attention.head_count_kv": 16,
+            "attention.head_count": 32,
+            "embedding_length": 5376,
+            "attention.key_length": 128,
+            "attention.value_length": 128,
+            "attention.sliding_window": 1024,
+            "full_attention_interval": 6,
+            "attention.kv_lora_rank": 512,
+            "attention.key_length_mla": 256,
+            "ssm.inner_size": 4096,
+            "ssm.state_size": 128,
+        }
+        b = _backend_from_gguf("testarch", fields)
+        assert b._context_length == 131072
+        assert b._n_layers == 62
+        assert b._n_kv_heads == 16
+        assert b._n_heads == 32
+        assert b._embedding_length == 5376
+        assert b._kv_key_length == 128
+        assert b._kv_value_length == 128
+        assert b._sliding_window == 1024
+        assert b._full_attention_interval == 6
+        assert b._kv_lora_rank == 512
+        assert b._key_length_mla == 256
+        assert b._ssm_inner_size == 4096
+        assert b._ssm_state_size == 128
+
+
+class TestGGUFParserReset:
+    """Verify that fields are properly reset between parses."""
+
+    def test_reset_between_parses(self):
+        # First parse with all fields
+        b = _backend_from_gguf(
+            "arch1",
+            {
+                "block_count": 32,
+                "attention.key_length": 128,
+                "attention.kv_lora_rank": 512,
+                "ssm.inner_size": 4096,
+            },
+        )
+        assert b._kv_key_length == 128
+        assert b._kv_lora_rank == 512
+        assert b._ssm_inner_size == 4096
+
+        # Second parse without those fields -- they should be None
+        kv = {"general.architecture": "arch2", "arch2.block_count": 64}
+        import tempfile, os
+
+        data = _make_gguf_bytes("arch2", kv)
+        fd, path = tempfile.mkstemp(suffix = ".gguf")
+        os.write(fd, data)
+        os.close(fd)
+        try:
+            b._read_gguf_metadata(path)
+        finally:
+            os.unlink(path)
+        assert b._kv_key_length is None
+        assert b._kv_lora_rank is None
+        assert b._ssm_inner_size is None
+        assert b._n_layers == 64
+
+
+# ---------------------------------------------------------------------------
+# B. _can_estimate_kv Gate Tests
+# ---------------------------------------------------------------------------
+
+
+class TestCanEstimateKV:
+    """Verify gate logic for all field combinations."""
+
+    def test_no_layers_returns_false(self):
+        b = LlamaCppBackend()
+        b._n_layers = None
+        b._kv_key_length = 128
+        assert not b._can_estimate_kv()
+
+    def test_explicit_both_dims_sufficient(self):
+        b = LlamaCppBackend()
+        b._n_layers = 32
+        b._kv_key_length = 128
+        b._kv_value_length = 128
+        assert b._can_estimate_kv()
+
+    def test_key_length_alone_insufficient(self):
+        """key_length without value_length should NOT be enough."""
+        b = LlamaCppBackend()
+        b._n_layers = 32
+        b._kv_key_length = 128
+        assert not b._can_estimate_kv()
+
+    def test_kv_lora_rank_sufficient(self):
+        b = LlamaCppBackend()
+        b._n_layers = 61
+        b._kv_lora_rank = 512
+        assert b._can_estimate_kv()
+
+    def test_legacy_embed_plus_heads(self):
+        b = LlamaCppBackend()
+        b._n_layers = 28
+        b._embedding_length = 1024
+        b._n_heads = 16
+        assert b._can_estimate_kv()
+
+    def test_legacy_embed_plus_kv_heads(self):
+        b = LlamaCppBackend()
+        b._n_layers = 28
+        b._embedding_length = 1024
+        b._n_kv_heads = 8
+        assert b._can_estimate_kv()
+
+    def test_legacy_no_embed_returns_false(self):
+        b = LlamaCppBackend()
+        b._n_layers = 28
+        b._n_heads = 16
+        # No embedding_length, no new-style fields
+        assert not b._can_estimate_kv()
+
+    def test_fresh_backend_returns_false(self):
+        b = LlamaCppBackend()
+        assert not b._can_estimate_kv()
+
+
+# ---------------------------------------------------------------------------
+# C. Path 1: MLA Estimation
+# ---------------------------------------------------------------------------
+
+
+class TestMLAEstimation:
+    """MLA: K-only cache using compressed KV latent + RoPE."""
+
+    def _mla_backend(self, **overrides):
+        defaults = {
+            "_n_layers": 61,
+            "_n_kv_heads": 1,
+            "_n_heads": 128,
+            "_embedding_length": 7168,
+            "_kv_key_length": 576,
+            "_kv_value_length": 512,
+            "_kv_lora_rank": 512,
+            "_key_length_mla": 192,
+        }
+        defaults.update(overrides)
+        b = LlamaCppBackend()
+        for k, v in defaults.items():
+            setattr(b, k, v)
+        return b
+
+    def test_deepseek_v3_f16(self):
+        b = self._mla_backend()
+        # 61 layers * 163840 ctx * 1 head * 576 key_len * 2 bpe
+        expected = 61 * 163840 * 1 * 576 * 2
+        assert b._estimate_kv_cache_bytes(163840, "f16") == expected
+
+    def test_mla_ignores_value_length(self):
+        """MLA should NOT add value_length -- V is reconstructed from the latent."""
+        b = self._mla_backend()
+        result = b._estimate_kv_cache_bytes(1000, "f16")
+        # Should be n_layers * ctx * 1 * key_len(576) * 2
+        expected = 61 * 1000 * 1 * 576 * 2
+        assert result == expected
+
+    def test_mla_fallback_when_no_key_length(self):
+        """If key_length is missing, fallback to kv_lora_rank + key_length_mla."""
+        b = self._mla_backend(_kv_key_length = None)
+        # _key_length_mla=192 in default, so rope_dim=192
+        result = b._estimate_kv_cache_bytes(1000, "f16")
+        expected = 61 * 1000 * 1 * (512 + 192) * 2  # 704
+        assert result == expected
+
+    def test_mla_fallback_no_key_length_mla(self):
+        """If both key_length and key_length_mla are missing, fallback to +64."""
+        b = self._mla_backend(_kv_key_length = None, _key_length_mla = None)
+        result = b._estimate_kv_cache_bytes(1000, "f16")
+        expected = 61 * 1000 * 1 * (512 + 64) * 2  # 576
+        assert result == expected
+
+    def test_mla_defaults_n_kv_to_1_when_heads_absent(self):
+        """MLA should use n_kv=1 even if n_kv_heads is None (not n_heads)."""
+        b = self._mla_backend(_n_kv_heads = None)  # n_heads=128 still set
+        result = b._estimate_kv_cache_bytes(1000, "f16")
+        # Should use n_kv_mla=1, NOT n_heads=128
+        expected = 61 * 1000 * 1 * 576 * 2
+        assert result == expected
+
+    def test_mla_q4_quantization(self):
+        b = self._mla_backend()
+        result_f16 = b._estimate_kv_cache_bytes(1000, "f16")
+        result_q4 = b._estimate_kv_cache_bytes(1000, "q4_0")
+        assert result_q4 < result_f16
+        # q4_0 bpe = 0.5625, f16 bpe = 2.0
+        assert result_q4 == int(61 * 1000 * 1 * 576 * 0.5625)
+
+
+# ---------------------------------------------------------------------------
+# D. Path 2: Hybrid Mamba Estimation
+# ---------------------------------------------------------------------------
+
+
+class TestHybridMambaEstimation:
+    """Hybrid Mamba: only attention layers (1 in N) need KV cache."""
+
+    def _hybrid_backend(self, **overrides):
+        defaults = {
+            "_n_layers": 64,
+            "_n_kv_heads": 4,
+            "_n_heads": 24,
+            "_embedding_length": 5120,
+            "_kv_key_length": 256,
+            "_kv_value_length": 256,
+            "_full_attention_interval": 4,
+            "_ssm_inner_size": 6144,
+            "_ssm_state_size": 128,
+        }
+        defaults.update(overrides)
+        b = LlamaCppBackend()
+        for k, v in defaults.items():
+            setattr(b, k, v)
+        return b
+
+    def test_qwen35_27b(self):
+        b = self._hybrid_backend()
+        # n_attn = 64 // 4 = 16
+        expected = 16 * 262144 * 4 * (256 + 256) * 2
+        assert b._estimate_kv_cache_bytes(262144, "f16") == expected
+
+    def test_qwen35_35b_a3b(self):
+        b = self._hybrid_backend(
+            _n_layers = 40,
+            _n_kv_heads = 2,
+            _n_heads = 16,
+            _embedding_length = 2048,
+            _ssm_inner_size = 4096,
+        )
+        # n_attn = 40 // 4 = 10
+        expected = 10 * 262144 * 2 * (256 + 256) * 2
+        assert b._estimate_kv_cache_bytes(262144, "f16") == expected
+
+    def test_hybrid_without_explicit_dims(self):
+        """Fallback to head_dim when key_length/value_length are missing."""
+        b = self._hybrid_backend(_kv_key_length = None, _kv_value_length = None)
+        head_dim = 5120 // 24  # 213
+        expected = 16 * 4096 * 4 * 2 * head_dim * 2
+        assert b._estimate_kv_cache_bytes(4096, "f16") == expected
+
+    def test_fai_zero_safety(self):
+        """full_attention_interval=0 should not cause ZeroDivisionError."""
+        b = self._hybrid_backend(_full_attention_interval = 0)
+        result = b._estimate_kv_cache_bytes(4096, "f16")
+        # fai=0 -> n_attn = n_layers (all layers)
+        expected = 64 * 4096 * 4 * (256 + 256) * 2
+        assert result == expected
+
+
+# ---------------------------------------------------------------------------
+# E. Path 3: Sliding Window Estimation
+# ---------------------------------------------------------------------------
+
+
+class TestSlidingWindowEstimation:
+    """SWA: half global (full ctx) + half sliding window."""
+
+    def _swa_backend(self, **overrides):
+        defaults = {
+            "_n_layers": 62,
+            "_n_kv_heads": 16,
+            "_n_heads": 32,
+            "_embedding_length": 5376,
+            "_kv_key_length": 128,
+            "_kv_value_length": 128,
+            "_sliding_window": 1024,
+        }
+        defaults.update(overrides)
+        b = LlamaCppBackend()
+        for k, v in defaults.items():
+            setattr(b, k, v)
+        return b
+
+    def test_gemma3(self):
+        b = self._swa_backend()
+        # 1/4 heuristic: 62 // 4 = 15 global, 47 SWA
+        n_global = max(1, 62 // 4)  # 15
+        n_swa = 62 - n_global  # 47
+        kv_per = 16 * (128 + 128) * 2
+        expected = int(n_global * 131072 * kv_per + n_swa * min(131072, 1024) * kv_per)
+        assert b._estimate_kv_cache_bytes(131072, "f16") == expected
+
+    def test_gpt_oss(self):
+        b = self._swa_backend(
+            _n_layers = 24,
+            _n_kv_heads = 8,
+            _n_heads = 64,
+            _embedding_length = 2880,
+            _kv_key_length = 64,
+            _kv_value_length = 64,
+            _sliding_window = 128,
+        )
+        # 1/4 heuristic: 24 // 4 = 6 global, 18 SWA
+        n_global = max(1, 24 // 4)  # 6
+        n_swa = 24 - n_global  # 18
+        kv_per = 8 * (64 + 64) * 2
+        expected = int(n_global * 131072 * kv_per + n_swa * min(131072, 128) * kv_per)
+        assert b._estimate_kv_cache_bytes(131072, "f16") == expected
+
+    def test_ctx_smaller_than_window(self):
+        """When context < sliding_window, SWA layers use full context anyway."""
+        b = self._swa_backend(_sliding_window = 8192)
+        n_global = max(1, 62 // 4)  # 15
+        n_swa = 62 - n_global  # 47
+        kv_per = 16 * (128 + 128) * 2
+        ctx = 4096
+        expected = int(n_global * ctx * kv_per + n_swa * min(ctx, 8192) * kv_per)
+        # min(4096, 8192) = 4096, so both pools use full ctx
+        assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
+
+    def test_odd_layer_count(self):
+        """Odd layer count: n_global = max(1, n//4), n_swa = n - n_global."""
+        b = self._swa_backend(_n_layers = 63)
+        n_global = max(1, 63 // 4)  # 15
+        n_swa = 63 - n_global  # 48
+        kv_per = 16 * (128 + 128) * 2
+        expected = int(n_global * 1000 * kv_per + n_swa * min(1000, 1024) * kv_per)
+        assert b._estimate_kv_cache_bytes(1000, "f16") == expected
+
+
+# ---------------------------------------------------------------------------
+# F. Path 4: Standard GQA Estimation
+# ---------------------------------------------------------------------------
+
+
+class TestStandardGQAEstimation:
+    """Standard GQA with explicit key_length/value_length."""
+
+    def _gqa_backend(self, **overrides):
+        defaults = {
+            "_n_layers": 28,
+            "_n_kv_heads": 8,
+            "_n_heads": 16,
+            "_embedding_length": 1024,
+            "_kv_key_length": 128,
+            "_kv_value_length": 128,
+        }
+        defaults.update(overrides)
+        b = LlamaCppBackend()
+        for k, v in defaults.items():
+            setattr(b, k, v)
+        return b
+
+    def test_qwen3_06b(self):
+        b = self._gqa_backend()
+        expected = 28 * 40960 * 8 * (128 + 128) * 2
+        assert b._estimate_kv_cache_bytes(40960, "f16") == expected
+
+    def test_asymmetric_kv_dims(self):
+        """key_length != value_length (some architectures have this)."""
+        b = self._gqa_backend(_kv_key_length = 192, _kv_value_length = 64)
+        expected = 28 * 4096 * 8 * (192 + 64) * 2
+        assert b._estimate_kv_cache_bytes(4096, "f16") == expected
+
+    def test_differs_from_legacy(self):
+        """GQA path should differ from legacy when key_length != embed//n_heads."""
+        b = self._gqa_backend()
+        head_dim = 1024 // 16  # 64
+        gqa_result = b._estimate_kv_cache_bytes(4096, "f16")
+        # Legacy would use: 2 * 8 * 64 * 28 * 4096 * 2
+        legacy_result = int(2 * 8 * head_dim * 28 * 4096 * 2)
+        # GQA: 28 * 4096 * 8 * (128+128) * 2 -- uses actual key_length=128
+        assert gqa_result != legacy_result
+        assert gqa_result > legacy_result  # key_length (128) > head_dim (64)
+
+
+# ---------------------------------------------------------------------------
+# G. Path 5: Legacy Fallback Estimation
+# ---------------------------------------------------------------------------
+
+
+class TestLegacyEstimation:
+    """Legacy: embed // n_heads, for old GGUFs without new fields."""
+
+    def _legacy_backend(self, **overrides):
+        defaults = {
+            "_n_layers": 32,
+            "_n_kv_heads": 8,
+            "_n_heads": 32,
+            "_embedding_length": 4096,
+        }
+        defaults.update(overrides)
+        b = LlamaCppBackend()
+        for k, v in defaults.items():
+            setattr(b, k, v)
+        return b
+
+    def test_basic_legacy(self):
+        b = self._legacy_backend()
+        head_dim = 4096 // 32  # 128
+        expected = int(2 * 8 * 128 * 32 * 4096 * 2)
+        assert b._estimate_kv_cache_bytes(4096, "f16") == expected
+
+    def test_legacy_with_only_n_heads(self):
+        """n_kv_heads is None, falls back to n_heads."""
+        b = self._legacy_backend(_n_kv_heads = None)
+        head_dim = 4096 // 32
+        expected = int(2 * 32 * head_dim * 32 * 4096 * 2)
+        assert b._estimate_kv_cache_bytes(4096, "f16") == expected
+
+    def test_legacy_identical_to_old_formula(self):
+        """Confirm legacy path produces the same result as the pre-PR formula."""
+        b = self._legacy_backend()
+        n_layers = 32
+        n_kv_heads = 8
+        head_dim = 4096 // 32
+        n_ctx = 8192
+        bpe = 2.0
+        old_formula = int(2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe)
+        assert b._estimate_kv_cache_bytes(n_ctx, "f16") == old_formula
+
+
+# ---------------------------------------------------------------------------
+# H. Path Priority (selection order)
+# ---------------------------------------------------------------------------
+
+
+class TestPathPriority:
+    """Confirm: MLA > Hybrid Mamba > SWA > GQA > Legacy."""
+
+    def test_mla_takes_priority_over_all(self):
+        """If kv_lora_rank is set, MLA path is used even if other fields are present."""
+        b = LlamaCppBackend()
+        b._n_layers = 61
+        b._n_kv_heads = 1
+        b._n_heads = 128
+        b._embedding_length = 7168
+        b._kv_key_length = 576
+        b._kv_value_length = 512
+        b._kv_lora_rank = 512
+        b._ssm_inner_size = 4096  # Would trigger Hybrid
+        b._full_attention_interval = 4
+        b._sliding_window = 1024  # Would trigger SWA
+
+        # MLA: 61 * 1000 * 1 * 576 * 2
+        expected_mla = int(61 * 1000 * 1 * 576 * 2)
+        assert b._estimate_kv_cache_bytes(1000, "f16") == expected_mla
+
+    def test_hybrid_over_swa(self):
+        """Hybrid takes priority over SWA when both fields present."""
+        b = LlamaCppBackend()
+        b._n_layers = 64
+        b._n_kv_heads = 4
+        b._n_heads = 24
+        b._embedding_length = 5120
+        b._kv_key_length = 256
+        b._kv_value_length = 256
+        b._ssm_inner_size = 6144
+        b._full_attention_interval = 4
+        b._sliding_window = 1024  # Would trigger SWA
+
+        n_attn = 64 // 4
+        expected_hybrid = int(n_attn * 1000 * 4 * (256 + 256) * 2)
+        assert b._estimate_kv_cache_bytes(1000, "f16") == expected_hybrid
+
+    def test_all_paths_produce_different_values(self):
+        """With carefully chosen params, each path should yield a distinct value."""
+        # Use embedding_length=768 so legacy head_dim (768//16=48) differs from
+        # key_length (256), and MLA key_len (256) != legacy K+V (2*48=96).
+        params = {
+            "_n_layers": 40,
+            "_n_kv_heads": 4,
+            "_n_heads": 16,
+            "_embedding_length": 768,
+            "_kv_key_length": 256,
+            "_kv_value_length": 256,
+        }
+        ctx = 4096
+
+        # Path 4: Standard GQA
+        b_gqa = LlamaCppBackend()
+        for k, v in params.items():
+            setattr(b_gqa, k, v)
+        gqa_val = b_gqa._estimate_kv_cache_bytes(ctx, "f16")
+
+        # Path 1: MLA
+        b_mla = LlamaCppBackend()
+        for k, v in params.items():
+            setattr(b_mla, k, v)
+        b_mla._kv_lora_rank = 512
+        mla_val = b_mla._estimate_kv_cache_bytes(ctx, "f16")
+
+        # Path 2: Hybrid Mamba
+        b_hybrid = LlamaCppBackend()
+        for k, v in params.items():
+            setattr(b_hybrid, k, v)
+        b_hybrid._ssm_inner_size = 4096
+        b_hybrid._full_attention_interval = 4
+        hybrid_val = b_hybrid._estimate_kv_cache_bytes(ctx, "f16")
+
+        # Path 3: SWA
+        b_swa = LlamaCppBackend()
+        for k, v in params.items():
+            setattr(b_swa, k, v)
+        b_swa._sliding_window = 512
+        swa_val = b_swa._estimate_kv_cache_bytes(ctx, "f16")
+
+        # Path 5: Legacy (no key_length/value_length)
+        b_legacy = LlamaCppBackend()
+        b_legacy._n_layers = 40
+        b_legacy._n_kv_heads = 4
+        b_legacy._n_heads = 16
+        b_legacy._embedding_length = 768
+        legacy_val = b_legacy._estimate_kv_cache_bytes(ctx, "f16")
+
+        values = [mla_val, hybrid_val, swa_val, gqa_val, legacy_val]
+        assert len(set(values)) == 5, f"Expected 5 distinct values, got {values}"
+
+
+# ---------------------------------------------------------------------------
+# I. KV Cache Quantization
+# ---------------------------------------------------------------------------
+
+
+class TestQuantization:
+    """Verify all supported cache_type_kv values produce correct scaling."""
+
+    @pytest.mark.parametrize(
+        "cache_type,expected_bpe",
+        [
+            ("f32", 4.0),
+            ("f16", 2.0),
+            ("bf16", 2.0),
+            ("q8_0", 34 / 32),
+            ("q5_1", 0.75),
+            ("q5_0", 0.6875),
+            ("q4_1", 0.625),
+            ("q4_0", 0.5625),
+            ("iq4_nl", 0.5625),
+            (None, 2.0),  # default is f16
+            ("unknown", 2.0),  # unknown falls back to f16
+        ],
+    )
+    def test_quantization_scaling(self, cache_type, expected_bpe):
+        b = LlamaCppBackend()
+        b._n_layers = 10
+        b._n_kv_heads = 1
+        b._n_heads = 8
+        b._embedding_length = 512
+        b._kv_key_length = 64
+        b._kv_value_length = 64
+        result = b._estimate_kv_cache_bytes(1000, cache_type)
+        expected = int(10 * 1000 * 1 * (64 + 64) * expected_bpe)
+        assert result == expected
+
+
+# ---------------------------------------------------------------------------
+# J. Edge Cases
+# ---------------------------------------------------------------------------
+
+
+class TestEdgeCases:
+    """Boundary conditions and degenerate inputs."""
+
+    def test_zero_context(self):
+        b = LlamaCppBackend()
+        b._n_layers = 32
+        b._kv_key_length = 128
+        assert b._estimate_kv_cache_bytes(0, "f16") == 0
+
+    def test_negative_context(self):
+        b = LlamaCppBackend()
+        b._n_layers = 32
+        b._kv_key_length = 128
+        assert b._estimate_kv_cache_bytes(-1, "f16") == 0
+
+    def test_context_of_one(self):
+        b = LlamaCppBackend()
+        b._n_layers = 10
+        b._n_kv_heads = 1
+        b._kv_key_length = 64
+        b._kv_value_length = 64
+        result = b._estimate_kv_cache_bytes(1, "f16")
+        assert result == int(10 * 1 * 1 * (64 + 64) * 2)
+
+    def test_very_large_context(self):
+        """1M context should not overflow or crash."""
+        b = LlamaCppBackend()
+        b._n_layers = 10
+        b._n_kv_heads = 1
+        b._kv_key_length = 128
+        b._kv_value_length = 128
+        result = b._estimate_kv_cache_bytes(1_000_000, "f16")
+        assert result > 0
+        assert isinstance(result, int)
+
+    def test_n_kv_heads_none_falls_to_n_heads(self):
+        b = LlamaCppBackend()
+        b._n_layers = 10
+        b._n_kv_heads = None
+        b._n_heads = 8
+        b._kv_key_length = 64
+        b._kv_value_length = 64
+        result = b._estimate_kv_cache_bytes(100, "f16")
+        expected = int(10 * 100 * 8 * (64 + 64) * 2)
+        assert result == expected
+
+    def test_both_heads_none_falls_to_one(self):
+        b = LlamaCppBackend()
+        b._n_layers = 10
+        b._n_kv_heads = None
+        b._n_heads = None
+        b._kv_key_length = 64
+        b._kv_value_length = 64
+        result = b._estimate_kv_cache_bytes(100, "f16")
+        expected = int(10 * 100 * 1 * (64 + 64) * 2)
+        assert result == expected
+
+
+# ---------------------------------------------------------------------------
+# K. Lifecycle Tests
+# ---------------------------------------------------------------------------
+
+
+class TestLifecycle:
+    """Init, unload, and reparse field management."""
+
+    def test_init_fields_none(self):
+        b = LlamaCppBackend()
+        for attr in [
+            "_kv_key_length",
+            "_kv_value_length",
+            "_sliding_window",
+            "_full_attention_interval",
+            "_kv_lora_rank",
+            "_key_length_mla",
+            "_ssm_inner_size",
+            "_ssm_state_size",
+        ]:
+            assert getattr(b, attr) is None
+
+    def test_unload_resets_fields(self):
+        b = LlamaCppBackend()
+        b._n_layers = 32
+        b._kv_key_length = 128
+        b._kv_lora_rank = 512
+        b._sliding_window = 1024
+        b._ssm_inner_size = 4096
+        b._full_attention_interval = 4
+        b.unload_model()
+        for attr in [
+            "_kv_key_length",
+            "_kv_value_length",
+            "_sliding_window",
+            "_full_attention_interval",
+            "_kv_lora_rank",
+            "_key_length_mla",
+            "_ssm_inner_size",
+            "_ssm_state_size",
+        ]:
+            assert getattr(b, attr) is None
+
+    def test_end_to_end_synthetic_mla(self):
+        """Full round-trip: write GGUF -> parse -> estimate."""
+        b = _backend_from_gguf(
+            "deepseek2",
+            {
+                "context_length": 163840,
+                "block_count": 61,
+                "attention.head_count_kv": 1,
+                "attention.head_count": 128,
+                "embedding_length": 7168,
+                "attention.key_length": 576,
+                "attention.value_length": 512,
+                "attention.kv_lora_rank": 512,
+                "attention.key_length_mla": 192,
+            },
+        )
+        assert b._can_estimate_kv()
+        result = b._estimate_kv_cache_bytes(163840, "f16")
+        expected = 61 * 163840 * 1 * 576 * 2
+        assert result == expected
+
+    def test_end_to_end_synthetic_hybrid(self):
+        b = _backend_from_gguf(
+            "qwen35",
+            {
+                "context_length": 262144,
+                "block_count": 64,
+                "attention.head_count_kv": 4,
+                "attention.head_count": 24,
+                "embedding_length": 5120,
+                "attention.key_length": 256,
+                "attention.value_length": 256,
+                "full_attention_interval": 4,
+                "ssm.inner_size": 6144,
+                "ssm.state_size": 128,
+            },
+        )
+        assert b._can_estimate_kv()
+        result = b._estimate_kv_cache_bytes(262144, "f16")
+        n_attn = 64 // 4
+        expected = n_attn * 262144 * 4 * (256 + 256) * 2
+        assert result == expected
+
+    def test_end_to_end_synthetic_swa(self):
+        b = _backend_from_gguf(
+            "gemma3",
+            {
+                "context_length": 131072,
+                "block_count": 62,
+                "attention.head_count_kv": 16,
+                "attention.head_count": 32,
+                "embedding_length": 5376,
+                "attention.key_length": 128,
+                "attention.value_length": 128,
+                "attention.sliding_window": 1024,
+            },
+        )
+        assert b._can_estimate_kv()
+        result = b._estimate_kv_cache_bytes(131072, "f16")
+        n_global = max(1, 62 // 4)  # 15
+        n_swa = 62 - n_global  # 47
+        kv_per = 16 * 256 * 2
+        expected = int(n_global * 131072 * kv_per + n_swa * 1024 * kv_per)
+        assert result == expected
+
+    def test_end_to_end_synthetic_gqa(self):
+        b = _backend_from_gguf(
+            "qwen3",
+            {
+                "context_length": 40960,
+                "block_count": 28,
+                "attention.head_count_kv": 8,
+                "attention.head_count": 16,
+                "embedding_length": 1024,
+                "attention.key_length": 128,
+                "attention.value_length": 128,
+            },
+        )
+        assert b._can_estimate_kv()
+        result = b._estimate_kv_cache_bytes(40960, "f16")
+        expected = 28 * 40960 * 8 * 256 * 2
+        assert result == expected
+
+    def test_end_to_end_synthetic_legacy(self):
+        b = _backend_from_gguf(
+            "llama",
+            {
+                "context_length": 4096,
+                "block_count": 32,
+                "attention.head_count_kv": 8,
+                "attention.head_count": 32,
+                "embedding_length": 4096,
+            },
+        )
+        assert b._can_estimate_kv()
+        result = b._estimate_kv_cache_bytes(4096, "f16")
+        head_dim = 4096 // 32
+        expected = int(2 * 8 * head_dim * 32 * 4096 * 2)
+        assert result == expected

studio/backend/tests/test_native_context_length.py

@@ -0,0 +1,518 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Tests for the native_context_length feature (PR #4746).
+
+Verifies that the new `native_context_length` property on LlamaCppBackend
+and the corresponding Pydantic model fields work correctly.  The raw GGUF
+`_context_length` must never be overwritten by VRAM-capping logic.
+
+Requires no GPU, network, or external libraries beyond pytest and pydantic.
+"""
+
+import io
+import json
+import struct
+import sys
+import types as _types
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# Stub heavy / unavailable external dependencies before importing the
+# module under test.  Same pattern as test_kv_cache_estimation.py.
+# ---------------------------------------------------------------------------
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+# loggers
+_loggers_stub = _types.ModuleType("loggers")
+_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
+sys.modules.setdefault("loggers", _loggers_stub)
+
+# structlog
+_structlog_stub = _types.ModuleType("structlog")
+sys.modules.setdefault("structlog", _structlog_stub)
+
+# httpx -- stub only the names referenced at import / class-definition time
+_httpx_stub = _types.ModuleType("httpx")
+for _exc_name in (
+    "ConnectError",
+    "TimeoutException",
+    "ReadTimeout",
+    "ReadError",
+    "RemoteProtocolError",
+    "CloseError",
+):
+    setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
+
+
+class _FakeTimeout:
+    def __init__(self, *a, **kw):
+        pass
+
+
+_httpx_stub.Timeout = _FakeTimeout
+_httpx_stub.Client = type(
+    "Client",
+    (),
+    {
+        "__init__": lambda self, **kw: None,
+        "__enter__": lambda self: self,
+        "__exit__": lambda self, *a: None,
+    },
+)
+sys.modules.setdefault("httpx", _httpx_stub)
+
+from core.inference.llama_cpp import LlamaCppBackend
+from models.inference import LoadResponse, InferenceStatusResponse
+
+
+# ── Helpers ──────────────────────────────────────────────────────────
+
+
+def _write_kv(buf: io.BytesIO, key: str, value, vtype: int) -> None:
+    """Append a single GGUF KV pair to *buf*."""
+    key_bytes = key.encode("utf-8")
+    buf.write(struct.pack("<Q", len(key_bytes)))
+    buf.write(key_bytes)
+    buf.write(struct.pack("<I", vtype))
+    if vtype == 4:  # UINT32
+        buf.write(struct.pack("<I", value))
+    elif vtype == 10:  # UINT64
+        buf.write(struct.pack("<Q", value))
+    elif vtype == 8:  # STRING
+        val_bytes = value.encode("utf-8")
+        buf.write(struct.pack("<Q", len(val_bytes)))
+        buf.write(val_bytes)
+    else:
+        raise ValueError(f"Unsupported vtype in test helper: {vtype}")
+
+
+def make_gguf(
+    tmp_path: Path,
+    arch: str,
+    kvs: list,
+    *,
+    arch_first: bool = True,
+    filename: str = "test.gguf",
+) -> str:
+    """Create a minimal valid GGUF v3 binary in *tmp_path*."""
+    buf = io.BytesIO()
+    buf.write(struct.pack("<I", 0x46554747))  # GGUF magic
+    buf.write(struct.pack("<I", 3))  # version 3
+    buf.write(struct.pack("<Q", 0))  # tensor count = 0
+
+    ordered = []
+    arch_entry = ("general.architecture", arch, 8)
+
+    if arch_first:
+        ordered.append(arch_entry)
+    for suffix, val, vt in kvs:
+        ordered.append((f"{arch}.{suffix}", val, vt))
+    if not arch_first:
+        ordered.append(arch_entry)
+
+    buf.write(struct.pack("<Q", len(ordered)))
+    for key, val, vt in ordered:
+        _write_kv(buf, key, val, vt)
+
+    path = tmp_path / filename
+    path.write_bytes(buf.getvalue())
+    return str(path)
+
+
+@pytest.fixture
+def backend():
+    """Create a fresh LlamaCppBackend with side effects disabled."""
+    with patch.object(LlamaCppBackend, "_kill_orphaned_servers"):
+        with patch("atexit.register"):
+            return LlamaCppBackend()
+
+
+# =====================================================================
+# A. TestNativeContextLengthProperty -- the new property
+# =====================================================================
+
+
+class TestNativeContextLengthProperty:
+    """Tests the new `native_context_length` property on LlamaCppBackend."""
+
+    def test_none_on_fresh_backend(self, backend):
+        """Returns None when no model loaded."""
+        assert backend.native_context_length is None
+
+    def test_returns_raw_gguf_value(self, backend):
+        """Directly returns _context_length when set."""
+        backend._context_length = 131072
+        assert backend.native_context_length == 131072
+
+    def test_not_capped_by_effective(self, backend):
+        """native_context_length ignores _effective_context_length."""
+        backend._context_length = 131072
+        backend._effective_context_length = 32768
+        assert backend.native_context_length == 131072
+
+    def test_not_capped_by_max(self, backend):
+        """native_context_length ignores _max_context_length."""
+        backend._context_length = 131072
+        backend._max_context_length = 65536
+        assert backend.native_context_length == 131072
+
+    def test_none_after_unload(self, backend):
+        """After unload_model(), returns None."""
+        backend._context_length = 131072
+        assert backend.native_context_length == 131072
+        backend.unload_model()
+        assert backend.native_context_length is None
+
+    def test_after_gguf_parse(self, tmp_path, backend):
+        """Synthetic GGUF with context_length=16384 populates the property."""
+        path = make_gguf(
+            tmp_path,
+            "llama",
+            [("context_length", 16384, 4)],
+        )
+        backend._read_gguf_metadata(path)
+        assert backend.native_context_length == 16384
+
+    def test_resets_between_parses(self, tmp_path, backend):
+        """Second GGUF without context_length resets native to None."""
+        path_a = make_gguf(
+            tmp_path,
+            "llama",
+            [("context_length", 16384, 4)],
+            filename = "a.gguf",
+        )
+        backend._read_gguf_metadata(path_a)
+        assert backend.native_context_length == 16384
+
+        path_b = make_gguf(
+            tmp_path,
+            "gpt2",
+            [("block_count", 12, 4)],
+            filename = "b.gguf",
+        )
+        backend._read_gguf_metadata(path_b)
+        assert backend.native_context_length is None
+
+
+# =====================================================================
+# B. TestContextValueSeparation -- core invariant
+# =====================================================================
+
+
+class TestContextValueSeparation:
+    """_context_length is never overwritten by VRAM logic."""
+
+    def test_preserved_after_effective_set(self, backend):
+        """Setting _effective_context_length does not change _context_length."""
+        backend._context_length = 131072
+        backend._effective_context_length = 32768
+        assert backend._context_length == 131072
+        assert backend.native_context_length == 131072
+
+    def test_ordering_when_capped(self, backend):
+        """native >= max >= effective holds when VRAM-capped."""
+        backend._context_length = 131072
+        backend._max_context_length = 65536
+        backend._effective_context_length = 32768
+        assert backend.native_context_length >= backend.max_context_length
+        assert backend.max_context_length >= backend.context_length
+
+    def test_all_equal_when_uncapped(self, backend):
+        """All three equal when no VRAM constraint."""
+        backend._context_length = 8192
+        # No effective or max set -- properties fall back to _context_length
+        assert backend.native_context_length == 8192
+        assert backend.max_context_length == 8192
+        assert backend.context_length == 8192
+
+    def test_fit_context_does_not_modify(self, backend):
+        """_fit_context_to_vram() does not touch _context_length."""
+        backend._context_length = 131072
+        backend._n_layers = 32
+        backend._n_kv_heads = 8
+        backend._n_heads = 32
+        backend._embedding_length = 4096
+        original = backend._context_length
+
+        # Simulate a very small VRAM budget that forces capping
+        result = backend._fit_context_to_vram(
+            requested_ctx = 131072,
+            available_mib = 512,  # very small
+            model_size_bytes = 0,
+        )
+        # _fit_context_to_vram returns the capped value, not modifying _context_length
+        assert backend._context_length == original
+        assert backend.native_context_length == original
+        # The returned capped value should be <= requested
+        assert result <= 131072
+
+    def test_native_gt_context_when_capped(self, backend):
+        """native_context_length > context_length after VRAM capping."""
+        backend._context_length = 131072
+        backend._effective_context_length = 16384
+        assert backend.native_context_length > backend.context_length
+
+
+# =====================================================================
+# C. TestPydanticModels -- LoadResponse & InferenceStatusResponse
+# =====================================================================
+
+
+class TestPydanticModels:
+    """Tests native_context_length field on Pydantic models."""
+
+    def test_load_response_has_field(self):
+        """Field exists in LoadResponse.model_fields."""
+        assert "native_context_length" in LoadResponse.model_fields
+
+    def test_load_response_defaults_none(self):
+        """Omitting native_context_length defaults to None."""
+        resp = LoadResponse(
+            status = "loaded",
+            model = "test",
+            display_name = "Test",
+            inference = {},
+        )
+        assert resp.native_context_length is None
+
+    def test_load_response_accepts_int(self):
+        """native_context_length=131072 stores correctly."""
+        resp = LoadResponse(
+            status = "loaded",
+            model = "test",
+            display_name = "Test",
+            inference = {},
+            native_context_length = 131072,
+        )
+        assert resp.native_context_length == 131072
+
+    def test_load_response_json_null(self):
+        """None serializes to JSON null."""
+        resp = LoadResponse(
+            status = "loaded",
+            model = "test",
+            display_name = "Test",
+            inference = {},
+        )
+        data = json.loads(resp.model_dump_json())
+        assert data["native_context_length"] is None
+
+    def test_load_response_json_int(self):
+        """131072 serializes to JSON number."""
+        resp = LoadResponse(
+            status = "loaded",
+            model = "test",
+            display_name = "Test",
+            inference = {},
+            native_context_length = 131072,
+        )
+        data = json.loads(resp.model_dump_json())
+        assert data["native_context_length"] == 131072
+
+    def test_status_response_has_field(self):
+        """Field exists in InferenceStatusResponse.model_fields."""
+        assert "native_context_length" in InferenceStatusResponse.model_fields
+
+    def test_status_response_defaults_none(self):
+        """Omitting native_context_length defaults to None."""
+        resp = InferenceStatusResponse()
+        assert resp.native_context_length is None
+
+    def test_roundtrip_preserves_value(self):
+        """model_validate_json(model_dump_json()) round-trips."""
+        resp = LoadResponse(
+            status = "loaded",
+            model = "test",
+            display_name = "Test",
+            inference = {},
+            native_context_length = 131072,
+        )
+        roundtripped = LoadResponse.model_validate_json(resp.model_dump_json())
+        assert roundtripped.native_context_length == 131072
+
+
+# =====================================================================
+# D. TestRouteCompleteness -- source-level verification
+# =====================================================================
+
+
+class TestRouteCompleteness:
+    """All response construction sites in routes/inference.py include native_context_length."""
+
+    @pytest.fixture(autouse = True)
+    def _load_source(self):
+        """Read routes/inference.py source once."""
+        routes_path = Path(__file__).resolve().parent.parent / "routes" / "inference.py"
+        self._source = routes_path.read_text()
+
+    def _find_construction_blocks(self, class_name: str) -> list[str]:
+        """Extract all code blocks that construct a given response class."""
+        blocks = []
+        idx = 0
+        while True:
+            start = self._source.find(f"{class_name}(", idx)
+            if start == -1:
+                break
+            # Find matching closing paren (simple depth counter)
+            depth = 0
+            end = start
+            for i, ch in enumerate(self._source[start:], start):
+                if ch == "(":
+                    depth += 1
+                elif ch == ")":
+                    depth -= 1
+                    if depth == 0:
+                        end = i + 1
+                        break
+            blocks.append(self._source[start:end])
+            idx = end
+        return blocks
+
+    def test_gguf_load_responses_have_field(self):
+        """Every GGUF LoadResponse (is_gguf = True) includes native_context_length."""
+        blocks = self._find_construction_blocks("LoadResponse")
+        gguf_blocks = [
+            b for b in blocks if "is_gguf = True" in b or "is_gguf=True" in b
+        ]
+        assert (
+            len(gguf_blocks) >= 2
+        ), f"Expected at least 2 GGUF LoadResponse blocks, found {len(gguf_blocks)}"
+        for i, block in enumerate(gguf_blocks):
+            assert (
+                "native_context_length" in block
+            ), f"GGUF LoadResponse block #{i} missing native_context_length:\n{block[:200]}"
+
+    def test_non_gguf_load_responses_omit_field(self):
+        """Non-GGUF LoadResponse blocks do not set native_context_length (defaults to None)."""
+        blocks = self._find_construction_blocks("LoadResponse")
+        non_gguf = [
+            b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b
+        ]
+        # Non-GGUF paths should not reference native_context_length
+        # (Pydantic defaults it to None, so not setting it is correct)
+        for block in non_gguf:
+            assert (
+                "native_context_length" not in block
+            ), f"Non-GGUF LoadResponse should not set native_context_length:\n{block[:200]}"
+
+    def test_status_path(self):
+        """InferenceStatusResponse construction with llama_backend has the field."""
+        blocks = self._find_construction_blocks("InferenceStatusResponse")
+        found = False
+        for block in blocks:
+            if "llama_backend" in block and "native_context_length" in block:
+                found = True
+                break
+        assert found, "No InferenceStatusResponse block with llama_backend has native_context_length"
+
+
+# =====================================================================
+# E. TestEdgeCases
+# =====================================================================
+
+
+class TestNativeContextEdgeCases:
+    """Edge cases for native_context_length."""
+
+    def test_context_length_zero(self, tmp_path, backend):
+        """GGUF context_length=0 returns 0, not None."""
+        path = make_gguf(tmp_path, "llama", [("context_length", 0, 4)])
+        backend._read_gguf_metadata(path)
+        assert backend.native_context_length == 0
+
+    def test_context_length_uint32_max(self, tmp_path, backend):
+        """2^32 - 1 survives without truncation."""
+        val = 2**32 - 1
+        path = make_gguf(tmp_path, "llama", [("context_length", val, 4)])
+        backend._read_gguf_metadata(path)
+        assert backend.native_context_length == val
+
+    def test_context_length_uint64(self, tmp_path, backend):
+        """UINT64 type context_length parsed correctly."""
+        val = 2**33  # exceeds UINT32 range
+        path = make_gguf(tmp_path, "llama", [("context_length", val, 10)])
+        backend._read_gguf_metadata(path)
+        assert backend.native_context_length == val
+
+    def test_no_context_length_in_gguf(self, tmp_path, backend):
+        """GGUF without context_length key yields None."""
+        path = make_gguf(tmp_path, "llama", [("block_count", 32, 4)])
+        backend._read_gguf_metadata(path)
+        assert backend.native_context_length is None
+
+    def test_native_equals_context_when_uncapped(self, backend):
+        """Both equal when no VRAM cap applied."""
+        backend._context_length = 8192
+        assert backend.native_context_length == backend.context_length
+
+    def test_native_survives_parse_then_cap(self, tmp_path, backend):
+        """Parse then set effective cap: native unchanged."""
+        path = make_gguf(
+            tmp_path,
+            "llama",
+            [
+                ("context_length", 131072, 4),
+                ("block_count", 32, 4),
+                ("attention.head_count", 32, 4),
+                ("attention.head_count_kv", 8, 4),
+                ("embedding_length", 4096, 4),
+            ],
+        )
+        backend._read_gguf_metadata(path)
+        assert backend.native_context_length == 131072
+
+        # Simulate VRAM capping by setting effective and max
+        backend._effective_context_length = 16384
+        backend._max_context_length = 32768
+        assert backend.native_context_length == 131072
+
+
+# =====================================================================
+# F. TestCrossPlatform -- binary I/O and serialization
+# =====================================================================
+
+
+class TestCrossPlatform:
+    """Binary I/O and serialization correctness across platforms."""
+
+    def test_le_uint32_context_length(self, tmp_path, backend):
+        """Little-endian UINT32 parsed correctly."""
+        path = make_gguf(tmp_path, "llama", [("context_length", 16384, 4)])
+        backend._read_gguf_metadata(path)
+        assert backend.native_context_length == 16384
+
+    def test_le_uint64_context_length(self, tmp_path, backend):
+        """Little-endian UINT64 parsed correctly."""
+        path = make_gguf(tmp_path, "llama", [("context_length", 16384, 10)])
+        backend._read_gguf_metadata(path)
+        assert backend.native_context_length == 16384
+
+    def test_gguf_magic_le_byte_order(self, tmp_path):
+        """Magic 0x46554747 matches GGUF spec (little-endian 'GGUF')."""
+        path = tmp_path / "magic_check.gguf"
+        buf = io.BytesIO()
+        buf.write(struct.pack("<I", 0x46554747))
+        raw = buf.getvalue()
+        # 'G' = 0x47, 'G' = 0x47, 'U' = 0x55, 'F' = 0x46
+        assert raw == b"GGUF"
+
+    def test_json_serialization_deterministic(self):
+        """model_dump_json() is consistent across calls."""
+        resp = LoadResponse(
+            status = "loaded",
+            model = "test",
+            display_name = "Test",
+            inference = {},
+            native_context_length = 131072,
+        )
+        json1 = resp.model_dump_json()
+        json2 = resp.model_dump_json()
+        assert json1 == json2
+        assert '"native_context_length":131072' in json1

studio/backend/utils/inference/inference_config.py

@@ -17,6 +17,7 @@ import structlog
 from loggers import get_logger
 
 from utils.models.model_config import load_model_defaults
+from utils.paths import is_local_path, normalize_path
 
 logger = get_logger(__name__)
 
@@ -93,8 +94,28 @@ def _has_specific_yaml(model_identifier: str) -> bool:
     if model_identifier.lower() in _REVERSE_MODEL_MAPPING:
         return True
 
-    # Check for exact filename match
-    model_filename = model_identifier.replace("/", "_") + ".yaml"
+    # For local filesystem paths (e.g. C:\Users\...\model on Windows),
+    # normalize backslashes so Path().parts splits correctly on POSIX/WSL,
+    # then try matching the last 1-2 path components against the registry
+    # (mirrors the logic in load_model_defaults).
+    _is_local = is_local_path(model_identifier)
+    _normalized = normalize_path(model_identifier) if _is_local else model_identifier
+
+    if _is_local:
+        parts = Path(_normalized).parts
+        for depth in (2, 1):
+            if len(parts) >= depth:
+                suffix = "/".join(parts[-depth:])
+                if suffix.lower() in _REVERSE_MODEL_MAPPING:
+                    return True
+        _lookup = Path(_normalized).name
+    else:
+        _lookup = model_identifier
+
+    # Check for exact filename match (basename for local paths to avoid
+    # passing absolute paths into rglob which raises
+    # "Non-relative patterns are unsupported" on Windows).
+    model_filename = _lookup.replace("/", "_") + ".yaml"
     for config_path in defaults_dir.rglob(model_filename):
         if config_path.is_file():
             return True

studio/backend/utils/models/__init__.py

@@ -19,6 +19,7 @@ from .model_config import (
     get_base_model_from_lora,
     load_model_config,
     list_gguf_variants,
+    extract_model_size_b,
     MODEL_NAME_MAPPING,
     UI_STATUS_INDICATORS,
 )
@@ -38,6 +39,7 @@ __all__ = [
     "get_base_model_from_lora",
     "load_model_config",
     "list_gguf_variants",
+    "extract_model_size_b",
     "MODEL_NAME_MAPPING",
     "UI_STATUS_INDICATORS",
     "scan_checkpoints",

studio/backend/utils/models/model_config.py

@@ -31,6 +31,37 @@ import yaml
 
 logger = get_logger(__name__)
 
+# ── Model size extraction ────────────────────────────────────
+import re as _re
+
+_MODEL_SIZE_RE = _re.compile(
+    r"(?:^|[-_/])(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE
+)
+# MoE active-parameter pattern: matches "A3B", "A3.5B", etc.
+_ACTIVE_SIZE_RE = _re.compile(
+    r"(?:^|[-_/])a(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE
+)
+
+
+def extract_model_size_b(model_id: str) -> float | None:
+    """Extract model size in billions from a model identifier.
+
+    Prefers MoE active-parameter notation (e.g. ``A3B`` in
+    ``Qwen3.5-35B-A3B``) over the total parameter count.
+    Handles both ``B`` (billions) and ``M`` (millions) suffixes.
+    """
+    mid = (model_id or "").lower()
+    active = _ACTIVE_SIZE_RE.search(mid)
+    if active:
+        val = float(active.group(1))
+        return val / 1000.0 if active.group(2).lower() == "m" else val
+    size = _MODEL_SIZE_RE.search(mid)
+    if not size:
+        return None
+    val = float(size.group(1))
+    return val / 1000.0 if size.group(2).lower() == "m" else val
+
+
 # Model name mapping: maps all equivalent model names to their canonical YAML config file
 # Format: "canonical_model_name.yaml": [list of all equivalent model names]
 # Based on the model mapper provided - canonical filename is based on the first model name in the mapper
@@ -1420,17 +1451,20 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
                         return config
 
         # If model_name is a local path (e.g. /home/.../Spark-TTS-0.5B/LLM from
-        # adapter_config.json), try matching the last 1-2 path components against
-        # the registry (e.g. "Spark-TTS-0.5B/LLM").
-        if model_name not in _REVERSE_MODEL_MAPPING and (
-            model_name.startswith("/") or model_name.startswith(".")
-        ):
-            parts = Path(model_name).parts
+        # adapter_config.json, or C:\Users\...\model on Windows), try matching
+        # the last 1-2 path components against the registry
+        # (e.g. "Spark-TTS-0.5B/LLM").
+        _is_local_path = is_local_path(model_name)
+        # Normalize Windows backslash paths so Path().parts splits correctly
+        # on POSIX/WSL hosts (pathlib treats backslashes as literals on Linux).
+        _normalized = normalize_path(model_name) if _is_local_path else model_name
+        if model_name.lower() not in _REVERSE_MODEL_MAPPING and _is_local_path:
+            parts = Path(_normalized).parts
             for depth in [2, 1]:
                 if len(parts) >= depth:
                     suffix = "/".join(parts[-depth:])
-                    if suffix in _REVERSE_MODEL_MAPPING:
-                        canonical_file = _REVERSE_MODEL_MAPPING[suffix]
+                    if suffix.lower() in _REVERSE_MODEL_MAPPING:
+                        canonical_file = _REVERSE_MODEL_MAPPING[suffix.lower()]
                         for config_path in defaults_dir.rglob(canonical_file):
                             if config_path.is_file():
                                 with open(config_path, "r", encoding = "utf-8") as f:
@@ -1440,8 +1474,12 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
                                     )
                                     return config
 
-        # Try exact model name match (for backward compatibility)
-        model_filename = model_name.replace("/", "_") + ".yaml"
+        # Try exact model name match (for backward compatibility).
+        # For local filesystem paths, use only the directory basename to
+        # avoid passing absolute paths (e.g. C:\...) into rglob which
+        # raises "Non-relative patterns are unsupported" on Windows.
+        _lookup_name = Path(_normalized).name if _is_local_path else model_name
+        model_filename = _lookup_name.replace("/", "_") + ".yaml"
         # Search in subfolders and root
         for config_path in defaults_dir.rglob(model_filename):
             if config_path.is_file():

studio/backend/utils/transformers_version.py

@@ -47,6 +47,8 @@ TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = (
     "qwen3.5",  # Qwen3.5 family (35B-A3B, etc.)
     "qwen3-next",  # Qwen3-Next and variants
     "tiny_qwen3_moe",  # imdatta0/tiny_qwen3_moe_2.8B_0.7B
+    "gemma-4",  # Gemma-4 (E2B-it, E4B-it, 31B-it, 26B-A4B-it)
+    "gemma4",  # Gemma-4 alternate naming
 )
 
 # Tokenizer classes that only exist in transformers>=5.x
@@ -58,7 +60,7 @@ _TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = {
 _tokenizer_class_cache: dict[str, bool] = {}
 
 # Versions
-TRANSFORMERS_5_VERSION = "5.3.0"
+TRANSFORMERS_5_VERSION = "5.5.0"
 TRANSFORMERS_DEFAULT_VERSION = "4.57.6"
 
 # Pre-installed directory for transformers 5.x — created by setup.sh / setup.ps1
@@ -258,7 +260,7 @@ def _purge_modules() -> int:
 
 _VENV_T5_PACKAGES = (
     f"transformers=={TRANSFORMERS_5_VERSION}",
-    "huggingface_hub==1.7.1",
+    "huggingface_hub==1.8.0",
     "hf_xet==1.4.2",
     "tiktoken",
 )

studio/frontend/src/components/assistant-ui/markdown-text.tsx

@@ -100,20 +100,27 @@ function getCodeFilename(language: string | null) {
 function isSvgFence(codeFence: CodeFence): boolean {
   const lang = codeFence.language?.toLowerCase() ?? "";
   if (lang === "svg") return true;
-  if ((lang === "xml" || lang === "html") && codeFence.source.trimStart().startsWith("<svg")) return true;
+  if (lang === "xml" || lang === "html") {
+    const trimmed = codeFence.source.trimStart();
+    // Match <svg directly or <?xml ...?> followed by <svg
+    if (trimmed.startsWith("<svg")) return true;
+    if (trimmed.startsWith("<?xml") && trimmed.includes("<svg")) return true;
+  }
   return false;
 }
 
 function isHtmlFence(codeFence: CodeFence): boolean {
   const lang = codeFence.language?.toLowerCase() ?? "";
-  return lang === "html" && !codeFence.source.trimStart().startsWith("<svg");
+  return lang === "html" && !isSvgFence(codeFence);
 }
 
 const UNSAFE_SVG_RE = /<script[\s>]|on\w+\s*=|javascript:|<foreignObject[\s>]|<iframe[\s>]|<embed[\s>]|<object[\s>]/i;
 
 function sanitizeSvg(source: string): string | null {
   if (UNSAFE_SVG_RE.test(source)) return null;
-  return source;
+  // Strip XML declaration (<?xml ...?>) -- not needed for data URI
+  // rendering and can cause issues with some renderers.
+  return source.replace(/^\s*<\?xml[^?]*\?>\s*/i, "");
 }
 
 function SvgPreview({ source }: { source: string }) {
@@ -403,7 +410,7 @@ const MarkdownTextImpl = () => {
   }
 
   return (
-    <div data-status={status.type}>
+    <div data-status={status.type} className="min-w-0 max-w-full">
       <Streamdown
         mode="streaming"
         isAnimating={status.type === "running"}

studio/frontend/src/components/assistant-ui/model-selector.tsx

@@ -34,6 +34,7 @@ interface ModelSelectorProps {
   activeGgufVariant?: string | null;
   onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void;
   onEject?: () => void;
+  onFoldersChange?: () => void;
   variant?: "outline" | "ghost" | "muted";
   size?: "sm" | "default" | "lg";
   className?: string;
@@ -100,6 +101,7 @@ function ModelSelectorContent({
   value,
   onSelect,
   onEject,
+  onFoldersChange,
   className,
   dataTour,
 }: {
@@ -108,6 +110,7 @@ function ModelSelectorContent({
   value?: string;
   onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
   onEject?: () => void;
+  onFoldersChange?: () => void;
   className?: string;
   dataTour?: string;
 }) {
@@ -124,7 +127,7 @@ function ModelSelectorContent({
       )}
     >
       {chatOnly ? (
-        <HubModelPicker models={models} value={value} onSelect={onSelect} />
+        <HubModelPicker models={models} value={value} onSelect={onSelect} onFoldersChange={onFoldersChange} />
       ) : (
         <Tabs defaultValue="hub" className="w-full">
           <TabsList className="mb-2 w-full">
@@ -133,7 +136,7 @@ function ModelSelectorContent({
           </TabsList>
 
           <TabsContent value="hub" className="m-0">
-            <HubModelPicker models={models} value={value} onSelect={onSelect} />
+            <HubModelPicker models={models} value={value} onSelect={onSelect} onFoldersChange={onFoldersChange} />
           </TabsContent>
 
           <TabsContent value="lora" className="m-0">
@@ -171,6 +174,7 @@ export function ModelSelector({
   activeGgufVariant,
   onValueChange,
   onEject,
+  onFoldersChange,
   variant = "outline",
   size = "default",
   className,
@@ -253,6 +257,7 @@ export function ModelSelector({
         value={selected}
         onSelect={handleSelect}
         onEject={onEject ? handleEject : undefined}
+        onFoldersChange={onFoldersChange}
         className={contentClassName}
         dataTour={contentDataTour}
       />

studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx

@@ -20,11 +20,15 @@ import {
 } from "@/components/ui/tooltip";
 import { usePlatformStore } from "@/config/env";
 import {
+  type ScanFolderInfo,
+  addScanFolder,
   deleteCachedModel,
   listCachedGguf,
   listCachedModels,
   listGgufVariants,
   listLocalModels,
+  listScanFolders,
+  removeScanFolder,
 } from "@/features/chat/api/chat-api";
 import type {
   CachedGgufRepo,
@@ -42,7 +46,7 @@ import {
 import { cn, formatCompact } from "@/lib/utils";
 import type { VramFitStatus } from "@/lib/vram";
 import { checkVramFit, estimateLoadingVram } from "@/lib/vram";
-import { Search01Icon } from "@hugeicons/core-free-icons";
+import { Add01Icon, Cancel01Icon, Folder02Icon, Search01Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { Trash2Icon } from "lucide-react";
 import {
@@ -123,23 +127,22 @@ function ModelRow({
       className={cn(
         "flex w-full items-center gap-2 rounded-md px-2.5 py-1.5 text-left text-sm transition-colors hover:bg-accent",
         selected && "bg-accent/60",
-        exceeds && "opacity-50",
       )}
     >
       <span
         className={cn(
           "block min-w-0 flex-1 truncate",
-          exceeds && "line-through decoration-muted-foreground/50",
+          exceeds && "!text-gray-500 dark:!text-gray-400",
         )}
       >
         {label}
       </span>
       <span className="ml-auto flex items-center gap-1.5 shrink-0">
         {vramStatus === "exceeds" && (
-          <span className="text-[9px] font-medium text-red-400">OOM</span>
+          <span className="text-[9px] font-medium !text-red-700 !bg-red-50 dark:!text-red-400 dark:!bg-red-950 px-1.5 py-0.5 rounded">OOM</span>
         )}
         {vramStatus === "tight" && (
-          <span className="text-[9px] font-medium text-amber-400">TIGHT</span>
+          <span className="text-[9px] font-medium !text-amber-400">TIGHT</span>
         )}
         {meta ? (
           <span className="text-[10px] text-muted-foreground">{meta}</span>
@@ -350,7 +353,7 @@ function GgufVariantExpander({
               )}
             >
               <span className="min-w-0 flex-1 truncate font-mono text-xs">
-                {v.quant}
+                <span className={cn(oom && "!text-gray-500 dark:!text-gray-400")}>{v.quant}</span>
                 {v.downloaded ? (
                   <span className="ml-1.5 text-[9px] font-sans font-medium text-green-400">
                     downloaded
@@ -363,12 +366,12 @@ function GgufVariantExpander({
               </span>
               <span className="flex items-center gap-1.5 shrink-0">
                 {oom && (
-                  <span className="text-[9px] font-medium text-red-400">
+                  <span className="text-[9px] font-medium !text-red-700 !bg-red-50 dark:!text-red-400 dark:!bg-red-950 px-1.5 py-0.5 rounded">
                     OOM
                   </span>
                 )}
                 {tight && (
-                  <span className="text-[9px] font-medium text-amber-400">
+                  <span className="text-[9px] font-medium !text-amber-400">
                     TIGHT
                   </span>
                 )}
@@ -415,6 +418,7 @@ let _cachedGgufCache: CachedGgufRepo[] = [];
 let _cachedModelsCache: CachedModelRepo[] = [];
 let _lmStudioCache: LocalModelInfo[] = [];
 let _customFolderCache: LocalModelInfo[] = [];
+let _scanFoldersCache: ScanFolderInfo[] = [];
 
 /** Sort LM Studio models with unsloth publisher first. */
 function sortLmStudio(models: LocalModelInfo[]): LocalModelInfo[] {
@@ -434,10 +438,12 @@ export function HubModelPicker({
   models,
   value,
   onSelect,
+  onFoldersChange,
 }: {
   models: ModelOption[];
   value?: string;
   onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
+  onFoldersChange?: () => void;
 }) {
   const gpu = useGpuInfo();
   const [query, setQuery] = useState("");
@@ -469,6 +475,13 @@ export function HubModelPicker({
   const [customFolderModels, setCustomFolderModels] =
     useState<LocalModelInfo[]>(_customFolderCache);
 
+  // Custom scan folders management
+  const [scanFolders, setScanFolders] = useState<ScanFolderInfo[]>(_scanFoldersCache);
+  const [folderInput, setFolderInput] = useState("");
+  const [folderError, setFolderError] = useState<string | null>(null);
+  const [showFolderInput, setShowFolderInput] = useState(false);
+  const [folderLoading, setFolderLoading] = useState(false);
+
   const refreshLocalModelsList = useCallback(() => {
     listLocalModels()
       .then((res) => {
@@ -484,6 +497,57 @@ export function HubModelPicker({
       .catch(() => {});
   }, []);
 
+  const refreshScanFolders = useCallback(() => {
+    listScanFolders()
+      .then((v) => {
+        _scanFoldersCache = v;
+        setScanFolders(v);
+      })
+      .catch(() => {});
+  }, []);
+
+  const handleAddFolder = useCallback(async () => {
+    const trimmed = folderInput.trim();
+    if (!trimmed || folderLoading) return;
+    setFolderError(null);
+    setFolderLoading(true);
+    try {
+      const created = await addScanFolder(trimmed);
+      // Backend returns existing row for duplicates, so deduplicate
+      const next = _scanFoldersCache.some((f) => f.id === created.id || f.path === created.path)
+        ? _scanFoldersCache
+        : [..._scanFoldersCache, created];
+      _scanFoldersCache = next;
+      setScanFolders(next);
+      setFolderInput("");
+      setShowFolderInput(false);
+      refreshLocalModelsList();
+      onFoldersChange?.();
+      // Background reconciliation with the server
+      void refreshScanFolders();
+    } catch (e) {
+      setFolderError(e instanceof Error ? e.message : "Failed to add folder");
+    } finally {
+      setFolderLoading(false);
+    }
+  }, [folderInput, folderLoading, refreshScanFolders, refreshLocalModelsList, onFoldersChange]);
+
+  const handleRemoveFolder = useCallback(async (id: number) => {
+    try {
+      await removeScanFolder(id);
+      // Optimistic update so the folder disappears immediately
+      const next = _scanFoldersCache.filter((f) => f.id !== id);
+      _scanFoldersCache = next;
+      setScanFolders(next);
+      refreshScanFolders();
+      refreshLocalModelsList();
+      onFoldersChange?.();
+    } catch (e) {
+      toast.error(e instanceof Error ? e.message : "Failed to remove folder");
+      refreshScanFolders();
+    }
+  }, [refreshScanFolders, refreshLocalModelsList, onFoldersChange]);
+
   const refreshCachedLists = useCallback(() => {
     listCachedGguf()
       .then((v) => {
@@ -503,6 +567,7 @@ export function HubModelPicker({
   useEffect(() => {
     // Always refresh LM Studio + custom folder models (not gated by alreadyCached)
     refreshLocalModelsList();
+    refreshScanFolders();
 
     if (alreadyCached) return;
     let done = 0;
@@ -523,7 +588,7 @@ export function HubModelPicker({
       })
       .catch(() => {})
       .finally(check);
-  }, [alreadyCached]);
+  }, [alreadyCached, refreshLocalModelsList, refreshScanFolders]);
 
   const handleDeleteConfirm = useCallback(async () => {
     if (!deleteTarget) return;
@@ -878,9 +943,95 @@ export function HubModelPicker({
             </>
           ) : null}
 
-          {!showHfSection && customFolderModels.length > 0 ? (
+          {!showHfSection ? (
             <>
-              <ListLabel>Custom Folders</ListLabel>
+              <div className="flex items-center justify-between px-2.5 py-1.5">
+                <span className="text-[10px] font-semibold uppercase tracking-wider text-muted-foreground">
+                  Custom Folders
+                </span>
+                <button
+                  type="button"
+                  aria-label={showFolderInput ? "Cancel adding folder" : "Add scan folder"}
+                  onClick={() => {
+                    setShowFolderInput((open) => {
+                      if (open) { setFolderInput(""); setFolderError(null); }
+                      return !open;
+                    });
+                  }}
+                  className="rounded p-0.5 text-muted-foreground/60 transition-colors hover:text-foreground"
+                >
+                  <HugeiconsIcon icon={showFolderInput ? Cancel01Icon : Add01Icon} className="size-3" />
+                </button>
+              </div>
+
+              {/* Folder paths */}
+              {scanFolders.map((f) => (
+                <div
+                  key={f.id}
+                  className="group flex items-center gap-1.5 px-3 py-0.5"
+                >
+                  <HugeiconsIcon icon={Folder02Icon} className="size-3 shrink-0 text-muted-foreground/40" />
+                  <span
+                    className="min-w-0 flex-1 truncate font-mono text-[10px] text-muted-foreground/70"
+                    title={f.path}
+                  >
+                    {f.path}
+                  </span>
+                  <button
+                    type="button"
+                    onClick={() => handleRemoveFolder(f.id)}
+                    aria-label={`Remove folder ${f.path}`}
+                    className="shrink-0 rounded p-0.5 text-muted-foreground/40 opacity-100 md:opacity-0 md:group-hover:opacity-100 focus-visible:opacity-100 transition-opacity hover:text-destructive"
+                  >
+                    <HugeiconsIcon icon={Cancel01Icon} className="size-2.5" />
+                  </button>
+                </div>
+              ))}
+
+              {/* Add folder input */}
+              {showFolderInput && (
+                <div className="px-2.5 pb-1 pt-0.5">
+                  <div className="flex items-center gap-1">
+                    <HugeiconsIcon icon={Folder02Icon} className="size-3 shrink-0 text-muted-foreground/40" />
+                    <input
+                      value={folderInput}
+                      onChange={(e) => { setFolderInput(e.target.value); setFolderError(null); }}
+                      onKeyDown={(e) => {
+                        if (e.key === "Enter") { e.preventDefault(); handleAddFolder(); }
+                        if (e.key === "Escape") { e.preventDefault(); e.stopPropagation(); setShowFolderInput(false); setFolderInput(""); setFolderError(null); }
+                      }}
+                      placeholder="/path/to/models"
+                      className="h-6 min-w-0 flex-1 rounded border border-border/50 bg-transparent px-1.5 font-mono text-[10px] text-foreground outline-none placeholder:text-muted-foreground/40 focus:border-foreground/20"
+                      disabled={folderLoading}
+                      autoFocus={true}
+                    />
+                    <button
+                      type="button"
+                      onClick={handleAddFolder}
+                      disabled={folderLoading || !folderInput.trim()}
+                      className="h-6 shrink-0 rounded border border-border/50 px-1.5 text-[10px] text-muted-foreground transition-colors hover:bg-accent disabled:opacity-40"
+                    >
+                      Add
+                    </button>
+                  </div>
+                  {folderError && (
+                    <p className="px-0.5 pt-0.5 text-[10px] text-destructive">{folderError}</p>
+                  )}
+                </div>
+              )}
+
+              {/* Empty state */}
+              {scanFolders.length === 0 && customFolderModels.length === 0 && !showFolderInput && (
+                <button
+                  type="button"
+                  onClick={() => setShowFolderInput(true)}
+                  className="px-2.5 pb-1.5 text-left text-[10px] text-muted-foreground/60 transition-colors hover:text-muted-foreground"
+                >
+                  + Add a folder to scan for local models
+                </button>
+              )}
+
+              {/* Models from custom folders */}
               {customFolderModels.map((m) => {
                 const isGguf =
                   isGgufRepo(m.id) ||

studio/frontend/src/components/assistant-ui/thread.tsx

@@ -73,7 +73,7 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
       }}
     >
       <ThreadPrimitive.Viewport
-        className="aui-thread-viewport relative flex flex-1 flex-col overflow-x-auto overflow-y-scroll scroll-smooth px-4 pt-4"
+        className="aui-thread-viewport relative flex min-w-0 flex-1 flex-col overflow-x-auto overflow-y-scroll scroll-smooth px-4 pt-4"
       >
         {!hideWelcome && (
           <AuiIf condition={({ thread }) => thread.isEmpty}>
@@ -89,7 +89,13 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
           }}
         />
 
-        <ThreadPrimitive.ViewportFooter className="aui-thread-viewport-footer sticky bottom-0 z-20 mt-auto flex w-full flex-col gap-4 overflow-visible bg-background pb-4 md:pb-4">
+        <ThreadPrimitive.ViewportFooter className={cn(
+          "aui-thread-viewport-footer sticky bottom-0 z-20 mt-auto flex w-full flex-col gap-4 overflow-visible pb-4 md:pb-4",
+          hideComposer ? "bg-background" : "relative bg-transparent",
+        )}>
+          {!hideComposer && (
+            <div className="pointer-events-none absolute inset-x-0 bottom-0 h-4 bg-background" aria-hidden />
+          )}
           <ThreadScrollToBottom />
           <AuiIf condition={({ thread }) => !thread.isEmpty}>
             {!hideComposer && <ComposerAnimated />}
@@ -118,7 +124,7 @@ const SUGGESTION_TOOLS: Record<string, Array<"thinking" | "search" | "code">> =
   "How do you fine-tune an audio model with Unsloth?": ["thinking", "search"],
   "Create a live weather dashboard in HTML using no API key. Show me the code": ["thinking", "code", "search"],
   "Solve the integral of x·sin(x), and verify it step by step": ["thinking", "code"],
-  "Draw an SVG of a cute sloth": ["thinking", "code", "search"],
+  "Draw an SVG of a cute sloth & show the code": ["thinking", "code", "search"],
 };
 
 const toolIconMap = {
@@ -220,14 +226,20 @@ const GeneratingSpinner: FC = () => {
 
 const ComposerAnimated: FC = () => {
   return (
-    <motion.div
-      layout={true}
-      layoutId="composer"
-      transition={{ type: "spring", bounce: 0.15, duration: 0.5 }}
-      className="mx-auto w-full max-w-(--thread-max-width)"
-    >
-      <Composer />
-    </motion.div>
+    <div className="relative mx-auto min-w-0 w-full max-w-(--thread-max-width)">
+      <div
+        className="pointer-events-none absolute inset-x-0 top-1/2 bottom-0 z-0 bg-background"
+        aria-hidden
+      />
+      <motion.div
+        layout={true}
+        layoutId="composer"
+        transition={{ type: "spring", bounce: 0.15, duration: 0.5 }}
+        className="relative z-10 w-full"
+      >
+        <Composer />
+      </motion.div>
+    </div>
   );
 };
 
@@ -262,7 +274,7 @@ const Composer: FC = () => {
         <ToolStatusDisplay />
         <ComposerPrimitive.Input
           placeholder="Send a message..."
-          className="aui-composer-input mb-1 max-h-32 min-h-12 w-full resize-none bg-transparent px-4 pt-2 pb-3 text-sm outline-none placeholder:text-muted-foreground focus-visible:ring-0"
+          className="aui-composer-input mb-1 max-h-32 min-h-12 w-full resize-none bg-transparent pl-5 pr-4 pt-2 pb-3 text-sm outline-none placeholder:text-muted-foreground focus-visible:ring-0"
           rows={1}
           autoFocus={true}
           aria-label="Message input"
@@ -437,21 +449,40 @@ const CodeToolsToggle: FC = () => {
 
 const ToolStatusDisplay: FC = () => {
   const toolStatus = useChatRuntimeStore((s) => s.toolStatus);
+  const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
   const [elapsed, setElapsed] = useState(0);
+  const [visible, setVisible] = useState(false);
 
   useEffect(() => {
     if (!toolStatus) {
       setElapsed(0);
+      if (!isThreadRunning) {
+        setVisible(false);
+      }
       return;
     }
+
     setElapsed(0);
+
+    // Debounce badge visibility by 300ms when the badge is not
+    // already on screen. Once visible from a prior tool, consecutive
+    // tools show immediately so the badge does not flicker. Fast
+    // tool calls that all complete under 300ms never show the badge.
+    let showTimer: ReturnType<typeof setTimeout> | undefined;
+    if (!visible) {
+      showTimer = setTimeout(() => setVisible(true), 300);
+    }
+
     const interval = setInterval(() => {
       setElapsed((prev) => prev + 1);
     }, 1000);
-    return () => clearInterval(interval);
-  }, [toolStatus]);
+    return () => {
+      clearInterval(interval);
+      if (showTimer) clearTimeout(showTimer);
+    };
+  }, [toolStatus, isThreadRunning]);
 
-  if (!toolStatus) return null;
+  if (!toolStatus || !visible) return null;
   const isRunning = toolStatus.startsWith("Running");
   const StatusIcon = isRunning ? TerminalIcon : GlobeIcon;
   return (
@@ -555,10 +586,10 @@ const GeneratingIndicator: FC = () => {
 const AssistantMessage: FC = () => {
   return (
     <MessagePrimitive.Root
-      className="aui-assistant-message-root fade-in slide-in-from-bottom-1 relative mx-auto w-full max-w-(--thread-max-width) animate-in py-3 duration-150"
+      className="aui-assistant-message-root fade-in slide-in-from-bottom-1 relative mx-auto min-w-0 w-full max-w-(--thread-max-width) animate-in py-3 duration-150"
       data-role="assistant"
     >
-      <div className="aui-assistant-message-content wrap-break-word px-2 text-foreground leading-relaxed">
+      <div className="aui-assistant-message-content wrap-break-word min-w-0 text-foreground leading-relaxed">
         <GeneratingIndicator />
         <MessagePrimitive.Parts
           components={{
@@ -581,7 +612,7 @@ const AssistantMessage: FC = () => {
         <MessageError />
       </div>
 
-      <div className="aui-assistant-message-footer mt-1 ml-2 flex">
+      <div className="aui-assistant-message-footer mt-1 flex">
         <BranchPicker />
         <AssistantActionBar />
       </div>

studio/frontend/src/components/assistant-ui/tool-fallback.tsx

@@ -156,12 +156,12 @@ function ToolFallbackTrigger({
       <span
         data-slot="tool-fallback-trigger-label"
         className={cn(
-          "aui-tool-fallback-trigger-label-wrapper relative inline-block grow text-left leading-none",
+          "aui-tool-fallback-trigger-label-wrapper relative inline-block grow text-left leading-none text-muted-foreground",
           isCancelled && "text-muted-foreground line-through",
         )}
       >
         <span>
-          {label}: <b>{toolName}</b>
+          {label}: <span className="font-medium text-foreground/85">{toolName}</span>
         </span>
         {isRunning && (
           <span
@@ -169,7 +169,7 @@ function ToolFallbackTrigger({
             data-slot="tool-fallback-trigger-shimmer"
             className="aui-tool-fallback-trigger-shimmer shimmer pointer-events-none absolute inset-0 motion-reduce:animate-none"
           >
-            {label}: <b>{toolName}</b>
+            {label}: <span className="font-medium text-foreground/85">{toolName}</span>
           </span>
         )}
       </span>

studio/frontend/src/components/assistant-ui/tool-ui-python.tsx

@@ -4,6 +4,7 @@
 "use client";
 
 import { copyToClipboard } from "@/lib/copy-to-clipboard";
+import { getAuthToken } from "@/features/auth/session";
 import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
 import { code as codePlugin } from "@streamdown/code";
 import { CheckIcon, CodeIcon, CopyIcon, LoaderIcon } from "lucide-react";
@@ -15,6 +16,12 @@ import {
   ToolFallbackTrigger,
 } from "./tool-fallback";
 
+interface StructuredResult {
+  text: string;
+  images: string[];
+  sessionId: string;
+}
+
 const MAX_DISPLAY = 10_000;
 const COPY_RESET_MS = 2000;
 const SHIKI_THEME = ["github-light", "github-dark"] as ["github-light", "github-dark"];
@@ -84,6 +91,16 @@ function HighlightedCode({ code: source, language }: { code: string; language: s
   );
 }
 
+function isStructuredResult(val: unknown): val is StructuredResult {
+  return (
+    typeof val === "object" &&
+    val !== null &&
+    "text" in val &&
+    "images" in val &&
+    "sessionId" in val
+  );
+}
+
 const PythonToolUIImpl: ToolCallMessagePartComponent = ({
   args,
   result,
@@ -92,12 +109,24 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({
   const code = (args as { code?: string })?.code ?? "";
   const firstLine = code.split("\n")[0]?.slice(0, 60) ?? "";
   const isRunning = status?.type === "running";
-  const output =
-    typeof result === "string"
-      ? result
-      : result
-        ? JSON.stringify(result, null, 2)
-        : "";
+
+  let output: string;
+  let images: string[] = [];
+  let sessionId = "";
+
+  if (isStructuredResult(result)) {
+    output = result.text;
+    images = result.images;
+    sessionId = result.sessionId;
+  } else if (typeof result === "string") {
+    output = result;
+  } else if (result) {
+    output = JSON.stringify(result, null, 2);
+  } else {
+    output = "";
+  }
+
+  const authToken = getAuthToken();
 
   return (
     <ToolFallbackRoot>
@@ -133,6 +162,21 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({
               </pre>
             </div>
           ) : null}
+
+          {/* Images from Python tool execution */}
+          {images.length > 0 && sessionId && (
+            <div className="mt-2 flex flex-col gap-2">
+              {images.map((filename) => (
+                <img
+                  key={filename}
+                  src={`/api/inference/sandbox/${encodeURIComponent(sessionId)}/${encodeURIComponent(filename)}${authToken ? `?token=${encodeURIComponent(authToken)}` : ""}`}
+                  alt={filename}
+                  loading="lazy"
+                  className="max-w-full rounded border border-border"
+                />
+              ))}
+            </div>
+          )}
         </div>
       </ToolFallbackContent>
     </ToolFallbackRoot>

studio/frontend/src/components/navbar.tsx

@@ -372,90 +372,94 @@ export function Navbar() {
           })}
         </nav>
 
-        {/* Right: docs/tour desktop */}
-        <div className="hidden items-center justify-self-end gap-2 md:flex">
-          <AnimatedThemeToggler
-            className="flex h-9 w-9 items-center justify-center rounded-md text-muted-foreground transition-colors hover:bg-accent hover:text-foreground [&_svg]:size-4"
-            title="Toggle theme"
-            aria-label="Toggle theme"
-          />
-          <HoverCard openDelay={200} closeDelay={100}>
-            <HoverCardTrigger asChild={true}>
-              <a
-                href="https://unsloth.ai/docs"
-                target="_blank"
-                rel="noopener noreferrer"
-                className="flex items-center gap-1.5 text-sm font-medium text-emerald-600 hover:text-emerald-700 transition-colors"
-              >
-                <HugeiconsIcon icon={Book03Icon} className="size-4" />
-                Learn more
-              </a>
-            </HoverCardTrigger>
-            <HoverCardContent align="end" className="w-80 p-0">
-              <a
-                href="https://unsloth.ai/docs"
-                target="_blank"
-                rel="noopener noreferrer"
-                className="group/card flex flex-col gap-1 p-4 no-underline"
-              >
-                <p className="text-sm font-semibold font-heading">
-                  Unsloth Documentation
-                </p>
-                <p className="text-xs text-muted-foreground leading-relaxed">
-                  Guides on fine-tuning LLMs 2x faster with 70% less memory.
-                  Covers LoRA, QLoRA, data formatting, and deployment.
-                </p>
-                <span className="mt-1 flex items-center gap-1 text-xs font-medium text-emerald-600 group-hover/card:underline">
-                  Visit docs
-                  <HugeiconsIcon icon={ArrowRight01Icon} className="size-3" />
-                </span>
-              </a>
-            </HoverCardContent>
-          </HoverCard>
-
-          <button
-            type="button"
-            onClick={tourId ? openTour : undefined}
-            className={cn(
-              "flex h-9 items-center gap-1.5 rounded-md px-3 text-muted-foreground transition-colors hover:bg-accent hover:text-foreground",
-              !tourId && "invisible pointer-events-none",
-            )}
-            title="Tour"
-            aria-hidden={!tourId}
-            tabIndex={tourId ? 0 : -1}
-          >
-            <HugeiconsIcon icon={CursorInfo02Icon} className="size-4" />
-            <span className="text-sm font-medium">Tour</span>
-          </button>
-
-          <HoverCard openDelay={200} closeDelay={100}>
-            <HoverCardTrigger asChild={true}>
+        {/* Right: docs/tour desktop — one wrapper per control so flex gap is even (HoverCard roots can confuse flex spacing). */}
+        <div className="hidden items-center justify-self-end gap-0 md:flex">
+          <div className="flex shrink-0 items-center">
+            <AnimatedThemeToggler
+              className="flex h-9 w-9 items-center justify-center rounded-md text-muted-foreground transition-colors hover:bg-accent hover:text-foreground [&_svg]:size-4"
+              title="Toggle theme"
+              aria-label="Toggle theme"
+            />
+          </div>
+          <div className="flex shrink-0 items-center">
+            <HoverCard openDelay={200} closeDelay={100}>
+              <HoverCardTrigger asChild={true}>
+                <a
+                  href="https://unsloth.ai/docs"
+                  target="_blank"
+                  rel="noopener noreferrer"
+                  className="inline-flex h-9 items-center gap-1.5 rounded-md px-3 text-sm font-medium text-emerald-600 transition-colors hover:bg-accent hover:text-emerald-700 dark:hover:text-emerald-400"
+                >
+                  <HugeiconsIcon icon={Book03Icon} className="size-4" />
+                  Learn more
+                </a>
+              </HoverCardTrigger>
+              <HoverCardContent align="end" className="w-80 p-0">
+                <a
+                  href="https://unsloth.ai/docs"
+                  target="_blank"
+                  rel="noopener noreferrer"
+                  className="group/card flex flex-col gap-1 p-4 no-underline"
+                >
+                  <p className="text-sm font-semibold font-heading">
+                    Unsloth Documentation
+                  </p>
+                  <p className="text-xs text-muted-foreground leading-relaxed">
+                    Guides on fine-tuning LLMs 2x faster with 70% less memory.
+                    Covers LoRA, QLoRA, data formatting, and deployment.
+                  </p>
+                  <span className="mt-1 flex items-center gap-1 text-xs font-medium text-emerald-600 group-hover/card:underline">
+                    Visit docs
+                    <HugeiconsIcon icon={ArrowRight01Icon} className="size-3" />
+                  </span>
+                </a>
+              </HoverCardContent>
+            </HoverCard>
+          </div>
+          {tourId ? (
+            <div className="flex shrink-0 items-center">
               <button
                 type="button"
-                className="flex h-9 items-center gap-1.5 rounded-md px-3 text-sm font-medium text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
-                aria-label="How to update Unsloth Studio"
+                onClick={openTour}
+                className="flex h-9 items-center gap-1.5 rounded-md px-3 text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
+                title="Tour"
               >
-                <HugeiconsIcon icon={ArrowReloadHorizontalIcon} className="size-4" />
-                Update
+                <HugeiconsIcon icon={CursorInfo02Icon} className="size-4" />
+                <span className="text-sm font-medium">Tour</span>
               </button>
-            </HoverCardTrigger>
-            <HoverCardContent align="end" className="w-[22.5rem] p-0">
-              <UpdateStudioInstructions
-                className="p-4"
-                defaultShell={defaultUpdateShell}
-              />
-            </HoverCardContent>
-          </HoverCard>
-
-          <button
-            type="button"
-            onClick={() => setShutdownOpen(true)}
-            className="-mr-1.5 flex h-9 w-9 items-center justify-center rounded-md text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
-            title="Shut down Unsloth Studio server"
-            aria-label="Shut down Unsloth Studio server"
-          >
-            <HugeiconsIcon icon={Cancel01Icon} className="size-5" />
-          </button>
+            </div>
+          ) : null}
+          <div className="flex shrink-0 items-center">
+            <HoverCard openDelay={200} closeDelay={100}>
+              <HoverCardTrigger asChild={true}>
+                <button
+                  type="button"
+                  className="flex h-9 items-center gap-1.5 rounded-md px-3 text-sm font-medium text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
+                  aria-label="How to update Unsloth Studio"
+                >
+                  <HugeiconsIcon icon={ArrowReloadHorizontalIcon} className="size-4" />
+                  Update
+                </button>
+              </HoverCardTrigger>
+              <HoverCardContent align="end" className="w-[22.5rem] p-0">
+                <UpdateStudioInstructions
+                  className="p-4"
+                  defaultShell={defaultUpdateShell}
+                />
+              </HoverCardContent>
+            </HoverCard>
+          </div>
+          <div className="flex shrink-0 items-center">
+            <button
+              type="button"
+              onClick={() => setShutdownOpen(true)}
+              className="flex h-9 w-9 items-center justify-center rounded-md text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
+              title="Shut down Unsloth Studio server"
+              aria-label="Shut down Unsloth Studio server"
+            >
+              <HugeiconsIcon icon={Cancel01Icon} className="size-5" />
+            </button>
+          </div>
         </div>
 
         {/* Right: mobile */}

studio/frontend/src/config/training.ts

@@ -141,6 +141,10 @@ export const MODEL_TYPE_TO_HF_TASK: Record<ModelType, PipelineType> = {
 
 
 export const PRIORITY_TRAINING_MODELS: readonly string[] = [
+  "unsloth/gemma-4-E2B-it",
+  "unsloth/gemma-4-E4B-it",
+  "unsloth/gemma-4-31B-it",
+  "unsloth/gemma-4-26B-A4B-it",
   "unsloth/Qwen3.5-2B",
   "unsloth/Qwen3.5-9B",
   "unsloth/gpt-oss-20b",

studio/frontend/src/features/chat/api/chat-adapter.ts

@@ -635,7 +635,23 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                 toolCallParts[toolCallParts.length - 1]?.toolCallId || "";
               const idx = toolCallParts.findIndex((p) => p.toolCallId === id);
               if (idx !== -1) {
-                toolCallParts[idx] = { ...toolCallParts[idx], result: toolEvent.result as string };
+                const rawResult = (toolEvent.result as string) ?? "";
+                const imgMarker = "\n__IMAGES__:";
+                const imgIdx = rawResult.lastIndexOf(imgMarker);
+                let parsedResult: string | { text: string; images: string[]; sessionId: string };
+                if (imgIdx !== -1) {
+                  const text = rawResult.slice(0, imgIdx);
+                  const sessionId = unstable_threadId || "";
+                  try {
+                    const images = JSON.parse(rawResult.slice(imgIdx + imgMarker.length)) as string[];
+                    parsedResult = { text, images, sessionId };
+                  } catch {
+                    parsedResult = rawResult;
+                  }
+                } else {
+                  parsedResult = rawResult;
+                }
+                toolCallParts[idx] = { ...toolCallParts[idx], result: parsedResult };
               }
             }
             // Yield cumulative state so tool UI updates (tools first, text after)

studio/frontend/src/features/chat/chat-page.tsx

@@ -162,10 +162,12 @@ const CompareContent = memo(function CompareContent({
   pairId,
   models,
   loraModels,
+  onFoldersChange,
 }: {
   pairId: string;
   models: ModelOption[];
   loraModels: LoraModelOption[];
+  onFoldersChange?: () => void;
 }): ReactElement {
   const isLoraCompare = useIsLoraCompare();
 
@@ -176,6 +178,7 @@ const CompareContent = memo(function CompareContent({
       pairId={pairId}
       models={models}
       loraModels={loraModels}
+      onFoldersChange={onFoldersChange}
     />
   );
 });
@@ -259,10 +262,12 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
   pairId,
   models,
   loraModels,
+  onFoldersChange,
 }: {
   pairId: string;
   models: ModelOption[];
   loraModels: LoraModelOption[];
+  onFoldersChange?: () => void;
 }): ReactElement {
   const handlesRef = useRef<Record<string, CompareHandle>>({});
   const [model1ThreadId, setModel1ThreadId] = useState<string>();
@@ -327,6 +332,7 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
                     ggufVariant: meta.ggufVariant,
                   })
                 }
+                onFoldersChange={onFoldersChange}
                 variant="ghost"
                 size="sm"
                 className="max-w-[50%]"
@@ -359,6 +365,7 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
                     ggufVariant: meta.ggufVariant,
                   })
                 }
+                onFoldersChange={onFoldersChange}
                 variant="ghost"
                 size="sm"
                 className="max-w-[50%]"
@@ -846,6 +853,7 @@ export function ChatPage(): ReactElement {
                 activeGgufVariant={activeGgufVariant}
                 onValueChange={handleCheckpointChange}
                 onEject={handleEject}
+                onFoldersChange={refreshLocalModels}
                 variant="ghost"
                 open={modelSelectorOpen}
                 onOpenChange={handleModelSelectorOpenChange}
@@ -858,14 +866,16 @@ export function ChatPage(): ReactElement {
                   label={
                     loadProgress?.phase === "starting"
                       ? "Starting model…"
-                      : loadingModel.isDownloaded
+                      : loadingModel.isDownloaded || loadingModel.isCachedLora
                         ? "Loading model…"
                         : "Downloading model…"
                   }
                   title={
                     loadingModel.isDownloaded
                       ? `Loading ${loadingModel.displayName} from cache.`
-                      : `Loading ${loadingModel.displayName}. This may include downloading.`
+                      : loadingModel.isCachedLora
+                        ? `Loading ${loadingModel.displayName} into memory.`
+                        : `Loading ${loadingModel.displayName}. This may include downloading.`
                   }
                   progressPercent={loadProgress?.percent}
                   progressLabel={loadProgress?.label}
@@ -911,6 +921,7 @@ export function ChatPage(): ReactElement {
               pairId={view.pairId}
               models={models}
               loraModels={loraModels}
+              onFoldersChange={refreshLocalModels}
             />
           )}
         </div>
@@ -934,7 +945,6 @@ export function ChatPage(): ReactElement {
               });
             }
           }}
-          onFoldersChange={refreshLocalModels}
         />
       </SidebarProvider>
     </div>

studio/frontend/src/features/chat/chat-settings-sheet.tsx

@@ -34,7 +34,6 @@ import {
   CodeIcon,
   Delete02Icon,
   FloppyDiskIcon,
-  FolderSearchIcon,
   PencilEdit01Icon,
   Settings02Icon,
   SlidersHorizontalIcon,
@@ -44,13 +43,7 @@ import {
 import { HugeiconsIcon } from "@hugeicons/react";
 import { AnimatePresence, motion } from "motion/react";
 import type { ReactNode } from "react";
-import { useCallback, useEffect, useMemo, useState } from "react";
-import {
-  type ScanFolderInfo,
-  addScanFolder,
-  listScanFolders,
-  removeScanFolder,
-} from "./api/chat-api";
+import { useEffect, useMemo, useState } from "react";
 import { useChatRuntimeStore } from "./stores/chat-runtime-store";
 import {
   DEFAULT_INFERENCE_PARAMS,
@@ -266,108 +259,6 @@ function CollapsibleSection({
   );
 }
 
-function ModelFoldersSection({
-  onFoldersChange,
-}: { onFoldersChange?: () => void }) {
-  const [folders, setFolders] = useState<ScanFolderInfo[]>([]);
-  const [input, setInput] = useState("");
-  const [error, setError] = useState<string | null>(null);
-  const [loading, setLoading] = useState(false);
-
-  const refresh = useCallback(() => {
-    listScanFolders()
-      .then(setFolders)
-      .catch(() => {});
-  }, []);
-
-  useEffect(() => {
-    refresh();
-  }, [refresh]);
-
-  const handleAdd = async () => {
-    const trimmed = input.trim();
-    if (!trimmed) return;
-    setError(null);
-    setLoading(true);
-    try {
-      await addScanFolder(trimmed);
-      setInput("");
-      refresh();
-      onFoldersChange?.();
-    } catch (e) {
-      setError(e instanceof Error ? e.message : "Failed to add folder");
-    } finally {
-      setLoading(false);
-    }
-  };
-
-  const handleRemove = async (id: number) => {
-    try {
-      await removeScanFolder(id);
-      onFoldersChange?.();
-    } catch (e) {
-      setError(e instanceof Error ? e.message : "Failed to remove folder");
-    } finally {
-      refresh();
-    }
-  };
-
-  return (
-    <CollapsibleSection icon={FolderSearchIcon} label="Model Folders">
-      <div className="flex flex-col gap-2 py-1">
-        {folders.length > 0 && (
-          <div className="flex flex-col gap-1">
-            {folders.map((f) => (
-              <div
-                key={f.id}
-                className="group flex items-center gap-1.5 rounded-md px-1.5 py-1 text-xs transition-colors hover:bg-accent"
-              >
-                <span
-                  className="min-w-0 flex-1 truncate text-muted-foreground"
-                  title={f.path}
-                >
-                  {f.path}
-                </span>
-                <button
-                  type="button"
-                  onClick={() => handleRemove(f.id)}
-                  className="shrink-0 rounded p-0.5 text-muted-foreground/50 opacity-0 transition-opacity group-hover:opacity-100 hover:text-destructive"
-                >
-                  <HugeiconsIcon icon={Delete02Icon} className="size-3" />
-                </button>
-              </div>
-            ))}
-          </div>
-        )}
-        <div className="flex gap-1.5">
-          <Input
-            value={input}
-            onChange={(e) => {
-              setInput(e.target.value);
-              setError(null);
-            }}
-            onKeyDown={(e) => {
-              if (e.key === "Enter") handleAdd();
-            }}
-            placeholder="/path/to/models"
-            className="h-7 flex-1 text-xs font-mono"
-            disabled={loading}
-          />
-          <button
-            type="button"
-            onClick={handleAdd}
-            disabled={loading || !input.trim()}
-            className="h-7 rounded-md border px-2 text-[11px] font-medium text-muted-foreground transition-colors hover:bg-accent disabled:cursor-not-allowed disabled:opacity-50"
-          >
-            Add
-          </button>
-        </div>
-        {error && <p className="text-[11px] text-destructive">{error}</p>}
-      </div>
-    </CollapsibleSection>
-  );
-}
-
 interface ChatSettingsPanelProps {
   open: boolean;
   onOpenChange?: (open: boolean) => void;
@@ -376,7 +267,6 @@ interface ChatSettingsPanelProps {
   autoTitle: boolean;
   onAutoTitleChange: (enabled: boolean) => void;
   onReloadModel?: () => void;
-  onFoldersChange?: () => void;
 }
 
 export function ChatSettingsPanel({
@@ -387,7 +277,6 @@ export function ChatSettingsPanel({
   autoTitle,
   onAutoTitleChange,
   onReloadModel,
-  onFoldersChange,
 }: ChatSettingsPanelProps) {
   const isMobile = useIsMobile();
   const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null;
@@ -395,6 +284,9 @@ export function ChatSettingsPanel({
   const ggufMaxContextLength = useChatRuntimeStore(
     (s) => s.ggufMaxContextLength,
   );
+  const ggufNativeContextLength = useChatRuntimeStore(
+    (s) => s.ggufNativeContextLength,
+  );
   const kvCacheDtype = useChatRuntimeStore((s) => s.kvCacheDtype);
   const setKvCacheDtype = useChatRuntimeStore((s) => s.setKvCacheDtype);
   const loadedKvCacheDtype = useChatRuntimeStore((s) => s.loadedKvCacheDtype);
@@ -404,7 +296,7 @@ export function ChatSettingsPanel({
   );
 
   const ctxDisplayValue = customContextLength ?? ggufContextLength ?? "";
-  const ctxMaxValue = ggufMaxContextLength ?? ggufContextLength ?? null;
+  const ctxMaxValue = ggufNativeContextLength ?? ggufContextLength ?? null;
   const kvDirty = kvCacheDtype !== loadedKvCacheDtype;
   const ctxDirty = customContextLength !== null;
   const modelSettingsDirty = kvDirty || ctxDirty;
@@ -655,6 +547,13 @@ export function ChatSettingsPanel({
                       );
                     }}
                   />
+                  {ggufMaxContextLength != null &&
+                    typeof ctxDisplayValue === "number" &&
+                    ctxDisplayValue > ggufMaxContextLength && (
+                      <p className="text-[11px] text-amber-500">
+                        Exceeds estimated VRAM capacity ({ggufMaxContextLength.toLocaleString()} tokens). The model may use system RAM.
+                      </p>
+                    )}
                 </div>
                 <div className="flex items-center justify-between gap-3">
                   <div className="min-w-0">
@@ -835,8 +734,6 @@ export function ChatSettingsPanel({
           </div>
         </CollapsibleSection>
 
-        <ModelFoldersSection onFoldersChange={onFoldersChange} />
-
         <ChatTemplateSection onReloadModel={onReloadModel} />
       </div>
       <Dialog

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

@@ -168,6 +168,7 @@ export function useChatModelRuntime() {
     id: string;
     displayName: string;
     isDownloaded?: boolean;
+    isCachedLora?: boolean;
   } | null>(null);
   const [loadToastDismissed, setLoadToastDismissed] = useState(false);
   const [loadProgress, setLoadProgress] = useState<{
@@ -246,12 +247,16 @@ export function useChatModelRuntime() {
         const ggufMaxContextLength = statusRes.is_gguf
           ? (statusRes.max_context_length ?? null)
           : null;
+        const ggufNativeContextLength = statusRes.is_gguf
+          ? (statusRes.native_context_length ?? null)
+          : null;
         useChatRuntimeStore.setState({
           supportsReasoning,
           reasoningAlwaysOn,
           supportsTools,
           ggufContextLength: currentGgufContextLength,
           ggufMaxContextLength,
+          ggufNativeContextLength,
         });
 
         // Set reasoning default for Qwen3.5 small models
@@ -290,8 +295,11 @@ export function useChatModelRuntime() {
     setLoadToastDismissedState(false);
     clearCheckpoint();
     if (tid != null) toast.dismiss(tid);
+    const isCachedOrLocal = model.isDownloaded || model.isCachedLora;
     toast.info("Stopped loading model", {
-      description: "The current download may still finish in the background.",
+      description: isCachedOrLocal
+        ? undefined
+        : "The current download may still finish in the background.",
     });
     // Fire-and-forget: tell backend to stop, don't block UI
     unloadModel({ model_path: model.id }).catch(() => {});
@@ -335,20 +343,24 @@ export function useChatModelRuntime() {
         : undefined;
       const previousIsLora =
         previousModel?.isLora ?? (previousLora ? true : false);
+      // Covers Unix absolute (/), relative (./  ../), tilde (~/), Windows drive (C:\), UNC (\\server)
+      const isLocal = /^(\/|\.{1,2}[\\\/]|~[\\\/]|[A-Za-z]:[\\\/]|\\\\)/.test(modelId);
+      const isCachedLora = isLora && isLocal;
       const loadingDescription = [
         currentCheckpoint ? "Switching models." : null,
         extraLoadingDescription ?? null,
         isDownloaded ? "Loading cached model into memory." : null,
+        !isDownloaded && isCachedLora ? "Loading trained model into memory." : null,
       ]
         .filter(Boolean)
         .join(" ");
       setModelsError(null);
       setLoadToastDismissedState(false);
-      const loadInfo = { id: modelId, displayName, isDownloaded };
+      const loadInfo = { id: modelId, displayName, isDownloaded, isCachedLora };
       setLoadingModel(loadInfo);
       useChatRuntimeStore.getState().setModelLoading(true);
       setLoadProgress(
-        isDownloaded
+        isDownloaded || isCachedLora
           ? { percent: null, label: null, phase: "starting" }
           : { percent: 0, label: "Preparing download", phase: "downloading" },
       );
@@ -425,6 +437,9 @@ export function useChatModelRuntime() {
             const reportedMaxCtx = loadResponse.is_gguf
               ? (loadResponse.max_context_length ?? null)
               : null;
+            const reportedNativeCtx = loadResponse.is_gguf
+              ? (loadResponse.native_context_length ?? null)
+              : null;
             // A successful reload has applied settings, so clear pending custom
             // context state and display the backend-reported effective context.
             const keepCustomCtx = null;
@@ -433,6 +448,7 @@ export function useChatModelRuntime() {
             useChatRuntimeStore.setState({
               ggufContextLength: nativeCtx,
               ggufMaxContextLength,
+              ggufNativeContextLength: reportedNativeCtx,
               supportsReasoning: loadResponse.supports_reasoning ?? false,
               reasoningAlwaysOn,
               reasoningEnabled: reasoningAlwaysOn ? true : reasoningDefault,
@@ -477,15 +493,16 @@ export function useChatModelRuntime() {
           }
         }
 
-        const toastTitle = isDownloaded ? "Starting model…" : "Downloading model…";
+        const isCachedLoad = isDownloaded || isCachedLora;
+        const toastTitle = isCachedLoad ? "Starting model…" : "Downloading model…";
         const toastId = toast(
           null,
           {
             description: renderLoadDescription(
               toastTitle,
               loadingDescription,
-              isDownloaded ? null : 0,
-              isDownloaded ? null : "Preparing download",
+              isCachedLoad ? null : 0,
+              isCachedLoad ? null : "Preparing download",
               cancelLoading,
             ),
             duration: Infinity,
@@ -503,7 +520,7 @@ export function useChatModelRuntime() {
 
         // Poll download progress for non-cached models (GGUF and non-GGUF)
         let progressInterval: ReturnType<typeof setInterval> | null = null;
-        if (!isDownloaded) {
+        if (!isDownloaded && !isCachedLora) {
           const expectedBytes =
             typeof selection !== "string" ? selection.expectedBytes ?? 0 : 0;
           let hasShownProgress = false;

studio/frontend/src/features/chat/runtime-provider.tsx

@@ -47,9 +47,9 @@ const DEFAULT_SUGGESTIONS = [
     prompt: "Solve the integral of x·sin(x), and verify it step by step",
   },
   {
-    title: "Draw an SVG of a cute sloth",
+    title: "Draw an SVG of a cute sloth & show the code",
     label: "SVG sloth",
-    prompt: "Draw an SVG of a cute sloth",
+    prompt: "Draw an SVG of a cute sloth & show the code",
   },
 ];
 
@@ -569,10 +569,32 @@ function ThreadHistoryProvider({
           store.setContextUsage(savedUsage);
         }
 
+        // If any message has a stored parentId, reconstruct the tree
+        // so retries/regenerations load as branches instead of being
+        // unrolled into a flat list.  For mixed legacy/new threads
+        // (old messages without parentId + new messages with), infer
+        // sequential parents for old messages to preserve the chain.
+        // Fall back to fromArray for fully legacy threads.
+        const hasParentIds = msgs.some((m) => "parentId" in m);
+        if (hasParentIds) {
+          let previousId: string | null = null;
+          return {
+            messages: msgs.map((m) => {
+              const parentId = "parentId" in m
+                ? (m.parentId ?? null)
+                : previousId;
+              previousId = m.id;
+              return {
+                parentId,
+                message: toThreadMessage(m),
+              };
+            }),
+          };
+        }
         return ExportedMessageRepository.fromArray(msgs.map(toThreadMessage));
       },
 
-      async append({ message }: ExportedMessageRepositoryItem) {
+      async append({ parentId, message }: ExportedMessageRepositoryItem) {
         const { remoteId } = await aui.threadListItem().initialize();
         const content = cloneContent(message.content);
         const attachments =
@@ -586,6 +608,7 @@ function ThreadHistoryProvider({
         await db.messages.put({
           id: message.id,
           threadId: remoteId,
+          parentId: parentId ?? null,
           role: message.role,
           content,
           ...(attachments.length > 0 && { attachments }),

studio/frontend/src/features/chat/shared-composer.tsx

@@ -473,7 +473,7 @@ export function SharedComposer({
         onChange={(e) => setText(e.target.value)}
         onKeyDown={onKeyDown}
         placeholder="Send to both models..."
-        className="mb-1 max-h-32 min-h-14 w-full resize-none bg-transparent px-4 pt-2 pb-3 text-sm outline-none placeholder:text-muted-foreground"
+        className="mb-1 max-h-32 min-h-14 w-full resize-none bg-transparent pl-5 pr-4 pt-2 pb-3 text-sm outline-none placeholder:text-muted-foreground"
         rows={1}
       />
       <div className="relative mx-2 mb-2 flex items-center justify-between">

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

@@ -151,6 +151,7 @@ type ChatRuntimeStore = {
   activeGgufVariant: string | null;
   ggufContextLength: number | null;
   ggufMaxContextLength: number | null;
+  ggufNativeContextLength: number | null;
   supportsReasoning: boolean;
   reasoningAlwaysOn: boolean;
   reasoningEnabled: boolean;
@@ -215,6 +216,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
   activeGgufVariant: null,
   ggufContextLength: null,
   ggufMaxContextLength: null,
+  ggufNativeContextLength: null,
   supportsReasoning: false,
   reasoningAlwaysOn: false,
   reasoningEnabled: true,
@@ -290,6 +292,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
       activeGgufVariant: null,
       ggufContextLength: null,
       ggufMaxContextLength: null,
+      ggufNativeContextLength: null,
       contextUsage: null,
       supportsReasoning: false,
       reasoningEnabled: true,

studio/frontend/src/features/chat/types.ts

@@ -20,6 +20,7 @@ export interface ThreadRecord {
 export interface MessageRecord {
   id: string;
   threadId: string;
+  parentId?: string | null;
   role: import("@assistant-ui/react").ThreadMessage["role"];
   content: import("@assistant-ui/react").ThreadMessage["content"];
   attachments?: import("@assistant-ui/react").ThreadMessage["attachments"];

studio/frontend/src/features/chat/types/api.ts

@@ -87,6 +87,7 @@ export interface LoadModelResponse {
   };
   context_length?: number | null;
   max_context_length?: number | null;
+  native_context_length?: number | null;
   supports_reasoning?: boolean;
   reasoning_always_on?: boolean;
   supports_tools?: boolean;
@@ -121,6 +122,7 @@ export interface InferenceStatusResponse {
   supports_tools?: boolean;
   context_length?: number | null;
   max_context_length?: number | null;
+  native_context_length?: number | null;
 }
 
 export interface AudioGenerationResponse {

studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx

@@ -268,12 +268,12 @@ export function ModelSelectionStep() {
                       <ComboboxItem
                         key={id}
                         value={id}
-                        className={`justify-between ${exceeds ? "opacity-50" : ""}`}
+                        className="justify-between"
                       >
                         <Tooltip>
                           <TooltipTrigger asChild={true}>
                             <span
-                              className={`min-w-0 flex-1 truncate ${exceeds ? "line-through decoration-muted-foreground/50" : ""}`}
+                              className={`min-w-0 flex-1 truncate ${exceeds ? "!text-gray-500 dark:!text-gray-400" : ""}`}
                             >
                               {id}
                             </span>
@@ -287,12 +287,12 @@ export function ModelSelectionStep() {
                         </Tooltip>
                         <span className="flex items-center gap-1.5 shrink-0">
                           {fitStatus === "exceeds" && (
-                            <span className="text-[9px] font-medium text-red-400">
+                            <span className="text-[9px] font-medium !text-red-700 !bg-red-50 dark:!text-red-400 dark:!bg-red-950 px-1.5 py-0.5 rounded">
                               OOM
                             </span>
                           )}
                           {fitStatus === "tight" && (
-                            <span className="text-[9px] font-medium text-amber-400">
+                            <span className="text-[9px] font-medium !text-amber-400">
                               TIGHT
                             </span>
                           )}

studio/frontend/src/features/studio/sections/model-section.tsx

@@ -489,12 +489,12 @@ export function ModelSection() {
                           <ComboboxItem
                             key={id}
                             value={id}
-                            className={`gap-2 ${exceeds ? "opacity-50" : ""}`}
+                            className="gap-2"
                           >
                             <Tooltip>
                               <TooltipTrigger asChild={true}>
                                 <span
-                                  className={`block min-w-0 flex-1 truncate ${exceeds ? "line-through decoration-muted-foreground/50" : ""}`}
+                                  className={`block min-w-0 flex-1 truncate ${exceeds ? "!text-gray-500 dark:!text-gray-400" : ""}`}
                                 >
                                   {id}
                                 </span>
@@ -519,12 +519,12 @@ export function ModelSection() {
                             </Tooltip>
                             <span className="ml-auto flex items-center gap-1.5 shrink-0">
                               {fitStatus === "exceeds" && (
-                                <span className="text-[9px] font-medium text-red-400">
+                                <span className="text-[9px] font-medium !text-red-700 !bg-red-50 dark:!text-red-400 dark:!bg-red-950 px-1.5 py-0.5 rounded">
                                   OOM
                                 </span>
                               )}
                               {fitStatus === "tight" && (
-                                <span className="text-[9px] font-medium text-amber-400">
+                                <span className="text-[9px] font-medium !text-amber-400">
                                   TIGHT
                                 </span>
                               )}

studio/frontend/src/hooks/use-hf-model-search.ts

@@ -104,6 +104,11 @@ function makeMapModel(excludeGguf: boolean) {
 
 /** Number of unsloth results to pull up-front before yielding general results. */
 const UNSLOTH_PREFETCH = 20;
+/** When the user searched for a specific publisher, show fewer unsloth results
+ *  before the pinned (original publisher) model. */
+const UNSLOTH_PINNED_PREFETCH = 4;
+/** Matches a valid "owner/repo" identifier (exactly two non-empty segments). */
+const PUBLISHER_RE = /^([^/\s]+)\/([^/\s]+)$/;
 
 /**
  * Prime the hf-cache from a listModels result. For public (non-gated,
@@ -131,6 +136,7 @@ async function* mergedModelIterator(
   query: string,
   task?: PipelineType,
   accessToken?: string,
+  pinnedId?: string,
 ): AsyncGenerator<unknown> {
   const common = {
     additionalFields: ["safetensors", "tags"] as ("safetensors" | "tags")[],
@@ -148,6 +154,18 @@ async function* mergedModelIterator(
     ...common,
   });
 
+  // Start pinned model lookup immediately so it can run in parallel with
+  // the Phase 1 unsloth iteration instead of blocking Phase 2.
+  const pinnedPromise = pinnedId
+    ? cachedModelInfo({
+        name: pinnedId,
+        additionalFields: ["safetensors", "tags"],
+        ...(accessToken ? { credentials: { accessToken } } : {}),
+      }).catch(() => null)
+    : null;
+
+  const limit = pinnedId ? UNSLOTH_PINNED_PREFETCH : UNSLOTH_PREFETCH;
+
   // Phase 1: pull & yield unsloth models first
   const seen = new Set<string>();
   let count = 0;
@@ -159,10 +177,26 @@ async function* mergedModelIterator(
     }
     yield model;
     count++;
-    if (count >= UNSLOTH_PREFETCH) break;
+    if (count >= limit) break;
   }
 
-  // Phase 2: yield general results, skipping already-seen unsloth models
+  // Phase 1b: yield the pinned (original publisher) model before general results
+  if (pinnedId && !seen.has(pinnedId) && pinnedPromise) {
+    const pinned = await pinnedPromise;
+    if (pinned) {
+      // Record both the raw input and the canonical name returned by HF
+      // so phase 2 deduplication works even when casing differs
+      // (e.g. user typed "OpenAI/gpt-oss-20b", HF returns "openai/gpt-oss-20b").
+      seen.add(pinnedId);
+      const canonicalName = (pinned as { name?: string }).name;
+      if (canonicalName && canonicalName !== pinnedId) {
+        seen.add(canonicalName);
+      }
+      yield pinned;
+    }
+  }
+
+  // Phase 2: yield general results, skipping already-seen models
   for await (const model of generalIter) {
     const m = model as { name?: string };
     if (m.name && seen.has(m.name)) continue;
@@ -235,11 +269,24 @@ export function useHfModelSearch(
 ) {
   const { task, accessToken, excludeGguf = false, priorityIds } = options ?? {};
 
+  // Parse publisher detection once and share between the iterator factory
+  // and the secondary sort gate (avoids duplicating the regex + logic).
+  const { isPublisherQuery, searchQuery, pinnedId, trimmed } = useMemo(() => {
+    const t = query.trim();
+    const m = PUBLISHER_RE.exec(t);
+    const is = !!m && m[1].toLowerCase() !== "unsloth";
+    return {
+      isPublisherQuery: is,
+      searchQuery: is ? m![2] : t,
+      pinnedId: is ? t : undefined,
+      trimmed: t,
+    };
+  }, [query]);
+
   const createIter = useCallback(
     () => {
-      const trimmed = query.trim();
       if (!trimmed) {
-        // No query → show priority models first (with full metadata), then general unsloth listing
+        // No query: show priority models first (with full metadata), then general unsloth listing
         if (priorityIds && priorityIds.length > 0) {
           return priorityThenListingIterator(priorityIds, task, accessToken) as AsyncGenerator<unknown>;
         }
@@ -250,24 +297,35 @@ export function useHfModelSearch(
           ...(accessToken ? { credentials: { accessToken } } : {}),
         }) as AsyncGenerator<unknown>;
       }
-      // Typed query: disable task filter so explicitly searched models still appear even if HF task metadata is wrong/missing.
-      return mergedModelIterator(trimmed, undefined, accessToken) as AsyncGenerator<unknown>;
+      // Typed query: disable task filter so explicitly searched models still
+      // appear even if HF task metadata is wrong/missing.
+      // If the query is a valid "owner/repo" identifier (exactly two non-empty,
+      // slash-free, space-free segments), strip the org prefix so unsloth
+      // variants surface, then pin the original publisher model after a small
+      // batch of unsloth results.  Queries for unsloth-owned models are left
+      // as-is so they get the full 20-result prefetch and secondary sort.
+      return mergedModelIterator(searchQuery, undefined, accessToken, pinnedId) as AsyncGenerator<unknown>;
     },
-    [query, task, accessToken, priorityIds],
+    [trimmed, searchQuery, pinnedId, task, accessToken, priorityIds],
   );
 
   const mapModel = useMemo(() => makeMapModel(excludeGguf), [excludeGguf]);
   const search = useHfPaginatedSearch(createIter, mapModel);
 
-  // Secondary sort guarantee: unsloth models always float to the top
+  // Secondary sort guarantee: unsloth models always float to the top.
+  // Skip when the user searched for a specific non-unsloth publisher
+  // (e.g. "openai/gpt-oss-20b") -- the iterator already handles the
+  // pinned ordering in that case.
   const results = useMemo(
     () =>
-      [...search.results].sort((a, b) => {
-        const aFirst = a.id.startsWith("unsloth/") ? 0 : 1;
-        const bFirst = b.id.startsWith("unsloth/") ? 0 : 1;
-        return aFirst - bFirst;
-      }),
-    [search.results],
+      isPublisherQuery
+        ? search.results
+        : [...search.results].sort((a, b) => {
+            const aFirst = a.id.startsWith("unsloth/") ? 0 : 1;
+            const bFirst = b.id.startsWith("unsloth/") ? 0 : 1;
+            return aFirst - bFirst;
+          }),
+    [search.results, isPublisherQuery],
   );
 
   return { ...search, results };

studio/frontend/src/index.css

@@ -12,6 +12,8 @@
 @plugin "@toolwind/corner-shape";
 @source "../node_modules/streamdown/dist/*.js";
 
+@custom-variant dark (&:is(.dark *));
+
 @font-face {
   font-family: "Hellix";
   src: url("/fonts/Hellix-SemiBold.woff2") format("woff2"),
@@ -21,8 +23,6 @@
   font-display: swap;
 }
 
-@custom-variant dark (&:is(.dark *));
-
 :root {
   /* Animation timing */
   --duration-micro: 100ms;
@@ -67,7 +67,7 @@
   --sidebar-ring: oklch(0.6929 0.1396 166.5513);
   --destructive-foreground: oklch(1 0 0);
   --font-sans: "Inter Variable", ui-sans-serif, sans-serif, system-ui;
-  --font-heading: "Hellix", "Space Grotesk Variable", ui-sans-serif, sans-serif;
+  --font-heading: "Hellix", "Space Grotesk Variable", var(--font-sans);
   --font-serif: Source Serif 4, serif;
   --font-mono: JetBrains Mono, monospace;
   --shadow-color: hsl(0 0% 0%);
@@ -159,7 +159,7 @@
 
 @theme inline {
   --font-sans: "Inter Variable", ui-sans-serif, sans-serif, system-ui;
-  --font-heading: "Hellix", "Space Grotesk Variable", ui-sans-serif, sans-serif;
+  --font-heading: "Hellix", "Space Grotesk Variable", var(--font-sans);
   --color-sidebar-ring: var(--sidebar-ring);
   --color-sidebar-border: var(--sidebar-border);
   --color-sidebar-accent-foreground: var(--sidebar-accent-foreground);
@@ -343,10 +343,58 @@
   [data-streamdown="code-block"] {
     gap: 0;
     padding: 0.5rem;
+    /* Wide lines must scroll inside the thread column, not widen past the composer (flex min-width:auto). */
+    max-width: 100%;
+    min-width: 0;
+    overflow-x: auto;
   }
   [data-streamdown="code-block-header"] {
     padding-left: 0.75rem;
   }
+
+  /* Chat thread: code slightly smaller by default; step up when the thread column is wide. */
+  .aui-thread-root [data-streamdown="code-block"] {
+    font-size: 0.8125rem;
+    line-height: 1.55;
+  }
+
+  .aui-thread-root [data-streamdown="code-block-header"] {
+    font-size: 0.6875rem;
+  }
+
+  @container (min-width: 36rem) {
+    .aui-thread-root [data-streamdown="code-block"] {
+      font-size: 0.875rem;
+    }
+
+    .aui-thread-root [data-streamdown="code-block-header"] {
+      font-size: 0.75rem;
+    }
+  }
+
+  /* Chat: use the app sans stack for UI + prose. */
+  .aui-thread-root {
+    --font-heading: var(--font-sans);
+    font-family: var(--font-sans);
+  }
+
+  /* Keep monospace for code fences and inline code (not KaTeX). */
+  .aui-thread-root [data-streamdown="code-block"] pre,
+  .aui-thread-root [data-streamdown="code-block"] code {
+    font-family: var(--font-mono), ui-monospace, monospace;
+  }
+
+  .aui-thread-root :where(p, li, td, th, blockquote, h1, h2, h3, h4, h5, h6) code {
+    font-family: var(--font-mono), ui-monospace, monospace;
+  }
+
+  /* Align fenced code blocks with the main chat column even when nested in lists. */
+  .aui-thread-root [data-streamdown="list-item"] > [data-streamdown="code-block"],
+  .aui-thread-root [data-streamdown="list-item"] [data-streamdown="code-block"] {
+    margin-left: -1.25rem;
+    width: calc(100% + 1.25rem);
+    max-width: calc(100% + 1.25rem);
+  }
 }
 
 /* Minimal scrollbar — thumb only, no track */

studio/setup.ps1

@@ -22,6 +22,19 @@ $ErrorActionPreference = "Stop"
 $ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path
 $PackageDir = Split-Path -Parent $ScriptDir
 
+# --------------------------------------------------------------------------
+#  Maintainer-editable defaults
+#  Change these in the GitHub-hosted script so users get updated defaults.
+#  User env vars always override these baked-in values.
+# --------------------------------------------------------------------------
+# Prefer "latest" over "master" -- "master" bypasses the prebuilt resolver
+# (no matching GitHub release), forces a source build, and causes HTTP 422
+# errors. Only use "master" temporarily when the latest release is missing
+# support for a new model architecture.
+$DefaultLlamaPrForce = ""
+$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp"
+$DefaultLlamaTag = "latest"
+
 # Verbose can be enabled either by CLI flag or by UNSLOTH_VERBOSE=1.
 $script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq '1')
 foreach ($a in $args) {
@@ -62,6 +75,12 @@ function Refresh-Environment {
     $env:Path = "$machinePath;$userPath"
 }
 
+# PowerShell 5.1 compatibility helper: avoid relying on New-TemporaryFile.
+function New-UnslothTemporaryFile {
+    $tempPath = [System.IO.Path]::GetTempFileName()
+    return Get-Item -LiteralPath $tempPath
+}
+
 # Find nvcc on PATH, CUDA_PATH, or standard toolkit dirs.
 # Returns the path to nvcc.exe, or $null if not found.
 function Find-Nvcc {
@@ -490,8 +509,7 @@ if (-not $HasNvidiaSmi) {
 if (-not $HasNvidiaSmi) {
     Write-Host ""
     step "gpu" "none (chat-only / GGUF)" "Yellow"
-    Write-Host "       Training and GPU inference require an NVIDIA GPU with drivers installed." -ForegroundColor Yellow
-    Write-Host "       https://www.nvidia.com/Download/index.aspx" -ForegroundColor Yellow
+    substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow"
     Write-Host ""
 } else {
     step "gpu" "NVIDIA GPU detected"
@@ -1544,7 +1562,7 @@ if (Test-Path $VenvT5Dir) { Remove-Item -Recurse -Force $VenvT5Dir }
 New-Item -ItemType Directory -Path $VenvT5Dir -Force | Out-Null
 $prevEAP_t5 = $ErrorActionPreference
 $ErrorActionPreference = "Continue"
-foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.7.1", "hf_xet==1.4.2")) {
+foreach ($pkg in @("transformers==5.5.0", "huggingface_hub==1.8.0", "hf_xet==1.4.2")) {
     if ($script:UnslothVerbose) {
         Fast-Install --target $VenvT5Dir --no-deps $pkg
         $t5PkgExit = $LASTEXITCODE
@@ -1590,43 +1608,156 @@ if (-not (Test-Path $UnslothHome)) { New-Item -ItemType Directory -Force $Unslot
 $LlamaCppDir = Join-Path $UnslothHome "llama.cpp"
 $NeedLlamaSourceBuild = $false
 $SkipPrebuiltInstall = $false
-$RequestedLlamaTag = if ($env:UNSLOTH_LLAMA_TAG) { $env:UNSLOTH_LLAMA_TAG } else { "latest" }
-$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO) { $env:UNSLOTH_LLAMA_RELEASE_REPO } else { "unslothai/llama.cpp" }
-$resolveOutput = & python "$PSScriptRoot\install_llama_prebuilt.py" --resolve-install-tag $RequestedLlamaTag --published-repo $HelperReleaseRepo 2>&1
-$resolveExit = $LASTEXITCODE
-$ResolvedLlamaTag = if ($resolveOutput) { ($resolveOutput | Select-Object -Last 1).ToString().Trim() } else { "" }
-if ($resolveExit -ne 0 -or [string]::IsNullOrWhiteSpace($ResolvedLlamaTag)) {
-    Write-Host ""
-    substep "Failed to resolve an installable prebuilt llama.cpp tag via $HelperReleaseRepo" "Yellow"
-    Write-LlamaFailureLog -Output ($resolveOutput | Out-String)
-    # Resolve the llama.cpp tag for source-build fallback. Pass --published-repo
-    # so the resolver prefers Unsloth's tested tag (e.g. b8508) over the upstream
-    # bleeding-edge tag (e.g. b8514) from ggml-org/llama.cpp.
-    $fallbackOutput = & python "$PSScriptRoot\install_llama_prebuilt.py" --resolve-llama-tag $RequestedLlamaTag --published-repo $HelperReleaseRepo 2>$null
-    $fallbackExit = $LASTEXITCODE
-    $ResolvedLlamaTag = if ($fallbackExit -eq 0 -and $fallbackOutput) {
-        ($fallbackOutput | Select-Object -Last 1).ToString().Trim()
-    } elseif ($RequestedLlamaTag -eq "latest") {
-        # Try Unsloth release repo first, then fall back to ggml-org upstream
-        $resolvedLatest = $null
-        try {
-            $latestRelease = Invoke-RestMethod -Uri "https://api.github.com/repos/$HelperReleaseRepo/releases/latest" -ErrorAction Stop
-            $resolvedLatest = $latestRelease.tag_name
-        } catch {}
-        if (-not $resolvedLatest) {
-            try {
-                $latestRelease = Invoke-RestMethod -Uri "https://api.github.com/repos/ggml-org/llama.cpp/releases/latest" -ErrorAction Stop
-                $resolvedLatest = $latestRelease.tag_name
-            } catch {}
-        }
-        if ($resolvedLatest) { $resolvedLatest } else { $RequestedLlamaTag }
-    } else {
-        $RequestedLlamaTag
-    }
+$RequestedLlamaTag = if ($env:UNSLOTH_LLAMA_TAG) { $env:UNSLOTH_LLAMA_TAG } else { $DefaultLlamaTag }
+$HelperReleaseRepo = "ggml-org/llama.cpp"
+$LlamaPr = if ($env:UNSLOTH_LLAMA_PR) { $env:UNSLOTH_LLAMA_PR.Trim() } else { "" }
+
+$LlamaPrForce = if ($env:UNSLOTH_LLAMA_PR_FORCE) { $env:UNSLOTH_LLAMA_PR_FORCE.Trim() } else { $DefaultLlamaPrForce }
+$LlamaSource = $DefaultLlamaSource
+if ($LlamaSource.EndsWith('.git')) { $LlamaSource = $LlamaSource.Substring(0, $LlamaSource.Length - 4) }
+$ResolvedSourceUrl = $LlamaSource
+$ResolvedSourceRef = $RequestedLlamaTag
+$ResolvedSourceRefKind = "tag"
+
+if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
     $NeedLlamaSourceBuild = $true
     $SkipPrebuiltInstall = $true
 }
 
+function Invoke-LlamaHelper {
+    param(
+        [string[]]$Arguments,
+        [string]$StderrPath = $null
+    )
+
+    $previousErrorActionPreference = $ErrorActionPreference
+    $previousNativeErrorPreference = $null
+    $restoreNativeErrorPreference = $false
+    $ErrorActionPreference = "Continue"
+    if ($PSVersionTable.PSVersion.Major -ge 7) {
+        $previousNativeErrorPreference = $PSNativeCommandUseErrorActionPreference
+        $PSNativeCommandUseErrorActionPreference = $false
+        $restoreNativeErrorPreference = $true
+    }
+
+    try {
+        # Capture all output (stdout + stderr) so that PowerShell does not
+        # convert stderr lines into visible ErrorRecord objects.  Separate
+        # stdout from stderr afterwards.
+        $allOutput = & python "$PSScriptRoot\install_llama_prebuilt.py" @Arguments 2>&1
+        $exitCode = $LASTEXITCODE
+        $stdoutLines = @()
+        $stderrLines = @()
+        foreach ($line in $allOutput) {
+            if ($line -is [System.Management.Automation.ErrorRecord]) {
+                $stderrLines += $line.ToString()
+            } else {
+                $stdoutLines += $line
+            }
+        }
+        if ($StderrPath -and $stderrLines.Count -gt 0) {
+            $stderrLines | Out-File -FilePath $StderrPath -Encoding utf8
+        }
+        return [pscustomobject]@{
+            Output = $stdoutLines
+            ExitCode = $exitCode
+        }
+    } finally {
+        $ErrorActionPreference = $previousErrorActionPreference
+        if ($restoreNativeErrorPreference) {
+            $PSNativeCommandUseErrorActionPreference = $previousNativeErrorPreference
+        }
+    }
+}
+
+if ($LlamaSource -ne "https://github.com/ggml-org/llama.cpp") {
+    step "llama.cpp" "custom source: $LlamaSource -- forcing source build" "Yellow"
+    $NeedLlamaSourceBuild = $true
+    $SkipPrebuiltInstall = $true
+}
+
+if (-not $LlamaPr -and $LlamaPrForce -and $LlamaPrForce -match '^\d+$' -and [int]$LlamaPrForce -gt 0) {
+    $LlamaPr = $LlamaPrForce
+    step "llama.cpp" "baked-in PR_FORCE=$LlamaPrForce" "Yellow"
+}
+
+if ($LlamaPr) {
+    if ($LlamaPr -notmatch '^\d+$' -or [int]$LlamaPr -le 0) {
+        Write-Host "[ERROR] UNSLOTH_LLAMA_PR=$LlamaPr is not a valid PR number" -ForegroundColor Red
+        exit 1
+    }
+    step "llama.cpp" "UNSLOTH_LLAMA_PR=$LlamaPr -- will build from PR head" "Yellow"
+    $ResolvedLlamaTag = "pr-$LlamaPr"
+    $ResolvedSourceUrl = $LlamaSource
+    $ResolvedSourceRef = "pr-$LlamaPr"
+    $ResolvedSourceRefKind = "pull"
+    $NeedLlamaSourceBuild = $true
+    $SkipPrebuiltInstall = $true
+} elseif ($SkipPrebuiltInstall) {
+    # Custom source or other override already forced source build; skip the
+    # prebuilt release resolution. When building from a custom fork, the fork
+    # may not carry upstream bNNNN tags.
+    if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
+        $ResolvedLlamaTag = $RequestedLlamaTag
+    } elseif ($LlamaSource -eq "https://github.com/ggml-org/llama.cpp") {
+        $resolveTagArgs = @("--resolve-llama-tag", $RequestedLlamaTag, "--published-repo", $HelperReleaseRepo, "--output-format", "json")
+        if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $resolveTagArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) }
+        $fallbackResult = Invoke-LlamaHelper -Arguments $resolveTagArgs
+        $fallbackOutput = $fallbackResult.Output
+        $fallbackExit = $fallbackResult.ExitCode
+        $ResolvedLlamaTag = if ($fallbackExit -eq 0 -and $fallbackOutput) {
+            try {
+                (($fallbackOutput | Out-String) | ConvertFrom-Json).llama_tag
+            } catch {
+                $RequestedLlamaTag
+            }
+        } else {
+            $RequestedLlamaTag
+        }
+    } else {
+        $ResolvedLlamaTag = $RequestedLlamaTag
+    }
+} else {
+    $resolveInstallArgs = @("--resolve-install-tag", $RequestedLlamaTag, "--published-repo", $HelperReleaseRepo, "--output-format", "json")
+    if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $resolveInstallArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) }
+    $resolveErrorLog = New-UnslothTemporaryFile
+    $resolveResult = Invoke-LlamaHelper -Arguments $resolveInstallArgs -StderrPath $resolveErrorLog
+    $resolveOutput = $resolveResult.Output
+    $resolveExit = $resolveResult.ExitCode
+    $ResolvedLlamaTag = if ($resolveOutput) {
+        try {
+            (($resolveOutput | Out-String) | ConvertFrom-Json).llama_tag
+        } catch {
+            ""
+        }
+    } else { "" }
+    if ($resolveExit -ne 0 -or [string]::IsNullOrWhiteSpace($ResolvedLlamaTag)) {
+        Write-Host ""
+        substep "Failed to resolve a published llama.cpp release via $HelperReleaseRepo" "Yellow"
+        Write-LlamaFailureLog -Output (Get-Content -Raw $resolveErrorLog)
+        # Resolve the llama.cpp tag for source-build fallback. Pass --published-repo
+        # so the resolver prefers the latest usable Unsloth-published upstream tag
+        # before falling back to the bleeding-edge ggml-org/llama.cpp tag.
+        $resolveFallbackArgs = @("--resolve-llama-tag", $RequestedLlamaTag, "--published-repo", $HelperReleaseRepo, "--output-format", "json")
+        if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $resolveFallbackArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) }
+        $fallbackResult = Invoke-LlamaHelper -Arguments $resolveFallbackArgs
+        $fallbackOutput = $fallbackResult.Output
+        $fallbackExit = $fallbackResult.ExitCode
+        $ResolvedLlamaTag = if ($fallbackExit -eq 0 -and $fallbackOutput) {
+            try {
+                (($fallbackOutput | Out-String) | ConvertFrom-Json).llama_tag
+            } catch {
+                $RequestedLlamaTag
+            }
+        } else {
+            $RequestedLlamaTag
+        }
+        $NeedLlamaSourceBuild = $true
+        $SkipPrebuiltInstall = $true
+    }
+    Remove-Item $resolveErrorLog -Force -ErrorAction SilentlyContinue
+}
+
 Write-Host ""
 substep "Resolved llama.cpp release tag: $ResolvedLlamaTag"
 
@@ -1646,7 +1777,7 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
         $prebuiltArgs = @(
             "$PSScriptRoot\install_llama_prebuilt.py",
             "--install-dir", $LlamaCppDir,
-            "--llama-tag", $ResolvedLlamaTag,
+            "--llama-tag", $RequestedLlamaTag,
             "--published-repo", $HelperReleaseRepo
         )
         if ($env:UNSLOTH_LLAMA_RELEASE_TAG) {
@@ -1654,21 +1785,46 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
         }
         $prevEAPPrebuilt = $ErrorActionPreference
         $ErrorActionPreference = "Continue"
-        if ($script:UnslothVerbose) {
-            # Show live output in verbose mode while still capturing for error log
-            $prebuiltLog = Join-Path $env:TEMP "unsloth-prebuilt-$PID.log"
-            & python @prebuiltArgs 2>&1 | Tee-Object -FilePath $prebuiltLog | Out-Host
-            $prebuiltExit = $LASTEXITCODE
-            $prebuiltOutput = if (Test-Path $prebuiltLog) { Get-Content $prebuiltLog -Raw } else { "" }
-            Remove-Item $prebuiltLog -ErrorAction SilentlyContinue
-        } else {
-            $prebuiltOutput = & python @prebuiltArgs 2>&1 | Out-String
-            $prebuiltExit = $LASTEXITCODE
+        $previousNativeErrorPreference = $null
+        $restoreNativeErrorPreference = $false
+        if ($PSVersionTable.PSVersion.Major -ge 7) {
+            $previousNativeErrorPreference = $PSNativeCommandUseErrorActionPreference
+            $PSNativeCommandUseErrorActionPreference = $false
+            $restoreNativeErrorPreference = $true
+        }
+        try {
+            if ($script:UnslothVerbose) {
+                # Show live output in verbose mode while still capturing for error log
+                $prebuiltLog = Join-Path $env:TEMP "unsloth-prebuilt-$PID.log"
+                & python @prebuiltArgs 2>&1 | Tee-Object -FilePath $prebuiltLog | Out-Host
+                $prebuiltExit = $LASTEXITCODE
+                $prebuiltOutput = if (Test-Path $prebuiltLog) { Get-Content $prebuiltLog -Raw } else { "" }
+                Remove-Item $prebuiltLog -ErrorAction SilentlyContinue
+            } else {
+                $prebuiltOutput = & python @prebuiltArgs 2>&1 | Out-String
+                $prebuiltExit = $LASTEXITCODE
+            }
+        } finally {
+            if ($restoreNativeErrorPreference) {
+                $PSNativeCommandUseErrorActionPreference = $previousNativeErrorPreference
+            }
         }
         $ErrorActionPreference = $prevEAPPrebuilt
 
         if ($prebuiltExit -eq 0) {
-            step "llama.cpp" "prebuilt installed and validated"
+            if ($prebuiltOutput -match "already matches") {
+                step "llama.cpp" "prebuilt up to date and validated"
+            } else {
+                step "llama.cpp" "prebuilt installed and validated"
+            }
+        } elseif ($prebuiltExit -eq 3) {
+            step "llama.cpp" "install blocked by active llama.cpp process" "Yellow"
+            Write-LlamaFailureLog -Output $prebuiltOutput
+            if (Test-Path $LlamaCppDir) {
+                substep "Existing install was restored" "Yellow"
+            }
+            substep "Close Studio or other llama.cpp users and retry" "Yellow"
+            exit 3
         } else {
             step "llama.cpp" "prebuilt install failed (continuing)" "Yellow"
             Write-LlamaFailureLog -Output $prebuiltOutput
@@ -1766,7 +1922,10 @@ if (Test-Path $LlamaServerBin) {
 if (-not $NeedLlamaSourceBuild) {
     Write-Host ""
     step "llama.cpp" "prebuilt (validated)"
-} elseif ((Test-Path $LlamaServerBin) -and -not $NeedRebuild) {
+} elseif ((Test-Path $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") {
+    # Skip rebuild only for pinned tags (e.g. b8635).  When the requested
+    # tag is "master" (a moving target), always rebuild so the binary picks
+    # up new model architecture support (e.g. Gemma 4).
     Write-Host ""
     step "llama.cpp" "already built"
 } elseif (-not $HasCmakeForBuild) {
@@ -1821,42 +1980,188 @@ if (-not $NeedLlamaSourceBuild) {
         [Environment]::SetEnvironmentVariable('CudaToolkitDir', "$CudaToolkitRoot\", 'Process')
     }
 
+    if (-not $LlamaPr) {
+        if ($LlamaSource -eq "https://github.com/ggml-org/llama.cpp") {
+            $resolveSourceArgs = @("--resolve-source-build", $RequestedLlamaTag, "--published-repo", $HelperReleaseRepo, "--output-format", "json")
+            if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $resolveSourceArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) }
+            $sourcePlanResult = Invoke-LlamaHelper -Arguments $resolveSourceArgs
+            $sourcePlanOutput = $sourcePlanResult.Output
+            $sourcePlanExit = $sourcePlanResult.ExitCode
+            if ($sourcePlanExit -eq 0 -and $sourcePlanOutput) {
+                try {
+                    $sourcePlan = ($sourcePlanOutput | Out-String) | ConvertFrom-Json
+                    $ResolvedSourceUrl = $sourcePlan.source_url
+                    $ResolvedSourceRefKind = $sourcePlan.source_ref_kind
+                    $ResolvedSourceRef = $sourcePlan.source_ref
+                } catch {
+                }
+            }
+        }
+        if ([string]::IsNullOrWhiteSpace($ResolvedSourceUrl)) { $ResolvedSourceUrl = $LlamaSource }
+        if ([string]::IsNullOrWhiteSpace($ResolvedSourceRef)) { $ResolvedSourceRef = $ResolvedLlamaTag }
+    }
+
     # -- Step A: Clone or pull llama.cpp --
 
-    $UseConcreteRef = ($ResolvedLlamaTag -ne "latest" -and -not [string]::IsNullOrWhiteSpace($ResolvedLlamaTag))
+    $UseConcreteRef = ($ResolvedSourceRef -ne "latest" -and -not [string]::IsNullOrWhiteSpace($ResolvedSourceRef))
 
     if (Test-Path (Join-Path $LlamaCppDir ".git")) {
-        Write-Host "   Syncing llama.cpp to $ResolvedLlamaTag..." -ForegroundColor Gray
-        if ($UseConcreteRef) {
-            $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin $ResolvedLlamaTag }
+        Write-Host "   Syncing llama.cpp to $ResolvedSourceRef..." -ForegroundColor Gray
+        # Always sync the remote URL so switching between default/fork sources works
+        Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir remote set-url origin "$ResolvedSourceUrl.git" } | Out-Null
+        if ($LlamaPr) {
+            $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin "pull/$LlamaPr/head" }
+            if ($gitFetchExit -ne 0) {
+                $BuildOk = $false
+                $FailedStep = "git fetch PR #$LlamaPr"
+            } else {
+                $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B "pr-$LlamaPr" FETCH_HEAD }
+                if ($gitCheckoutExit -ne 0) {
+                    $BuildOk = $false
+                    $FailedStep = "git checkout PR #$LlamaPr"
+                } else {
+                    Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null
+                }
+            }
+        } elseif ($ResolvedSourceRefKind -eq "pull") {
+            $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin $ResolvedSourceRef }
+            if ($gitFetchExit -ne 0) {
+                substep "git fetch failed -- using existing source" "Yellow"
+            } else {
+                $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD }
+                if ($gitCheckoutExit -ne 0) {
+                    $BuildOk = $false
+                    $FailedStep = "git checkout"
+                } else {
+                    Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null
+                }
+            }
+        } elseif ($ResolvedSourceRefKind -eq "commit") {
+            $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin $ResolvedSourceRef }
+            if ($gitFetchExit -ne 0) {
+                substep "git fetch failed -- using existing source" "Yellow"
+            } else {
+                $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD }
+                if ($gitCheckoutExit -ne 0) {
+                    $BuildOk = $false
+                    $FailedStep = "git checkout"
+                } else {
+                    Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null
+                }
+            }
+        } elseif ($UseConcreteRef) {
+            $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin $ResolvedSourceRef }
+            if ($gitFetchExit -ne 0) {
+                substep "git fetch failed -- using existing source" "Yellow"
+            } else {
+                $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD }
+                if ($gitCheckoutExit -ne 0) {
+                    $BuildOk = $false
+                    $FailedStep = "git checkout"
+                } else {
+                    Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null
+                }
+            }
         } else {
             $gitFetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir fetch --depth 1 origin }
-        }
-        if ($gitFetchExit -ne 0) {
-            substep "git fetch failed -- using existing source" "Yellow"
-        } else {
-            $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD }
-            if ($gitCheckoutExit -ne 0) {
-                $BuildOk = $false
-                $FailedStep = "git checkout"
+            if ($gitFetchExit -ne 0) {
+                substep "git fetch failed -- using existing source" "Yellow"
             } else {
-                Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null
+                $gitCheckoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir checkout -B unsloth-llama-build FETCH_HEAD }
+                if ($gitCheckoutExit -ne 0) {
+                    $BuildOk = $false
+                    $FailedStep = "git checkout"
+                } else {
+                    Invoke-SetupCommand -AlwaysQuiet { git -C $LlamaCppDir clean -fdx } | Out-Null
+                }
             }
         }
     } else {
-        Write-Host "   Cloning llama.cpp @ $ResolvedLlamaTag..." -ForegroundColor Gray
+        Write-Host "   Cloning llama.cpp @ $ResolvedSourceRef..." -ForegroundColor Gray
         $buildTmp = "$LlamaCppDir.build.$PID"
+        $null = New-Item -ItemType Directory -Force -Path (Split-Path $LlamaCppDir -Parent)
         if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
-        $cloneArgs = @("clone", "--depth", "1")
-        if ($UseConcreteRef) {
-            $cloneArgs += @("--branch", $ResolvedLlamaTag)
-        }
-        $cloneArgs += @("https://github.com/ggml-org/llama.cpp.git", $buildTmp)
-        $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git @cloneArgs }
-        if ($cloneExit -ne 0) {
-            $BuildOk = $false
-            $FailedStep = "git clone"
-            if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
+        if ($LlamaPr) {
+            $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git clone --depth 1 "$LlamaSource.git" $buildTmp }
+            if ($cloneExit -ne 0) {
+                $BuildOk = $false
+                $FailedStep = "git clone"
+                if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
+            }
+            if ($BuildOk) {
+                $fetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp fetch --depth 1 origin "pull/$LlamaPr/head:pr-$LlamaPr" }
+                if ($fetchExit -ne 0) {
+                    $BuildOk = $false
+                    $FailedStep = "git fetch PR #$LlamaPr"
+                    if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
+                }
+            }
+            if ($BuildOk) {
+                $checkoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp checkout "pr-$LlamaPr" }
+                if ($checkoutExit -ne 0) {
+                    $BuildOk = $false
+                    $FailedStep = "git checkout PR #$LlamaPr"
+                    if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
+                }
+            }
+        } elseif ($ResolvedSourceRefKind -eq "pull") {
+            $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git clone --depth 1 "$ResolvedSourceUrl.git" $buildTmp }
+            if ($cloneExit -ne 0) {
+                $BuildOk = $false
+                $FailedStep = "git clone"
+                if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
+            }
+            if ($BuildOk) {
+                $fetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp fetch --depth 1 origin $ResolvedSourceRef }
+                if ($fetchExit -ne 0) {
+                    $BuildOk = $false
+                    $FailedStep = "git fetch source PR ref"
+                    if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
+                }
+            }
+            if ($BuildOk) {
+                $checkoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp checkout -B unsloth-llama-build FETCH_HEAD }
+                if ($checkoutExit -ne 0) {
+                    $BuildOk = $false
+                    $FailedStep = "git checkout source PR ref"
+                    if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
+                }
+            }
+        } elseif ($ResolvedSourceRefKind -eq "commit") {
+            $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git clone --depth 1 "$ResolvedSourceUrl.git" $buildTmp }
+            if ($cloneExit -ne 0) {
+                $BuildOk = $false
+                $FailedStep = "git clone"
+                if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
+            }
+            if ($BuildOk) {
+                $fetchExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp fetch --depth 1 origin $ResolvedSourceRef }
+                if ($fetchExit -ne 0) {
+                    $BuildOk = $false
+                    $FailedStep = "git fetch source commit"
+                    if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
+                }
+            }
+            if ($BuildOk) {
+                $checkoutExit = Invoke-SetupCommand -AlwaysQuiet { git -C $buildTmp checkout -B unsloth-llama-build FETCH_HEAD }
+                if ($checkoutExit -ne 0) {
+                    $BuildOk = $false
+                    $FailedStep = "git checkout source commit"
+                    if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
+                }
+            }
+        } else {
+            $cloneArgs = @("clone", "--depth", "1")
+            if ($UseConcreteRef) {
+                $cloneArgs += @("--branch", $ResolvedSourceRef)
+            }
+            $cloneArgs += @("$ResolvedSourceUrl.git", $buildTmp)
+            $cloneExit = Invoke-SetupCommand -AlwaysQuiet { git @cloneArgs }
+            if ($cloneExit -ne 0) {
+                $BuildOk = $false
+                $FailedStep = "git clone"
+                if (Test-Path $buildTmp) { Remove-Item -Recurse -Force $buildTmp }
+            }
         }
         # Use temp dir for build; swap into $LlamaCppDir only after build succeeds
         if ($BuildOk) {

studio/setup.sh

@@ -8,6 +8,24 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
 RULE=$(printf '\342\224\200%.0s' {1..52})
 
+# ── Maintainer-editable defaults ──────────────────────────────────────────
+# Change these in the GitHub-hosted script so all users get updated defaults.
+# User environment variables always override these baked-in values.
+#
+#   _DEFAULT_LLAMA_PR_FORCE : PR number to build by default ("" = normal path)
+#   _DEFAULT_LLAMA_SOURCE   : git clone URL for source builds
+#   _DEFAULT_LLAMA_TAG      : llama.cpp ref to build ("latest" = newest release,
+#                             "master" = bleeding-edge, "bNNNN" = specific tag)
+#                             Prefer "latest" over "master" -- "master" bypasses
+#                             the prebuilt resolver (no matching GitHub release),
+#                             forces a source build, and causes HTTP 422 errors.
+#                             Only use "master" temporarily when the latest release
+#                             is missing support for a new model architecture.
+# ──────────────────────────────────────────────────────────────────────────
+_DEFAULT_LLAMA_PR_FORCE=""
+_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp"
+_DEFAULT_LLAMA_TAG="latest"
+
 # ── Colors (same palette as startup_banner / install_python_stack) ──
 if [ -n "${NO_COLOR:-}" ]; then
     C_TITLE= C_DIM= C_OK= C_WARN= C_ERR= C_RST=
@@ -108,6 +126,10 @@ echo ""
 printf "  ${C_TITLE}%s${C_RST}\n" "🦥 Unsloth Studio Setup"
 printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
 verbose_substep "verbose diagnostics enabled"
+_LLAMA_ONLY="${UNSLOTH_STUDIO_LLAMA_ONLY:-0}"
+if [ "$_LLAMA_ONLY" = "1" ]; then
+    substep "llama.cpp only mode"
+fi
 # ── Clean up stale caches ──
 rm -rf "$REPO_ROOT/unsloth_compiled_cache"
 rm -rf "$SCRIPT_DIR/backend/unsloth_compiled_cache"
@@ -120,6 +142,7 @@ if [[ "$keynames" == *$'\nCOLAB_'* ]]; then
     IS_COLAB=true
 fi
 
+if [ "$_LLAMA_ONLY" != "1" ]; then
 # ── Frontend ──
 _NEED_FRONTEND_BUILD=true
 if [ -d "$SCRIPT_DIR/frontend/dist" ]; then
@@ -444,8 +467,8 @@ if [ "$_SKIP_PYTHON_DEPS" = false ]; then
     # at runtime (slow, ~10-15s), we pre-install into a separate directory.
     # The training subprocess just prepends .venv_t5/ to sys.path -- instant switch.
     mkdir -p "$VENV_T5_DIR"
-    run_quiet "install transformers 5.x" fast_install --target "$VENV_T5_DIR" --no-deps "transformers==5.3.0"
-    run_quiet "install huggingface_hub for t5" fast_install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.7.1"
+    run_quiet "install transformers 5.x" fast_install --target "$VENV_T5_DIR" --no-deps "transformers==5.5.0"
+    run_quiet "install huggingface_hub for t5" fast_install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.8.0"
     run_quiet "install hf_xet for t5" fast_install --target "$VENV_T5_DIR" --no-deps "hf_xet==1.4.2"
     run_quiet "install tiktoken for t5" fast_install --target "$VENV_T5_DIR" "tiktoken"
     step "transformers" "5.x pre-installed"
@@ -453,6 +476,7 @@ else
     step "python" "dependencies up to date"
     verbose_substep "python deps check: installed=$_PKG_NAME@${INSTALLED_VER:-unknown} latest=${LATEST_VER:-unknown}"
 fi
+fi
 
 # ── 7. Prefer prebuilt llama.cpp bundles before any source build path ──
 UNSLOTH_HOME="$HOME/.unsloth"
@@ -462,46 +486,131 @@ LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server"
 _NEED_LLAMA_SOURCE_BUILD=false
 _LLAMA_CPP_DEGRADED=false
 _LLAMA_FORCE_COMPILE="${UNSLOTH_LLAMA_FORCE_COMPILE:-0}"
-_REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-latest}"
-_HELPER_RELEASE_REPO="${UNSLOTH_LLAMA_RELEASE_REPO:-unslothai/llama.cpp}"
-_RESOLVE_LLAMA_LOG="$(mktemp)"
-set +e
-python "$SCRIPT_DIR/install_llama_prebuilt.py" \
-    --resolve-install-tag "$_REQUESTED_LLAMA_TAG" \
-    --published-repo "$_HELPER_RELEASE_REPO" >"$_RESOLVE_LLAMA_LOG" 2>&1
-_RESOLVE_LLAMA_STATUS=$?
-set -e
-if [ "$_RESOLVE_LLAMA_STATUS" -eq 0 ]; then
-    _RESOLVED_LLAMA_TAG="$(tail -n 1 "$_RESOLVE_LLAMA_LOG" | tr -d '\r')"
-else
-    _RESOLVED_LLAMA_TAG=""
+_REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-${_DEFAULT_LLAMA_TAG}}"
+# Force all installs to use mainline llama.cpp from ggml-org.
+_HELPER_RELEASE_REPO="ggml-org/llama.cpp"
+_LLAMA_PR="${UNSLOTH_LLAMA_PR:-}"
+
+_LLAMA_PR_FORCE="${UNSLOTH_LLAMA_PR_FORCE:-${_DEFAULT_LLAMA_PR_FORCE}}"
+# Force mainline source -- no env var override for now.
+_LLAMA_SOURCE="${_DEFAULT_LLAMA_SOURCE}"
+_LLAMA_SOURCE="${_LLAMA_SOURCE%.git}"  # normalize: strip trailing .git
+_RESOLVED_SOURCE_URL="$_LLAMA_SOURCE"
+_RESOLVED_SOURCE_REF="$_REQUESTED_LLAMA_TAG"
+_RESOLVED_SOURCE_REF_KIND="tag"
+
+if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then
+    _NEED_LLAMA_SOURCE_BUILD=true
+    _SKIP_PREBUILT_INSTALL=true
 fi
-if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
-    step "llama.cpp" "failed to resolve prebuilt tag via $_HELPER_RELEASE_REPO" "$C_WARN"
-    print_llama_error_log "$_RESOLVE_LLAMA_LOG"
-    set +e
-    # Resolve the llama.cpp tag for source-build fallback. Pass --published-repo
-    # so the resolver prefers Unsloth's tested tag (e.g. b8508) over the upstream
-    # bleeding-edge tag (e.g. b8514) from ggml-org/llama.cpp.
-    _RESOLVED_LLAMA_TAG="$(python "$SCRIPT_DIR/install_llama_prebuilt.py" --resolve-llama-tag "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO" 2>/dev/null)"
-    _RESOLVE_UPSTREAM_STATUS=$?
-    set -e
-    if [ "$_RESOLVE_UPSTREAM_STATUS" -ne 0 ] || [ -z "$_RESOLVED_LLAMA_TAG" ]; then
-        if [ "$_REQUESTED_LLAMA_TAG" = "latest" ]; then
-            # Try Unsloth release repo first, then fall back to ggml-org upstream
-            _RESOLVED_LLAMA_TAG="$(curl -fsSL "https://api.github.com/repos/${_HELPER_RELEASE_REPO}/releases/latest" 2>/dev/null | python -c "import sys,json; print(json.load(sys.stdin)['tag_name'])" 2>/dev/null)" || _RESOLVED_LLAMA_TAG=""
-            if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
-                _RESOLVED_LLAMA_TAG="$(curl -fsSL https://api.github.com/repos/ggml-org/llama.cpp/releases/latest 2>/dev/null | python -c "import sys,json; print(json.load(sys.stdin)['tag_name'])" 2>/dev/null)" || _RESOLVED_LLAMA_TAG=""
-            fi
+
+# Non-default source URL forces source build (fork has different code than prebuilt).
+if [ "$_LLAMA_SOURCE" != "https://github.com/ggml-org/llama.cpp" ]; then
+    step "llama.cpp" "custom source: $_LLAMA_SOURCE -- forcing source build" "$C_WARN"
+    _NEED_LLAMA_SOURCE_BUILD=true
+    _SKIP_PREBUILT_INSTALL=true
+fi
+
+# Baked-in PR_FORCE promotes to _LLAMA_PR when user hasn't set one.
+if [ -z "$_LLAMA_PR" ] && [ -n "$_LLAMA_PR_FORCE" ] && \
+   [[ "$_LLAMA_PR_FORCE" =~ ^[0-9]+$ ]] && [ "$_LLAMA_PR_FORCE" -gt 0 ]; then
+    _LLAMA_PR="$_LLAMA_PR_FORCE"
+    step "llama.cpp" "baked-in PR_FORCE=$_LLAMA_PR_FORCE" "$C_WARN"
+fi
+
+if [ -n "$_LLAMA_PR" ]; then
+    if ! [[ "$_LLAMA_PR" =~ ^[0-9]+$ ]] || [ "$_LLAMA_PR" -le 0 ]; then
+        step "llama.cpp" "UNSLOTH_LLAMA_PR=$_LLAMA_PR is not a valid PR number" "$C_ERR"
+        exit 1
+    fi
+    step "llama.cpp" "UNSLOTH_LLAMA_PR=$_LLAMA_PR -- will build from PR head" "$C_WARN"
+    _RESOLVED_LLAMA_TAG="pr-$_LLAMA_PR"
+    _RESOLVED_SOURCE_URL="$_LLAMA_SOURCE"
+    _RESOLVED_SOURCE_REF="pr-$_LLAMA_PR"
+    _RESOLVED_SOURCE_REF_KIND="pull"
+    _NEED_LLAMA_SOURCE_BUILD=true
+    _SKIP_PREBUILT_INSTALL=true
+elif [ "${_SKIP_PREBUILT_INSTALL:-false}" = true ]; then
+    # Custom source or other override already forced source build; skip
+    # the prebuilt release resolution entirely. When building from a custom
+    # fork, the fork may not carry upstream bNNNN tags, so resolve the tag
+    # only when the source is the default ggml-org repo.
+    if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then
+        _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"
+    elif [ "$_LLAMA_SOURCE" = "https://github.com/ggml-org/llama.cpp" ]; then
+        _RESOLVE_TAG_ARGS=(--resolve-llama-tag "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO")
+        _RESOLVE_TAG_ARGS+=(--output-format json)
+        if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
+            _RESOLVE_TAG_ARGS+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG")
+        fi
+        set +e
+        _RESOLVE_TAG_JSON="$(python "$SCRIPT_DIR/install_llama_prebuilt.py" "${_RESOLVE_TAG_ARGS[@]}" 2>/dev/null)"
+        _RESOLVE_UPSTREAM_STATUS=$?
+        set -e
+        if [ "$_RESOLVE_UPSTREAM_STATUS" -eq 0 ] && [ -n "${_RESOLVE_TAG_JSON:-}" ]; then
+            _RESOLVED_LLAMA_TAG="$(
+                printf '%s' "$_RESOLVE_TAG_JSON" | python -c 'import json,sys; print(json.load(sys.stdin).get("llama_tag",""))' 2>/dev/null || true
+            )"
+        else
+            _RESOLVED_LLAMA_TAG=""
         fi
         if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
             _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"
         fi
+    else
+        _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"
     fi
-    _NEED_LLAMA_SOURCE_BUILD=true
-    _SKIP_PREBUILT_INSTALL=true
+else
+    _RESOLVE_INSTALL_ARGS=(--resolve-install-tag "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO")
+    _RESOLVE_INSTALL_ARGS+=(--output-format json)
+    if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
+        _RESOLVE_INSTALL_ARGS+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG")
+    fi
+    _RESOLVE_LLAMA_LOG="$(mktemp)"
+    set +e
+    _RESOLVE_INSTALL_JSON="$(
+        python "$SCRIPT_DIR/install_llama_prebuilt.py" \
+            "${_RESOLVE_INSTALL_ARGS[@]}" 2>"$_RESOLVE_LLAMA_LOG"
+    )"
+    _RESOLVE_LLAMA_STATUS=$?
+    set -e
+    if [ "$_RESOLVE_LLAMA_STATUS" -eq 0 ]; then
+        _RESOLVED_LLAMA_TAG="$(
+            printf '%s' "${_RESOLVE_INSTALL_JSON:-}" | python -c 'import json,sys; print(json.load(sys.stdin).get("llama_tag",""))' 2>/dev/null || true
+        )"
+    else
+        _RESOLVED_LLAMA_TAG=""
+    fi
+    if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
+        step "llama.cpp" "failed to resolve a published llama.cpp release via $_HELPER_RELEASE_REPO" "$C_WARN"
+        print_llama_error_log "$_RESOLVE_LLAMA_LOG"
+        set +e
+        # Resolve the llama.cpp tag for source-build fallback. Pass --published-repo
+        # so the resolver prefers the latest usable Unsloth-published upstream tag
+        # before falling back to the bleeding-edge ggml-org/llama.cpp tag.
+        _RESOLVE_FALLBACK_ARGS=(--resolve-llama-tag "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO")
+        _RESOLVE_FALLBACK_ARGS+=(--output-format json)
+        if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
+            _RESOLVE_FALLBACK_ARGS+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG")
+        fi
+        _RESOLVE_FALLBACK_JSON="$(python "$SCRIPT_DIR/install_llama_prebuilt.py" "${_RESOLVE_FALLBACK_ARGS[@]}" 2>/dev/null)"
+        _RESOLVE_UPSTREAM_STATUS=$?
+        set -e
+        if [ "$_RESOLVE_UPSTREAM_STATUS" -eq 0 ] && [ -n "${_RESOLVE_FALLBACK_JSON:-}" ]; then
+            _RESOLVED_LLAMA_TAG="$(
+                printf '%s' "$_RESOLVE_FALLBACK_JSON" | python -c 'import json,sys; print(json.load(sys.stdin).get("llama_tag",""))' 2>/dev/null || true
+            )"
+        else
+            _RESOLVED_LLAMA_TAG=""
+        fi
+        if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
+            _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"
+        fi
+        _NEED_LLAMA_SOURCE_BUILD=true
+        _SKIP_PREBUILT_INSTALL=true
+    fi
+    rm -f "$_RESOLVE_LLAMA_LOG"
 fi
-rm -f "$_RESOLVE_LLAMA_LOG"
 
 substep "resolved llama.cpp tag: $_RESOLVED_LLAMA_TAG"
 verbose_substep "requested llama.cpp tag: $_REQUESTED_LLAMA_TAG (repo: $_HELPER_RELEASE_REPO)"
@@ -520,7 +629,7 @@ else
         _PREBUILT_CMD=(
             python "$SCRIPT_DIR/install_llama_prebuilt.py"
             --install-dir "$LLAMA_CPP_DIR"
-            --llama-tag "$_RESOLVED_LLAMA_TAG"
+            --llama-tag "$_REQUESTED_LLAMA_TAG"
             --published-repo "$_HELPER_RELEASE_REPO"
         )
         if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
@@ -538,9 +647,22 @@ else
         set -e
 
         if [ "$_PREBUILT_STATUS" -eq 0 ]; then
-            step "llama.cpp" "prebuilt installed and validated"
+            if grep -Fq "already matches" "$_PREBUILT_LOG"; then
+                step "llama.cpp" "prebuilt up to date and validated"
+            else
+                step "llama.cpp" "prebuilt installed and validated"
+            fi
             verbose_substep "llama.cpp install dir: $LLAMA_CPP_DIR"
             rm -f "$_PREBUILT_LOG"
+        elif [ "$_PREBUILT_STATUS" -eq 3 ]; then
+            step "llama.cpp" "install blocked by active llama.cpp process" "$C_WARN"
+            print_llama_error_log "$_PREBUILT_LOG"
+            rm -f "$_PREBUILT_LOG"
+            if [ -d "$LLAMA_CPP_DIR" ]; then
+                substep "existing install was restored"
+            fi
+            substep "close Studio or other llama.cpp users and retry"
+            exit 3
         else
             step "llama.cpp" "prebuilt install failed (continuing)" "$C_WARN"
             print_llama_error_log "$_PREBUILT_LOG"
@@ -623,21 +745,99 @@ else
         step "llama.cpp" "skipped (git not found)" "$C_WARN"
         [ -f "$LLAMA_SERVER_BIN" ] || _LLAMA_CPP_DEGRADED=true
     else
-        BUILD_OK=true
-        _CLONE_BRANCH_ARGS=()
-        if [ "$_RESOLVED_LLAMA_TAG" != "latest" ] && [ -n "$_RESOLVED_LLAMA_TAG" ]; then
-            _CLONE_BRANCH_ARGS=(--branch "$_RESOLVED_LLAMA_TAG")
+        if [ -z "$_LLAMA_PR" ]; then
+            if [ "$_LLAMA_SOURCE" = "https://github.com/ggml-org/llama.cpp" ]; then
+                _RESOLVE_SOURCE_ARGS=(--resolve-source-build "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO")
+                _RESOLVE_SOURCE_ARGS+=(--output-format json)
+                if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
+                    _RESOLVE_SOURCE_ARGS+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG")
+                fi
+                set +e
+                _SOURCE_BUILD_PLAN="$(python "$SCRIPT_DIR/install_llama_prebuilt.py" "${_RESOLVE_SOURCE_ARGS[@]}" 2>/dev/null)"
+                _RESOLVE_SOURCE_STATUS=$?
+                set -e
+                if [ "$_RESOLVE_SOURCE_STATUS" -eq 0 ] && [ -n "$_SOURCE_BUILD_PLAN" ]; then
+                    _RESOLVED_SOURCE_URL="$(
+                        printf '%s' "$_SOURCE_BUILD_PLAN" | python -c 'import json,sys; print(json.load(sys.stdin).get("source_url",""))' 2>/dev/null || true
+                    )"
+                    _RESOLVED_SOURCE_REF_KIND="$(
+                        printf '%s' "$_SOURCE_BUILD_PLAN" | python -c 'import json,sys; print(json.load(sys.stdin).get("source_ref_kind",""))' 2>/dev/null || true
+                    )"
+                    _RESOLVED_SOURCE_REF="$(
+                        printf '%s' "$_SOURCE_BUILD_PLAN" | python -c 'import json,sys; print(json.load(sys.stdin).get("source_ref",""))' 2>/dev/null || true
+                    )"
+                fi
+            fi
+            if [ -z "$_RESOLVED_SOURCE_URL" ]; then
+                _RESOLVED_SOURCE_URL="$_LLAMA_SOURCE"
+            fi
+            if [ -z "$_RESOLVED_SOURCE_REF" ]; then
+                _RESOLVED_SOURCE_REF="$_RESOLVED_LLAMA_TAG"
+            fi
         fi
+        verbose_substep "source build repo: $_RESOLVED_SOURCE_URL"
+        verbose_substep "source build ref: ${_RESOLVED_SOURCE_REF:-latest} (${_RESOLVED_SOURCE_REF_KIND})"
+        BUILD_OK=true
+        mkdir -p "$(dirname "$LLAMA_CPP_DIR")"
         _BUILD_TMP="${LLAMA_CPP_DIR}.build.$$"
         rm -rf "$_BUILD_TMP"
-        run_quiet_no_exit "clone llama.cpp" git clone --depth 1 "${_CLONE_BRANCH_ARGS[@]}" https://github.com/ggml-org/llama.cpp.git "$_BUILD_TMP" || BUILD_OK=false
+        if [ -n "$_LLAMA_PR" ]; then
+            run_quiet_no_exit "clone llama.cpp" \
+                git clone --depth 1 "${_LLAMA_SOURCE}.git" "$_BUILD_TMP" || BUILD_OK=false
+            if [ "$BUILD_OK" = true ]; then
+                run_quiet_no_exit "fetch PR #$_LLAMA_PR" \
+                    git -C "$_BUILD_TMP" fetch --depth 1 origin "pull/$_LLAMA_PR/head:pr-$_LLAMA_PR" || BUILD_OK=false
+            fi
+            if [ "$BUILD_OK" = true ]; then
+                run_quiet_no_exit "checkout PR #$_LLAMA_PR" \
+                    git -C "$_BUILD_TMP" checkout "pr-$_LLAMA_PR" || BUILD_OK=false
+            fi
+        elif [ "$_RESOLVED_SOURCE_REF_KIND" = "pull" ] && [ -n "$_RESOLVED_SOURCE_REF" ]; then
+            run_quiet_no_exit "clone llama.cpp" \
+                git clone --depth 1 "${_RESOLVED_SOURCE_URL}.git" "$_BUILD_TMP" || BUILD_OK=false
+            if [ "$BUILD_OK" = true ]; then
+                run_quiet_no_exit "fetch source PR ref" \
+                    git -C "$_BUILD_TMP" fetch --depth 1 origin "$_RESOLVED_SOURCE_REF" || BUILD_OK=false
+            fi
+            if [ "$BUILD_OK" = true ]; then
+                run_quiet_no_exit "checkout source PR ref" \
+                    git -C "$_BUILD_TMP" checkout -B unsloth-llama-build FETCH_HEAD || BUILD_OK=false
+            fi
+        elif [ "$_RESOLVED_SOURCE_REF_KIND" = "commit" ] && [ -n "$_RESOLVED_SOURCE_REF" ]; then
+            run_quiet_no_exit "clone llama.cpp" \
+                git clone --depth 1 "${_RESOLVED_SOURCE_URL}.git" "$_BUILD_TMP" || BUILD_OK=false
+            if [ "$BUILD_OK" = true ]; then
+                run_quiet_no_exit "fetch source commit" \
+                    git -C "$_BUILD_TMP" fetch --depth 1 origin "$_RESOLVED_SOURCE_REF" || BUILD_OK=false
+            fi
+            if [ "$BUILD_OK" = true ]; then
+                run_quiet_no_exit "checkout source commit" \
+                    git -C "$_BUILD_TMP" checkout -B unsloth-llama-build FETCH_HEAD || BUILD_OK=false
+            fi
+        else
+            _CLONE_ARGS=(git clone --depth 1)
+            if [ "$_RESOLVED_SOURCE_REF" != "latest" ] && [ -n "$_RESOLVED_SOURCE_REF" ]; then
+                _CLONE_ARGS+=(--branch "$_RESOLVED_SOURCE_REF")
+            fi
+            _CLONE_ARGS+=("${_RESOLVED_SOURCE_URL}.git" "$_BUILD_TMP")
+            run_quiet_no_exit "clone llama.cpp" \
+                "${_CLONE_ARGS[@]}" || BUILD_OK=false
+        fi
 
         if [ "$BUILD_OK" = true ]; then
             CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=ON -DGGML_NATIVE=ON"
+            _TRY_METAL_CPU_FALLBACK=false
+            _HOST_SYSTEM="$(uname -s 2>/dev/null || true)"
+            _HOST_MACHINE="$(uname -m 2>/dev/null || true)"
+            _IS_MACOS_ARM64=false
+            if [ "$_HOST_SYSTEM" = "Darwin" ] && { [ "$_HOST_MACHINE" = "arm64" ] || [ "$_HOST_MACHINE" = "aarch64" ]; }; then
+                _IS_MACOS_ARM64=true
+            fi
 
             if command -v ccache &>/dev/null; then
                 CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache"
             fi
+            CPU_FALLBACK_CMAKE_ARGS="$CMAKE_ARGS"
 
             GPU_BACKEND=""
             NVCC_PATH=""
@@ -673,7 +873,13 @@ else
             fi
 
             _BUILD_DESC="building"
-            if [ -n "$NVCC_PATH" ]; then
+            if [ "$_IS_MACOS_ARM64" = true ]; then
+                # Metal takes precedence on Apple Silicon (CUDA/ROCm not functional on macOS)
+                _BUILD_DESC="building (Metal)"
+                CMAKE_ARGS="$CMAKE_ARGS -DGGML_METAL=ON -DGGML_METAL_EMBED_LIBRARY=ON -DGGML_METAL_USE_BF16=ON -DCMAKE_INSTALL_RPATH=@loader_path -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON"
+                CPU_FALLBACK_CMAKE_ARGS="$CPU_FALLBACK_CMAKE_ARGS -DGGML_METAL=OFF"
+                _TRY_METAL_CPU_FALLBACK=true
+            elif [ -n "$NVCC_PATH" ]; then
                 CMAKE_ARGS="$CMAKE_ARGS -DGGML_CUDA=ON"
 
                 CUDA_ARCHS=""
@@ -755,11 +961,37 @@ else
                 CMAKE_GENERATOR_ARGS="-G Ninja"
             fi
 
-            run_quiet_no_exit "cmake llama.cpp" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CMAKE_ARGS || BUILD_OK=false
+            if ! run_quiet_no_exit "cmake llama.cpp" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CMAKE_ARGS; then
+                if [ "$_TRY_METAL_CPU_FALLBACK" = true ]; then
+                    _TRY_METAL_CPU_FALLBACK=false
+                    substep "Metal configure failed; retrying CPU build..." "$C_WARN"
+                    rm -rf "$_BUILD_TMP/build"
+                    run_quiet_no_exit "cmake llama.cpp (cpu fallback)" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS || BUILD_OK=false
+                    if [ "$BUILD_OK" = true ]; then
+                        _BUILD_DESC="building (CPU fallback)"
+                    fi
+                else
+                    BUILD_OK=false
+                fi
+            fi
         fi
 
         if [ "$BUILD_OK" = true ]; then
-            run_quiet_no_exit "build llama-server" cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false
+            if ! run_quiet_no_exit "build llama-server" cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU"; then
+                if [ "$_TRY_METAL_CPU_FALLBACK" = true ]; then
+                    _TRY_METAL_CPU_FALLBACK=false
+                    substep "Metal build failed; retrying CPU build..." "$C_WARN"
+                    rm -rf "$_BUILD_TMP/build"
+                    if run_quiet_no_exit "cmake llama.cpp (cpu fallback)" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS; then
+                        _BUILD_DESC="building (CPU fallback)"
+                        run_quiet_no_exit "build llama-server (cpu fallback)" cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false
+                    else
+                        BUILD_OK=false
+                    fi
+                else
+                    BUILD_OK=false
+                fi
+            fi
         fi
 
         if [ "$BUILD_OK" = true ]; then
@@ -794,7 +1026,16 @@ else
 fi  # end _SKIP_GGUF_BUILD check
 
 # ── Footer ──
-if [ "$IS_COLAB" = true ]; then
+if [ "$_LLAMA_ONLY" = "1" ]; then
+    echo ""
+    printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
+    if [ "$_LLAMA_CPP_DEGRADED" = true ]; then
+        printf "  ${C_WARN}%s${C_RST}\n" "llama.cpp update finished (limited: llama.cpp unavailable)"
+    else
+        printf "  ${C_TITLE}%s${C_RST}\n" "llama.cpp update finished"
+    fi
+    printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
+elif [ "$IS_COLAB" = true ]; then
     echo ""
     printf "  ${C_DIM}%s${C_RST}\n" "$RULE"
     if [ "$_LLAMA_CPP_DEGRADED" = true ]; then

tests/python/conftest.py

@@ -5,3 +5,6 @@ def pytest_configure(config):
     config.addinivalue_line(
         "markers", "server: heavyweight tests requiring studio venv"
     )
+    config.addinivalue_line(
+        "markers", "e2e: end-to-end tests requiring network and venv creation"
+    )

tests/python/test_tokenizers_and_torch_constraint.py

@@ -0,0 +1,572 @@
+"""
+Tests for two install fixes:
+  1. tokenizers added to no-torch-runtime.txt  (prevents AutoConfig crash)
+  2. TORCH_CONSTRAINT variable in install.sh    (arm64 macOS + py313+ -> torch>=2.6)
+"""
+
+from __future__ import annotations
+
+import pathlib
+import re
+import subprocess
+import textwrap
+
+import pytest
+
+# ── Locate source files relative to this test ──────────────────────────
+_TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent  # tests/
+_REPO_ROOT = _TESTS_DIR.parent  # unsloth/
+_INSTALL_SH = _REPO_ROOT / "install.sh"
+_INSTALL_PS1 = _REPO_ROOT / "install.ps1"
+_NO_TORCH_RT = (
+    _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt"
+)
+
+
+def _read(path: pathlib.Path) -> str:
+    return path.read_text(encoding = "utf-8")
+
+
+def _lines(path: pathlib.Path) -> list[str]:
+    """Return non-comment, non-blank lines stripped."""
+    return [
+        ln.strip()
+        for ln in _read(path).splitlines()
+        if ln.strip() and not ln.strip().startswith("#")
+    ]
+
+
+# ======================================================================
+# Group 1 -- Structural checks (no network, instant)
+# ======================================================================
+class TestStructuralTokenizers:
+    """Verify tokenizers presence and ordering in no-torch-runtime.txt."""
+
+    def test_tokenizers_present(self):
+        """tokenizers must be a standalone package line."""
+        pkgs = _lines(_NO_TORCH_RT)
+        bare_names = [
+            p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs
+        ]
+        assert "tokenizers" in bare_names
+
+    def test_tokenizers_before_transformers(self):
+        """tokenizers should appear before transformers (install order intent)."""
+        pkgs = _lines(_NO_TORCH_RT)
+        bare_names = [
+            p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs
+        ]
+        idx_tok = bare_names.index("tokenizers")
+        idx_tf = bare_names.index("transformers")
+        assert idx_tok < idx_tf, (
+            f"tokenizers at index {idx_tok} should appear before "
+            f"transformers at index {idx_tf}"
+        )
+
+    def test_torch_not_in_no_torch_file(self):
+        """torch itself must NOT be listed in the no-torch requirements."""
+        pkgs = _lines(_NO_TORCH_RT)
+        bare_names = [
+            p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs
+        ]
+        assert "torch" not in bare_names
+
+
+class TestStructuralTorchConstraint:
+    """Verify TORCH_CONSTRAINT wiring in install.sh."""
+
+    _sh = _read(_INSTALL_SH)
+
+    def test_default_assignment_exists(self):
+        assert 'TORCH_CONSTRAINT="torch>=2.4,<2.11.0"' in self._sh
+
+    def test_tightened_assignment_exists(self):
+        assert 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' in self._sh
+
+    def test_variable_used_in_pip_install(self):
+        """$TORCH_CONSTRAINT must appear in a uv pip install line."""
+        assert '"$TORCH_CONSTRAINT"' in self._sh
+
+    def test_hardcoded_torch_constraint_only_once(self):
+        """The hard-coded torch>=2.4,<2.11.0 string should appear exactly once
+        in install.sh (the default assignment), not in pip install lines."""
+        count = self._sh.count('"torch>=2.4,<2.11.0"')
+        assert count == 1, f"Expected 1, found {count}"
+
+    def test_tightening_guarded_by_skip_torch(self):
+        """The block must check SKIP_TORCH=false."""
+        # Find the tightening if-block
+        m = re.search(
+            r"if\s.*SKIP_TORCH.*=\s*false.*&&.*OS.*=.*macos.*&&.*_ARCH.*=.*arm64",
+            self._sh,
+        )
+        assert m is not None, "Guard not found: SKIP_TORCH + macos + arm64"
+
+    def test_tightening_guarded_by_arch(self):
+        m = re.search(r"_ARCH.*=.*arm64", self._sh)
+        assert m is not None
+
+    def test_tightening_guarded_by_os(self):
+        m = re.search(r"OS.*=.*macos", self._sh)
+        assert m is not None
+
+
+class TestStructuralInstallPs1Unchanged:
+    """install.ps1 should NOT have TORCH_CONSTRAINT variable."""
+
+    _ps1 = _read(_INSTALL_PS1)
+
+    def test_no_torch_constraint_variable(self):
+        assert "TORCH_CONSTRAINT" not in self._ps1
+        assert "$TorchConstraint" not in self._ps1
+
+    def test_hardcoded_torch_constraint_present(self):
+        assert '"torch>=2.4,<2.11.0"' in self._ps1
+
+
+# ======================================================================
+# Group 2 -- Shell snippet tests (bash subprocess, mocked python)
+# ======================================================================
+class TestTorchConstraintShell:
+    """Test the TORCH_CONSTRAINT block using bash subprocesses with
+    mocked python binaries that return controlled minor versions."""
+
+    # The extracted snippet we test in isolation.  We override OS, _ARCH,
+    # SKIP_TORCH, and provide a mock python at $VENV_DIR/bin/python.
+    _SNIPPET_TEMPLATE = textwrap.dedent(r"""
+        #!/bin/bash
+        set -e
+        SKIP_TORCH={skip_torch}
+        OS="{os}"
+        _ARCH="{arch}"
+        VENV_DIR="{venv_dir}"
+
+        TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
+        if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
+            _PY_MINOR=$("$VENV_DIR/bin/python" -c \
+                "import sys; print(sys.version_info.minor)" 2>/dev/null || echo "0")
+            if [ "$_PY_MINOR" -ge 13 ] 2>/dev/null; then
+                TORCH_CONSTRAINT="torch>=2.6,<2.11.0"
+            fi
+        fi
+        echo "$TORCH_CONSTRAINT"
+    """).strip()
+
+    @staticmethod
+    def _make_mock_python(tmp_path: pathlib.Path, minor: int) -> pathlib.Path:
+        """Create a mock python that prints a controlled minor version."""
+        venv = tmp_path / "venv"
+        bin_dir = venv / "bin"
+        bin_dir.mkdir(parents = True, exist_ok = True)
+        mock_py = bin_dir / "python"
+        mock_py.write_text(
+            textwrap.dedent(f"""\
+            #!/bin/bash
+            # Mock python: always report minor={minor}
+            if echo "$@" | grep -q "sys.version_info.minor"; then
+                echo "{minor}"
+            else
+                echo "0"
+            fi
+        """)
+        )
+        mock_py.chmod(0o755)
+        return venv
+
+    def _run(
+        self,
+        tmp_path: pathlib.Path,
+        *,
+        py_minor: int = 12,
+        os_val: str = "macos",
+        arch: str = "arm64",
+        skip_torch: str = "false",
+    ) -> str:
+        venv = self._make_mock_python(tmp_path, py_minor)
+        script = self._SNIPPET_TEMPLATE.format(
+            skip_torch = skip_torch,
+            os = os_val,
+            arch = arch,
+            venv_dir = str(venv),
+        )
+        script_file = tmp_path / "test_snippet.sh"
+        script_file.write_text(script)
+        script_file.chmod(0o755)
+        result = subprocess.run(
+            ["bash", str(script_file)],
+            capture_output = True,
+            text = True,
+            timeout = 10,
+        )
+        assert result.returncode == 0, f"Script failed: {result.stderr}"
+        return result.stdout.strip()
+
+    # -- arm64 macOS tightening cases --
+
+    def test_arm64_macos_py313_tightened(self, tmp_path):
+        out = self._run(tmp_path, py_minor = 13, os_val = "macos", arch = "arm64")
+        assert out == "torch>=2.6,<2.11.0"
+
+    def test_arm64_macos_py314_tightened(self, tmp_path):
+        out = self._run(tmp_path, py_minor = 14, os_val = "macos", arch = "arm64")
+        assert out == "torch>=2.6,<2.11.0"
+
+    # -- arm64 macOS default (older python) --
+
+    def test_arm64_macos_py312_default(self, tmp_path):
+        out = self._run(tmp_path, py_minor = 12, os_val = "macos", arch = "arm64")
+        assert out == "torch>=2.4,<2.11.0"
+
+    def test_arm64_macos_py311_default(self, tmp_path):
+        out = self._run(tmp_path, py_minor = 11, os_val = "macos", arch = "arm64")
+        assert out == "torch>=2.4,<2.11.0"
+
+    # -- Linux (unaffected) --
+
+    def test_linux_x86_py313_default(self, tmp_path):
+        out = self._run(tmp_path, py_minor = 13, os_val = "linux", arch = "x86_64")
+        assert out == "torch>=2.4,<2.11.0"
+
+    def test_linux_aarch64_py313_default(self, tmp_path):
+        out = self._run(tmp_path, py_minor = 13, os_val = "linux", arch = "aarch64")
+        assert out == "torch>=2.4,<2.11.0"
+
+    # -- Intel Mac (arch mismatch) --
+
+    def test_intel_mac_x86_py313_default(self, tmp_path):
+        out = self._run(tmp_path, py_minor = 13, os_val = "macos", arch = "x86_64")
+        assert out == "torch>=2.4,<2.11.0"
+
+    # -- SKIP_TORCH bypass --
+
+    def test_skip_torch_arm64_macos_py313_default(self, tmp_path):
+        out = self._run(
+            tmp_path,
+            py_minor = 13,
+            os_val = "macos",
+            arch = "arm64",
+            skip_torch = "true",
+        )
+        assert out == "torch>=2.4,<2.11.0"
+
+    # -- WSL --
+
+    def test_wsl_py313_default(self, tmp_path):
+        out = self._run(tmp_path, py_minor = 13, os_val = "wsl", arch = "x86_64")
+        assert out == "torch>=2.4,<2.11.0"
+
+    # -- Edge cases --
+
+    def test_py_minor_0_fallback_default(self, tmp_path):
+        """If python query fails (returns 0), should stay at default."""
+        out = self._run(tmp_path, py_minor = 0, os_val = "macos", arch = "arm64")
+        assert out == "torch>=2.4,<2.11.0"
+
+    def test_boundary_py_minor_12_not_tightened(self, tmp_path):
+        out = self._run(tmp_path, py_minor = 12, os_val = "macos", arch = "arm64")
+        assert out == "torch>=2.4,<2.11.0"
+
+    def test_boundary_py_minor_13_tightened(self, tmp_path):
+        out = self._run(tmp_path, py_minor = 13, os_val = "macos", arch = "arm64")
+        assert out == "torch>=2.6,<2.11.0"
+
+    def test_mock_uv_receives_correct_constraint(self, tmp_path):
+        """Verify a mock uv would receive the correct constraint string."""
+        venv = self._make_mock_python(tmp_path, minor = 13)
+
+        # Create a mock uv that logs its arguments
+        mock_uv = tmp_path / "mock_uv"
+        log_file = tmp_path / "uv_log.txt"
+        mock_uv.write_text(
+            textwrap.dedent(f"""\
+            #!/bin/bash
+            echo "$@" >> {log_file}
+        """)
+        )
+        mock_uv.chmod(0o755)
+
+        script = textwrap.dedent(f"""\
+            #!/bin/bash
+            set -e
+            SKIP_TORCH=false
+            OS="macos"
+            _ARCH="arm64"
+            VENV_DIR="{venv}"
+
+            TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
+            if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
+                _PY_MINOR=$("$VENV_DIR/bin/python" -c \\
+                    "import sys; print(sys.version_info.minor)" 2>/dev/null || echo "0")
+                if [ "$_PY_MINOR" -ge 13 ] 2>/dev/null; then
+                    TORCH_CONSTRAINT="torch>=2.6,<2.11.0"
+                fi
+            fi
+            # Simulate the uv pip install line
+            {mock_uv} pip install --python "$VENV_DIR/bin/python" "$TORCH_CONSTRAINT" torchvision torchaudio
+        """)
+        script_file = tmp_path / "test_uv.sh"
+        script_file.write_text(script)
+        script_file.chmod(0o755)
+
+        result = subprocess.run(
+            ["bash", str(script_file)],
+            capture_output = True,
+            text = True,
+            timeout = 10,
+        )
+        assert result.returncode == 0, f"Script failed: {result.stderr}"
+        logged = log_file.read_text()
+        assert "torch>=2.6,<2.11.0" in logged, f"uv log: {logged}"
+
+    def test_mock_uv_receives_default_constraint(self, tmp_path):
+        """On py3.12 arm64 macOS, uv should receive the default constraint."""
+        venv = self._make_mock_python(tmp_path, minor = 12)
+        mock_uv = tmp_path / "mock_uv"
+        log_file = tmp_path / "uv_log.txt"
+        mock_uv.write_text(
+            textwrap.dedent(f"""\
+            #!/bin/bash
+            echo "$@" >> {log_file}
+        """)
+        )
+        mock_uv.chmod(0o755)
+
+        script = textwrap.dedent(f"""\
+            #!/bin/bash
+            set -e
+            SKIP_TORCH=false
+            OS="macos"
+            _ARCH="arm64"
+            VENV_DIR="{venv}"
+
+            TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
+            if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
+                _PY_MINOR=$("$VENV_DIR/bin/python" -c \\
+                    "import sys; print(sys.version_info.minor)" 2>/dev/null || echo "0")
+                if [ "$_PY_MINOR" -ge 13 ] 2>/dev/null; then
+                    TORCH_CONSTRAINT="torch>=2.6,<2.11.0"
+                fi
+            fi
+            {mock_uv} pip install --python "$VENV_DIR/bin/python" "$TORCH_CONSTRAINT" torchvision torchaudio
+        """)
+        script_file = tmp_path / "test_uv.sh"
+        script_file.write_text(script)
+        script_file.chmod(0o755)
+
+        result = subprocess.run(
+            ["bash", str(script_file)],
+            capture_output = True,
+            text = True,
+            timeout = 10,
+        )
+        assert result.returncode == 0, f"Script failed: {result.stderr}"
+        logged = log_file.read_text()
+        assert "torch>=2.4,<2.11.0" in logged, f"uv log: {logged}"
+
+
+# ======================================================================
+# Group 3 -- E2E tokenizers fix (requires network, ~2-5 min)
+# ======================================================================
+@pytest.mark.e2e
+class TestE2ETokenizersFix:
+    """Creates real uv venvs to verify tokenizers + transformers work
+    without torch installed."""
+
+    @staticmethod
+    def _create_venv(tmp_path: pathlib.Path, name: str, py: str) -> pathlib.Path:
+        venv = tmp_path / name
+        result = subprocess.run(
+            ["uv", "venv", str(venv), "--python", py],
+            capture_output = True,
+            text = True,
+            timeout = 120,
+        )
+        if result.returncode != 0:
+            pytest.skip(f"uv venv creation failed for {py}: {result.stderr}")
+        return venv
+
+    @staticmethod
+    def _pip_install(venv: pathlib.Path, *args: str) -> subprocess.CompletedProcess:
+        py = str(venv / "bin" / "python")
+        cmd = ["uv", "pip", "install", "--python", py, *args]
+        return subprocess.run(cmd, capture_output = True, text = True, timeout = 300)
+
+    @staticmethod
+    def _run_python(venv: pathlib.Path, code: str) -> subprocess.CompletedProcess:
+        py = str(venv / "bin" / "python")
+        return subprocess.run(
+            [py, "-c", code],
+            capture_output = True,
+            text = True,
+            timeout = 60,
+        )
+
+    @pytest.mark.parametrize("py_version", ["3.12", "3.13"])
+    def test_autoconfig_works_with_no_torch_runtime(self, tmp_path, py_version):
+        """Install from no-torch-runtime.txt with --no-deps (matching the
+        real install.sh path), then verify AutoConfig imports successfully."""
+        venv = self._create_venv(tmp_path, f"tok-{py_version}", py_version)
+        r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
+        assert r.returncode == 0, f"Install failed: {r.stderr}"
+
+        result = self._run_python(
+            venv, "from transformers import AutoConfig; print('OK')"
+        )
+        assert (
+            result.returncode == 0
+        ), f"AutoConfig import failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
+        assert "OK" in result.stdout
+
+    @pytest.mark.parametrize("py_version", ["3.12", "3.13"])
+    def test_tokenizers_directly_importable(self, tmp_path, py_version):
+        venv = self._create_venv(tmp_path, f"tok-imp-{py_version}", py_version)
+        r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
+        assert r.returncode == 0, f"Install failed: {r.stderr}"
+        result = self._run_python(venv, "import tokenizers; print('OK')")
+        assert result.returncode == 0, f"Failed: {result.stderr}"
+
+    @pytest.mark.parametrize("py_version", ["3.12", "3.13"])
+    def test_torch_not_importable(self, tmp_path, py_version):
+        """In the no-torch scenario, torch should not be available."""
+        venv = self._create_venv(tmp_path, f"no-torch-{py_version}", py_version)
+        r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
+        assert r.returncode == 0, f"Install failed: {r.stderr}"
+        result = self._run_python(venv, "import torch")
+        assert result.returncode != 0, "torch should NOT be importable"
+
+    def test_negative_control_no_tokenizers(self, tmp_path):
+        """Without tokenizers, AutoConfig should fail. We create a copy of
+        no-torch-runtime.txt with the tokenizers line removed."""
+        venv = self._create_venv(tmp_path, "neg-ctrl", "3.12")
+        req_no_tokenizers = tmp_path / "no-tokenizers.txt"
+        req_no_tokenizers.write_text(
+            "\n".join(
+                line
+                for line in _read(_NO_TORCH_RT).splitlines()
+                if line.strip() != "tokenizers"
+            ),
+            encoding = "utf-8",
+        )
+        r = self._pip_install(venv, "--no-deps", "-r", str(req_no_tokenizers))
+        assert r.returncode == 0, f"Install failed: {r.stderr}"
+        result = self._run_python(venv, "from transformers import AutoConfig")
+        assert (
+            result.returncode != 0
+        ), "AutoConfig should fail without tokenizers installed"
+        assert (
+            "tokenizers" in result.stderr.lower()
+            or "ModuleNotFoundError" in result.stderr
+        )
+
+
+# ======================================================================
+# Group 4 -- Integration: install.sh reads no-torch-runtime.txt correctly
+# ======================================================================
+class TestInstallShNoTorchIntegration:
+    """Verify install.sh has the correct no-torch-runtime.txt wiring."""
+
+    _sh = _read(_INSTALL_SH)
+
+    def test_find_no_torch_runtime_exists(self):
+        assert "_find_no_torch_runtime()" in self._sh
+
+    def test_no_deps_invocation_for_migrated(self):
+        """Migrated path should use --no-deps -r."""
+        assert '--no-deps -r "$_NO_TORCH_RT"' in self._sh
+
+    def test_no_deps_invocation_for_fresh(self):
+        """Fresh install path should also use --no-deps -r."""
+        # Count occurrences of the no-deps -r pattern
+        count = self._sh.count('--no-deps -r "$_NO_TORCH_RT"')
+        assert count >= 2, f"Expected >=2 no-deps -r invocations, found {count}"
+
+    def test_mock_uv_skip_torch_reads_requirements(self, tmp_path):
+        """When SKIP_TORCH=true, the _find_no_torch_runtime path should be used."""
+        # We test this structurally: verify the SKIP_TORCH=true blocks contain
+        # _find_no_torch_runtime calls
+        skip_blocks = re.findall(
+            r'if \[ "\$SKIP_TORCH" = true \].*?(?=\n    (?:else|elif|fi))',
+            self._sh,
+            re.DOTALL,
+        )
+        found = any("_find_no_torch_runtime" in block for block in skip_blocks)
+        assert found, "SKIP_TORCH=true block should call _find_no_torch_runtime"
+
+
+# ======================================================================
+# Group 5 -- Full no-torch sandbox (requires network, ~5 min)
+# ======================================================================
+@pytest.mark.e2e
+class TestE2EFullNoTorchSandbox:
+    """Creates venvs and installs the actual no-torch-runtime.txt."""
+
+    @staticmethod
+    def _create_venv(tmp_path: pathlib.Path, name: str) -> pathlib.Path:
+        venv = tmp_path / name
+        result = subprocess.run(
+            ["uv", "venv", str(venv), "--python", "3.12"],
+            capture_output = True,
+            text = True,
+            timeout = 120,
+        )
+        if result.returncode != 0:
+            pytest.skip(f"uv venv creation failed: {result.stderr}")
+        return venv
+
+    @staticmethod
+    def _pip_install(venv: pathlib.Path, *args: str) -> subprocess.CompletedProcess:
+        py = str(venv / "bin" / "python")
+        cmd = ["uv", "pip", "install", "--python", py, *args]
+        return subprocess.run(cmd, capture_output = True, text = True, timeout = 600)
+
+    @staticmethod
+    def _run_python(venv: pathlib.Path, code: str) -> subprocess.CompletedProcess:
+        py = str(venv / "bin" / "python")
+        return subprocess.run(
+            [py, "-c", code],
+            capture_output = True,
+            text = True,
+            timeout = 60,
+        )
+
+    def test_autoconfig_succeeds(self, tmp_path):
+        """The real bug fix: install with --no-deps (matching install.sh)
+        and verify from transformers import AutoConfig works."""
+        venv = self._create_venv(tmp_path, "full-no-torch")
+        r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
+        assert r.returncode == 0, f"Install failed: {r.stderr}"
+        result = self._run_python(
+            venv, "from transformers import AutoConfig; print('OK')"
+        )
+        assert (
+            result.returncode == 0
+        ), f"AutoConfig failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
+
+    def test_torch_not_importable(self, tmp_path):
+        """With --no-deps (as install.sh uses), torch must not be pulled in."""
+        venv = self._create_venv(tmp_path, "no-torch-check")
+        r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
+        assert r.returncode == 0, f"Install failed: {r.stderr}"
+        result = self._run_python(venv, "import torch")
+        assert result.returncode != 0, "torch should NOT be importable"
+
+    def test_tokenizers_importable(self, tmp_path):
+        venv = self._create_venv(tmp_path, "tok-check")
+        r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
+        assert r.returncode == 0, f"Install failed: {r.stderr}"
+        result = self._run_python(venv, "import tokenizers; print('OK')")
+        assert result.returncode == 0, f"tokenizers import failed: {result.stderr}"
+
+    def test_safetensors_importable(self, tmp_path):
+        venv = self._create_venv(tmp_path, "st-check")
+        r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
+        assert r.returncode == 0, f"Install failed: {r.stderr}"
+        result = self._run_python(venv, "import safetensors; print('OK')")
+        assert result.returncode == 0, f"safetensors import failed: {result.stderr}"
+
+    def test_huggingface_hub_importable(self, tmp_path):
+        venv = self._create_venv(tmp_path, "hfhub-check")
+        r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
+        assert r.returncode == 0, f"Install failed: {r.stderr}"
+        result = self._run_python(venv, "import huggingface_hub; print('OK')")
+        assert result.returncode == 0, f"huggingface_hub import failed: {result.stderr}"

tests/run_all.sh

@@ -7,6 +7,7 @@ TESTS_DIR="$(cd "$(dirname "$0")" && pwd)"
 echo "=== Bash tests ==="
 sh "$TESTS_DIR/sh/test_get_torch_index_url.sh"
 sh "$TESTS_DIR/sh/test_mac_intel_compat.sh"
+sh "$TESTS_DIR/sh/test_torch_constraint.sh"
 
 echo ""
 echo "=== Python tests ==="
@@ -14,6 +15,7 @@ python -m pytest "$TESTS_DIR/python/test_install_python_stack.py" -v
 python -m pytest "$TESTS_DIR/python/test_cross_platform_parity.py" -v
 python -m pytest "$TESTS_DIR/python/test_no_torch_filtering.py" -v
 python -m pytest "$TESTS_DIR/python/test_studio_import_no_torch.py" -v
+python -m pytest "$TESTS_DIR/python/test_tokenizers_and_torch_constraint.py" -v -k "not e2e"
 
 echo ""
 echo "All tests passed."

tests/saving/test_save_shell_injection.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+
+SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py"
+
+
+def _function_calls(source: str, function_name: str) -> list[ast.Call]:
+    tree = ast.parse(source, filename = str(SAVE_PY))
+    for node in tree.body:
+        if isinstance(node, ast.FunctionDef) and node.name == function_name:
+            return [child for child in ast.walk(node) if isinstance(child, ast.Call)]
+    raise AssertionError(f"Function {function_name} not found in save.py")
+
+
+def _assert_safe_ggml_calls(calls: list[ast.Call]) -> None:
+    popen_calls = []
+    for call in calls:
+        if isinstance(call.func, ast.Attribute) and call.func.attr == "Popen":
+            if (
+                isinstance(call.func.value, ast.Name)
+                and call.func.value.id == "subprocess"
+            ):
+                popen_calls.append(call)
+
+    assert popen_calls, "Expected at least one subprocess.Popen call"
+
+    ggml_calls = []
+    for call in popen_calls:
+        if not call.args:
+            continue
+        argv = call.args[0]
+        if isinstance(argv, ast.List) and len(argv.elts) >= 2:
+            second_arg = argv.elts[1]
+            if (
+                isinstance(second_arg, ast.Constant)
+                and second_arg.value == "llama.cpp/convert-lora-to-ggml.py"
+            ):
+                ggml_calls.append(call)
+
+    assert ggml_calls, "Expected the GGML conversion subprocess call"
+
+    for call in ggml_calls:
+        shell_kwargs = [
+            keyword
+            for keyword in call.keywords
+            if keyword.arg == "shell"
+            and isinstance(keyword.value, ast.Constant)
+            and keyword.value.value is True
+        ]
+        assert not shell_kwargs, "subprocess.Popen must not use shell=True"
+
+        assert call.args, "subprocess.Popen must receive argv as a positional argument"
+        argv = call.args[0]
+        assert isinstance(
+            argv, ast.List
+        ), "subprocess.Popen must be called with an argv list"
+        assert len(argv.elts) == 5, "GGML conversion argv should have five elements"
+
+        second_arg = argv.elts[1]
+        assert isinstance(second_arg, ast.Constant)
+        assert second_arg.value == "llama.cpp/convert-lora-to-ggml.py"
+
+
+def test_ggml_conversion_paths_do_not_use_shell() -> None:
+    source = SAVE_PY.read_text(encoding = "utf-8")
+    for function_name in (
+        "unsloth_convert_lora_to_ggml_and_push_to_hub",
+        "unsloth_convert_lora_to_ggml_and_save_locally",
+    ):
+        calls = _function_calls(source, function_name)
+        _assert_safe_ggml_calls(calls)

tests/sh/test_torch_constraint.sh

@@ -0,0 +1,266 @@
+#!/bin/bash
+# Tests for TORCH_CONSTRAINT variable in install.sh and tokenizers in no-torch-runtime.txt.
+# Follows the same assertion pattern as test_mac_intel_compat.sh.
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+INSTALL_SH="$SCRIPT_DIR/../../install.sh"
+INSTALL_PS1="$SCRIPT_DIR/../../install.ps1"
+NO_TORCH_RT="$SCRIPT_DIR/../../studio/backend/requirements/no-torch-runtime.txt"
+PASS=0
+FAIL=0
+
+assert_eq() {
+    _label="$1"; _expected="$2"; _actual="$3"
+    if [ "$_actual" = "$_expected" ]; then
+        echo "  PASS: $_label"
+        PASS=$((PASS + 1))
+    else
+        echo "  FAIL: $_label (expected '$_expected', got '$_actual')"
+        FAIL=$((FAIL + 1))
+    fi
+}
+
+assert_contains() {
+    _label="$1"; _haystack="$2"; _needle="$3"
+    if echo "$_haystack" | grep -qF "$_needle"; then
+        echo "  PASS: $_label"
+        PASS=$((PASS + 1))
+    else
+        echo "  FAIL: $_label (expected to find '$_needle')"
+        FAIL=$((FAIL + 1))
+    fi
+}
+
+assert_not_contains() {
+    _label="$1"; _haystack="$2"; _needle="$3"
+    if echo "$_haystack" | grep -qF "$_needle"; then
+        echo "  FAIL: $_label (found '$_needle' but should not)"
+        FAIL=$((FAIL + 1))
+    else
+        echo "  PASS: $_label"
+        PASS=$((PASS + 1))
+    fi
+}
+
+# ── Helper: create a mock python that reports a given minor version ──
+make_mock_python() {
+    _minor="$1"
+    _venv_dir="$2"
+    mkdir -p "$_venv_dir/bin"
+    cat > "$_venv_dir/bin/python" <<MOCK_EOF
+#!/bin/bash
+if echo "\$@" | grep -q "sys.version_info.minor"; then
+    echo "$_minor"
+else
+    echo "0"
+fi
+MOCK_EOF
+    chmod +x "$_venv_dir/bin/python"
+}
+
+# ── Helper: run the TORCH_CONSTRAINT snippet with given params ──
+run_constraint_snippet() {
+    _skip_torch="$1"
+    _os="$2"
+    _arch="$3"
+    _py_minor="$4"
+    _venv_dir="$5"
+
+    make_mock_python "$_py_minor" "$_venv_dir"
+
+    bash -c "
+        SKIP_TORCH=\"$_skip_torch\"
+        OS=\"$_os\"
+        _ARCH=\"$_arch\"
+        VENV_DIR=\"$_venv_dir\"
+        TORCH_CONSTRAINT=\"torch>=2.4,<2.11.0\"
+        if [ \"\$SKIP_TORCH\" = false ] && [ \"\$OS\" = \"macos\" ] && [ \"\$_ARCH\" = \"arm64\" ]; then
+            _PY_MINOR=\$(\"\$VENV_DIR/bin/python\" -c \"import sys; print(sys.version_info.minor)\" 2>/dev/null || echo \"0\")
+            if [ \"\$_PY_MINOR\" -ge 13 ] 2>/dev/null; then
+                TORCH_CONSTRAINT=\"torch>=2.6,<2.11.0\"
+            fi
+        fi
+        echo \"\$TORCH_CONSTRAINT\"
+    " 2>/dev/null
+}
+
+# ======================================================================
+# Structural checks
+# ======================================================================
+echo "=== Structural: TORCH_CONSTRAINT in install.sh ==="
+
+_SH_CONTENT=$(cat "$INSTALL_SH")
+
+_count=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<2.11.0"' "$INSTALL_SH" || true)
+assert_eq "default TORCH_CONSTRAINT assignment exists" "1" "$_count"
+
+_count=$(grep -c 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' "$INSTALL_SH" || true)
+assert_eq "tightened TORCH_CONSTRAINT assignment exists" "1" "$_count"
+
+_count=$(grep -c '"\$TORCH_CONSTRAINT"' "$INSTALL_SH" || true)
+_has_var=$([ "$_count" -ge 1 ] && echo "yes" || echo "no")
+assert_eq "\$TORCH_CONSTRAINT used in pip install" "yes" "$_has_var"
+
+# Hardcoded torch>=2.4,<2.11.0 should only appear once (the default assignment)
+_hardcoded=$(grep -c '"torch>=2.4,<2.11.0"' "$INSTALL_SH" || true)
+assert_eq "hardcoded torch>=2.4 appears exactly once" "1" "$_hardcoded"
+
+echo ""
+echo "=== Structural: tokenizers in no-torch-runtime.txt ==="
+
+_has_tokenizers=$(grep -c '^tokenizers$' "$NO_TORCH_RT" || true)
+assert_eq "tokenizers present as standalone line" "1" "$_has_tokenizers"
+
+# tokenizers before transformers
+_tok_line=$(grep -n '^tokenizers$' "$NO_TORCH_RT" | head -1 | cut -d: -f1)
+_tf_line=$(grep -n '^transformers' "$NO_TORCH_RT" | head -1 | cut -d: -f1)
+_tok_first=$([ "$_tok_line" -lt "$_tf_line" ] && echo "yes" || echo "no")
+assert_eq "tokenizers before transformers" "yes" "$_tok_first"
+
+# torch itself NOT in no-torch file
+_has_torch=$(grep -c '^torch$' "$NO_TORCH_RT" || true)
+assert_eq "torch not in no-torch-runtime.txt" "0" "$_has_torch"
+
+echo ""
+echo "=== Structural: install.ps1 unchanged ==="
+
+_PS1_CONTENT=$(cat "$INSTALL_PS1")
+_ps1_has_var=$(echo "$_PS1_CONTENT" | grep -c 'TORCH_CONSTRAINT\|TorchConstraint' || true)
+assert_eq "install.ps1 has no TORCH_CONSTRAINT variable" "0" "$_ps1_has_var"
+
+_ps1_hardcoded=$(echo "$_PS1_CONTENT" | grep -c '"torch>=2.4,<2.11.0"' || true)
+_ps1_has_hc=$([ "$_ps1_hardcoded" -ge 1 ] && echo "yes" || echo "no")
+assert_eq "install.ps1 has hardcoded torch constraint" "yes" "$_ps1_has_hc"
+
+# ======================================================================
+# Runtime: mocked platform/version combos
+# ======================================================================
+echo ""
+echo "=== Runtime: TORCH_CONSTRAINT with mocked inputs ==="
+
+TMPDIR_BASE=$(mktemp -d)
+trap 'rm -rf "$TMPDIR_BASE"' EXIT
+
+# 1. arm64 macOS py3.13 -> tightened
+_result=$(run_constraint_snippet false macos arm64 13 "$TMPDIR_BASE/v1")
+assert_eq "arm64+macos+py313 -> tightened" "torch>=2.6,<2.11.0" "$_result"
+
+# 2. arm64 macOS py3.14 -> tightened (future-proofed)
+_result=$(run_constraint_snippet false macos arm64 14 "$TMPDIR_BASE/v2")
+assert_eq "arm64+macos+py314 -> tightened" "torch>=2.6,<2.11.0" "$_result"
+
+# 3. arm64 macOS py3.12 -> default
+_result=$(run_constraint_snippet false macos arm64 12 "$TMPDIR_BASE/v3")
+assert_eq "arm64+macos+py312 -> default" "torch>=2.4,<2.11.0" "$_result"
+
+# 4. arm64 macOS py3.11 -> default
+_result=$(run_constraint_snippet false macos arm64 11 "$TMPDIR_BASE/v4")
+assert_eq "arm64+macos+py311 -> default" "torch>=2.4,<2.11.0" "$_result"
+
+# 5. Linux x86_64 py3.13 -> default (Linux unaffected)
+_result=$(run_constraint_snippet false linux x86_64 13 "$TMPDIR_BASE/v5")
+assert_eq "linux+x86_64+py313 -> default" "torch>=2.4,<2.11.0" "$_result"
+
+# 6. Linux aarch64 py3.13 -> default (guard checks OS=macos)
+_result=$(run_constraint_snippet false linux aarch64 13 "$TMPDIR_BASE/v6")
+assert_eq "linux+aarch64+py313 -> default" "torch>=2.4,<2.11.0" "$_result"
+
+# 7. Intel Mac x86_64 py3.12 -> default (arch mismatch)
+_result=$(run_constraint_snippet false macos x86_64 12 "$TMPDIR_BASE/v7")
+assert_eq "macos+x86_64+py312 -> default" "torch>=2.4,<2.11.0" "$_result"
+
+# 8. SKIP_TORCH=true arm64 macOS py3.13 -> block skipped, default
+_result=$(run_constraint_snippet true macos arm64 13 "$TMPDIR_BASE/v8")
+assert_eq "SKIP_TORCH=true -> default" "torch>=2.4,<2.11.0" "$_result"
+
+# 9. WSL py3.13 -> default
+_result=$(run_constraint_snippet false wsl x86_64 13 "$TMPDIR_BASE/v9")
+assert_eq "wsl+py313 -> default" "torch>=2.4,<2.11.0" "$_result"
+
+# 10. py_minor=0 (failed query fallback) -> default
+_result=$(run_constraint_snippet false macos arm64 0 "$TMPDIR_BASE/v10")
+assert_eq "py_minor=0 fallback -> default" "torch>=2.4,<2.11.0" "$_result"
+
+# 11. Boundary: py_minor=12 -> NOT tightened
+_result=$(run_constraint_snippet false macos arm64 12 "$TMPDIR_BASE/v11")
+assert_eq "boundary py_minor=12 -> default" "torch>=2.4,<2.11.0" "$_result"
+
+# 12. Boundary: py_minor=13 -> tightened
+_result=$(run_constraint_snippet false macos arm64 13 "$TMPDIR_BASE/v12")
+assert_eq "boundary py_minor=13 -> tightened" "torch>=2.6,<2.11.0" "$_result"
+
+# 13. Intel Mac py3.13 -> default (arch=x86_64, not arm64)
+_result=$(run_constraint_snippet false macos x86_64 13 "$TMPDIR_BASE/v13")
+assert_eq "macos+x86_64+py313 -> default" "torch>=2.4,<2.11.0" "$_result"
+
+# ======================================================================
+# Mock uv integration
+# ======================================================================
+echo ""
+echo "=== Mock uv: verify constraint passed to uv ==="
+
+# arm64 + py313 -> uv receives torch>=2.6
+_UV_LOG="$TMPDIR_BASE/uv_log_tight.txt"
+make_mock_python 13 "$TMPDIR_BASE/uv_venv1"
+cat > "$TMPDIR_BASE/mock_uv_tight" <<UVEOF
+#!/bin/bash
+echo "\$@" >> $_UV_LOG
+UVEOF
+chmod +x "$TMPDIR_BASE/mock_uv_tight"
+
+bash -c "
+    SKIP_TORCH=false
+    OS=\"macos\"
+    _ARCH=\"arm64\"
+    VENV_DIR=\"$TMPDIR_BASE/uv_venv1\"
+    TORCH_CONSTRAINT=\"torch>=2.4,<2.11.0\"
+    if [ \"\$SKIP_TORCH\" = false ] && [ \"\$OS\" = \"macos\" ] && [ \"\$_ARCH\" = \"arm64\" ]; then
+        _PY_MINOR=\$(\"\$VENV_DIR/bin/python\" -c \"import sys; print(sys.version_info.minor)\" 2>/dev/null || echo \"0\")
+        if [ \"\$_PY_MINOR\" -ge 13 ] 2>/dev/null; then
+            TORCH_CONSTRAINT=\"torch>=2.6,<2.11.0\"
+        fi
+    fi
+    \"$TMPDIR_BASE/mock_uv_tight\" pip install --python \"\$VENV_DIR/bin/python\" \"\$TORCH_CONSTRAINT\" torchvision torchaudio
+" 2>/dev/null
+_uv_got=$(cat "$_UV_LOG" 2>/dev/null || echo "")
+assert_contains "mock uv arm64+py313 receives torch>=2.6" "$_uv_got" "torch>=2.6,<2.11.0"
+
+# arm64 + py312 -> uv receives torch>=2.4
+_UV_LOG2="$TMPDIR_BASE/uv_log_default.txt"
+make_mock_python 12 "$TMPDIR_BASE/uv_venv2"
+cat > "$TMPDIR_BASE/mock_uv_default" <<UVEOF
+#!/bin/bash
+echo "\$@" >> $_UV_LOG2
+UVEOF
+chmod +x "$TMPDIR_BASE/mock_uv_default"
+
+bash -c "
+    SKIP_TORCH=false
+    OS=\"macos\"
+    _ARCH=\"arm64\"
+    VENV_DIR=\"$TMPDIR_BASE/uv_venv2\"
+    TORCH_CONSTRAINT=\"torch>=2.4,<2.11.0\"
+    if [ \"\$SKIP_TORCH\" = false ] && [ \"\$OS\" = \"macos\" ] && [ \"\$_ARCH\" = \"arm64\" ]; then
+        _PY_MINOR=\$(\"\$VENV_DIR/bin/python\" -c \"import sys; print(sys.version_info.minor)\" 2>/dev/null || echo \"0\")
+        if [ \"\$_PY_MINOR\" -ge 13 ] 2>/dev/null; then
+            TORCH_CONSTRAINT=\"torch>=2.6,<2.11.0\"
+        fi
+    fi
+    \"$TMPDIR_BASE/mock_uv_default\" pip install --python \"\$VENV_DIR/bin/python\" \"\$TORCH_CONSTRAINT\" torchvision torchaudio
+" 2>/dev/null
+_uv_got2=$(cat "$_UV_LOG2" 2>/dev/null || echo "")
+assert_contains "mock uv arm64+py312 receives torch>=2.4" "$_uv_got2" "torch>=2.4,<2.11.0"
+
+# ======================================================================
+# Summary
+# ======================================================================
+echo ""
+echo "=== Results ==="
+echo "  PASS: $PASS"
+echo "  FAIL: $FAIL"
+if [ "$FAIL" -gt 0 ]; then
+    echo "FAILED"
+    exit 1
+fi
+echo "ALL PASSED"

tests/studio/install/smoke_test_llama_prebuilt.py

@@ -39,7 +39,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--llama-tag",
         default = "latest",
-        help = "llama.cpp tag to resolve. Defaults to the approved prebuilt tag for this host.",
+        help = "llama.cpp tag to resolve. Defaults to the latest usable published Unsloth release.",
     )
     parser.add_argument(
         "--published-repo",

tests/studio/install/test_llama_pr_force_and_source.py

@@ -0,0 +1,635 @@
+"""
+Tests for the current llama.cpp wrapper policy in setup.sh / setup.ps1.
+
+Tests cover:
+  - Bash subprocess: PR_FORCE promotion, user-override, zero/empty/invalid ignored
+  - Bash subprocess: source remains pinned to ggml-org even if env source is set
+  - Static source checks: mainline repo/source are hardcoded for now
+  - PowerShell subprocess: PR_FORCE promotion and fixed-source parity
+
+Run: pytest tests/studio/install/test_llama_pr_force_and_source.py -v
+"""
+
+import os
+import shlex
+import subprocess
+import textwrap
+from pathlib import Path
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# Paths
+# ---------------------------------------------------------------------------
+PACKAGE_ROOT = Path(__file__).resolve().parents[3]
+SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
+SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
+
+BASH = "/bin/bash"
+PWSH = "/usr/bin/pwsh"
+PWSH_AVAILABLE = os.path.isfile(PWSH) and os.access(PWSH, os.X_OK)
+requires_pwsh = pytest.mark.skipif(not PWSH_AVAILABLE, reason = "pwsh not available")
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+def run_bash(
+    script: str, *, timeout: int = 10, env: dict | None = None
+) -> subprocess.CompletedProcess:
+    """Run a bash script fragment and return the CompletedProcess."""
+    run_env = os.environ.copy()
+    if env:
+        run_env.update(env)
+    return subprocess.run(
+        [BASH, "-c", script],
+        capture_output = True,
+        text = True,
+        timeout = timeout,
+        env = run_env,
+    )
+
+
+def run_pwsh(
+    script: str, *, timeout: int = 10, env: dict | None = None
+) -> subprocess.CompletedProcess:
+    """Run a PowerShell script fragment and return the CompletedProcess."""
+    run_env = os.environ.copy()
+    run_env["NO_COLOR"] = "1"
+    if env:
+        run_env.update(env)
+    return subprocess.run(
+        [PWSH, "-NoProfile", "-Command", script],
+        capture_output = True,
+        text = True,
+        timeout = timeout,
+        env = run_env,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Shared bash stubs
+# ---------------------------------------------------------------------------
+BASH_STUBS = textwrap.dedent("""\
+    step()    { echo "step:$1:$2"; }
+    substep() { :; }
+    verbose_substep() { :; }
+    print_llama_error_log() { :; }
+    C_ERR= C_WARN= C_OK= C_RST= C_TITLE= C_DIM=
+""")
+
+RUN_QUIET_STUB = textwrap.dedent("""\
+    run_quiet_no_exit() { local _label="$1"; shift; "$@"; return $?; }
+""")
+
+
+def make_mock_git(tmp_path: Path, *, fail_on: str = "") -> tuple[Path, Path]:
+    """Create a mock git binary that logs calls. Returns (mock_bin, log_file)."""
+    mock_bin = tmp_path / "mock_bin"
+    mock_bin.mkdir(exist_ok = True)
+    log_file = tmp_path / "git_calls.log"
+
+    if fail_on:
+        script = (
+            f'#!/bin/bash\necho "$*" >> {log_file}\n'
+            f'_args=("$@")\n'
+            f"_i=0\n"
+            f'while [ "${{_args[$_i]:-}}" = "-C" ]; do _i=$((_i+2)); done\n'
+            f'_subcmd="${{_args[$_i]:-}}"\n'
+            f'if [ "$_subcmd" = "{fail_on}" ]; then exit 1; fi\n'
+            f"exit 0\n"
+        )
+    else:
+        script = f'#!/bin/bash\necho "$*" >> {log_file}\nexit 0\n'
+
+    git_bin = mock_bin / "git"
+    git_bin.write_text(script)
+    git_bin.chmod(0o755)
+    return mock_bin, log_file
+
+
+# =========================================================================
+# Bash fragment that exercises PR_FORCE and fixed _LLAMA_SOURCE resolution
+# =========================================================================
+def _bash_resolution_fragment(
+    llama_pr: str = "",
+    llama_pr_force: str = "",
+    llama_source: str = "",
+    default_pr_force: str = "",
+    default_source: str = "https://github.com/ggml-org/llama.cpp",
+) -> str:
+    """Build the bash fragment that mirrors setup.sh resolution logic."""
+    return BASH_STUBS + textwrap.dedent(f"""\
+        _LLAMA_PR={shlex.quote(llama_pr) if llama_pr else '""'}
+        _DEFAULT_LLAMA_PR_FORCE={shlex.quote(default_pr_force) if default_pr_force else '""'}
+        _DEFAULT_LLAMA_SOURCE={shlex.quote(default_source)}
+
+        _LLAMA_PR_FORCE={shlex.quote(llama_pr_force) if llama_pr_force else '"$_DEFAULT_LLAMA_PR_FORCE"'}
+        export UNSLOTH_LLAMA_SOURCE={shlex.quote(llama_source) if llama_source else '""'}
+        _LLAMA_SOURCE="$_DEFAULT_LLAMA_SOURCE"
+        _LLAMA_SOURCE="${{_LLAMA_SOURCE%.git}}"
+
+        _NEED_LLAMA_SOURCE_BUILD=false
+        _SKIP_PREBUILT_INSTALL=false
+
+        if [ "$_LLAMA_SOURCE" != "https://github.com/ggml-org/llama.cpp" ]; then
+            step "llama.cpp" "custom source: $_LLAMA_SOURCE -- forcing source build"
+            _NEED_LLAMA_SOURCE_BUILD=true
+            _SKIP_PREBUILT_INSTALL=true
+        fi
+
+        if [ -z "$_LLAMA_PR" ] && [ -n "$_LLAMA_PR_FORCE" ] && \\
+           [[ "$_LLAMA_PR_FORCE" =~ ^[0-9]+$ ]] && [ "$_LLAMA_PR_FORCE" -gt 0 ]; then
+            _LLAMA_PR="$_LLAMA_PR_FORCE"
+            step "llama.cpp" "baked-in PR_FORCE=$_LLAMA_PR_FORCE"
+        fi
+
+        echo "LLAMA_PR=$_LLAMA_PR"
+        echo "LLAMA_SOURCE=$_LLAMA_SOURCE"
+        echo "NEED_SOURCE=$_NEED_LLAMA_SOURCE_BUILD"
+        echo "SKIP_PREBUILT=$_SKIP_PREBUILT_INSTALL"
+    """)
+
+
+# =========================================================================
+# TEST GROUP A: Bash PR_FORCE promotion (subprocess)
+# =========================================================================
+class TestBashPrForcePromotion:
+    """PR_FORCE promotes to _LLAMA_PR when user hasn't set one."""
+
+    def test_baked_in_pr_force_promotes(self):
+        script = _bash_resolution_fragment(default_pr_force = "12345")
+        r = run_bash(script)
+        assert r.returncode == 0
+        assert "LLAMA_PR=12345" in r.stdout
+        assert "baked-in PR_FORCE=12345" in r.stdout
+
+    def test_env_pr_force_promotes(self):
+        script = _bash_resolution_fragment(llama_pr_force = "999")
+        r = run_bash(script)
+        assert r.returncode == 0
+        assert "LLAMA_PR=999" in r.stdout
+
+    def test_user_pr_overrides_pr_force(self):
+        """UNSLOTH_LLAMA_PR takes priority over PR_FORCE."""
+        script = _bash_resolution_fragment(
+            llama_pr = "100",
+            llama_pr_force = "200",
+        )
+        r = run_bash(script)
+        assert r.returncode == 0
+        assert "LLAMA_PR=100" in r.stdout
+        assert "baked-in PR_FORCE" not in r.stdout
+
+    def test_user_pr_overrides_baked_in(self):
+        script = _bash_resolution_fragment(
+            llama_pr = "100",
+            default_pr_force = "200",
+        )
+        r = run_bash(script)
+        assert r.returncode == 0
+        assert "LLAMA_PR=100" in r.stdout
+        assert "baked-in PR_FORCE" not in r.stdout
+
+    def test_pr_force_zero_ignored(self):
+        script = _bash_resolution_fragment(llama_pr_force = "0")
+        r = run_bash(script)
+        assert r.returncode == 0
+        assert "LLAMA_PR=" in r.stdout
+        assert "baked-in PR_FORCE" not in r.stdout
+
+    def test_pr_force_empty_ignored(self):
+        script = _bash_resolution_fragment(default_pr_force = "")
+        r = run_bash(script)
+        assert r.returncode == 0
+        assert "LLAMA_PR=" in r.stdout
+        assert "baked-in PR_FORCE" not in r.stdout
+
+    def test_pr_force_alpha_ignored(self):
+        script = _bash_resolution_fragment(llama_pr_force = "abc")
+        r = run_bash(script)
+        assert r.returncode == 0
+        assert "LLAMA_PR=" in r.stdout
+        assert "baked-in PR_FORCE" not in r.stdout
+
+    def test_pr_force_negative_ignored(self):
+        script = _bash_resolution_fragment(llama_pr_force = "-5")
+        r = run_bash(script)
+        assert r.returncode == 0
+        assert "LLAMA_PR=" in r.stdout
+
+    def test_pr_force_decimal_ignored(self):
+        script = _bash_resolution_fragment(llama_pr_force = "12.34")
+        r = run_bash(script)
+        assert r.returncode == 0
+        assert "LLAMA_PR=" in r.stdout
+
+
+# =========================================================================
+# TEST GROUP B: Bash fixed mainline source (subprocess)
+# =========================================================================
+class TestBashFixedMainlineSource:
+    """Source remains pinned to ggml-org while the temporary policy is active."""
+
+    def test_default_source_no_force(self):
+        script = _bash_resolution_fragment()
+        r = run_bash(script)
+        assert r.returncode == 0
+        assert "NEED_SOURCE=false" in r.stdout
+        assert "SKIP_PREBUILT=false" in r.stdout
+        assert "custom source:" not in r.stdout
+
+    def test_env_source_override_is_ignored(self):
+        script = _bash_resolution_fragment(
+            llama_source = "https://github.com/unslothai/llama.cpp.git",
+        )
+        r = run_bash(script)
+        assert r.returncode == 0
+        assert "LLAMA_SOURCE=https://github.com/ggml-org/llama.cpp" in r.stdout
+        assert "NEED_SOURCE=false" in r.stdout
+        assert "SKIP_PREBUILT=false" in r.stdout
+
+    def test_baked_in_source_stays_mainline(self):
+        script = _bash_resolution_fragment(
+            default_source = "https://github.com/ggml-org/llama.cpp",
+        )
+        r = run_bash(script)
+        assert r.returncode == 0
+        assert "LLAMA_SOURCE=https://github.com/ggml-org/llama.cpp" in r.stdout
+
+
+# =========================================================================
+# TEST GROUP C: Bash clone URL parameterization (subprocess with mock git)
+# =========================================================================
+class TestBashCloneUrlParameterized:
+    """Verify git clone uses _LLAMA_SOURCE instead of hardcoded URL."""
+
+    @staticmethod
+    def _clone_script(
+        mock_bin: Path,
+        build_tmp: str,
+        llama_pr: str = "",
+        llama_source: str = "https://github.com/ggml-org/llama.cpp",
+        resolved_tag: str = "b8508",
+    ) -> str:
+        return RUN_QUIET_STUB + textwrap.dedent(f"""\
+            export PATH="{mock_bin}:$PATH"
+            _LLAMA_PR={shlex.quote(llama_pr) if llama_pr else '""'}
+            _LLAMA_SOURCE={shlex.quote(llama_source)}
+            _RESOLVED_LLAMA_TAG={shlex.quote(resolved_tag)}
+            _BUILD_TMP={shlex.quote(build_tmp)}
+            BUILD_OK=true
+
+            if [ -n "$_LLAMA_PR" ]; then
+                run_quiet_no_exit "clone llama.cpp" \\
+                    git clone --depth 1 "${{_LLAMA_SOURCE}}.git" "$_BUILD_TMP" || BUILD_OK=false
+            else
+                _CLONE_ARGS=(git clone --depth 1)
+                if [ "$_RESOLVED_LLAMA_TAG" != "latest" ] && [ -n "$_RESOLVED_LLAMA_TAG" ]; then
+                    _CLONE_ARGS+=(--branch "$_RESOLVED_LLAMA_TAG")
+                fi
+                _CLONE_ARGS+=("${{_LLAMA_SOURCE}}.git" "$_BUILD_TMP")
+                run_quiet_no_exit "clone llama.cpp" \\
+                    "${{_CLONE_ARGS[@]}}" || BUILD_OK=false
+            fi
+            echo "BUILD_OK=$BUILD_OK"
+        """)
+
+    def test_pr_path_uses_custom_source(self, tmp_path: Path):
+        mock_bin, log_file = make_mock_git(tmp_path)
+        build_tmp = str(tmp_path / "build_tmp")
+        script = self._clone_script(
+            mock_bin,
+            build_tmp,
+            llama_pr = "123",
+            llama_source = "https://github.com/unslothai/llama.cpp",
+        )
+        r = run_bash(script)
+        assert r.returncode == 0
+        log = log_file.read_text()
+        assert "unslothai/llama.cpp.git" in log
+        assert "ggml-org" not in log
+
+    def test_non_pr_path_uses_custom_source(self, tmp_path: Path):
+        mock_bin, log_file = make_mock_git(tmp_path)
+        build_tmp = str(tmp_path / "build_tmp")
+        script = self._clone_script(
+            mock_bin,
+            build_tmp,
+            llama_source = "https://github.com/unslothai/llama.cpp",
+        )
+        r = run_bash(script)
+        assert r.returncode == 0
+        log = log_file.read_text()
+        assert "unslothai/llama.cpp.git" in log
+        assert "ggml-org" not in log
+
+    def test_default_source_unchanged(self, tmp_path: Path):
+        mock_bin, log_file = make_mock_git(tmp_path)
+        build_tmp = str(tmp_path / "build_tmp")
+        script = self._clone_script(mock_bin, build_tmp)
+        r = run_bash(script)
+        assert r.returncode == 0
+        log = log_file.read_text()
+        assert "ggml-org/llama.cpp.git" in log
+
+    def test_latest_tag_omits_branch_flag(self, tmp_path: Path):
+        """resolved_tag='latest' should not pass --branch to git clone."""
+        mock_bin, log_file = make_mock_git(tmp_path)
+        build_tmp = str(tmp_path / "build_tmp")
+        script = self._clone_script(
+            mock_bin,
+            build_tmp,
+            resolved_tag = "latest",
+        )
+        r = run_bash(script)
+        assert r.returncode == 0
+        log = log_file.read_text()
+        assert "--branch" not in log
+        assert "ggml-org/llama.cpp.git" in log
+
+    def test_empty_tag_omits_branch_flag(self, tmp_path: Path):
+        """resolved_tag='' (empty) should not pass --branch to git clone."""
+        mock_bin, log_file = make_mock_git(tmp_path)
+        build_tmp = str(tmp_path / "build_tmp")
+        script = self._clone_script(
+            mock_bin,
+            build_tmp,
+            resolved_tag = "",
+        )
+        r = run_bash(script)
+        assert r.returncode == 0
+        log = log_file.read_text()
+        assert "--branch" not in log
+        assert "ggml-org/llama.cpp.git" in log
+
+
+# =========================================================================
+# TEST GROUP D: Static source patterns -- setup.sh
+# =========================================================================
+class TestSourcePatternsSh:
+    """Verify setup.sh keeps the temporary mainline-only llama.cpp policy."""
+
+    @pytest.fixture(autouse = True)
+    def _load_source(self):
+        self.content = SETUP_SH.read_text()
+
+    def test_has_default_pr_force(self):
+        assert '_DEFAULT_LLAMA_PR_FORCE=""' in self.content
+
+    def test_has_default_source(self):
+        assert (
+            '_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp"'
+            in self.content
+        )
+
+    def test_has_pr_force_env_read(self):
+        assert "UNSLOTH_LLAMA_PR_FORCE" in self.content
+
+    def test_source_env_override_removed(self):
+        assert "UNSLOTH_LLAMA_SOURCE:-${_DEFAULT_LLAMA_SOURCE}" not in self.content
+        assert '_LLAMA_SOURCE="${_DEFAULT_LLAMA_SOURCE}"' in self.content
+
+    def test_release_repo_override_removed(self):
+        assert "UNSLOTH_LLAMA_RELEASE_REPO:-unslothai/llama.cpp" not in self.content
+        assert '_HELPER_RELEASE_REPO="ggml-org/llama.cpp"' in self.content
+
+    def test_force_compile_skips_prebuilt_resolution_early(self):
+        assert 'if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then' in self.content
+        assert "_SKIP_PREBUILT_INSTALL=true" in self.content
+
+    def test_force_compile_uses_requested_tag_without_helper(self):
+        assert 'if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then' in self.content
+        assert '_RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"' in self.content
+
+    def test_pr_force_resolution_block(self):
+        assert '_LLAMA_PR="$_LLAMA_PR_FORCE"' in self.content
+
+    def test_source_trailing_git_strip(self):
+        assert "${_LLAMA_SOURCE%.git}" in self.content
+
+    def test_clone_urls_parameterized_pr_path(self):
+        """PR clone path uses ${_LLAMA_SOURCE}.git, not hardcoded URL."""
+        pr_clone_idx = self.content.index(
+            'if [ -n "$_LLAMA_PR" ]; then\n'
+            '            run_quiet_no_exit "clone llama.cpp"'
+        )
+        else_idx = self.content.index("else\n", pr_clone_idx)
+        pr_block = self.content[pr_clone_idx:else_idx]
+        assert '"${_LLAMA_SOURCE}.git"' in pr_block
+        assert "ggml-org/llama.cpp.git" not in pr_block
+
+    def test_clone_urls_parameterized_tag_path(self):
+        """Non-PR clone path uses the resolved source URL, not a hardcoded URL."""
+        # Find the non-PR clone line (after _CLONE_ARGS)
+        idx = self.content.index("_CLONE_ARGS=(git clone --depth 1)")
+        block = self.content[idx : idx + 400]
+        assert '"${_RESOLVED_SOURCE_URL}.git"' in block
+        assert "ggml-org/llama.cpp.git" not in block
+
+    def test_no_hardcoded_clone_urls(self):
+        """No remaining hardcoded ggml-org clone URLs in clone commands."""
+        lines = self.content.splitlines()
+        for i, line in enumerate(lines, 1):
+            if "git clone" in line and "ggml-org/llama.cpp.git" in line:
+                pytest.fail(
+                    f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}"
+                )
+
+
+# =========================================================================
+# TEST GROUP E: Static source patterns -- setup.ps1
+# =========================================================================
+class TestSourcePatternsPs1:
+    """Verify setup.ps1 keeps the temporary mainline-only llama.cpp policy."""
+
+    @pytest.fixture(autouse = True)
+    def _load_source(self):
+        self.content = SETUP_PS1.read_text()
+
+    def test_has_default_pr_force(self):
+        assert '$DefaultLlamaPrForce = ""' in self.content
+
+    def test_has_default_source(self):
+        assert (
+            '$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp"'
+            in self.content
+        )
+
+    def test_has_pr_force_env_read(self):
+        assert "$env:UNSLOTH_LLAMA_PR_FORCE" in self.content
+
+    def test_source_env_override_removed(self):
+        assert "$LlamaSource = if ($env:UNSLOTH_LLAMA_SOURCE)" not in self.content
+        assert "$LlamaSource = $DefaultLlamaSource" in self.content
+
+    def test_release_repo_override_removed(self):
+        assert (
+            "$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO)"
+            not in self.content
+        )
+        assert '$HelperReleaseRepo = "ggml-org/llama.cpp"' in self.content
+
+    def test_force_compile_skips_prebuilt_resolution_early(self):
+        assert 'if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {' in self.content
+        assert "$SkipPrebuiltInstall = $true" in self.content
+
+    def test_force_compile_uses_requested_tag_without_helper(self):
+        assert 'if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {' in self.content
+        assert "$ResolvedLlamaTag = $RequestedLlamaTag" in self.content
+
+    def test_pr_force_promotion_block(self):
+        assert "$LlamaPr = $LlamaPrForce" in self.content
+
+    def test_source_trailing_git_strip(self):
+        assert ".EndsWith('.git')" in self.content
+
+    def test_clone_urls_parameterized_pr_path(self):
+        """PR clone path uses $LlamaSource.git, not hardcoded URL."""
+        pr_idx = self.content.index(
+            "if ($LlamaPr) {\n", self.content.index("Cloning llama.cpp")
+        )
+        else_idx = self.content.index("} else {", pr_idx)
+        pr_block = self.content[pr_idx:else_idx]
+        assert '"$LlamaSource.git"' in pr_block
+        assert "ggml-org/llama.cpp.git" not in pr_block
+
+    def test_clone_urls_parameterized_tag_path(self):
+        """Non-PR clone path uses the resolved source URL, not a hardcoded URL."""
+        clone_args_idx = self.content.index('$cloneArgs = @("clone"')
+        block = self.content[clone_args_idx : clone_args_idx + 400]
+        assert '"$ResolvedSourceUrl.git"' in block
+        assert "ggml-org/llama.cpp.git" not in block
+
+    def test_no_hardcoded_clone_urls(self):
+        """No remaining hardcoded ggml-org clone URLs in clone commands."""
+        lines = self.content.splitlines()
+        for i, line in enumerate(lines, 1):
+            if "git clone" in line and "ggml-org/llama.cpp.git" in line:
+                pytest.fail(
+                    f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}"
+                )
+
+
+# =========================================================================
+# TEST GROUP F: PowerShell PR_FORCE promotion (subprocess)
+# =========================================================================
+@requires_pwsh
+class TestPwshPrForcePromotion:
+    """PR_FORCE promotion and fixed-source logic via pwsh subprocess."""
+
+    FRAGMENT_TEMPLATE = textwrap.dedent("""\
+        function step($a, $b, $c) { Write-Output "step:$a`:$b" }
+
+        $DefaultLlamaPrForce = "%%DEFAULT_PR_FORCE%%"
+        $DefaultLlamaSource = "%%DEFAULT_SOURCE%%"
+
+        $LlamaPr = if ($env:UNSLOTH_LLAMA_PR) { $env:UNSLOTH_LLAMA_PR.Trim() } else { "" }
+        $LlamaPrForce = if ($env:UNSLOTH_LLAMA_PR_FORCE) { $env:UNSLOTH_LLAMA_PR_FORCE.Trim() } else { $DefaultLlamaPrForce }
+        $LlamaSource = $DefaultLlamaSource
+        if ($LlamaSource.EndsWith('.git')) { $LlamaSource = $LlamaSource.Substring(0, $LlamaSource.Length - 4) }
+
+        $NeedLlamaSourceBuild = $false
+        $SkipPrebuiltInstall = $false
+
+        if ($LlamaSource -ne "https://github.com/ggml-org/llama.cpp") {
+            step "llama.cpp" "custom source: $LlamaSource -- forcing source build" "Yellow"
+            $NeedLlamaSourceBuild = $true
+            $SkipPrebuiltInstall = $true
+        }
+
+        if (-not $LlamaPr -and $LlamaPrForce -and $LlamaPrForce -match '^\\d+$' -and [int]$LlamaPrForce -gt 0) {
+            $LlamaPr = $LlamaPrForce
+            step "llama.cpp" "baked-in PR_FORCE=$LlamaPrForce" "Yellow"
+        }
+
+        Write-Output "LLAMA_PR=$LlamaPr"
+        Write-Output "LLAMA_SOURCE=$LlamaSource"
+        Write-Output "NEED_SOURCE=$NeedLlamaSourceBuild"
+        Write-Output "SKIP_PREBUILT=$SkipPrebuiltInstall"
+    """)
+
+    def _run(
+        self,
+        default_pr_force: str = "",
+        default_source: str = "https://github.com/ggml-org/llama.cpp",
+        env: dict | None = None,
+    ) -> subprocess.CompletedProcess:
+        script = self.FRAGMENT_TEMPLATE.replace(
+            "%%DEFAULT_PR_FORCE%%",
+            default_pr_force,
+        ).replace(
+            "%%DEFAULT_SOURCE%%",
+            default_source,
+        )
+        run_env = {}
+        # Ensure env vars are unset by default
+        run_env["UNSLOTH_LLAMA_PR"] = ""
+        run_env["UNSLOTH_LLAMA_PR_FORCE"] = ""
+        if env:
+            run_env.update(env)
+        return run_pwsh(script, env = run_env)
+
+    def test_baked_in_pr_force_promotes(self):
+        r = self._run(default_pr_force = "12345")
+        assert r.returncode == 0
+        assert "LLAMA_PR=12345" in r.stdout
+        assert "baked-in PR_FORCE=12345" in r.stdout
+
+    def test_env_pr_force_promotes(self):
+        r = self._run(env = {"UNSLOTH_LLAMA_PR_FORCE": "999"})
+        assert r.returncode == 0
+        assert "LLAMA_PR=999" in r.stdout
+
+    def test_user_pr_overrides_pr_force(self):
+        r = self._run(
+            env = {
+                "UNSLOTH_LLAMA_PR": "100",
+                "UNSLOTH_LLAMA_PR_FORCE": "200",
+            }
+        )
+        assert r.returncode == 0
+        assert "LLAMA_PR=100" in r.stdout
+        assert "baked-in PR_FORCE" not in r.stdout
+
+    def test_pr_force_zero_ignored(self):
+        r = self._run(env = {"UNSLOTH_LLAMA_PR_FORCE": "0"})
+        assert r.returncode == 0
+        assert "LLAMA_PR=" in r.stdout
+        assert "baked-in PR_FORCE" not in r.stdout
+
+    def test_pr_force_alpha_ignored(self):
+        r = self._run(env = {"UNSLOTH_LLAMA_PR_FORCE": "abc"})
+        assert r.returncode == 0
+        assert "baked-in PR_FORCE" not in r.stdout
+
+    def test_env_source_override_is_ignored(self):
+        r = self._run(
+            env = {
+                "UNSLOTH_LLAMA_SOURCE": "https://github.com/unslothai/llama.cpp",
+            }
+        )
+        assert r.returncode == 0
+        assert "LLAMA_SOURCE=https://github.com/ggml-org/llama.cpp" in r.stdout
+        assert "NEED_SOURCE=False" in r.stdout
+        assert "SKIP_PREBUILT=False" in r.stdout
+
+    def test_default_source_no_force(self):
+        r = self._run()
+        assert r.returncode == 0
+        assert "NEED_SOURCE=False" in r.stdout
+        assert "SKIP_PREBUILT=False" in r.stdout
+
+    def test_trailing_git_override_is_ignored(self):
+        r = self._run(
+            env = {
+                "UNSLOTH_LLAMA_SOURCE": "https://github.com/unslothai/llama.cpp.git",
+            }
+        )
+        assert r.returncode == 0
+        assert "LLAMA_SOURCE=https://github.com/ggml-org/llama.cpp" in r.stdout
+
+    def test_baked_in_source_stays_mainline(self):
+        r = self._run(default_source = "https://github.com/ggml-org/llama.cpp")
+        assert r.returncode == 0
+        assert "LLAMA_SOURCE=https://github.com/ggml-org/llama.cpp" in r.stdout

tests/studio/install/test_pr4562_bugfixes.py

@@ -6,7 +6,7 @@ Tests cover:
   - Bug 2: Source-build fallback ignores pinned tag (both .sh and .ps1)
   - Bug 3: Unix fallback deletes install before checking prerequisites
   - Bug 4: Linux LD_LIBRARY_PATH missing build/bin
-  - "latest" tag resolution fallback chain (Unsloth -> ggml-org -> raw)
+  - "latest" tag resolution fallback chain (helper only)
   - Cross-platform binary_env (Linux, macOS, Windows)
   - Edge cases: malformed JSON, empty responses, env overrides
 
@@ -14,13 +14,11 @@ Run: pytest tests/studio/install/test_pr4562_bugfixes.py -v
 """
 
 import importlib.util
-import json
 import os
 import subprocess
 import sys
 import textwrap
 from pathlib import Path
-from unittest.mock import patch
 
 import pytest
 
@@ -40,6 +38,10 @@ SPEC.loader.exec_module(MOD)
 binary_env = MOD.binary_env
 HostInfo = MOD.HostInfo
 resolve_requested_llama_tag = MOD.resolve_requested_llama_tag
+PublishedReleaseBundle = MOD.PublishedReleaseBundle
+ApprovedArtifactHash = MOD.ApprovedArtifactHash
+ApprovedReleaseChecksums = MOD.ApprovedReleaseChecksums
+source_archive_logical_name = MOD.source_archive_logical_name
 
 SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
 SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
@@ -82,6 +84,9 @@ def run_bash(script: str, *, timeout: int = 10, env: dict | None = None) -> str:
         timeout = timeout,
         env = run_env,
     )
+    assert (
+        result.returncode == 0
+    ), f"bash script failed (exit {result.returncode}):\n{result.stderr}"
     return result.stdout.strip()
 
 
@@ -240,6 +245,107 @@ class TestResolveRequestedLlamaTag:
         monkeypatch.setattr(MOD, "latest_upstream_release_tag", lambda: "b5555")
         assert resolve_requested_llama_tag("") == "b5555"
 
+    def test_latest_with_published_repo_uses_latest_valid_published_release(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        invalid = PublishedReleaseBundle(
+            repo = "unslothai/llama.cpp",
+            release_tag = "v2.0",
+            upstream_tag = "b9000",
+            assets = {},
+            manifest_asset_name = "llama-prebuilt-manifest.json",
+            artifacts = [],
+            selection_log = [],
+        )
+        valid = PublishedReleaseBundle(
+            repo = "unslothai/llama.cpp",
+            release_tag = "v1.0",
+            upstream_tag = "b8999",
+            assets = {},
+            manifest_asset_name = "llama-prebuilt-manifest.json",
+            artifacts = [],
+            selection_log = [],
+        )
+
+        monkeypatch.setattr(
+            MOD,
+            "iter_published_release_bundles",
+            lambda repo, published_release_tag = "": iter([invalid, valid]),
+        )
+
+        def fake_load(repo, release_tag):
+            if release_tag == "v2.0":
+                raise MOD.PrebuiltFallback("checksum asset missing")
+            return ApprovedReleaseChecksums(
+                repo = repo,
+                release_tag = release_tag,
+                upstream_tag = "b8999",
+                source_commit = None,
+                artifacts = {
+                    source_archive_logical_name("b8999"): ApprovedArtifactHash(
+                        asset_name = source_archive_logical_name("b8999"),
+                        sha256 = "a" * 64,
+                        repo = "ggml-org/llama.cpp",
+                        kind = "upstream-source",
+                    )
+                },
+            )
+
+        monkeypatch.setattr(MOD, "load_approved_release_checksums", fake_load)
+        monkeypatch.setattr(MOD, "latest_upstream_release_tag", lambda: "b7777")
+
+        assert resolve_requested_llama_tag("latest", "unslothai/llama.cpp") == "b8999"
+
+    def test_latest_with_published_release_tag_passes_pin_through(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        captured = {}
+
+        def fake_resolve(requested_tag, published_repo, published_release_tag = ""):
+            captured["requested_tag"] = requested_tag
+            captured["published_repo"] = published_repo
+            captured["published_release_tag"] = published_release_tag
+            return MOD.ResolvedPublishedRelease(
+                bundle = PublishedReleaseBundle(
+                    repo = published_repo,
+                    release_tag = published_release_tag,
+                    upstream_tag = "b9001",
+                    assets = {},
+                    manifest_asset_name = "llama-prebuilt-manifest.json",
+                    artifacts = [],
+                    selection_log = [],
+                ),
+                checksums = ApprovedReleaseChecksums(
+                    repo = published_repo,
+                    release_tag = published_release_tag,
+                    upstream_tag = "b9001",
+                    artifacts = {
+                        source_archive_logical_name("b9001"): ApprovedArtifactHash(
+                            asset_name = source_archive_logical_name("b9001"),
+                            sha256 = "a" * 64,
+                            repo = "ggml-org/llama.cpp",
+                            kind = "upstream-source",
+                        )
+                    },
+                ),
+            )
+
+        monkeypatch.setattr(MOD, "resolve_published_release", fake_resolve)
+
+        assert (
+            resolve_requested_llama_tag(
+                "latest",
+                "unslothai/llama.cpp",
+                "llama-prebuilt-main",
+            )
+            == "b9001"
+        )
+        assert captured == {
+            "requested_tag": "latest",
+            "published_repo": "unslothai/llama.cpp",
+            "published_release_tag": "llama-prebuilt-main",
+        }
+
 
 # =========================================================================
 # TEST GROUP C: setup.sh logic (bash subprocess tests)
@@ -429,140 +535,61 @@ class TestSetupShLogic:
 # TEST GROUP D: "latest" tag resolution (bash subprocess)
 # =========================================================================
 class TestLatestTagResolution:
-    """Test the fallback chain: Unsloth API -> ggml-org API -> raw."""
+    """Test the fallback chain: helper resolver -> raw."""
 
     RESOLVE_TEMPLATE = textwrap.dedent("""\
-        export PATH="{mock_bin}:$PATH"
         _REQUESTED_LLAMA_TAG="{requested_tag}"
         _RESOLVED_LLAMA_TAG=""
-        _RESOLVE_UPSTREAM_STATUS=1
-        _HELPER_RELEASE_REPO="unslothai/llama.cpp"
-        if [ "$_RESOLVE_UPSTREAM_STATUS" -ne 0 ] || [ -z "$_RESOLVED_LLAMA_TAG" ]; then
-            if [ "$_REQUESTED_LLAMA_TAG" = "latest" ]; then
-                _RESOLVED_LLAMA_TAG="$(curl -fsSL "https://api.github.com/repos/${{_HELPER_RELEASE_REPO}}/releases/latest" 2>/dev/null | python -c "import sys,json; print(json.load(sys.stdin)['tag_name'])" 2>/dev/null)" || _RESOLVED_LLAMA_TAG=""
-                if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
-                    _RESOLVED_LLAMA_TAG="$(curl -fsSL https://api.github.com/repos/ggml-org/llama.cpp/releases/latest 2>/dev/null | python -c "import sys,json; print(json.load(sys.stdin)['tag_name'])" 2>/dev/null)" || _RESOLVED_LLAMA_TAG=""
-                fi
-            fi
-            if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
-                _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"
-            fi
+        _RESOLVE_UPSTREAM_STATUS={resolve_status}
+        if [ "$_RESOLVE_UPSTREAM_STATUS" -eq 0 ] && [ -n "{resolved_tag}" ]; then
+            _RESOLVED_LLAMA_TAG="{resolved_tag}"
+        else
+            _RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"
         fi
         echo "$_RESOLVED_LLAMA_TAG"
     """)
 
-    @staticmethod
-    def _make_curl_mock(
-        mock_bin: Path, unsloth_response: str | None, ggml_response: str | None
-    ):
-        """Create a curl mock that returns different responses per repo."""
-        lines = ["#!/bin/bash"]
-        if unsloth_response is not None:
-            lines.append(
-                f'if echo "$*" | grep -q "unslothai/llama.cpp"; then echo \'{unsloth_response}\'; exit 0; fi'
-            )
-        else:
-            lines.append(
-                'if echo "$*" | grep -q "unslothai/llama.cpp"; then exit 1; fi'
-            )
-        if ggml_response is not None:
-            lines.append(
-                f'if echo "$*" | grep -q "ggml-org/llama.cpp"; then echo \'{ggml_response}\'; exit 0; fi'
-            )
-        else:
-            lines.append('if echo "$*" | grep -q "ggml-org/llama.cpp"; then exit 1; fi')
-        lines.append("exit 1")
-        curl_path = mock_bin / "curl"
-        curl_path.write_text("\n".join(lines) + "\n")
-        curl_path.chmod(0o755)
-
     def _run_resolve(
         self,
         tmp_path: Path,
         requested_tag: str,
-        unsloth_resp: str | None,
-        ggml_resp: str | None,
+        resolved_tag: str,
+        resolve_status: int,
     ) -> str:
-        mock_bin = tmp_path / "mock_bin"
-        mock_bin.mkdir(exist_ok = True)
-        self._make_curl_mock(mock_bin, unsloth_resp, ggml_resp)
         script = self.RESOLVE_TEMPLATE.format(
-            mock_bin = mock_bin, requested_tag = requested_tag
+            requested_tag = requested_tag,
+            resolved_tag = resolved_tag,
+            resolve_status = resolve_status,
         )
         return run_bash(script)
 
-    def test_unsloth_succeeds(self, tmp_path: Path):
+    def test_helper_resolution_succeeds(self, tmp_path: Path):
         output = self._run_resolve(
             tmp_path,
             "latest",
-            unsloth_resp = '{"tag_name":"b8508"}',
-            ggml_resp = '{"tag_name":"b9000"}',
+            resolved_tag = "b8508",
+            resolve_status = 0,
         )
         assert output == "b8508"
 
-    def test_unsloth_fails_ggml_succeeds(self, tmp_path: Path):
+    def test_helper_resolution_falls_back_to_raw_requested_tag(self, tmp_path: Path):
         output = self._run_resolve(
             tmp_path,
             "latest",
-            unsloth_resp = None,
-            ggml_resp = '{"tag_name":"b9000"}',
-        )
-        assert output == "b9000"
-
-    def test_both_fail_raw_fallback(self, tmp_path: Path):
-        output = self._run_resolve(
-            tmp_path,
-            "latest",
-            unsloth_resp = None,
-            ggml_resp = None,
+            resolved_tag = "",
+            resolve_status = 1,
         )
         assert output == "latest"
 
-    def test_concrete_tag_passes_through(self, tmp_path: Path):
+    def test_concrete_tag_passes_through_when_helper_fails(self, tmp_path: Path):
         output = self._run_resolve(
             tmp_path,
             "b7777",
-            unsloth_resp = '{"tag_name":"b8508"}',
-            ggml_resp = '{"tag_name":"b9000"}',
+            resolved_tag = "",
+            resolve_status = 1,
         )
         assert output == "b7777"
 
-    def test_unsloth_malformed_json_falls_through(self, tmp_path: Path):
-        output = self._run_resolve(
-            tmp_path,
-            "latest",
-            unsloth_resp = '{"bad_key":"no_tag"}',
-            ggml_resp = '{"tag_name":"b9001"}',
-        )
-        assert output == "b9001"
-
-    def test_both_malformed_json_raw_fallback(self, tmp_path: Path):
-        output = self._run_resolve(
-            tmp_path,
-            "latest",
-            unsloth_resp = '{"bad":"data"}',
-            ggml_resp = '{"also":"bad"}',
-        )
-        assert output == "latest"
-
-    def test_unsloth_empty_body_falls_through(self, tmp_path: Path):
-        output = self._run_resolve(
-            tmp_path,
-            "latest",
-            unsloth_resp = "",
-            ggml_resp = '{"tag_name":"b7000"}',
-        )
-        assert output == "b7000"
-
-    def test_unsloth_empty_tag_name_falls_through(self, tmp_path: Path):
-        output = self._run_resolve(
-            tmp_path,
-            "latest",
-            unsloth_resp = '{"tag_name":""}',
-            ggml_resp = '{"tag_name":"b6000"}',
-        )
-        assert output == "b6000"
-
     def test_env_override_unsloth_llama_tag(self):
         output = run_bash(
             'echo "${UNSLOTH_LLAMA_TAG:-latest}"',
@@ -593,10 +620,10 @@ class TestSourceCodePatterns:
     def test_setup_sh_no_rm_before_prereq_check(self):
         """rm -rf must appear AFTER cmake/git checks, not before."""
         content = SETUP_SH.read_text()
-        # Find the source-build block
-        idx_else = content.find("# Check prerequisites")
-        assert idx_else != -1
-        block = content[idx_else:]
+        # Anchor on the source-build cmake check block.
+        idx_block = content.find("command -v cmake")
+        assert idx_block != -1
+        block = content[idx_block:]
         # rm -rf should appear after the cmake/git checks
         idx_cmake = block.find("command -v cmake")
         idx_git = block.find("command -v git")
@@ -605,28 +632,102 @@ class TestSourceCodePatterns:
         assert idx_rm > idx_git, "rm -rf should come after git check"
 
     def test_setup_sh_clone_uses_branch_tag(self):
-        """git clone in source-build should use --branch via _CLONE_BRANCH_ARGS."""
+        """git clone in source-build should use --branch via the clone args array."""
         content = SETUP_SH.read_text()
-        # The clone line should use _CLONE_BRANCH_ARGS (which conditionally includes --branch)
+        assert "_CLONE_ARGS=(git clone --depth 1)" in content
         assert (
-            "_CLONE_BRANCH_ARGS" in content
-        ), "Clone should use _CLONE_BRANCH_ARGS array"
-        assert (
-            '--branch "$_RESOLVED_LLAMA_TAG"' in content
-        ), "_CLONE_BRANCH_ARGS should be set to --branch $_RESOLVED_LLAMA_TAG"
+            '_CLONE_ARGS+=(--branch "$_RESOLVED_SOURCE_REF")' in content
+        ), "_CLONE_ARGS should be extended with --branch $_RESOLVED_SOURCE_REF"
         # Verify the guard: --branch is only used when tag is not "latest"
         assert (
-            '_RESOLVED_LLAMA_TAG" != "latest"' in content
+            '_RESOLVED_SOURCE_REF" != "latest"' in content
         ), "Should guard against literal 'latest' tag"
 
-    def test_setup_sh_latest_resolution_queries_unsloth_first(self):
-        """The Unsloth repo should be queried before ggml-org."""
+    def test_setup_sh_source_build_uses_helper_resolution(self):
+        """Shell source fallback should consult the helper for repo/ref planning."""
         content = SETUP_SH.read_text()
-        idx_unsloth = content.find("_HELPER_RELEASE_REPO}/releases/latest")
-        idx_ggml = content.find("ggml-org/llama.cpp/releases/latest")
-        assert idx_unsloth != -1, "Unsloth API query not found"
-        assert idx_ggml != -1, "ggml-org API query not found"
-        assert idx_unsloth < idx_ggml, "Unsloth should be queried before ggml-org"
+        assert "--resolve-source-build" in content
+        assert "--output-format json" in content
+        assert "_RESOLVED_SOURCE_URL" in content
+        assert "_RESOLVED_SOURCE_REF_KIND" in content
+        assert "_RESOLVED_SOURCE_REF" in content
+
+    def test_setup_sh_latest_resolution_uses_helper_only(self):
+        """Shell fallback should rely on helper output, not raw GitHub API tag_name."""
+        content = SETUP_SH.read_text()
+        assert "--resolve-install-tag" in content
+        assert "--resolve-llama-tag" in content
+        assert 'tail -n 1 "$_RESOLVE_LLAMA_LOG"' not in content
+        assert "json.load" in content
+        assert "_HELPER_RELEASE_REPO}/releases/latest" not in content
+        assert "ggml-org/llama.cpp/releases/latest" not in content
+
+    def test_setup_sh_macos_arm64_uses_metal_flags(self):
+        """Apple Silicon source builds should explicitly enable Metal like upstream."""
+        content = SETUP_SH.read_text()
+        assert "_IS_MACOS_ARM64=true" in content
+        assert 'if [ "$_IS_MACOS_ARM64" = true ]; then' in content
+        assert "-DGGML_METAL=ON" in content
+        assert "-DGGML_METAL_EMBED_LIBRARY=ON" in content
+        assert "-DGGML_METAL_USE_BF16=ON" in content
+        assert "-DCMAKE_INSTALL_RPATH=@loader_path" in content
+        assert "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" in content
+
+    def test_setup_sh_macos_metal_configure_has_cpu_fallback(self):
+        """If Metal configure or build fails, setup should retry with CPU fallback."""
+        content = SETUP_SH.read_text()
+        assert "_TRY_METAL_CPU_FALLBACK=true" in content
+        assert (
+            'substep "Metal configure failed; retrying CPU build..." "$C_WARN"'
+            in content
+        )
+        assert (
+            'substep "Metal build failed; retrying CPU build..." "$C_WARN"' in content
+        )
+        assert 'run_quiet_no_exit "cmake llama.cpp (cpu fallback)"' in content
+        assert "-DGGML_METAL=OFF" in content
+        # _TRY_METAL_CPU_FALLBACK must be reset to false in both fallback branches
+        # (1 init + 2 resets = at least 3 occurrences of =false)
+        assert content.count("_TRY_METAL_CPU_FALLBACK=false") >= 3, (
+            "_TRY_METAL_CPU_FALLBACK=false should appear at least 3 times "
+            "(init + configure fallback + build fallback)"
+        )
+
+    def test_macos_arm64_cpu_fallback_args_exclude_rpath(self):
+        """CPU fallback args must NOT contain Metal-only RPATH flags at runtime."""
+        script = (
+            '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n'
+            + _GPU_BACKEND_FRAGMENT
+        )
+        output = run_bash(script)
+        fallback_line = next(
+            line
+            for line in output.splitlines()
+            if line.startswith("CPU_FALLBACK_CMAKE_ARGS=")
+        )
+        assert "-DGGML_METAL=OFF" in fallback_line
+        assert (
+            "@loader_path" not in fallback_line
+        ), "CPU fallback args should not contain RPATH flags"
+        assert (
+            "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" not in fallback_line
+        ), "CPU fallback args should not contain RPATH build flag"
+
+    def test_setup_sh_does_not_enable_metal_for_intel_macos(self):
+        """Intel macOS should stay on the existing non-Metal path in this patch."""
+        content = SETUP_SH.read_text()
+        assert 'if [ "$_IS_MACOS_ARM64" = true ]; then' in content
+        assert (
+            'Darwin" ] && { [ "$_HOST_MACHINE" = "arm64" ] || [ "$_HOST_MACHINE" = "aarch64" ]; }'
+            in content
+        )
+        assert (
+            "x86_64"
+            not in content[
+                content.find("-DGGML_METAL=ON") - 200 : content.find("-DGGML_METAL=ON")
+                + 200
+            ]
+        )
 
     def test_setup_ps1_uses_checkout_b(self):
         """PS1 should use checkout -B, not checkout --force FETCH_HEAD."""
@@ -637,7 +738,7 @@ class TestSourceCodePatterns:
     def test_setup_ps1_clone_uses_branch_tag(self):
         """PS1 clone should use --branch with the resolved tag."""
         content = SETUP_PS1.read_text()
-        assert "--branch" in content and "$ResolvedLlamaTag" in content
+        assert "--branch" in content and "$ResolvedSourceRef" in content
         # The old commented-out line should be gone
         assert "# git clone --depth 1 --branch" not in content
 
@@ -658,14 +759,52 @@ class TestSourceCodePatterns:
                         f"Found 'git pull' in llama.cpp build section at line {i+1}"
                     )
 
-    def test_setup_ps1_latest_resolution_queries_unsloth_first(self):
-        """PS1 should query Unsloth repo before ggml-org."""
+    def test_setup_ps1_latest_resolution_uses_helper_only(self):
+        """PS1 fallback should rely on helper output, not raw GitHub API tag_name."""
         content = SETUP_PS1.read_text()
-        idx_unsloth = content.find("$HelperReleaseRepo/releases/latest")
-        idx_ggml = content.find("ggml-org/llama.cpp/releases/latest")
-        assert idx_unsloth != -1, "Unsloth API query not found in PS1"
-        assert idx_ggml != -1, "ggml-org API query not found in PS1"
-        assert idx_unsloth < idx_ggml, "Unsloth should be queried before ggml-org"
+        assert "--resolve-install-tag" in content
+        assert "--resolve-llama-tag" in content
+        assert '--output-format", "json"' in content
+        assert "ConvertFrom-Json" in content
+        assert "$HelperReleaseRepo/releases/latest" not in content
+        assert "ggml-org/llama.cpp/releases/latest" not in content
+
+    def test_setup_ps1_source_build_uses_helper_resolution(self):
+        """PS1 source fallback should consult the helper for repo/ref planning."""
+        content = SETUP_PS1.read_text()
+        assert "--resolve-source-build" in content
+        assert '--output-format", "json"' in content
+        assert "$ResolvedSourceUrl" in content
+        assert "$ResolvedSourceRefKind" in content
+        assert "$ResolvedSourceRef" in content
+
+    def test_setup_ps1_prebuilt_install_disables_native_error_abort(self):
+        """PS1 prebuilt install should not abort setup on helper stderr."""
+        content = SETUP_PS1.read_text()
+        install_idx = content.index("& python @prebuiltArgs 2>&1")
+        block = content[max(0, install_idx - 800) : install_idx + 800]
+        assert "$PSNativeCommandUseErrorActionPreference = $false" in block
+        assert "$restoreNativeErrorPreference = $true" in block
+        assert (
+            "$PSNativeCommandUseErrorActionPreference = $previousNativeErrorPreference"
+            in block
+        )
+
+    def test_setup_ps1_helper_disables_error_action_abort(self):
+        """Helper resolution should suppress terminating NativeCommandError on PS 5.1."""
+        content = SETUP_PS1.read_text()
+        helper_idx = content.index("function Invoke-LlamaHelper")
+        block = content[helper_idx : helper_idx + 1200]
+        assert "$previousErrorActionPreference = $ErrorActionPreference" in block
+        assert '$ErrorActionPreference = "Continue"' in block
+        assert "$ErrorActionPreference = $previousErrorActionPreference" in block
+
+    def test_setup_ps1_uses_local_tempfile_helper(self):
+        """PS1 should not depend on New-TemporaryFile being available."""
+        content = SETUP_PS1.read_text()
+        assert "function New-UnslothTemporaryFile" in content
+        assert "$resolveErrorLog = New-UnslothTemporaryFile" in content
+        assert "$resolveErrorLog = New-TemporaryFile" not in content
 
     def test_binary_env_linux_has_binary_parent(self):
         """The Linux branch of binary_env should include binary_path.parent."""
@@ -685,3 +824,274 @@ class TestSourceCodePatterns:
                 found = True
                 break
         assert found, "binary_path.parent not found in Linux branch of binary_env"
+
+
+# =========================================================================
+# TEST GROUP F: macOS Metal build logic (bash subprocess tests)
+# =========================================================================
+
+# Minimal bash fragment that mirrors setup.sh's GPU backend decision chain.
+# Variables _IS_MACOS_ARM64, NVCC_PATH, GPU_BACKEND are injected by tests.
+_GPU_BACKEND_FRAGMENT = textwrap.dedent("""\
+    CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF"
+    _TRY_METAL_CPU_FALLBACK=false
+    CPU_FALLBACK_CMAKE_ARGS="$CMAKE_ARGS"
+
+    _BUILD_DESC="building"
+    if [ "$_IS_MACOS_ARM64" = true ]; then
+        _BUILD_DESC="building (Metal)"
+        CMAKE_ARGS="$CMAKE_ARGS -DGGML_METAL=ON -DGGML_METAL_EMBED_LIBRARY=ON -DGGML_METAL_USE_BF16=ON -DCMAKE_INSTALL_RPATH=@loader_path -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON"
+        CPU_FALLBACK_CMAKE_ARGS="$CPU_FALLBACK_CMAKE_ARGS -DGGML_METAL=OFF"
+        _TRY_METAL_CPU_FALLBACK=true
+    elif [ -n "$NVCC_PATH" ]; then
+        CMAKE_ARGS="$CMAKE_ARGS -DGGML_CUDA=ON"
+        _BUILD_DESC="building (CUDA)"
+    elif [ "$GPU_BACKEND" = "rocm" ]; then
+        CMAKE_ARGS="$CMAKE_ARGS -DGGML_HIP=ON"
+        _BUILD_DESC="building (ROCm)"
+    else
+        _BUILD_DESC="building (CPU)"
+    fi
+
+    echo "CMAKE_ARGS=$CMAKE_ARGS"
+    echo "CPU_FALLBACK_CMAKE_ARGS=$CPU_FALLBACK_CMAKE_ARGS"
+    echo "BUILD_DESC=$_BUILD_DESC"
+    echo "TRY_METAL_CPU_FALLBACK=$_TRY_METAL_CPU_FALLBACK"
+""")
+
+
+class TestMacOSMetalBuildLogic:
+    """Behavioral bash subprocess tests for the Metal GPU backend logic."""
+
+    def test_macos_arm64_cmake_args_contain_metal_flags(self):
+        """macOS arm64 should enable Metal, not CUDA."""
+        script = (
+            '_IS_MACOS_ARM64=true\nNVCC_PATH=""\nGPU_BACKEND=""\n'
+            + _GPU_BACKEND_FRAGMENT
+        )
+        output = run_bash(script)
+        assert "-DGGML_METAL=ON" in output
+        assert "-DGGML_CUDA=ON" not in output
+        assert "BUILD_DESC=building (Metal)" in output
+
+    def test_intel_macos_no_metal_flags(self):
+        """Intel macOS (not arm64) should not get Metal flags."""
+        script = (
+            '_IS_MACOS_ARM64=false\nNVCC_PATH=""\nGPU_BACKEND=""\n'
+            + _GPU_BACKEND_FRAGMENT
+        )
+        output = run_bash(script)
+        assert "-DGGML_METAL=ON" not in output
+        assert "BUILD_DESC=building (CPU)" in output
+
+    def test_macos_arm64_metal_precedes_nvcc(self):
+        """Even with nvcc in PATH, macOS arm64 should use Metal, not CUDA."""
+        script = (
+            '_IS_MACOS_ARM64=true\nNVCC_PATH="/usr/local/cuda/bin/nvcc"\n'
+            'GPU_BACKEND="cuda"\n' + _GPU_BACKEND_FRAGMENT
+        )
+        output = run_bash(script)
+        assert "-DGGML_METAL=ON" in output
+        assert "-DGGML_CUDA=ON" not in output
+        assert "BUILD_DESC=building (Metal)" in output
+
+    def test_metal_cpu_fallback_triggers_on_cmake_failure(self, tmp_path: Path):
+        """When cmake fails on Metal, the fallback should retry with -DGGML_METAL=OFF."""
+        mock_bin = tmp_path / "mock_bin"
+        mock_bin.mkdir()
+        calls_file = tmp_path / "cmake_calls.log"
+        # cmake that logs args and fails on first call (Metal), succeeds on second (CPU fallback)
+        cmake_script = mock_bin / "cmake"
+        cmake_script.write_text(
+            textwrap.dedent(f"""\
+            #!/bin/bash
+            echo "$*" >> "{calls_file}"
+            COUNTER_FILE="{tmp_path}/cmake_counter"
+            if [ ! -f "$COUNTER_FILE" ]; then
+                echo 1 > "$COUNTER_FILE"
+                exit 1
+            fi
+            exit 0
+        """)
+        )
+        cmake_script.chmod(0o755)
+
+        script = textwrap.dedent(f"""\
+            export PATH="{mock_bin}:$PATH"
+            _IS_MACOS_ARM64=true
+            NVCC_PATH=""
+            GPU_BACKEND=""
+            CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF"
+            _TRY_METAL_CPU_FALLBACK=false
+            CPU_FALLBACK_CMAKE_ARGS="$CMAKE_ARGS"
+
+            _BUILD_DESC="building"
+            if [ "$_IS_MACOS_ARM64" = true ]; then
+                _BUILD_DESC="building (Metal)"
+                CMAKE_ARGS="$CMAKE_ARGS -DGGML_METAL=ON -DGGML_METAL_EMBED_LIBRARY=ON -DGGML_METAL_USE_BF16=ON -DCMAKE_INSTALL_RPATH=@loader_path -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON"
+                CPU_FALLBACK_CMAKE_ARGS="$CPU_FALLBACK_CMAKE_ARGS -DGGML_METAL=OFF"
+                _TRY_METAL_CPU_FALLBACK=true
+            fi
+
+            BUILD_OK=true
+            _BUILD_TMP="{tmp_path}/build_tmp"
+            mkdir -p "$_BUILD_TMP"
+            if ! cmake -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CMAKE_ARGS; then
+                if [ "$_TRY_METAL_CPU_FALLBACK" = true ]; then
+                    _TRY_METAL_CPU_FALLBACK=false
+                    echo "FALLBACK_TRIGGERED"
+                    rm -rf "$_BUILD_TMP/build"
+                    cmake -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS || BUILD_OK=false
+                    if [ "$BUILD_OK" = true ]; then
+                        _BUILD_DESC="building (CPU fallback)"
+                    fi
+                else
+                    BUILD_OK=false
+                fi
+            fi
+
+            echo "BUILD_OK=$BUILD_OK"
+            echo "BUILD_DESC=$_BUILD_DESC"
+            echo "TRY_METAL_CPU_FALLBACK=$_TRY_METAL_CPU_FALLBACK"
+        """)
+        output = run_bash(script)
+        assert "FALLBACK_TRIGGERED" in output
+        assert "BUILD_OK=true" in output
+        assert "BUILD_DESC=building (CPU fallback)" in output
+        assert (
+            "TRY_METAL_CPU_FALLBACK=false" in output
+        ), "Fallback flag should be reset to false after configure fallback"
+
+        # Verify cmake args: first call has Metal ON, second has Metal OFF
+        calls = calls_file.read_text().splitlines()
+        assert len(calls) >= 2, f"Expected >= 2 cmake calls, got {len(calls)}"
+        assert (
+            "-DGGML_METAL=ON" in calls[0]
+        ), f"First cmake call should have Metal ON: {calls[0]}"
+        assert (
+            "-DGGML_METAL=OFF" in calls[1]
+        ), f"Second cmake call should have Metal OFF: {calls[1]}"
+        assert (
+            "-DGGML_METAL=ON" not in calls[1]
+        ), f"Second cmake call should NOT have Metal ON: {calls[1]}"
+        assert (
+            "@loader_path" not in calls[1]
+        ), f"CPU fallback should not have RPATH: {calls[1]}"
+        assert (
+            "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" not in calls[1]
+        ), f"CPU fallback should not have RPATH build flag: {calls[1]}"
+
+    def test_metal_build_failure_retries_cpu_fallback(self, tmp_path: Path):
+        """When cmake --build fails on Metal, the fallback should re-configure and rebuild with CPU."""
+        mock_bin = tmp_path / "mock_bin"
+        mock_bin.mkdir()
+        calls_file = tmp_path / "cmake_calls.log"
+        # cmake mock: configure always succeeds; first --build fails, rest succeed
+        cmake_script = mock_bin / "cmake"
+        cmake_script.write_text(
+            textwrap.dedent(f"""\
+            #!/bin/bash
+            echo "$*" >> "{calls_file}"
+            if [ "$1" = "--build" ]; then
+                BUILD_COUNTER_FILE="{tmp_path}/build_counter"
+                if [ ! -f "$BUILD_COUNTER_FILE" ]; then
+                    echo 1 > "$BUILD_COUNTER_FILE"
+                    exit 1
+                fi
+            fi
+            exit 0
+        """)
+        )
+        cmake_script.chmod(0o755)
+
+        script = textwrap.dedent(f"""\
+            export PATH="{mock_bin}:$PATH"
+            _IS_MACOS_ARM64=true
+            NVCC_PATH=""
+            GPU_BACKEND=""
+            CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF"
+            _TRY_METAL_CPU_FALLBACK=false
+            CPU_FALLBACK_CMAKE_ARGS="$CMAKE_ARGS"
+            CMAKE_GENERATOR_ARGS=""
+            NCPU=2
+
+            _BUILD_DESC="building"
+            if [ "$_IS_MACOS_ARM64" = true ]; then
+                _BUILD_DESC="building (Metal)"
+                CMAKE_ARGS="$CMAKE_ARGS -DGGML_METAL=ON -DGGML_METAL_EMBED_LIBRARY=ON -DGGML_METAL_USE_BF16=ON -DCMAKE_INSTALL_RPATH=@loader_path -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON"
+                CPU_FALLBACK_CMAKE_ARGS="$CPU_FALLBACK_CMAKE_ARGS -DGGML_METAL=OFF"
+                _TRY_METAL_CPU_FALLBACK=true
+            fi
+
+            BUILD_OK=true
+            _BUILD_TMP="{tmp_path}/build_tmp"
+            mkdir -p "$_BUILD_TMP"
+
+            # Configure (succeeds)
+            if ! cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CMAKE_ARGS; then
+                if [ "$_TRY_METAL_CPU_FALLBACK" = true ]; then
+                    _TRY_METAL_CPU_FALLBACK=false
+                    echo "CONFIGURE_FALLBACK"
+                    rm -rf "$_BUILD_TMP/build"
+                    cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS || BUILD_OK=false
+                    if [ "$BUILD_OK" = true ]; then
+                        _BUILD_DESC="building (CPU fallback)"
+                    fi
+                else
+                    BUILD_OK=false
+                fi
+            fi
+
+            # Build (first --build fails, triggers fallback)
+            if [ "$BUILD_OK" = true ]; then
+                if ! cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU"; then
+                    if [ "$_TRY_METAL_CPU_FALLBACK" = true ]; then
+                        _TRY_METAL_CPU_FALLBACK=false
+                        echo "BUILD_FALLBACK_TRIGGERED"
+                        rm -rf "$_BUILD_TMP/build"
+                        if cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS; then
+                            _BUILD_DESC="building (CPU fallback)"
+                            cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false
+                        else
+                            BUILD_OK=false
+                        fi
+                    else
+                        BUILD_OK=false
+                    fi
+                fi
+            fi
+
+            echo "BUILD_OK=$BUILD_OK"
+            echo "BUILD_DESC=$_BUILD_DESC"
+            echo "TRY_METAL_CPU_FALLBACK=$_TRY_METAL_CPU_FALLBACK"
+        """)
+        output = run_bash(script)
+        assert "CONFIGURE_FALLBACK" not in output, "Configure should have succeeded"
+        assert "BUILD_FALLBACK_TRIGGERED" in output
+        assert "BUILD_OK=true" in output
+        assert "BUILD_DESC=building (CPU fallback)" in output
+        assert (
+            "TRY_METAL_CPU_FALLBACK=false" in output
+        ), "Fallback flag should be reset to false after build fallback"
+
+        # Verify: configure with Metal ON, build fails, re-configure with Metal OFF, rebuild
+        calls = calls_file.read_text().splitlines()
+        assert len(calls) >= 4, f"Expected >= 4 cmake calls, got {len(calls)}: {calls}"
+        # First call: configure with Metal ON
+        assert "-DGGML_METAL=ON" in calls[0]
+        # Second call: build (fails)
+        assert "--build" in calls[1]
+        # Third call: re-configure with Metal OFF and no RPATH flags
+        assert "-DGGML_METAL=OFF" in calls[2]
+        assert "-DGGML_METAL=ON" not in calls[2]
+        assert (
+            "@loader_path" not in calls[2]
+        ), f"CPU fallback should not have RPATH: {calls[2]}"
+        assert (
+            "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" not in calls[2]
+        ), f"CPU fallback should not have RPATH build flag: {calls[2]}"
+        assert (
+            "-DLLAMA_BUILD_TESTS=OFF" in calls[2]
+        ), f"CPU fallback should preserve baseline flags: {calls[2]}"
+        # Fourth call: rebuild (succeeds)
+        assert "--build" in calls[3]

unsloth/chat_templates.py

@@ -863,6 +863,114 @@ DEFAULT_SYSTEM_MESSAGE["gemma-3n"] = None # No system message in Gemma-3n
 CHAT_TEMPLATES["gemma3n"] = (gemma3n_template, gemma3n_template_eos_token, False, gemma3n_ollama,)
 DEFAULT_SYSTEM_MESSAGE["gemma3n"] = None # No system message in Gemma-3n
 
+# =========================================== Gemma-4
+# Gemma-4 uses <|turn>role\n...<turn|>\n format
+gemma4_template = \
+"""{%- if messages[0]['role'] == 'system' -%}
+    {%- set first_user_prefix = messages[0]['content'] + '\n\n' -%}
+    {%- set loop_messages = messages[1:] -%}
+{%- else -%}
+    {%- set first_user_prefix = "" -%}
+    {%- set loop_messages = messages -%}
+{%- endif -%}
+{%- for message in loop_messages -%}
+    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}
+        {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }}
+    {%- endif -%}
+    {%- if (message['role'] == 'assistant') -%}
+        {%- set role = "model" -%}
+    {%- else -%}
+        {%- set role = message['role'] -%}
+    {%- endif -%}
+    {{ '<|turn>' + role + '\n' + (first_user_prefix if loop.first else "") }}
+    {%- if message['content'] is string -%}
+        {{ message['content'] | trim }}
+    {%- elif message['content'] is iterable -%}
+        {%- for item in message['content'] -%}
+            {%- if item['type'] == 'audio' -%}
+                {{ '<|audio|>' }}
+            {%- elif item['type'] == 'image' -%}
+                {{ '<|image|>' }}
+            {%- elif item['type'] == 'video' -%}
+                {{ '<|video|>' }}
+            {%- elif item['type'] == 'text' -%}
+                {{ item['text'] | trim }}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- else -%}
+        {{ raise_exception("Invalid content type") }}
+    {%- endif -%}
+    {{ '<turn|>\n' }}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{'<|turn>model\n'}}
+{%- endif -%}
+"""
+
+try:
+    gemma4_ollama = _ollama_template("gemma-4")
+except KeyError:
+    gemma4_ollama = ""
+gemma4_template_eos_token = "<turn|>"
+CHAT_TEMPLATES["gemma-4"] = (gemma4_template, gemma4_template_eos_token, False, gemma4_ollama,)
+DEFAULT_SYSTEM_MESSAGE["gemma-4"] = None
+
+CHAT_TEMPLATES["gemma4"] = (gemma4_template, gemma4_template_eos_token, False, gemma4_ollama,)
+DEFAULT_SYSTEM_MESSAGE["gemma4"] = None
+
+# Gemma-4 with empty thought channel (required for larger models like 31B, 26B-A4B)
+# Injects <|channel>thought\n<channel|> at the start of each model response during training
+gemma4_thinking_template = \
+"""{%- if messages[0]['role'] == 'system' -%}
+    {%- set first_user_prefix = messages[0]['content'] + '\n\n' -%}
+    {%- set loop_messages = messages[1:] -%}
+{%- else -%}
+    {%- set first_user_prefix = "" -%}
+    {%- set loop_messages = messages -%}
+{%- endif -%}
+{%- for message in loop_messages -%}
+    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) -%}
+        {{ raise_exception("Conversation roles must alternate user/assistant/user/assistant/...") }}
+    {%- endif -%}
+    {%- if (message['role'] == 'assistant') -%}
+        {%- set role = "model" -%}
+    {%- else -%}
+        {%- set role = message['role'] -%}
+    {%- endif -%}
+    {{ '<|turn>' + role + '\n' + (first_user_prefix if loop.first else "") }}
+    {%- if role == "model" -%}
+        {{ '<|channel>thought\n<channel|>' }}
+    {%- endif -%}
+    {%- if message['content'] is string -%}
+        {{ message['content'] | trim }}
+    {%- elif message['content'] is iterable -%}
+        {%- for item in message['content'] -%}
+            {%- if item['type'] == 'audio' -%}
+                {{ '<|audio|>' }}
+            {%- elif item['type'] == 'image' -%}
+                {{ '<|image|>' }}
+            {%- elif item['type'] == 'video' -%}
+                {{ '<|video|>' }}
+            {%- elif item['type'] == 'text' -%}
+                {{ item['text'] | trim }}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- else -%}
+        {{ raise_exception("Invalid content type") }}
+    {%- endif -%}
+    {{ '<turn|>\n' }}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{'<|turn>model\n'}}
+{%- endif -%}
+"""
+
+CHAT_TEMPLATES["gemma-4-thinking"] = (gemma4_thinking_template, gemma4_template_eos_token, False, gemma4_ollama,)
+DEFAULT_SYSTEM_MESSAGE["gemma-4-thinking"] = None
+
+CHAT_TEMPLATES["gemma4-thinking"] = (gemma4_thinking_template, gemma4_template_eos_token, False, gemma4_ollama,)
+DEFAULT_SYSTEM_MESSAGE["gemma4-thinking"] = None
+
 # =========================================== GPT-OSS
 # Obtained via
 # print(tokenizer.chat_template.replace("}\n", "####").replace("\n", "\\n").replace("####", "}\n"))

unsloth/models/_utils.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "2026.3.18"
+__version__ = "2026.4.1"
 
 __all__ = [
     "SUPPORTS_BFLOAT16",
@@ -64,7 +64,8 @@ __all__ = [
     "patch_compiled_autograd",
     "process_vision_info",
     "unsloth_compile_transformers",
-    "prefer_flex_attn_if_supported",
+    "determine_attention_implementation",
+    "_set_attn_impl",
     "patch_fast_lora",
     "validate_loftq_config",
     "RaiseUninitialized",
@@ -222,44 +223,76 @@ def apply_unsloth_gradient_checkpointing(
     return use_gradient_checkpointing
 
 
-def prefer_flex_attn_if_supported(model_class, config):
-    if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") == "0":
-        return None
-    try:
-        from transformers.utils.import_utils import is_torch_flex_attn_available
+# Models that don't work with flex_attention:
+# GPT-OSS: left padding issues cause incorrect outputs.
+# Mllama: BlockMask Q_LEN!=KV_LEN ValueError on decode.
+# NemotronH: hybrid Mamba-2 + Transformer, raises NotImplementedError.
+# Gemma3N: timm vision wrappers don't support flex_attention.
+# ModernBERT: create_block_mask with _compile=True hits CUDA illegal memory
+# access on some GPU architectures (B200). Falls back to eager safely.
+_FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h", "modernbert")
+_EAGER_ONLY_PREFIXES = ("gemma3n",)
 
-        if not is_torch_flex_attn_available():
-            return None
-        if model_class is None or not getattr(
-            model_class, "_supports_flex_attn", False
-        ):
-            return None
 
-        attention_dropout = getattr(config, "attention_dropout", 0) or 0
-        if attention_dropout > 0:
-            return None
-        # GPT-OSS, Mllama and Gemma3N use eager/sdpa attention during
-        # inference since flex attention returns incorrect results or errors out.
-        # GPT-OSS: left padding issues cause incorrect outputs.
-        # Mllama: _update_causal_mask uses make_flex_block_causal_mask which
-        # creates BlockMask with Q_LEN=KV_LEN=total_seq_len, but during
-        # decode q_len=1, causing ValueError. Needs transformers update.
-        # Gemma3N: timm vision wrappers (eg Gemma3nVisionConfig) do not
-        # support flex_attention.
-        # NemotronH: hybrid Mamba-2 + Transformer model that does not
-        # support flex_attention (raises NotImplementedError from transformers).
-        model_type = getattr(config, "model_type", "") if config else ""
-        if model_type in ("gpt_oss", "mllama", "nemotron_h") or str(
-            model_type
-        ).startswith("gemma3n"):
-            return None
-        if config is not None:
-            setattr(config, "_attn_implementation", "flex_attention")
-            if hasattr(config, "attn_implementation"):
-                setattr(config, "attn_implementation", "flex_attention")
-        return "flex_attention"
-    except Exception:
-        return None
+def _is_flex_excluded(model_type):
+    return model_type in _FLEX_EXCLUDED_MODELS
+
+
+def _is_eager_only(model_type):
+    return any(model_type.startswith(p) for p in _EAGER_ONLY_PREFIXES)
+
+
+def _set_attn_impl(config, impl):
+    """Helper function to set attention implementation on config and return it."""
+    if config is not None:
+        setattr(config, "_attn_implementation", impl)
+        if hasattr(config, "attn_implementation"):
+            setattr(config, "attn_implementation", impl)
+    return impl
+
+
+def determine_attention_implementation(model_class, config):
+    model_type = getattr(config, "model_type", "").lower()
+
+    # Eager-only models (e.g. gemma3n timm vision towers)
+    if _is_eager_only(model_type):
+        _set_attn_impl(config, "eager")
+        return "eager"
+
+    # Flash Attention 2
+    if HAS_FLASH_ATTENTION and model_class is not None:
+        supports_fa2 = getattr(model_class, "_supports_flash_attn_2", False) or getattr(
+            model_class, "_supports_flash_attn", False
+        )
+        if supports_fa2:
+            _set_attn_impl(config, "flash_attention_2")
+            return "flash_attention_2"
+
+    # Flex Attention
+    if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "1") != "0":
+        try:
+            from transformers.utils.import_utils import is_torch_flex_attn_available
+
+            if (
+                is_torch_flex_attn_available()
+                and model_class is not None
+                and getattr(model_class, "_supports_flex_attn", False)
+                and not _is_flex_excluded(model_type)
+            ):
+                attention_dropout = getattr(config, "attention_dropout", 0) or 0
+                if attention_dropout == 0:
+                    _set_attn_impl(config, "flex_attention")
+                    return "flex_attention"
+        except Exception:
+            pass
+
+    # SDPA
+    if model_class is not None and getattr(model_class, "_supports_sdpa", False):
+        _set_attn_impl(config, "sdpa")
+        return "sdpa"
+
+    _set_attn_impl(config, "eager")
+    return "eager"
 
 
 def _run_temporary_patches(phase):
@@ -504,6 +537,15 @@ try:
 except:
     pass
 
+# Gemma4 It is strongly recommended to train Gemma4 models with the `eager`
+try:
+    from transformers.models.gemma4.modeling_gemma4 import logger as gemma4_logger
+
+    gemma4_logger.addFilter(HideLoggingMessage("strongly recommended"))
+    del gemma4_logger
+except:
+    pass
+
 # Xet Storage is enabled for this repo, but the 'hf_xet' package is not installed.
 try:
     from huggingface_hub.file_download import logger as hub_logger
@@ -765,7 +807,16 @@ model_architectures = [
     "falcon_h1",
 ]
 
+# Transformers 5.x uses class-level annotations with @strict, @auto_docstring,
+# and interval() in config classes. exec(inspect.getsource(...)) fails because
+# those symbols are not in scope. Skip the exec-based config patching for 5.x
+# since those configs already use rope_parameters (the v5 replacement for
+# rope_scaling).
+_skip_config_exec_patch = Version(transformers_version) >= Version("5.0.0")
+
 for model_name in model_architectures:
+    if _skip_config_exec_patch:
+        break
     config_filepath = f"transformers.models.{model_name}.configuration_{model_name}"
     model_filepath = f"transformers.models.{model_name}.modeling_{model_name}"
     config_filename = f"{model_name.title().replace('_','')}Config"  # qwen3 arch folder is qwen3_moe but config is Qwen3Config. Need to remove underscore(_) for now
@@ -799,9 +850,12 @@ for model_name in model_architectures:
         if Version(transformers_version) <= Version("4.42.4"):
             config = patch_mistral_nemo_config(config)
 
-    exec(config, globals())
-    exec(f"import {config_filepath}", globals())
-    exec(f"{config_filepath}.{config_filename} = {config_filename}", globals())
+    try:
+        exec(config, globals())
+        exec(f"import {config_filepath}", globals())
+        exec(f"{config_filepath}.{config_filename} = {config_filename}", globals())
+    except Exception:
+        continue
 # =============================================
 
 # =============================================
@@ -1885,6 +1939,18 @@ def _unsloth_pre_compute_loss(self, model, inputs, *args, **kwargs):
             _has_ccm = _mod is not None and hasattr(_mod, "create_causal_mask_mapping")
             if _has_ccm and _inner.training:
                 inputs["token_type_ids"] = torch.zeros_like(inputs["input_ids"])
+    # Gemma4 uses mm_token_type_ids (not token_type_ids) for VLM masking
+    if "mm_token_type_ids" not in inputs and "input_ids" in inputs:
+        _inner = model
+        for _attr in ("base_model", "model", "model"):
+            _inner = getattr(_inner, _attr, _inner)
+        if getattr(getattr(_inner, "config", None), "model_type", "") in ("gemma4",):
+            import sys as _sys
+
+            _mod = _sys.modules.get(type(_inner).__module__)
+            _has_ccm = _mod is not None and hasattr(_mod, "create_causal_mask_mapping")
+            if _has_ccm and _inner.training:
+                inputs["mm_token_type_ids"] = torch.zeros_like(inputs["input_ids"])
 
     outputs = self._old_compute_loss(model, inputs, *args, **kwargs)
     return outputs

unsloth/models/cohere.py

@@ -357,6 +357,9 @@ def CohereAttention_fast_forward_inference(
     # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len)
     # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)
     cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index)
+    # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last
+    if position_ids.dim() >= 2 and position_ids.shape[-1] > 1:
+        position_ids = position_ids[:, -1:]
     cos = cos[position_ids].unsqueeze(1)
     sin = sin[position_ids].unsqueeze(1)
     h = self.half_head_dim

unsloth/models/falcon_h1.py

@@ -313,6 +313,9 @@ def FalconH1Attention_fast_forward_inference(
     # or else error
     self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2)
     cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index)
+    # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last
+    if position_ids.dim() >= 2 and position_ids.shape[-1] > 1:
+        position_ids = position_ids[:, -1:]
     cos = cos[position_ids].unsqueeze(1)
     sin = sin[position_ids].unsqueeze(1)
     h = self.half_head_dim

unsloth/models/gemma2.py

@@ -394,6 +394,9 @@ def Gemma2Attention_fast_forward_inference(
     # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len)
     # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)
     cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index)
+    # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last
+    if position_ids.dim() >= 2 and position_ids.shape[-1] > 1:
+        position_ids = position_ids[:, -1:]
     cos = cos[position_ids].unsqueeze(1)
     sin = sin[position_ids].unsqueeze(1)
     h = self.half_head_dim

unsloth/models/granite.py

@@ -355,6 +355,9 @@ def GraniteAttention_fast_forward_inference(
     # cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len)
     # Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)
     cos, sin = position_embeddings
+    # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last
+    if position_ids.dim() >= 2 and position_ids.shape[-1] > 1:
+        position_ids = position_ids[:, -1:]
     cos, sin = cos[position_ids], sin[position_ids]
     h = self.half_head_dim
 

unsloth/models/llama.py

@@ -496,6 +496,10 @@ def LlamaAttention_fast_forward_inference(
     # ensure correct shape
     if position_ids.dim() == 1:
         position_ids = position_ids[:, None]
+    # Transformers 5.x generate() accumulates position_ids as [batch, full_seq_len]
+    # across decode steps. In single-token inference we only need the last position.
+    if position_ids.shape[-1] > 1:
+        position_ids = position_ids[:, -1:]
     position_ids = position_ids.to(Qn.device)
 
     if rotary_seq_len is None:
@@ -2341,8 +2345,8 @@ class FastLlamaModel:
         model_function = MODEL_FOR_CAUSAL_LM_MAPPING[model_config.__class__]
         IS_FALCON_H1 = model_config.model_type.startswith("falcon_h1")
 
-        preferred_attn_impl = (
-            prefer_flex_attn_if_supported(model_function, model_config) or "eager"
+        preferred_attn_impl = determine_attention_implementation(
+            model_function, model_config
         )
 
         has_rope_scaling = False
@@ -2414,14 +2418,24 @@ class FastLlamaModel:
 
         raise_handler = RaiseUninitialized()
         if num_labels is not None:
+            # Transformers 5.x @strict config classes reject unexpected kwargs
+            # like num_labels and max_position_embeddings. Set on the config
+            # object directly and pass config= instead.
+            model_config.num_labels = num_labels
+            if max_position_embeddings is not None:
+                model_config.max_position_embeddings = max_position_embeddings
+            # Pop config-level attrs that would be rejected by @strict model init
+            for _cfg_key in ("id2label", "label2id", "rope_scaling"):
+                _cfg_val = kwargs.pop(_cfg_key, None)
+                if _cfg_val is not None:
+                    setattr(model_config, _cfg_key, _cfg_val)
             model = AutoModelForSequenceClassification.from_pretrained(
                 model_name,
+                config = model_config,
                 device_map = device_map,
                 # torch_dtype             = dtype, # transformers changed torch_dtype to dtype
-                num_labels = num_labels,
                 # quantization_config     = bnb_config,
                 token = token,
-                max_position_embeddings = max_position_embeddings,
                 trust_remote_code = trust_remote_code,
                 attn_implementation = preferred_attn_impl,
                 **kwargs,

unsloth/models/loader.py

@@ -78,6 +78,7 @@ SUPPORTS_QWEN3_MOE = transformers_version >= Version("4.50.3")
 SUPPORTS_FALCON_H1 = transformers_version >= Version("4.53.0")
 SUPPORTS_GEMMA3N = transformers_version >= Version("4.53.0")
 SUPPORTS_GPTOSS = transformers_version >= Version("4.55.0")
+SUPPORTS_GEMMA4 = transformers_version >= Version("5.5.0")
 # Transformers v5 meta-device loading corrupts non-persistent buffers (inv_freq).
 # See _fix_rope_inv_freq() below for details.
 _NEEDS_ROPE_FIX = transformers_version >= Version("5.0.0")
@@ -107,6 +108,8 @@ FORCE_FLOAT32 = [
     "gemma3n",
     "gpt_oss",
     "qwen3_5",  # Qwen3.5 GDN layers produce NaN grad norms in float16 training
+    "gemma4,",  # Add comma bc gemma4 will match gemma4_text
+    "gemma4_text",
 ]
 
 global DISABLE_COMPILE_MODEL_NAMES
@@ -1130,6 +1133,17 @@ class FastModel(FastBaseModel):
             raise RuntimeError(
                 "Unsloth: Qwen 2.5 only works on transformers >= 4.49.0." + LATEST
             )
+        # Gemma 4 must be before Gemma 3N and Gemma 3
+        elif "gemma4" in model_types_all:
+            if not SUPPORTS_GEMMA4:
+                raise RuntimeError(
+                    "Unsloth: Gemma 4 requires transformers >= 5.5.0" + LATEST
+                )
+            os.environ["UNSLOTH_DISABLE_STATIC_GENERATION"] = "1"
+            os.environ["UNSLOTH_HIGH_PRECISION_LAYERNORM"] = "1"
+            # Disable flex_attention for Gemma-4: flex compile overhead is 2.7x slower
+            # than SDPA. Our attention patch ensures Q/K/V dtype alignment for SDPA.
+            os.environ["UNSLOTH_ENABLE_FLEX_ATTENTION"] = "0"
         # Gemma 3N must be before Gemma 3
         elif "gemma3n" in model_types_all:
             if transformers_version < Version("4.53.0"):
@@ -1407,8 +1421,14 @@ class FastModel(FastBaseModel):
             architectures = []
         is_vlm = any(x.endswith("ForConditionalGeneration") for x in architectures)
         is_vlm = is_vlm or hasattr(model_config, "vision_config")
+        # If num_labels is set, use AutoModelForSequenceClassification
+        _num_labels = kwargs.get("num_labels", None)
         if auto_model is None:
-            if is_vlm:
+            if _num_labels is not None:
+                from transformers import AutoModelForSequenceClassification
+
+                auto_model = AutoModelForSequenceClassification
+            elif is_vlm:
                 # Check if the model's auto_map supports the VLM auto class.
                 # Some VL models (e.g. Nemotron-VL) only register AutoModelForCausalLM
                 # in their auto_map, not AutoModelForImageTextToText/AutoModelForVision2Seq.

unsloth/models/qwen3.py

@@ -302,6 +302,9 @@ def Qwen3Attention_fast_forward_inference(
     # or else error
     self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2)
     cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index)
+    # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last
+    if position_ids.dim() >= 2 and position_ids.shape[-1] > 1:
+        position_ids = position_ids[:, -1:]
     cos = cos[position_ids].unsqueeze(1)
     sin = sin[position_ids].unsqueeze(1)
     h = self.half_head_dim

unsloth/models/rl_replacements.py

@@ -542,6 +542,37 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
 
         function = patched
 
+    # Transformers 5.x: Extend mm_token_type_ids for completion tokens (Qwen3VL M-RoPE).
+    # TRL handles token_type_ids but not mm_token_type_ids.
+    _tt_search = (
+        'if "token_type_ids" in forward_kwargs:\n'
+        '            token_type_ids = forward_kwargs["token_type_ids"]\n'
+        '            forward_kwargs["token_type_ids"] = torch.cat(\n'
+        "                [token_type_ids, token_type_ids.new_zeros(completion_ids.shape)], dim=1\n"
+        "            )"
+    )
+    _tt_replace = (
+        _tt_search + "\n"
+        '        if "mm_token_type_ids" in forward_kwargs:\n'
+        '            mm_tti = forward_kwargs["mm_token_type_ids"]\n'
+        '            forward_kwargs["mm_token_type_ids"] = torch.cat(\n'
+        "                [mm_tti, mm_tti.new_zeros(completion_ids.shape)], dim=1\n"
+        "            )"
+    )
+    function = function.replace(_tt_search, _tt_replace)
+
+    # Save mm_token_type_ids to output dict alongside token_type_ids
+    _save_search = (
+        'if "token_type_ids" in forward_kwargs:\n'
+        '            output["token_type_ids"] = forward_kwargs["token_type_ids"]'
+    )
+    _save_replace = (
+        _save_search + "\n"
+        '        if "mm_token_type_ids" in forward_kwargs:\n'
+        '            output["mm_token_type_ids"] = forward_kwargs["mm_token_type_ids"]'
+    )
+    function = function.replace(_save_search, _save_replace)
+
     return function
 
 
@@ -714,6 +745,9 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
                 kwargs.get("pixel_attention_mask", None),
                 kwargs.get("image_sizes", None),
             )
+            # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models
+            token_type_ids = kwargs.get("token_type_ids", None)
+            mm_token_type_ids = kwargs.get("mm_token_type_ids", None)
 
             unwrapped_model = self.accelerator.unwrap_model(
                 model, keep_fp32_wrapper = False
@@ -831,6 +865,10 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
             if logit_scale_divide is None:
                 logit_scale_divide = 0
 
+            # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models
+            token_type_ids_chunks = chunk_optional(token_type_ids, B)
+            mm_token_type_ids_chunks = chunk_optional(mm_token_type_ids, B)
+
             zipped_inputs = zip(
                 input_ids_chunks,
                 attention_mask_chunks,
@@ -838,6 +876,8 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
                 image_grid_thw_chunks,
                 pixel_attention_mask_chunks,
                 image_sizes_chunks,
+                token_type_ids_chunks,
+                mm_token_type_ids_chunks,
             )
             os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
 
@@ -849,7 +889,16 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
                     image_grid_thw_chunk,
                     pixel_attention_mask_chunk,
                     image_sizes_chunk,
+                    token_type_ids_chunk,
+                    mm_token_type_ids_chunk,
                 ) in zipped_inputs:
+                    _extra_vision_kwargs = {}
+                    if token_type_ids_chunk is not None:
+                        _extra_vision_kwargs["token_type_ids"] = token_type_ids_chunk
+                    if mm_token_type_ids_chunk is not None:
+                        _extra_vision_kwargs["mm_token_type_ids"] = (
+                            mm_token_type_ids_chunk
+                        )
                     with torch.amp.autocast(
                         device_type = "cuda", dtype = self._autocast_dtype
                     ):
@@ -861,6 +910,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
                                 image_grid_thw = image_grid_thw_chunk,
                                 pixel_attention_mask = pixel_attention_mask_chunk,
                                 image_sizes = image_sizes_chunk,
+                                **_extra_vision_kwargs,
                             ).logits
 
                             completion_input_ids_chunk = input_ids_chunk[
@@ -893,6 +943,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
                                 pixel_attention_mask = pixel_attention_mask_chunk,
                                 image_sizes = image_sizes_chunk,
                                 logits_to_keep = logits_to_keep + 1,
+                                **_extra_vision_kwargs,
                             ).logits
 
                             logits_chunk = logits_chunk[:, :-1, :]
@@ -993,6 +1044,9 @@ def grpo_trainer_compute_loss(function_name, function):
             inputs.get("pixel_attention_mask", None),
             inputs.get("image_sizes", None),
         )
+        # Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models
+        token_type_ids = inputs.get("token_type_ids", None)
+        mm_token_type_ids = inputs.get("mm_token_type_ids", None)
         num_items_in_batch = inputs.get("num_items_in_batch", None)
         sampling_per_token_logps = inputs.get("sampling_per_token_logps", None)
         current_gradient_accumulation_steps = self.current_gradient_accumulation_steps
@@ -1136,6 +1190,8 @@ def grpo_trainer_compute_loss(function_name, function):
                     current_gradient_accumulation_steps = current_gradient_accumulation_steps,
                     num_processes = num_processes,
                     sampling_per_token_logps = sampling_per_token_logps,
+                    token_type_ids = token_type_ids,
+                    mm_token_type_ids = mm_token_type_ids,
                 )
             else:
                 # to ensure backwards compatibility with trl 0.15.2 and maybe even 0.17
@@ -1154,6 +1210,8 @@ def grpo_trainer_compute_loss(function_name, function):
                         logit_scale_multiply = logit_scale_multiply,
                         logit_scale_divide = logit_scale_divide,
                         attention_mask = attention_mask,
+                        token_type_ids = token_type_ids,
+                        mm_token_type_ids = mm_token_type_ids,
                     )
                 )
         if "train" in self._metrics:

unsloth/models/vision.py

@@ -216,6 +216,10 @@ def unsloth_base_fast_generate(
         kwargs["pixel_values"] = kwargs["pixel_values"].to(dtype)
     except:
         pass
+    try:
+        kwargs["pixel_values_videos"] = kwargs["pixel_values_videos"].to(dtype)
+    except:
+        pass
 
     # Mixed precision autocast
     if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
@@ -597,8 +601,6 @@ class FastBaseModel:
                 custom_datatype = None
                 correct_dtype = None
 
-        # Stop SDPA for some archs like Pixtral / Mistral3
-        flex_attn_impl = None
         if auto_config is None:
             auto_config = AutoConfig.from_pretrained(
                 model_name,
@@ -609,7 +611,14 @@ class FastBaseModel:
             model_class = auto_model._model_mapping[auto_config.__class__]
         except Exception:
             model_class = None
-        flex_attn_impl = prefer_flex_attn_if_supported(model_class, auto_config)
+        if model_class is None:
+            # When model_class cannot be resolved (remote-code or unmapped
+            # configs), preserve the old fallback of sdpa when supported.
+            attn_impl = _set_attn_impl(
+                auto_config, "sdpa" if supports_sdpa else "eager"
+            )
+        else:
+            attn_impl = determine_attention_implementation(model_class, auto_config)
 
         # Handle FP8 models: get_model_name has already redirected this to BF16 sibling if the model ships with
         # FP8 weights. We just need to update it here for sanity.
@@ -620,21 +629,15 @@ class FastBaseModel:
         except Exception:
             model_class = None
 
-        model_type = str(getattr(auto_config, "model_type", "")).lower()
-        if model_type.startswith("gemma3n"):
-            # Gemma3N variants initialize timm-based vision towers which do
-            # not support flex_attention, so default to eager unless overridden.
-            default_attn_impl = "eager"
-        else:
-            default_attn_impl = "flex_attention" if flex_attn_impl else "sdpa"
         if not ("attn_implementation" in kwargs):
-            kwargs["attn_implementation"] = default_attn_impl
+            kwargs["attn_implementation"] = attn_impl
         if not supports_sdpa and kwargs.get("attn_implementation") == "sdpa":
-            if os.environ.get("UNSLOTH_ENABLE_FLEX_ATTENTION", "0") == "0":
-                print(
-                    f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager."
-                )
+            print(
+                f"Unsloth: {model_type_arch.title()} does not support SDPA - switching to fast eager."
+            )
             del kwargs["attn_implementation"]
+            # Re-stamp config so it stays consistent with the actual impl
+            _set_attn_impl(auto_config, "eager")
 
         bnb_config = None
         user_quantization_config = kwargs.get("quantization_config", None)
@@ -788,6 +791,15 @@ class FastBaseModel:
         if not fast_inference:
             # Prevent load_in_fp8 from being forwarded into HF internal model loading
             load_in_fp8 = kwargs.pop("load_in_fp8", None)
+            # Transformers 5.x @strict config classes reject unexpected kwargs.
+            # Move config-level attributes onto the config object directly.
+            _num_labels = kwargs.pop("num_labels", None)
+            if _num_labels is not None:
+                model_config.num_labels = _num_labels
+            for _cfg_key in ("id2label", "label2id", "max_position_embeddings"):
+                _cfg_val = kwargs.pop(_cfg_key, None)
+                if _cfg_val is not None:
+                    setattr(model_config, _cfg_key, _cfg_val)
             model = auto_model.from_pretrained(
                 model_name,
                 config = model_config,
@@ -1021,6 +1033,15 @@ class FastBaseModel:
                     f"Unsloth: Warning - VLM processor fallback returned None for model_type={model_type_arch}",
                     file = sys.stderr,
                 )
+        # Backwards compat: if processor has no chat_template (e.g. old saves without
+        # chat_template.jinja) but the inner tokenizer does, copy it to the processor.
+        if (
+            hasattr(tokenizer, "tokenizer")
+            and getattr(tokenizer, "chat_template", None) is None
+            and getattr(tokenizer.tokenizer, "chat_template", None) is not None
+        ):
+            tokenizer.chat_template = tokenizer.tokenizer.chat_template
+
         if hasattr(tokenizer, "tokenizer"):
             __tokenizer = tokenizer.tokenizer
             # Add padding side as well
@@ -1277,7 +1298,59 @@ class FastBaseModel:
             model,
             use_gradient_checkpointing = use_gradient_checkpointing,
         )
+        # Gemma4 ClippableLinear wraps nn.Linear -- PEFT can't inject LoRA on it directly.
+        # Monkey-patch PEFT to target the inner .linear child instead.
+        _clippable_linear_cls = None
+        try:
+            from transformers.models.gemma4.modeling_gemma4 import (
+                Gemma4ClippableLinear as _clippable_linear_cls,
+            )
+        except ImportError:
+            pass
+        if _clippable_linear_cls is not None:
+            from peft.tuners.lora.model import LoraModel as _LoraModel
+
+            _original_car = _LoraModel._create_and_replace
+
+            def _patched_car(
+                self,
+                peft_config,
+                adapter_name,
+                target,
+                target_name,
+                parent,
+                current_key = None,
+                **kwargs,
+            ):
+                if isinstance(target, _clippable_linear_cls):
+                    return _original_car(
+                        self,
+                        peft_config,
+                        adapter_name,
+                        target.linear,
+                        "linear",
+                        target,
+                        current_key = current_key,
+                        **kwargs,
+                    )
+                return _original_car(
+                    self,
+                    peft_config,
+                    adapter_name,
+                    target,
+                    target_name,
+                    parent,
+                    current_key = current_key,
+                    **kwargs,
+                )
+
+            _LoraModel._create_and_replace = _patched_car
+
         model = _get_peft_model(model, lora_config)
+
+        # Restore original PEFT method
+        if _clippable_linear_cls is not None:
+            _LoraModel._create_and_replace = _original_car
         # Apply QAT + LoRA if specified
         if qat_scheme is not None:
             print("Unsloth: Applying QAT to mitigate quantization degradation")
@@ -1375,7 +1448,7 @@ class FastBaseModel:
         # after this point, so we intercept gradient_checkpointing_enable
         # to always force use_reentrant=True for Gemma3N.
         _model_type = getattr(getattr(model, "config", None), "model_type", "") or ""
-        if "gemma3n" in _model_type.lower():
+        if "gemma3n" in _model_type.lower() or "gemma4" in _model_type.lower():
             _original_gc_enable = model.gradient_checkpointing_enable
 
             def _gc_enable_reentrant(**kwargs):

unsloth/ollama_template_mappers.py

@@ -1199,6 +1199,21 @@ TEMPLATE """{{- range $i, $_ := .Messages }}
 OLLAMA_TEMPLATES["gemma-3n"] = gemma3n_ollama
 OLLAMA_TEMPLATES["gemma3n"] = gemma3n_ollama
 
+# =========================================== Gemma-4
+gemma4_ollama = '''
+FROM {__FILE_LOCATION__}
+TEMPLATE """{{- range $i, $_ := .Messages }}
+{{- $last := eq (len (slice $.Messages $i)) 1 }}
+<|turn>{{ .Role }}
+{{ .Content }}{{ if not $last }}<turn|>
+{{ end }}
+{{- end }}<turn|>
+<|turn>model
+"""
+'''
+OLLAMA_TEMPLATES["gemma-4"] = gemma4_ollama
+OLLAMA_TEMPLATES["gemma4"] = gemma4_ollama
+
 # =========================================== GPT-OSS
 
 # Ollama from https://ollama.com/library/gpt-oss:latest/blobs/fa6710a93d78
@@ -1961,6 +1976,16 @@ OLLAMA_TEMPLATE_TO_MODEL_MAPPER = {
         "google/medgemma-27b-text-it",
         "unsloth/medgemma-27b-text-it-bnb-4bit",
     ),
+    "gemma4": (
+        "unsloth/gemma-4-E2B-it",
+        "unsloth/gemma-4-E2B",
+        "unsloth/gemma-4-E4B-it",
+        "unsloth/gemma-4-E4B",
+        "unsloth/gemma-4-31B-it",
+        "unsloth/gemma-4-31B",
+        "unsloth/gemma-4-26B-A4B-it",
+        "unsloth/gemma-4-26B-A4B",
+    ),
     "gemma3n": (
         "unsloth/gemma-3n-E4B-it-unsloth-bnb-4bit",
         "unsloth/gemma-3n-E4B-it",

unsloth/save.py

@@ -2533,12 +2533,15 @@ def unsloth_convert_lora_to_ggml_and_push_to_hub(
     )
     print(f"The output file will be {output_file}")
 
-    command = f"python3 llama.cpp/convert-lora-to-ggml.py {lora_directory_push} {output_file} llama"
-
     try:
         with subprocess.Popen(
-            command,
-            shell = True,
+            [
+                sys.executable,
+                "llama.cpp/convert-lora-to-ggml.py",
+                lora_directory_push,
+                output_file,
+                "llama",
+            ],
             stdout = subprocess.PIPE,
             stderr = subprocess.PIPE,
             bufsize = 1,
@@ -2550,7 +2553,7 @@ def unsloth_convert_lora_to_ggml_and_push_to_hub(
                 print(line, end = "", flush = True)
             sp.wait()
             if sp.returncode != 0:
-                raise subprocess.CalledProcessError(sp.returncode, command)
+                raise subprocess.CalledProcessError(sp.returncode, sp.args)
     except subprocess.CalledProcessError as e:
         print(f"Error: Conversion failed with return code {e.returncode}")
         return
@@ -2612,12 +2615,15 @@ def unsloth_convert_lora_to_ggml_and_save_locally(
     )
     print(f"The output file will be {output_file}")
 
-    command = f"python3 llama.cpp/convert-lora-to-ggml.py {save_directory} {output_file} llama"
-
     try:
         with subprocess.Popen(
-            command,
-            shell = True,
+            [
+                sys.executable,
+                "llama.cpp/convert-lora-to-ggml.py",
+                save_directory,
+                output_file,
+                "llama",
+            ],
             stdout = subprocess.PIPE,
             stderr = subprocess.PIPE,
             bufsize = 1,
@@ -2629,7 +2635,7 @@ def unsloth_convert_lora_to_ggml_and_save_locally(
                 print(line, end = "", flush = True)
             sp.wait()
             if sp.returncode != 0:
-                raise subprocess.CalledProcessError(sp.returncode, command)
+                raise subprocess.CalledProcessError(sp.returncode, sp.args)
     except subprocess.CalledProcessError as e:
         print(f"Error: Conversion failed with return code {e.returncode}")
         return
@@ -2777,19 +2783,79 @@ def unsloth_generic_save(
     elif save_method == "merged_4bit_forced":
         save_method = "merged_4bit"
 
-    merge_and_overwrite_lora(
-        get_model_name,
-        model = model,
-        tokenizer = tokenizer,
-        save_directory = save_directory,
-        push_to_hub = push_to_hub,
-        private = private,
-        token = token,
-        save_method = save_method,
-        output_dtype = None,
-        low_disk_space_usage = True,
-        use_temp_file = False,
-    )
+    # Full-finetuned models (no LoRA) cannot use merge_and_overwrite_lora
+    # since there are no adapters to merge. Fall back to save_pretrained.
+    # This mirrors the non-PeftModel handling in save_pretrained_torchao
+    # and the GGUF save path.
+    _is_peft = isinstance(model, PeftModel)
+    if not _is_peft:
+        if not is_main_process:
+            return
+
+        # Honor merged_16bit by casting to the target dtype if needed
+        _save_kwargs = dict(
+            safe_serialization = safe_serialization,
+            max_shard_size = max_shard_size,
+            variant = variant,
+        )
+        if "16bit" in save_method:
+            _target_dtype = (
+                torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+            )
+            _save_kwargs["state_dict"] = {
+                k: v.to(dtype = _target_dtype) if v.is_floating_point() else v
+                for k, v in model.state_dict().items()
+            }
+
+        if push_to_hub:
+            print(f"Unsloth: Pushing full fine-tuned model to '{save_directory}' ...")
+            model.push_to_hub(
+                repo_id = save_directory,
+                token = token,
+                private = private,
+                commit_message = commit_message,
+                create_pr = create_pr,
+                revision = revision,
+                commit_description = commit_description,
+                tags = tags,
+                **_save_kwargs,
+            )
+            if tokenizer is not None:
+                old_padding_side = tokenizer.padding_side
+                tokenizer.padding_side = "left"
+                tokenizer.push_to_hub(
+                    save_directory,
+                    token = token,
+                    private = private,
+                    commit_message = commit_message,
+                    create_pr = create_pr,
+                    revision = revision,
+                )
+                tokenizer.padding_side = old_padding_side
+        else:
+            print(f"Unsloth: Saving full fine-tuned model to '{save_directory}' ...")
+            model.save_pretrained(save_directory, **_save_kwargs)
+            if tokenizer is not None:
+                old_padding_side = tokenizer.padding_side
+                tokenizer.padding_side = "left"
+                tokenizer.save_pretrained(save_directory)
+                tokenizer.padding_side = old_padding_side
+
+        print(f"Unsloth: Model saved successfully to '{save_directory}'")
+    else:
+        merge_and_overwrite_lora(
+            get_model_name,
+            model = model,
+            tokenizer = tokenizer,
+            save_directory = save_directory,
+            push_to_hub = push_to_hub,
+            private = private,
+            token = token,
+            save_method = save_method,
+            output_dtype = None,
+            low_disk_space_usage = True,
+            use_temp_file = False,
+        )
 
     if push_to_hub and datasets:
         try: