Merge main into feature/rocm-support-v2

install.ps1

@@ -819,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.18" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.4.2" unsloth-zoo }
             if ($baseInstallExit -eq 0) {
                 $NoTorchReq = Find-NoTorchRuntimeFile
                 if ($NoTorchReq) {
@@ -827,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.18" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.4.2" unsloth-zoo }
         }
         if ($baseInstallExit -ne 0) {
             Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@@ -857,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.18" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.4.2" unsloth-zoo }
             if ($baseInstallExit -eq 0) {
                 $NoTorchReq = Find-NoTorchRuntimeFile
                 if ($NoTorchReq) {
@@ -865,7 +865,7 @@ shell.Run cmd, 0, False
                 }
             }
         } elseif ($StudioLocalInstall) {
-            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.18" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.4.2" unsloth-zoo }
         } else {
             $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName" }
         }
@@ -886,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.18" --torch-backend=auto }
+            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.4.2" --torch-backend=auto }
             if ($baseInstallExit -ne 0) {
                 Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
                 return

install.sh

@@ -1206,7 +1206,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.18" unsloth-zoo
+            "unsloth>=2026.4.2" 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"
@@ -1214,7 +1214,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.18" unsloth-zoo
+            "unsloth>=2026.4.2" unsloth-zoo
     fi
     if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
         substep "overlaying local repo (editable)..."
@@ -1284,7 +1284,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.18" unsloth-zoo
+            "unsloth>=2026.4.2" 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"
@@ -1295,7 +1295,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.18" unsloth-zoo
+            --upgrade-package unsloth "unsloth>=2026.4.2" 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
@@ -1306,7 +1306,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.18" --torch-backend=auto
+        run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.4.2" --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

studio/backend/assets/configs/inference_defaults.json

@@ -93,6 +93,14 @@
       "min_p": 0.0,
       "repetition_penalty": 1.0
     },
+    "gemma-4": {
+      "temperature": 1.0,
+      "top_p": 0.95,
+      "top_k": 64,
+      "min_p": 0.0,
+      "repetition_penalty": 1.0,
+      "presence_penalty": 0.0
+    },
     "gemma-3n": {
       "temperature": 1.0,
       "top_p": 0.95,
@@ -366,7 +374,7 @@
     "qwen2.5-coder", "qwen2.5-vl", "qwen2.5-omni", "qwen2.5-math", "qwen2.5",
     "qwen2-vl", "qwen2",
     "qwq",
-    "gemma-3n", "gemma-3", "medgemma", "gemma-2",
+    "gemma-4", "gemma-3n", "gemma-3", "medgemma", "gemma-2",
     "llama-4", "llama-3.3", "llama-3.2", "llama-3.1", "llama-3",
     "phi-4", "phi-3",
     "mistral-nemo", "mistral-small", "mistral-large", "magistral", "ministral",

studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml

@@ -0,0 +1,47 @@
+# Model defaults for unsloth/gemma-4-26B-A4B-it
+# Also applies to: google/gemma-4-26B-A4B-it, unsloth/gemma-4-26B-A4B-it-GGUF
+
+training:
+  trust_remote_code: false
+  max_seq_length: 2048
+  num_epochs: 0
+  learning_rate: 2e-4
+  batch_size: 2
+  gradient_accumulation_steps: 4
+  warmup_steps: 5
+  max_steps: 30
+  save_steps: 30
+  weight_decay: 0.001
+  random_seed: 3407
+  packing: false
+  train_on_completions: true
+  gradient_checkpointing: "unsloth"
+  optim: "adamw_8bit"
+  lr_scheduler_type: "linear"
+
+lora:
+  lora_r: 8
+  lora_alpha: 8
+  lora_dropout: 0.0
+  target_modules:
+    - "all-linear"
+  use_rslora: false
+  use_loftq: false
+  finetune_vision_layers: true
+  finetune_language_layers: true
+  finetune_attention_modules: true
+  finetune_mlp_modules: true
+
+logging:
+  enable_wandb: false
+  wandb_project: "llm-finetuning"
+  enable_tensorboard: false
+  tensorboard_dir: "runs"
+  log_frequency: 10
+
+inference:
+  trust_remote_code: false
+  temperature: 1.0
+  top_p: 0.95
+  top_k: 64
+  min_p: 0.0

studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml

@@ -0,0 +1,47 @@
+# Model defaults for unsloth/gemma-4-26B-A4B (base/pretrained)
+# Also applies to: google/gemma-4-26B-A4B
+
+training:
+  trust_remote_code: false
+  max_seq_length: 2048
+  num_epochs: 0
+  learning_rate: 2e-4
+  batch_size: 2
+  gradient_accumulation_steps: 4
+  warmup_steps: 5
+  max_steps: 30
+  save_steps: 30
+  weight_decay: 0.001
+  random_seed: 3407
+  packing: false
+  train_on_completions: true
+  gradient_checkpointing: "unsloth"
+  optim: "adamw_8bit"
+  lr_scheduler_type: "linear"
+
+lora:
+  lora_r: 8
+  lora_alpha: 8
+  lora_dropout: 0.0
+  target_modules:
+    - "all-linear"
+  use_rslora: false
+  use_loftq: false
+  finetune_vision_layers: true
+  finetune_language_layers: true
+  finetune_attention_modules: true
+  finetune_mlp_modules: true
+
+logging:
+  enable_wandb: false
+  wandb_project: "llm-finetuning"
+  enable_tensorboard: false
+  tensorboard_dir: "runs"
+  log_frequency: 10
+
+inference:
+  trust_remote_code: false
+  temperature: 1.0
+  top_p: 0.95
+  top_k: 64
+  min_p: 0.0

studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml

@@ -0,0 +1,47 @@
+# Model defaults for unsloth/gemma-4-31B-it
+# Also applies to: google/gemma-4-31B-it, unsloth/gemma-4-31B-it-GGUF
+
+training:
+  trust_remote_code: false
+  max_seq_length: 2048
+  num_epochs: 0
+  learning_rate: 2e-4
+  batch_size: 2
+  gradient_accumulation_steps: 4
+  warmup_steps: 5
+  max_steps: 30
+  save_steps: 30
+  weight_decay: 0.001
+  random_seed: 3407
+  packing: false
+  train_on_completions: true
+  gradient_checkpointing: "unsloth"
+  optim: "adamw_8bit"
+  lr_scheduler_type: "linear"
+
+lora:
+  lora_r: 8
+  lora_alpha: 8
+  lora_dropout: 0.0
+  target_modules:
+    - "all-linear"
+  use_rslora: false
+  use_loftq: false
+  finetune_vision_layers: true
+  finetune_language_layers: true
+  finetune_attention_modules: true
+  finetune_mlp_modules: true
+
+logging:
+  enable_wandb: false
+  wandb_project: "llm-finetuning"
+  enable_tensorboard: false
+  tensorboard_dir: "runs"
+  log_frequency: 10
+
+inference:
+  trust_remote_code: false
+  temperature: 1.0
+  top_p: 0.95
+  top_k: 64
+  min_p: 0.0

studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml

@@ -0,0 +1,47 @@
+# Model defaults for unsloth/gemma-4-31B (base/pretrained)
+# Also applies to: google/gemma-4-31B
+
+training:
+  trust_remote_code: false
+  max_seq_length: 2048
+  num_epochs: 0
+  learning_rate: 2e-4
+  batch_size: 2
+  gradient_accumulation_steps: 4
+  warmup_steps: 5
+  max_steps: 30
+  save_steps: 30
+  weight_decay: 0.001
+  random_seed: 3407
+  packing: false
+  train_on_completions: true
+  gradient_checkpointing: "unsloth"
+  optim: "adamw_8bit"
+  lr_scheduler_type: "linear"
+
+lora:
+  lora_r: 8
+  lora_alpha: 8
+  lora_dropout: 0.0
+  target_modules:
+    - "all-linear"
+  use_rslora: false
+  use_loftq: false
+  finetune_vision_layers: true
+  finetune_language_layers: true
+  finetune_attention_modules: true
+  finetune_mlp_modules: true
+
+logging:
+  enable_wandb: false
+  wandb_project: "llm-finetuning"
+  enable_tensorboard: false
+  tensorboard_dir: "runs"
+  log_frequency: 10
+
+inference:
+  trust_remote_code: false
+  temperature: 1.0
+  top_p: 0.95
+  top_k: 64
+  min_p: 0.0

studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml

@@ -0,0 +1,47 @@
+# Model defaults for unsloth/gemma-4-E2B-it
+# Also applies to: google/gemma-4-E2B-it, unsloth/gemma-4-E2B-it-GGUF
+
+training:
+  trust_remote_code: false
+  max_seq_length: 2048
+  num_epochs: 0
+  learning_rate: 2e-4
+  batch_size: 2
+  gradient_accumulation_steps: 4
+  warmup_steps: 5
+  max_steps: 30
+  save_steps: 30
+  weight_decay: 0.001
+  random_seed: 3407
+  packing: false
+  train_on_completions: true
+  gradient_checkpointing: "unsloth"
+  optim: "adamw_8bit"
+  lr_scheduler_type: "linear"
+
+lora:
+  lora_r: 8
+  lora_alpha: 8
+  lora_dropout: 0.0
+  target_modules:
+    - "all-linear"
+  use_rslora: false
+  use_loftq: false
+  finetune_vision_layers: true
+  finetune_language_layers: true
+  finetune_attention_modules: true
+  finetune_mlp_modules: true
+
+logging:
+  enable_wandb: false
+  wandb_project: "llm-finetuning"
+  enable_tensorboard: false
+  tensorboard_dir: "runs"
+  log_frequency: 10
+
+inference:
+  trust_remote_code: false
+  temperature: 1.0
+  top_p: 0.95
+  top_k: 64
+  min_p: 0.0

studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml

@@ -0,0 +1,47 @@
+# Model defaults for unsloth/gemma-4-E2B (base/pretrained)
+# Also applies to: google/gemma-4-E2B
+
+training:
+  trust_remote_code: false
+  max_seq_length: 2048
+  num_epochs: 0
+  learning_rate: 2e-4
+  batch_size: 2
+  gradient_accumulation_steps: 4
+  warmup_steps: 5
+  max_steps: 30
+  save_steps: 30
+  weight_decay: 0.001
+  random_seed: 3407
+  packing: false
+  train_on_completions: true
+  gradient_checkpointing: "unsloth"
+  optim: "adamw_8bit"
+  lr_scheduler_type: "linear"
+
+lora:
+  lora_r: 8
+  lora_alpha: 8
+  lora_dropout: 0.0
+  target_modules:
+    - "all-linear"
+  use_rslora: false
+  use_loftq: false
+  finetune_vision_layers: true
+  finetune_language_layers: true
+  finetune_attention_modules: true
+  finetune_mlp_modules: true
+
+logging:
+  enable_wandb: false
+  wandb_project: "llm-finetuning"
+  enable_tensorboard: false
+  tensorboard_dir: "runs"
+  log_frequency: 10
+
+inference:
+  trust_remote_code: false
+  temperature: 1.0
+  top_p: 0.95
+  top_k: 64
+  min_p: 0.0

studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml

@@ -0,0 +1,47 @@
+# Model defaults for unsloth/gemma-4-E4B-it
+# Also applies to: google/gemma-4-E4B-it, unsloth/gemma-4-E4B-it-GGUF
+
+training:
+  trust_remote_code: false
+  max_seq_length: 2048
+  num_epochs: 0
+  learning_rate: 2e-4
+  batch_size: 2
+  gradient_accumulation_steps: 4
+  warmup_steps: 5
+  max_steps: 30
+  save_steps: 30
+  weight_decay: 0.001
+  random_seed: 3407
+  packing: false
+  train_on_completions: true
+  gradient_checkpointing: "unsloth"
+  optim: "adamw_8bit"
+  lr_scheduler_type: "linear"
+
+lora:
+  lora_r: 8
+  lora_alpha: 8
+  lora_dropout: 0.0
+  target_modules:
+    - "all-linear"
+  use_rslora: false
+  use_loftq: false
+  finetune_vision_layers: true
+  finetune_language_layers: true
+  finetune_attention_modules: true
+  finetune_mlp_modules: true
+
+logging:
+  enable_wandb: false
+  wandb_project: "llm-finetuning"
+  enable_tensorboard: false
+  tensorboard_dir: "runs"
+  log_frequency: 10
+
+inference:
+  trust_remote_code: false
+  temperature: 1.0
+  top_p: 0.95
+  top_k: 64
+  min_p: 0.0

studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml

@@ -0,0 +1,47 @@
+# Model defaults for unsloth/gemma-4-E4B (base/pretrained)
+# Also applies to: google/gemma-4-E4B
+
+training:
+  trust_remote_code: false
+  max_seq_length: 2048
+  num_epochs: 0
+  learning_rate: 2e-4
+  batch_size: 2
+  gradient_accumulation_steps: 4
+  warmup_steps: 5
+  max_steps: 30
+  save_steps: 30
+  weight_decay: 0.001
+  random_seed: 3407
+  packing: false
+  train_on_completions: true
+  gradient_checkpointing: "unsloth"
+  optim: "adamw_8bit"
+  lr_scheduler_type: "linear"
+
+lora:
+  lora_r: 8
+  lora_alpha: 8
+  lora_dropout: 0.0
+  target_modules:
+    - "all-linear"
+  use_rslora: false
+  use_loftq: false
+  finetune_vision_layers: true
+  finetune_language_layers: true
+  finetune_attention_modules: true
+  finetune_mlp_modules: true
+
+logging:
+  enable_wandb: false
+  wandb_project: "llm-finetuning"
+  enable_tensorboard: false
+  tensorboard_dir: "runs"
+  log_frequency: 10
+
+inference:
+  trust_remote_code: false
+  temperature: 1.0
+  top_p: 0.95
+  top_k: 64
+  min_p: 0.0

studio/backend/core/data_recipe/service.py

@@ -167,12 +167,7 @@ def _validate_recipe_runtime_support(
     recipe: dict[str, Any],
     model_providers: list[Any],
 ) -> None:
-    if not _recipe_has_llm_columns(recipe):
-        raise ValueError(
-            "Recipe Studio currently requires at least one AI generation step."
-        )
-
-    if not model_providers:
+    if _recipe_has_llm_columns(recipe) and not model_providers:
         raise ValueError("Add a Provider connection block before running this recipe.")
 
 
@@ -266,6 +261,21 @@ def create_data_designer(
     model_providers = build_model_providers(recipe)
     _validate_recipe_runtime_support(recipe, model_providers)
 
+    # DataDesigner requires at least one model provider in its registry even
+    # when the pipeline contains no LLM columns.  Supply a lightweight stub
+    # so sampler/expression-only recipes can run without a real provider.
+    if not model_providers:
+        from data_designer.config.models import ModelProvider
+
+        model_providers = [
+            ModelProvider(
+                name = "_unused",
+                endpoint = "http://localhost",
+                provider_type = "openai",
+                api_key = None,
+            )
+        ]
+
     return DataDesigner(
         artifact_path = artifact_path,
         model_providers = model_providers,

studio/backend/core/inference/llama_cpp.py

@@ -22,6 +22,7 @@ import threading
 import time
 from pathlib import Path
 from typing import Generator, Optional
+from urllib.parse import urlparse
 
 import httpx
 
@@ -108,6 +109,7 @@ class LlamaCppBackend:
         self._supports_tools: bool = False
         self._cache_type_kv: Optional[str] = None
         self._reasoning_default: bool = True
+        self._speculative_type: Optional[str] = None
         # KV-cache estimation fields (populated by _read_gguf_metadata)
         self._n_layers: Optional[int] = None
         self._n_kv_heads: Optional[int] = None
@@ -197,6 +199,10 @@ class LlamaCppBackend:
     def cache_type_kv(self) -> Optional[str]:
         return self._cache_type_kv
 
+    @property
+    def speculative_type(self) -> Optional[str]:
+        return self._speculative_type
+
     # ── Binary discovery ──────────────────────────────────────────
 
     @staticmethod
@@ -1054,6 +1060,7 @@ class LlamaCppBackend:
         n_ctx: int = 4096,
         chat_template_override: Optional[str] = None,
         cache_type_kv: Optional[str] = None,
+        speculative_type: Optional[str] = None,
         n_threads: Optional[int] = None,
         n_gpu_layers: Optional[int] = None,  # Accepted for caller compat, unused
     ) -> bool:
@@ -1314,6 +1321,46 @@ class LlamaCppBackend:
             else:
                 self._cache_type_kv = None
 
+            # Speculative decoding (n-gram self-speculation, zero VRAM cost)
+            # ngram-mod: ~16 MB shared hash pool, constant memory/complexity,
+            # variable draft lengths.  Helps most when the model repeats
+            # existing text (code refactoring, summarization, reasoning).
+            # For general chat with low repetition, overhead is ~5 ms.
+            #
+            # Benchmarks from llama.cpp PRs #18471, #19164:
+            #   Scenario                        | Without | With    | Speedup
+            #   gpt-oss-120b code refactor      | 181 t/s | 446 t/s | 2.5x
+            #   Qwen3-235B offloaded            |  12 t/s |  21 t/s | 1.8x
+            #   gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x
+            #
+            # Params from llama.cpp docs (docs/speculative.md):
+            #   --spec-ngram-size-n 24  (small n not recommended)
+            #   --draft-min 48 --draft-max 64 (MoEs need long drafts;
+            #     dense models can reduce these)
+            # ref: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md
+            # ref: https://github.com/ggml-org/llama.cpp/pull/19164
+            # ref: https://github.com/ggml-org/llama.cpp/pull/18471
+            _valid_spec_types = {"ngram-simple", "ngram-mod"}
+            if speculative_type and speculative_type in _valid_spec_types:
+                if not is_vision:  # spec decoding disabled for vision models
+                    cmd.extend(["--spec-type", speculative_type])
+                    if speculative_type == "ngram-mod":
+                        cmd.extend(
+                            [
+                                "--spec-ngram-size-n",
+                                "24",
+                                "--draft-min",
+                                "48",
+                                "--draft-max",
+                                "64",
+                            ]
+                        )
+                    self._speculative_type = speculative_type
+                else:
+                    self._speculative_type = None
+            else:
+                self._speculative_type = None
+
             # Apply custom chat template override if provided
             if chat_template_override:
                 import tempfile
@@ -1552,6 +1599,7 @@ class LlamaCppBackend:
             self._reasoning_always_on = False
             self._supports_tools = False
             self._cache_type_kv = None
+            self._speculative_type = None
             self._n_layers = None
             self._n_kv_heads = None
             self._n_heads = None
@@ -2270,7 +2318,7 @@ class LlamaCppBackend:
         Agentic loop: let the model call tools, execute them, and continue.
 
         Yields dicts with:
-          {"type": "status", "text": "Searching: ..."}   -- tool status updates
+          {"type": "status", "text": "Searching: ..."/"Reading: ..."}   -- tool status updates
           {"type": "content", "text": "token"}            -- streamed content tokens (cumulative)
           {"type": "reasoning", "text": "token"}          -- streamed reasoning tokens (cumulative)
         """
@@ -2837,7 +2885,18 @@ class LlamaCppBackend:
                         arguments = raw_args
 
                     if tool_name == "web_search":
-                        status_text = f"Searching: {arguments.get('query', '')}"
+                        _ws_url = (arguments.get("url") or "").strip()
+                        if _ws_url:
+                            _parsed = urlparse(_ws_url)
+                            if _parsed.scheme in ("http", "https") and _parsed.hostname:
+                                _ws_host = _parsed.hostname
+                                if _ws_host.startswith("www."):
+                                    _ws_host = _ws_host[4:]
+                                status_text = f"Reading: {_ws_host}"
+                            else:
+                                status_text = "Reading page..."
+                        else:
+                            status_text = f"Searching: {arguments.get('query', '')}"
                     elif tool_name == "python":
                         preview = (
                             (arguments.get("code") or "").strip().split("\n")[0][:60]

studio/backend/core/inference/tools.py

@@ -14,6 +14,8 @@ import os
 os.environ["UNSLOTH_IS_PRESENT"] = "1"
 
 import random
+import re
+import shlex
 import ssl
 import subprocess
 import sys
@@ -27,14 +29,239 @@ logger = get_logger(__name__)
 
 _EXEC_TIMEOUT = 300  # 5 minutes
 
+# Pre-import modules used in _sandbox_preexec at module level so that
+# the preexec_fn closure does not trigger the import machinery in the
+# forked child (which can deadlock in multi-threaded servers).
+_libc = None
+if sys.platform == "linux":
+    try:
+        import ctypes
+        import ctypes.util
+
+        _libc_name = ctypes.util.find_library("c")
+        if _libc_name:
+            _libc = ctypes.CDLL(_libc_name, use_errno = True)
+    except (OSError, AttributeError):
+        pass
+
+_resource = None
+if sys.platform != "win32":
+    try:
+        import resource as _resource
+    except ImportError:
+        pass
+
 # 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"}
+_BLOCKED_COMMANDS_COMMON = frozenset(
+    {
+        "rm",
+        "sudo",
+        "su",
+        "dd",
+        "chmod",
+        "chown",
+        "mkfs",
+        "shutdown",
+        "reboot",
+        "passwd",
+        "mount",
+        "umount",
+        "fdisk",
+        "kill",
+        "killall",
+        "pkill",
+    }
+)
+_BLOCKED_COMMANDS_WIN = frozenset(
+    {
+        "rmdir",
+        "takeown",
+        "icacls",
+        "runas",
+        "powershell",
+        "pwsh",
+    }
+)
+_BLOCKED_COMMANDS = (
+    _BLOCKED_COMMANDS_COMMON | _BLOCKED_COMMANDS_WIN
+    if sys.platform == "win32"
+    else _BLOCKED_COMMANDS_COMMON
+)
+
+
+def _find_blocked_commands(command: str) -> set[str]:
+    """Detect blocked commands using shlex tokenization and regex scanning.
+
+    Catches: full paths (/usr/bin/sudo), quoted strings ("sudo"),
+    split-quotes (su""do), backslash escapes (\\rm), and command-position
+    words after ;, |, &&, $().
+    """
+    blocked = set()
+
+    # 1. shlex tokenization (handles quotes, escapes, concatenation)
+    try:
+        tokens = (
+            shlex.split(command)
+            if sys.platform != "win32"
+            else shlex.split(command, posix = False)
+        )
+    except ValueError:
+        tokens = command.split()
+
+    for token in tokens:
+        base = os.path.basename(token).lower()
+        # Strip common Windows executable extensions so that
+        # runas.exe, shutdown.bat, etc. match the blocklist.
+        stem, ext = os.path.splitext(base)
+        if ext in {".exe", ".com", ".bat", ".cmd"}:
+            base = stem
+        if base in _BLOCKED_COMMANDS:
+            blocked.add(base)
+
+    # 2. Regex: catch blocked words at shell command boundaries
+    #    (semicolons, pipes, &&, ||, backticks, $(), <(), subshells, newlines)
+    #    Uses a single combined pattern for all blocked words.
+    #    Handles optional Unix path prefix (/usr/bin/) and Windows drive
+    #    letter prefix (C:\Windows\...\).
+    lowered = command.lower()
+    if _BLOCKED_COMMANDS:
+        words_alt = "|".join(re.escape(w) for w in sorted(_BLOCKED_COMMANDS))
+        pattern = (
+            rf"(?:^|[;&|`\n(]\s*|[$]\(\s*|<\(\s*)"
+            rf"(?:[\w./\\-]*/|[a-zA-Z]:[/\\][\w./\\-]*)?"
+            rf"({words_alt})(?:\.(?:exe|com|bat|cmd))?\b"
+        )
+        blocked.update(re.findall(pattern, lowered))
+
+    # 3. Check for nested shell invocations (bash -c 'sudo whoami',
+    #    bash -lc '...', bash --login -c '...', cmd /c '...').
+    #    When a -c or /c flag is found, look backwards for a shell name
+    #    (skipping intermediate flags like --login, -l, -x) and recursively
+    #    scan the nested command string.
+    _SHELLS = {"bash", "sh", "zsh", "dash", "ksh", "csh", "tcsh", "fish"}
+    _SHELLS_WIN = {"cmd", "cmd.exe"}
+    for i, token in enumerate(tokens):
+        tok_lower = token.lower()
+        # Match -c exactly, or combined flags ending in c (e.g. -lc, -xc)
+        is_unix_c = tok_lower == "-c" or (
+            tok_lower.startswith("-")
+            and tok_lower.endswith("c")
+            and not tok_lower.startswith("--")
+        )
+        is_win_c = tok_lower == "/c"
+        if not (is_unix_c or is_win_c) or i < 1 or i + 1 >= len(tokens):
+            continue
+        # Look backwards past any flags to find the shell binary.
+        # On Unix, flags start with - (skip those). On Windows, flags
+        # start with / but so do absolute paths, so only skip short
+        # single-char /X flags (not /bin/bash style paths).
+        for j in range(i - 1, -1, -1):
+            prev = tokens[j]
+            if prev.startswith("-"):
+                continue  # skip Unix flags like --login, -l
+            if is_win_c and prev.startswith("/") and len(prev) <= 3:
+                continue  # skip Windows flags like /s, /q (not /bin/bash)
+            prev_base = os.path.basename(prev).lower()
+            if is_unix_c and prev_base in _SHELLS:
+                blocked |= _find_blocked_commands(tokens[i + 1])
+            elif is_win_c and prev_base in _SHELLS_WIN:
+                blocked |= _find_blocked_commands(tokens[i + 1])
+            break  # stop at first non-flag token
+
+    return blocked
+
+
+def _build_safe_env(workdir: str) -> dict[str, str]:
+    """Build a minimal, credential-free environment for sandboxed subprocesses.
+
+    Strips HF_TOKEN, WANDB_API_KEY, AWS_*, GH_TOKEN, LD_PRELOAD, DYLD_*, etc.
+    Preserves the active Python interpreter and virtualenv directories in PATH
+    so that pip, uv, and packages installed in the Studio runtime remain
+    accessible.
+    """
+    # Start with the directory containing the running Python interpreter
+    # so that subprocess calls to 'python', 'pip', etc. resolve to the
+    # same environment the Studio server is running in.
+    exe_dir = os.path.dirname(sys.executable)
+    path_entries = [exe_dir] if exe_dir else []
+
+    # If a virtualenv is active, include its bin/Scripts directory.
+    venv = os.environ.get("VIRTUAL_ENV")
+    if venv:
+        venv_bin = os.path.join(venv, "Scripts" if sys.platform == "win32" else "bin")
+        if venv_bin not in path_entries:
+            path_entries.append(venv_bin)
+
+    if sys.platform == "win32":
+        sysroot = os.environ.get("SystemRoot", r"C:\Windows")
+        path_entries.extend([os.path.join(sysroot, "System32"), sysroot])
+    else:
+        path_entries.extend(["/usr/local/bin", "/usr/bin", "/bin"])
+
+    # Deduplicate while preserving order
+    deduped = list(dict.fromkeys(p for p in path_entries if p))
+
+    env = {
+        "PATH": os.pathsep.join(deduped),
+        "HOME": workdir,
+        "TMPDIR": workdir,
+        "LANG": os.environ.get("LANG", "C.UTF-8"),
+        "TERM": "dumb",
+        "PYTHONIOENCODING": "utf-8",
+    }
+    if venv:
+        env["VIRTUAL_ENV"] = venv
+    # Windows needs SystemRoot for Python/subprocess to work
+    if sys.platform == "win32":
+        env["SystemRoot"] = os.environ.get("SystemRoot", r"C:\Windows")
+    return env
+
+
+def _sandbox_preexec():
+    """Pre-exec hook: drop privilege escalation ability and set resource limits.
+
+    On Linux, applies PR_SET_NO_NEW_PRIVS so sudo/su/pkexec fail at the
+    kernel level. On Linux and macOS, sets RLIMIT_FSIZE.
+    No-op on Windows (use creationflags instead).
+
+    Note: RLIMIT_NPROC is intentionally NOT set because Linux enforces it
+    per real UID, not per process tree, so it would starve the Studio
+    server and other sessions sharing the same user account.
+
+    All modules and handles are resolved at import time (module level) so
+    this function does not trigger Python imports in the forked child,
+    avoiding potential deadlocks in multi-threaded servers.
+    """
+    if _libc is not None:
+        try:
+            # PR_SET_NO_NEW_PRIVS = 38, arg2 = 1 (enable)
+            _libc.prctl(38, 1, 0, 0, 0)
+        except (OSError, AttributeError):
+            pass  # Not available (container, old kernel, etc.)
+
+    if _resource is not None:
+        try:
+            # Limit file size to 100MB (prevents disk filling)
+            _resource.setrlimit(
+                _resource.RLIMIT_FSIZE, (100 * 1024 * 1024, 100 * 1024 * 1024)
+            )
+        except (ValueError, OSError):
+            pass
+
+
+def _get_shell_cmd(command: str) -> list[str]:
+    """Return the platform-appropriate shell invocation for a command string."""
+    if sys.platform == "win32":
+        return ["cmd", "/c", command]
+    return ["bash", "-c", command]
+
 
 # Per-session working directories so each chat thread gets its own sandbox.
-# Falls back to a shared ~/studio_sandbox/ for API callers without a session_id.
+# Falls back to a shared ~/studio_sandbox/_default for API callers without a
+# session_id.
 _workdirs: dict[str, str] = {}
 
 
@@ -55,7 +282,7 @@ def _get_workdir(session_id: str | None = None) -> str:
             if not os.path.realpath(workdir).startswith(os.path.realpath(sandbox_root)):
                 workdir = os.path.join(sandbox_root, "_invalid")
         else:
-            workdir = sandbox_root
+            workdir = os.path.join(sandbox_root, "_default")
         os.makedirs(workdir, exist_ok = True)
         _workdirs[key] = workdir
     return _workdirs[key]
@@ -428,6 +655,7 @@ def _check_signal_escape_patterns(code: str):
 
     signal_tampering = []
     exception_catching = []
+    shell_escapes = []
     warnings = []
 
     def _ast_name_matches(node, names):
@@ -445,10 +673,84 @@ def _check_signal_escape_patterns(code: str):
             return full_name in names
         return False
 
+    # Dangerous os/subprocess functions that can execute shell commands
+    _SHELL_EXEC_FUNCS = frozenset(
+        {
+            "os.system",
+            "os.popen",
+            "os.popen2",
+            "os.popen3",
+            "os.popen4",
+            "os.execl",
+            "os.execle",
+            "os.execlp",
+            "os.execlpe",
+            "os.execv",
+            "os.execve",
+            "os.execvp",
+            "os.execvpe",
+            "os.spawnl",
+            "os.spawnle",
+            "os.spawnlp",
+            "os.spawnlpe",
+            "os.spawnv",
+            "os.spawnve",
+            "os.spawnvp",
+            "os.spawnvpe",
+            "os.posix_spawn",
+            "os.posix_spawnp",
+            "subprocess.run",
+            "subprocess.call",
+            "subprocess.check_call",
+            "subprocess.check_output",
+            "subprocess.Popen",
+            "subprocess.getoutput",
+            "subprocess.getstatusoutput",
+        }
+    )
+
+    def _extract_string_from_node(node):
+        """Extract a plain string value from an AST node, if it is a constant."""
+        if isinstance(node, ast.Constant) and isinstance(node.value, str):
+            return node.value
+        return None
+
+    def _extract_strings_from_list(node):
+        """Extract string elements from an AST List or Tuple node."""
+        if isinstance(node, (ast.List, ast.Tuple)):
+            parts = []
+            for elt in node.elts:
+                s = _extract_string_from_node(elt)
+                if s is not None:
+                    parts.append(s)
+            return parts
+        return []
+
+    # Keyword argument names that carry command content (as opposed to
+    # control flags like check=True, text=True, capture_output=True).
+    _CMD_KWARGS = frozenset({"args", "command", "executable", "path", "file"})
+
+    def _check_args_for_blocked(args_nodes):
+        """Check if any call arguments contain blocked commands."""
+        found = set()
+        for arg in args_nodes:
+            s = _extract_string_from_node(arg)
+            if s is not None:
+                found |= _find_blocked_commands(s)
+            strs = _extract_strings_from_list(arg)
+            for s in strs:
+                found |= _find_blocked_commands(s)
+        return found
+
     class SignalEscapeVisitor(ast.NodeVisitor):
         def __init__(self):
             self.imports_signal = False
             self.signal_aliases = {"signal"}
+            self.os_aliases = {"os"}
+            self.subprocess_aliases = {"subprocess"}
+            # Maps bare function names to their fully-qualified form
+            # for from-import tracking (e.g. "system" -> "os.system")
+            self.shell_exec_aliases: dict[str, str] = {}
             self.loop_depth = 0
 
         def visit_Import(self, node):
@@ -457,6 +759,10 @@ def _check_signal_escape_patterns(code: str):
                     self.imports_signal = True
                     if alias.asname:
                         self.signal_aliases.add(alias.asname)
+                elif alias.name == "os":
+                    self.os_aliases.add(alias.asname or "os")
+                elif alias.name == "subprocess":
+                    self.subprocess_aliases.add(alias.asname or "subprocess")
             self.generic_visit(node)
 
         def visit_ImportFrom(self, node):
@@ -474,6 +780,16 @@ def _check_signal_escape_patterns(code: str):
                         "alarm",
                     ):
                         self.signal_aliases.add(alias.asname or alias.name)
+            elif node.module in ("os", "subprocess"):
+                if node.module == "os":
+                    self.os_aliases.add("os")
+                else:
+                    self.subprocess_aliases.add("subprocess")
+                # Track from-imports of dangerous functions
+                for alias in node.names:
+                    fq = f"{node.module}.{alias.name}"
+                    if fq in _SHELL_EXEC_FUNCS:
+                        self.shell_exec_aliases[alias.asname or alias.name] = fq
             self.generic_visit(node)
 
         def visit_While(self, node):
@@ -538,6 +854,111 @@ def _check_signal_escape_patterns(code: str):
                             "description": "Modifies signal mask (may block SIGALRM)",
                         }
                     )
+
+            # --- Shell escape detection ---
+            # Resolve the fully qualified function name for os.*/subprocess.*
+            shell_func = None
+            if isinstance(func, ast.Attribute):
+                if isinstance(func.value, ast.Name):
+                    if func.value.id in self.os_aliases:
+                        shell_func = f"os.{func.attr}"
+                    elif func.value.id in self.subprocess_aliases:
+                        shell_func = f"subprocess.{func.attr}"
+            elif isinstance(func, ast.Name):
+                # Check from-import aliases: from os import system; system(...)
+                shell_func = self.shell_exec_aliases.get(func.id)
+
+            if shell_func and shell_func in _SHELL_EXEC_FUNCS:
+                # Expand **kwargs dicts to inspect their keys
+                expanded_kwargs: dict[str, ast.AST] = {}
+                has_opaque_kwargs = False
+                for kw in node.keywords:
+                    if kw.arg is not None:
+                        expanded_kwargs[kw.arg] = kw.value
+                    elif isinstance(kw.value, ast.Dict):
+                        for k, v in zip(kw.value.keys, kw.value.values):
+                            key = _extract_string_from_node(k) if k else None
+                            if key is not None:
+                                expanded_kwargs[key] = v
+                    else:
+                        has_opaque_kwargs = True
+
+                cmd_kw_values = [
+                    v for k, v in expanded_kwargs.items() if k in _CMD_KWARGS
+                ]
+                all_call_args = list(node.args) + cmd_kw_values
+                blocked_in_args = _check_args_for_blocked(all_call_args)
+
+                if has_opaque_kwargs:
+                    # Can't inspect dynamic **kwargs -- flag as unsafe
+                    shell_escapes.append(
+                        {
+                            "type": "shell_escape_dynamic",
+                            "line": node.lineno,
+                            "description": (
+                                f"{shell_func}() called with dynamic **kwargs"
+                            ),
+                        }
+                    )
+                elif blocked_in_args:
+                    shell_escapes.append(
+                        {
+                            "type": "shell_escape",
+                            "line": node.lineno,
+                            "description": (
+                                f"{shell_func}() invokes blocked command(s): "
+                                f"{', '.join(sorted(blocked_in_args))}"
+                            ),
+                        }
+                    )
+                else:
+                    # Only flag dynamic args for functions that interpret
+                    # strings as shell commands, or when shell= might be
+                    # enabled.  Treat any non-literal-False shell= value
+                    # as potentially True (conservative).
+                    _STRING_SHELL_FUNCS = frozenset(
+                        {
+                            "os.system",
+                            "os.popen",
+                            "os.popen2",
+                            "os.popen3",
+                            "os.popen4",
+                            "subprocess.getoutput",
+                            "subprocess.getstatusoutput",
+                        }
+                    )
+                    shell_node = expanded_kwargs.get("shell")
+                    shell_safe = shell_node is None or (
+                        isinstance(shell_node, ast.Constant)
+                        and shell_node.value is False
+                    )
+                    if shell_func in _STRING_SHELL_FUNCS or not shell_safe:
+
+                        def _is_safe_literal(n):
+                            if _extract_string_from_node(n) is not None:
+                                return True
+                            if isinstance(n, (ast.List, ast.Tuple)):
+                                return all(
+                                    _extract_string_from_node(e) is not None
+                                    for e in n.elts
+                                )
+                            return False
+
+                        has_non_literal = any(
+                            not _is_safe_literal(a) for a in all_call_args
+                        )
+                        if has_non_literal:
+                            shell_escapes.append(
+                                {
+                                    "type": "shell_escape_dynamic",
+                                    "line": node.lineno,
+                                    "description": (
+                                        f"{shell_func}() called with non-literal "
+                                        f"shell command (potential shell escape)"
+                                    ),
+                                }
+                            )
+
             self.generic_visit(node)
 
         def visit_ExceptHandler(self, node):
@@ -553,7 +974,12 @@ def _check_signal_escape_patterns(code: str):
                     }
                 )
             elif isinstance(node.type, ast.Name):
-                if node.type.id in ("TimeoutError", "BaseException", "Exception"):
+                # Only flag BaseException and TimeoutError, NOT Exception.
+                # except Exception does not catch SystemExit or
+                # KeyboardInterrupt, so it cannot suppress timeout
+                # enforcement.  Flagging Exception causes false positives
+                # on normal error-handling patterns.
+                if node.type.id in ("TimeoutError", "BaseException"):
                     exception_catching.append(
                         {
                             "type": f"catches_{node.type.id}_in_loop",
@@ -564,7 +990,7 @@ def _check_signal_escape_patterns(code: str):
             elif isinstance(node.type, ast.Tuple):
                 for elt in node.type.elts:
                     if isinstance(elt, ast.Name):
-                        if elt.id in ("TimeoutError", "BaseException", "Exception"):
+                        if elt.id in ("TimeoutError", "BaseException"):
                             exception_catching.append(
                                 {
                                     "type": f"catches_{elt.id}_in_loop",
@@ -580,10 +1006,15 @@ def _check_signal_escape_patterns(code: str):
     if visitor.imports_signal and not signal_tampering:
         warnings.append("Code imports 'signal' module - review manually for safety")
 
-    is_safe = len(signal_tampering) == 0 and len(exception_catching) == 0
+    is_safe = (
+        len(signal_tampering) == 0
+        and len(exception_catching) == 0
+        and len(shell_escapes) == 0
+    )
     return is_safe, {
         "signal_tampering": signal_tampering,
         "exception_catching": exception_catching,
+        "shell_escapes": shell_escapes,
         "warnings": warnings,
     }
 
@@ -604,10 +1035,18 @@ def _check_code_safety(code: str) -> str | None:
         reasons = [
             item.get("description", "") for item in info.get("signal_tampering", [])
         ]
-        return (
-            f"Error: unsafe code detected ({'; '.join(reasons)}). "
-            f"Please remove signal manipulation from your code."
-        )
+        shell_reasons = [
+            item.get("description", "") for item in info.get("shell_escapes", [])
+        ]
+        exception_reasons = [
+            item.get("description", "") for item in info.get("exception_catching", [])
+        ]
+        all_reasons = [r for r in reasons + shell_reasons + exception_reasons if r]
+        if all_reasons:
+            return (
+                f"Error: unsafe code detected ({'; '.join(all_reasons)}). "
+                f"Please remove unsafe patterns from your code."
+            )
 
     return None
 
@@ -662,13 +1101,20 @@ def _python_exec(
         with os.fdopen(fd, "w") as f:
             f.write(code)
 
-        proc = subprocess.Popen(
-            [sys.executable, tmp_path],
+        safe_env = _build_safe_env(workdir)
+        popen_kwargs = dict(
             stdout = subprocess.PIPE,
             stderr = subprocess.STDOUT,
             text = True,
             cwd = workdir,
+            env = safe_env,
         )
+        if sys.platform != "win32":
+            popen_kwargs["preexec_fn"] = _sandbox_preexec
+        else:
+            popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
+
+        proc = subprocess.Popen([sys.executable, tmp_path], **popen_kwargs)
 
         # Spawn cancel watcher if we have a cancel event
         if cancel_event is not None:
@@ -734,21 +1180,27 @@ def _bash_exec(
     if not command or not command.strip():
         return "No command provided."
 
-    # Block dangerous commands
-    tokens = set(command.lower().split())
-    blocked = tokens & _BASH_BLOCKED_WORDS
+    # Block dangerous commands (shlex + regex based)
+    blocked = _find_blocked_commands(command)
     if blocked:
         return f"Blocked command(s) for safety: {', '.join(sorted(blocked))}"
 
     try:
         workdir = _get_workdir(session_id)
-        proc = subprocess.Popen(
-            ["bash", "-c", command],
+        safe_env = _build_safe_env(workdir)
+        popen_kwargs = dict(
             stdout = subprocess.PIPE,
             stderr = subprocess.STDOUT,
             text = True,
             cwd = workdir,
+            env = safe_env,
         )
+        if sys.platform != "win32":
+            popen_kwargs["preexec_fn"] = _sandbox_preexec
+        else:
+            popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
+
+        proc = subprocess.Popen(_get_shell_cmd(command), **popen_kwargs)
 
         if cancel_event is not None:
             watcher = threading.Thread(

studio/backend/core/inference/worker.py

@@ -145,6 +145,8 @@ def _get_hf_download_state(
         blobs_dirs: list[Path] = []
 
         if model_names:
+            from utils.paths import resolve_cached_repo_id_case
+
             for name in model_names:
                 if not name:
                     continue
@@ -154,6 +156,7 @@ def _get_hf_download_state(
                 # relative paths, and Windows paths.
                 if name.startswith(("/", ".", "~")) or "\\" in name:
                     continue
+                name = resolve_cached_repo_id_case(name)
                 # HF cache dir format: models--org--name (slashes -> --)
                 cache_dir_name = "models--" + name.replace("/", "--")
                 blobs_dir = cache / cache_dir_name / "blobs"

studio/backend/models/inference.py

@@ -48,6 +48,10 @@ class LoadRequest(BaseModel):
         None,
         description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.",
     )
+    speculative_type: Optional[str] = Field(
+        None,
+        description = "Speculative decoding mode for GGUF models (e.g. 'ngram-simple', 'ngram-mod'). Ignored for non-GGUF and vision models.",
+    )
 
 
 class UnloadRequest(BaseModel):
@@ -163,6 +167,10 @@ class LoadResponse(BaseModel):
         None,
         description = "Jinja2 chat template string (from GGUF metadata or tokenizer)",
     )
+    speculative_type: Optional[str] = Field(
+        None,
+        description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled",
+    )
 
 
 class UnloadResponse(BaseModel):
@@ -225,6 +233,10 @@ class InferenceStatusResponse(BaseModel):
         None,
         description = "Model's native context length from GGUF metadata (not capped by VRAM)",
     )
+    speculative_type: Optional[str] = Field(
+        None,
+        description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled",
+    )
 
 
 # =====================================================================

studio/backend/routes/inference.py

@@ -179,6 +179,7 @@ async def load_model(
                     supports_reasoning = llama_backend.supports_reasoning,
                     reasoning_always_on = llama_backend.reasoning_always_on,
                     chat_template = llama_backend.chat_template,
+                    speculative_type = llama_backend.speculative_type,
                 )
         else:
             if (
@@ -263,6 +264,7 @@ async def load_model(
                     n_ctx = request.max_seq_length,
                     chat_template_override = request.chat_template_override,
                     cache_type_kv = request.cache_type_kv,
+                    speculative_type = request.speculative_type,
                 )
             else:
                 # Local mode: llama-server loads via -m <path>
@@ -275,6 +277,7 @@ async def load_model(
                     n_ctx = request.max_seq_length,
                     chat_template_override = request.chat_template_override,
                     cache_type_kv = request.cache_type_kv,
+                    speculative_type = request.speculative_type,
                 )
 
             if not success:
@@ -317,6 +320,7 @@ async def load_model(
                 supports_tools = llama_backend.supports_tools,
                 cache_type_kv = llama_backend.cache_type_kv,
                 chat_template = llama_backend.chat_template,
+                speculative_type = llama_backend.speculative_type,
             )
 
         # ── Standard path: load via Unsloth/transformers ──────────
@@ -652,6 +656,7 @@ async def get_status(
                 context_length = llama_backend.context_length,
                 max_context_length = llama_backend.max_context_length,
                 native_context_length = llama_backend.native_context_length,
+                speculative_type = llama_backend.speculative_type,
             )
 
         # Otherwise, report Unsloth backend status

studio/backend/routes/models.py

@@ -49,8 +49,10 @@ try:
     )
     from core.inference import get_inference_backend
     from utils.paths import (
+        is_local_path,
         outputs_root,
         exports_root,
+        resolve_cached_repo_id_case,
         resolve_output_dir,
         resolve_export_dir,
     )
@@ -77,8 +79,10 @@ except ImportError:
     )
     from core.inference import get_inference_backend
     from utils.paths import (
+        is_local_path,
         outputs_root,
         exports_root,
+        resolve_cached_repo_id_case,
         resolve_output_dir,
         resolve_export_dir,
     )
@@ -597,10 +601,15 @@ async def get_model_config(
     This endpoint wraps the backend load_model_defaults function.
     """
     try:
-        from utils.models.model_config import is_local_path
-
         if not is_local_path(model_name):
-            model_name = model_name.lower()
+            resolved = resolve_cached_repo_id_case(model_name)
+            if resolved != model_name:
+                logger.info(
+                    "Using cached repo_id casing '%s' for requested '%s'",
+                    resolved,
+                    model_name,
+                )
+            model_name = resolved
 
         logger.info(f"Getting model config for: {model_name}")
         from utils.models.model_config import detect_audio_type

studio/backend/tests/test_cache_case_resolution.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+from pathlib import Path
+import sys
+import types
+
+# Keep this test runnable in lightweight environments where optional logging
+# deps are not installed.
+if "structlog" not in sys.modules:
+
+    class _DummyLogger:
+        def __getattr__(self, _name):
+            return lambda *args, **kwargs: None
+
+    sys.modules["structlog"] = types.SimpleNamespace(
+        BoundLogger = _DummyLogger,
+        get_logger = lambda *args, **kwargs: _DummyLogger(),
+    )
+
+from utils.paths.path_utils import (
+    resolve_cached_repo_id_case,
+    get_cache_case_resolution_stats,
+    reset_cache_case_resolution_state,
+)
+import utils.paths.path_utils as path_utils
+
+
+def _mk_cache_repo(cache_root: Path, repo_id: str) -> Path:
+    repo_dir = cache_root / f"models--{repo_id.replace('/', '--')}"
+    repo_dir.mkdir(parents = True, exist_ok = True)
+    return repo_dir
+
+
+def test_resolve_cached_repo_id_case_exact_hit(tmp_path, monkeypatch):
+    reset_cache_case_resolution_state()
+    _mk_cache_repo(tmp_path, "Org/Model")
+    monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path)
+
+    resolved = resolve_cached_repo_id_case("Org/Model")
+
+    assert resolved == "Org/Model"
+    stats = get_cache_case_resolution_stats()
+    assert stats["calls"] == 1
+    assert stats["exact_hits"] == 1
+    assert stats["variant_hits"] == 0
+
+
+def test_resolve_cached_repo_id_case_variant_hit(tmp_path, monkeypatch):
+    reset_cache_case_resolution_state()
+    _mk_cache_repo(tmp_path, "Org/Model")
+    monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path)
+
+    resolved = resolve_cached_repo_id_case("org/model")
+
+    assert resolved == "Org/Model"
+    stats = get_cache_case_resolution_stats()
+    assert stats["variant_hits"] == 1
+    assert stats["tie_breaks"] == 0
+
+
+def test_resolve_cached_repo_id_case_tie_break_deterministic(tmp_path, monkeypatch):
+    reset_cache_case_resolution_state()
+    _mk_cache_repo(tmp_path, "Org/Model")
+    _mk_cache_repo(tmp_path, "org/model")
+    monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path)
+
+    resolved = resolve_cached_repo_id_case("oRg/mOdEl")
+
+    # Deterministic rule: lexical sort of candidate repo ids.
+    assert resolved == "Org/Model"
+    stats = get_cache_case_resolution_stats()
+    assert stats["variant_hits"] == 1
+    assert stats["tie_breaks"] == 1
+
+
+def test_resolve_cached_repo_id_case_no_cache_fallback(tmp_path, monkeypatch):
+    reset_cache_case_resolution_state()
+    monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path)
+
+    resolved = resolve_cached_repo_id_case("Org/Missing")
+
+    assert resolved == "Org/Missing"
+    stats = get_cache_case_resolution_stats()
+    assert stats["fallbacks"] == 1
+    assert stats["variant_hits"] == 0
+    assert stats["exact_hits"] == 0
+
+
+def test_resolve_cached_repo_id_case_memoization(tmp_path, monkeypatch):
+    reset_cache_case_resolution_state()
+    _mk_cache_repo(tmp_path, "Org/Model")
+    monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path)
+
+    first = resolve_cached_repo_id_case("org/model")
+    second = resolve_cached_repo_id_case("org/model")
+
+    assert first == "Org/Model"
+    assert second == "Org/Model"
+    stats = get_cache_case_resolution_stats()
+    assert stats["calls"] == 2
+    assert stats["variant_hits"] == 1
+    assert stats["memo_hits"] == 1
+
+
+def test_resolve_cached_repo_id_case_late_cache_population(tmp_path, monkeypatch):
+    """Regression guard: memoized fallback should not hide a later cache variant."""
+    reset_cache_case_resolution_state()
+    monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path)
+
+    first = resolve_cached_repo_id_case("org/model")
+    assert first == "org/model"
+
+    # Simulate cache being populated after first miss (e.g. another code path/download).
+    _mk_cache_repo(tmp_path, "Org/Model")
+
+    second = resolve_cached_repo_id_case("org/model")
+
+    # Desired behavior: second lookup should pick up the now-existing variant.
+    assert second == "Org/Model"

studio/backend/tests/test_models_get_model_config_case_resolution.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import asyncio
+import sys
+import types
+
+# Keep this test runnable in lightweight environments where optional logging
+# deps are not installed.
+if "structlog" not in sys.modules:
+
+    class _DummyLogger:
+        def __getattr__(self, _name):
+            return lambda *args, **kwargs: None
+
+    sys.modules["structlog"] = types.SimpleNamespace(
+        BoundLogger = _DummyLogger,
+        get_logger = lambda *args, **kwargs: _DummyLogger(),
+    )
+
+import routes.models as models_route
+import utils.models.model_config as model_config_module
+
+
+def test_get_model_config_resolves_cached_case_before_model_checks(monkeypatch):
+    calls: dict[str, str] = {}
+
+    class _DummyModelConfig:
+        is_lora = False
+        base_model = None
+
+    def _record_load(model_name):
+        calls["load_model_defaults"] = model_name
+        return {}
+
+    def _record_vision(model_name, hf_token = None):
+        calls["is_vision_model"] = model_name
+        return False
+
+    def _record_embedding(model_name, hf_token = None):
+        calls["is_embedding_model"] = model_name
+        return False
+
+    def _record_audio(model_name, hf_token = None):
+        calls["detect_audio_type"] = model_name
+        return None
+
+    def _record_from_identifier(cls, model_name):
+        calls["from_identifier"] = model_name
+        return _DummyModelConfig()
+
+    monkeypatch.setattr(models_route, "is_local_path", lambda _: False)
+    monkeypatch.setattr(
+        models_route, "resolve_cached_repo_id_case", lambda _: "Org/Model"
+    )
+    monkeypatch.setattr(models_route, "load_model_defaults", _record_load)
+    monkeypatch.setattr(models_route, "is_vision_model", _record_vision)
+    monkeypatch.setattr(models_route, "is_embedding_model", _record_embedding)
+    monkeypatch.setattr(model_config_module, "detect_audio_type", _record_audio)
+    monkeypatch.setattr(
+        models_route.ModelConfig,
+        "from_identifier",
+        classmethod(_record_from_identifier),
+    )
+    monkeypatch.setattr(models_route, "_get_max_position_embeddings", lambda _: 4096)
+    monkeypatch.setattr(models_route, "_get_model_size_bytes", lambda *_args, **_kw: 0)
+
+    result = asyncio.run(
+        models_route.get_model_config(
+            model_name = "org/model",
+            hf_token = None,
+            current_subject = "test-subject",
+        )
+    )
+
+    assert result.model_name == "Org/Model"
+    assert calls["load_model_defaults"] == "Org/Model"
+    assert calls["is_vision_model"] == "Org/Model"
+    assert calls["is_embedding_model"] == "Org/Model"
+    assert calls["detect_audio_type"] == "Org/Model"
+    assert calls["from_identifier"] == "Org/Model"

studio/backend/utils/models/model_config.py

@@ -11,6 +11,8 @@ from utils.paths import (
     normalize_path,
     is_local_path,
     is_model_cached,
+    get_cache_path,
+    resolve_cached_repo_id_case,
     outputs_root,
     exports_root,
     resolve_output_dir,
@@ -157,6 +159,38 @@ MODEL_NAME_MAPPING = {
         "unsloth/gemma-3n-E4B-unsloth-bnb-4bit",
         "google/gemma-3n-E4B",
     ],
+    "unsloth_gemma-4-31B-it.yaml": [
+        "unsloth/gemma-4-31B-it",
+        "google/gemma-4-31B-it",
+    ],
+    "unsloth_gemma-4-26B-A4B-it.yaml": [
+        "unsloth/gemma-4-26B-A4B-it",
+        "google/gemma-4-26B-A4B-it",
+    ],
+    "unsloth_gemma-4-E2B-it.yaml": [
+        "unsloth/gemma-4-E2B-it",
+        "google/gemma-4-E2B-it",
+    ],
+    "unsloth_gemma-4-E4B-it.yaml": [
+        "unsloth/gemma-4-E4B-it",
+        "google/gemma-4-E4B-it",
+    ],
+    "unsloth_gemma-4-31B.yaml": [
+        "unsloth/gemma-4-31B",
+        "google/gemma-4-31B",
+    ],
+    "unsloth_gemma-4-26B-A4B.yaml": [
+        "unsloth/gemma-4-26B-A4B",
+        "google/gemma-4-26B-A4B",
+    ],
+    "unsloth_gemma-4-E2B.yaml": [
+        "unsloth/gemma-4-E2B",
+        "google/gemma-4-E2B",
+    ],
+    "unsloth_gemma-4-E4B.yaml": [
+        "unsloth/gemma-4-E4B",
+        "google/gemma-4-E4B",
+    ],
     "unsloth_gpt-oss-20b.yaml": [
         "openai/gpt-oss-20b",
         "unsloth/gpt-oss-20b-unsloth-bnb-4bit",
@@ -711,12 +745,8 @@ def _detect_audio_from_tokenizer(
 
     # 1) Check local HF cache first (works for gated/offline models)
     try:
-        from huggingface_hub.constants import HF_HUB_CACHE
-
-        cache_dir = Path(HF_HUB_CACHE)
-        repo_dir_name = f"models--{model_name.replace('/', '--')}"
-        repo_dir = cache_dir / repo_dir_name
-        if repo_dir.exists():
+        repo_dir = get_cache_path(model_name)
+        if repo_dir is not None and repo_dir.exists():
             snapshots_dir = repo_dir / "snapshots"
             if snapshots_dir.exists():
                 for snapshot in snapshots_dir.iterdir():
@@ -1627,11 +1657,18 @@ class ModelConfig:
             identifier = f"unsloth/{identifier}"
             path = identifier
 
-        # Enforce lowercase for remote Hugging Face identifiers to prevent cache duplication
-        # Hugging Face Hub APIs are case-insensitive remotely, but case-sensitive locally (repo_folder_name).
+        # Preserve requested casing, but if a case-variant already exists in local HF cache,
+        # reuse that exact repo_id spelling to avoid one-time re-downloads after #2592.
         if not is_local:
-            identifier = identifier.lower()
-            path = path.lower()
+            resolved_identifier = resolve_cached_repo_id_case(identifier)
+            if resolved_identifier != identifier:
+                logger.info(
+                    "Using cached repo_id casing '%s' for requested '%s'",
+                    resolved_identifier,
+                    identifier,
+                )
+                identifier = resolved_identifier
+                path = resolved_identifier
 
         # Auto-detect GGUF models (check before LoRA/vision detection)
         if is_local:
@@ -1852,6 +1889,12 @@ class ModelConfig:
             identifier = f"unsloth/{identifier}"
             path = identifier
 
+        if not is_local:
+            resolved_identifier = resolve_cached_repo_id_case(identifier)
+            if resolved_identifier != identifier:
+                identifier = resolved_identifier
+                path = resolved_identifier
+
         # --- Logic for Base Model and Vision Detection ---
         base_model = None
         is_vision = False

studio/backend/utils/paths/__init__.py

@@ -5,7 +5,15 @@
 Path utilities for model and dataset handling
 """
 
-from .path_utils import normalize_path, is_local_path, is_model_cached, get_cache_path
+from .path_utils import (
+    normalize_path,
+    is_local_path,
+    is_model_cached,
+    get_cache_path,
+    resolve_cached_repo_id_case,
+    get_cache_case_resolution_stats,
+    reset_cache_case_resolution_state,
+)
 from .storage_roots import (
     studio_root,
     assets_root,
@@ -40,6 +48,9 @@ __all__ = [
     "is_local_path",
     "is_model_cached",
     "get_cache_path",
+    "resolve_cached_repo_id_case",
+    "get_cache_case_resolution_stats",
+    "reset_cache_case_resolution_state",
     "studio_root",
     "assets_root",
     "datasets_root",

studio/backend/utils/paths/path_utils.py

@@ -14,6 +14,20 @@ from loggers import get_logger
 
 logger = get_logger(__name__)
 
+# Per-process cache to avoid repeated cache-dir scans for the same identifier.
+_CACHE_CASE_RESOLUTION_MEMO: dict[str, str] = {}
+
+# Lightweight instrumentation counters for operational visibility.
+_CACHE_CASE_RESOLUTION_STATS: dict[str, int] = {
+    "calls": 0,
+    "memo_hits": 0,
+    "exact_hits": 0,
+    "variant_hits": 0,
+    "tie_breaks": 0,
+    "fallbacks": 0,
+    "errors": 0,
+}
+
 
 def _is_wsl() -> bool:
     """Detect if we are running inside WSL (Windows Subsystem for Linux)."""
@@ -94,8 +108,9 @@ def is_local_path(path: str) -> bool:
 
 def get_cache_path(model_name: str) -> Optional[Path]:
     """Get HuggingFace cache path for a model if it exists."""
-    cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
-    model_cache_name = model_name.replace("/", "--")
+    cache_dir = _hf_hub_cache_dir()
+    resolved_name = resolve_cached_repo_id_case(model_name)
+    model_cache_name = resolved_name.replace("/", "--")
     model_cache_path = cache_dir / f"models--{model_cache_name}"
 
     return model_cache_path if model_cache_path.exists() else None
@@ -113,3 +128,102 @@ def is_model_cached(model_name: str) -> bool:
             return True
 
     return False
+
+
+def _hf_hub_cache_dir() -> Path:
+    """Return HF cache root honoring HF_HUB_CACHE when available."""
+    try:
+        from huggingface_hub.constants import HF_HUB_CACHE
+
+        return Path(HF_HUB_CACHE)
+    except Exception as exc:
+        logger.debug(
+            "Could not read huggingface_hub HF_HUB_CACHE, using default hub path: %s",
+            exc,
+        )
+        return Path.home() / ".cache" / "huggingface" / "hub"
+
+
+def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str:
+    """Resolve repo_id to the exact casing already present in local HF cache.
+
+    Policy: prefer the requested/canonical repo_id, but if a case-variant already
+    exists in local HF cache, reuse that exact cached spelling. This avoids
+    duplicate downloads while preserving user intent whenever possible.
+    """
+    _CACHE_CASE_RESOLUTION_STATS["calls"] += 1
+
+    if not model_name or "/" not in model_name:
+        _CACHE_CASE_RESOLUTION_STATS["fallbacks"] += 1
+        return model_name
+
+    cache_dir = _hf_hub_cache_dir()
+    if not cache_dir.exists():
+        _CACHE_CASE_RESOLUTION_STATS["fallbacks"] += 1
+        return model_name
+
+    expected_dir = f"models--{model_name.replace('/', '--')}"
+
+    # Always check the exact-case path first so a newly-appeared exact match
+    # wins over any previously memoized variant.
+    exact_path = cache_dir / expected_dir
+    if exact_path.is_dir():
+        if use_memo:
+            _CACHE_CASE_RESOLUTION_MEMO[model_name] = model_name
+        _CACHE_CASE_RESOLUTION_STATS["exact_hits"] += 1
+        return model_name
+
+    # Validate memoized entries still exist on disk before returning them.
+    # This prevents stale results when cache dirs are deleted/recreated.
+    if use_memo:
+        cached = _CACHE_CASE_RESOLUTION_MEMO.get(model_name)
+        if cached is not None:
+            cached_path = cache_dir / f"models--{cached.replace('/', '--')}"
+            if cached_path.is_dir():
+                _CACHE_CASE_RESOLUTION_STATS["memo_hits"] += 1
+                return cached
+            # Stale entry -- drop it and re-scan below.
+            _CACHE_CASE_RESOLUTION_MEMO.pop(model_name, None)
+
+    expected_lower = expected_dir.lower()
+    try:
+        candidates: list[str] = []
+        for entry in cache_dir.iterdir():
+            if not entry.is_dir():
+                continue
+            if entry.name.lower() != expected_lower:
+                continue
+            if not entry.name.startswith("models--"):
+                continue
+            repo_part = entry.name[len("models--") :]
+            if not repo_part:
+                continue
+            candidates.append(repo_part.replace("--", "/"))
+
+        if candidates:
+            # Deterministic tie-break if multiple case variants coexist.
+            resolved = sorted(candidates)[0]
+            if len(candidates) > 1:
+                _CACHE_CASE_RESOLUTION_STATS["tie_breaks"] += 1
+            _CACHE_CASE_RESOLUTION_STATS["variant_hits"] += 1
+            if use_memo:
+                _CACHE_CASE_RESOLUTION_MEMO[model_name] = resolved
+            return resolved
+    except Exception as exc:
+        _CACHE_CASE_RESOLUTION_STATS["errors"] += 1
+        logger.debug(f"Could not resolve cached repo_id case for '{model_name}': {exc}")
+
+    _CACHE_CASE_RESOLUTION_STATS["fallbacks"] += 1
+    return model_name
+
+
+def get_cache_case_resolution_stats() -> dict[str, int]:
+    """Return a copy of case-resolution instrumentation counters."""
+    return dict(_CACHE_CASE_RESOLUTION_STATS)
+
+
+def reset_cache_case_resolution_state() -> None:
+    """Clear resolver memo and counters (primarily for tests)."""
+    _CACHE_CASE_RESOLUTION_MEMO.clear()
+    for key in _CACHE_CASE_RESOLUTION_STATS:
+        _CACHE_CASE_RESOLUTION_STATS[key] = 0

studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx

@@ -52,6 +52,18 @@ const WebSearchToolUIImpl: ToolCallMessagePartComponent = ({
   status,
 }) => {
   const query = (args as { query?: string })?.query ?? "";
+  const url = ((args as { url?: string })?.url ?? "").trim();
+  const isUrlFetch = !!url;
+  const displayDomain = (() => {
+    if (!url) return "";
+    try {
+      const parsed = new URL(url);
+      if (parsed.protocol !== "http:" && parsed.protocol !== "https:") return "";
+      return parsed.hostname.replace(/^www\./, "");
+    } catch {
+      return "";
+    }
+  })();
   const isRunning = status?.type === "running";
   const sources = result
     ? parseSearchResults(
@@ -75,7 +87,13 @@ const WebSearchToolUIImpl: ToolCallMessagePartComponent = ({
   return (
     <ToolFallbackRoot open={open} onOpenChange={setOpen}>
       <ToolFallbackTrigger
-        toolName={query ? `Searched "${query}"` : "Web Search"}
+        toolName={
+          isUrlFetch
+            ? displayDomain ? `Read ${displayDomain}` : "Read page"
+            : query
+              ? `Searched "${query}"`
+              : "Web Search"
+        }
         status={status}
         icon={GlobeIcon}
       />
@@ -83,7 +101,12 @@ const WebSearchToolUIImpl: ToolCallMessagePartComponent = ({
         {isRunning ? (
           <div className="flex items-center gap-2 text-sm text-muted-foreground">
             <LoaderIcon className="size-3.5 animate-spin" />
-            <span>Searching for &ldquo;{query}&rdquo;&hellip;</span>
+            <span>
+              {isUrlFetch
+                ? <>Reading {displayDomain || "page"}&hellip;</>
+                : <>Searching for &ldquo;{query}&rdquo;&hellip;</>
+              }
+            </span>
           </div>
         ) : sources.length > 0 ? (
           <div className="flex flex-wrap gap-1.5">

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

@@ -421,6 +421,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
   return {
     async *run({ messages, abortSignal, unstable_threadId }) {
       let runtime = useChatRuntimeStore.getState();
+      // Capture the thread ID once at the start so it stays stable even if
+      // the user switches chats while waiting for model load / auto-load.
+      const resolvedThreadId =
+        (unstable_threadId ?? runtime.activeThreadId) || undefined;
 
       // Wait for in-progress model load to finish before inferring
       if (runtime.modelLoading) {
@@ -473,14 +477,14 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         }
         runtime.clearPendingAudio();
       }
-      const useAdapter = await resolveUseAdapter(unstable_threadId);
+      const useAdapter = await resolveUseAdapter(resolvedThreadId);
 
       // ── Audio model path (non-streaming) ─────────────────────
       const activeModel = runtime.models.find(
         (m) => m.id === params.checkpoint,
       );
       if (activeModel?.isAudio && !activeModel?.hasAudioInput) {
-        const threadKey = unstable_threadId || "__default";
+        const threadKey = resolvedThreadId || "__default";
         runtime.setThreadRunning(threadKey, true);
         try {
           yield {
@@ -527,7 +531,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         return;
       }
 
-      const threadKey = unstable_threadId || "__default";
+      const threadKey = resolvedThreadId || "__default";
       let waitingFirstChunk = true;
       let firstTokenSettled = false;
       const streamStartTime = Date.now();
@@ -600,7 +604,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                     const mins = useChatRuntimeStore.getState().toolCallTimeout;
                     return mins >= 9999 ? 9999 : mins * 60;
                   })(),
-                  session_id: unstable_threadId || undefined,
+                  session_id: resolvedThreadId,
                 }
               : {}),
           },
@@ -641,7 +645,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                 let parsedResult: string | { text: string; images: string[]; sessionId: string };
                 if (imgIdx !== -1) {
                   const text = rawResult.slice(0, imgIdx);
-                  const sessionId = unstable_threadId || "";
+                  // Fall back to "_default" to match the backend sandbox directory
+                  // used when no session_id is provided (see tools.py _get_workdir).
+                  const sessionId = resolvedThreadId || "_default";
                   try {
                     const images = JSON.parse(rawResult.slice(imgIdx + imgMarker.length)) as string[];
                     parsedResult = { text, images, sessionId };

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

@@ -592,6 +592,9 @@ export function ChatPage(): ReactElement {
   }, []);
   const handleNewCompare = useCallback(() => {
     setView({ mode: "compare", pairId: crypto.randomUUID() });
+    // Clear activeThreadId so compare panes do not inherit the single-chat
+    // thread ID as a fallback for session_id routing.
+    useChatRuntimeStore.getState().setActiveThreadId(null);
     useChatRuntimeStore.getState().setContextUsage(null);
   }, []);
 
@@ -619,6 +622,9 @@ export function ChatPage(): ReactElement {
   const enterCompare = useCallback(() => {
     setViewBeforeCompare((prev) => prev ?? view);
     setView({ mode: "compare", pairId: crypto.randomUUID() });
+    // Clear activeThreadId so compare panes do not inherit the single-chat
+    // thread ID as a fallback for session_id routing.
+    useChatRuntimeStore.getState().setActiveThreadId(null);
     useChatRuntimeStore.getState().setContextUsage(null);
   }, [view]);
 
@@ -626,9 +632,13 @@ export function ChatPage(): ReactElement {
     if (!viewBeforeCompare) return;
     setView(viewBeforeCompare);
     setViewBeforeCompare(null);
-    // Restore context usage from the active thread's last assistant message
+    // Restore context usage from the active thread's last assistant message.
+    // Use the thread ID from the saved view rather than the store, because
+    // activeThreadId may have been cleared on compare entry.
     const store = useChatRuntimeStore.getState();
-    const threadId = store.activeThreadId;
+    const threadId =
+      ("threadId" in viewBeforeCompare ? viewBeforeCompare.threadId : null) ??
+      store.activeThreadId;
     if (threadId) {
       void db.messages
         .where("threadId")
@@ -735,6 +745,7 @@ export function ChatPage(): ReactElement {
           await selectModelRef.current({ id: targetLora.id, isLora: true });
           if (canceled) return;
           setView({ mode: "compare", pairId: crypto.randomUUID() });
+          useChatRuntimeStore.getState().setActiveThreadId(null);
           useChatRuntimeStore.getState().setContextUsage(null);
           clearHandoff();
           console.info("[chat-handoff] loaded lora + opened compare");

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

@@ -280,6 +280,15 @@ export function ChatSettingsPanel({
 }: ChatSettingsPanelProps) {
   const isMobile = useIsMobile();
   const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null;
+  const speculativeType = useChatRuntimeStore((s) => s.speculativeType);
+  const setSpeculativeType = useChatRuntimeStore((s) => s.setSpeculativeType);
+  const loadedSpeculativeType = useChatRuntimeStore(
+    (s) => s.loadedSpeculativeType,
+  );
+  const currentModels = useChatRuntimeStore((s) => s.models);
+  const currentCheckpoint = params.checkpoint;
+  const currentModelIsVision =
+    currentModels.find((m) => m.id === currentCheckpoint)?.isVision ?? false;
   const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
   const ggufMaxContextLength = useChatRuntimeStore(
     (s) => s.ggufMaxContextLength,
@@ -299,7 +308,8 @@ export function ChatSettingsPanel({
   const ctxMaxValue = ggufNativeContextLength ?? ggufContextLength ?? null;
   const kvDirty = kvCacheDtype !== loadedKvCacheDtype;
   const ctxDirty = customContextLength !== null;
-  const modelSettingsDirty = kvDirty || ctxDirty;
+  const specDirty = speculativeType !== loadedSpeculativeType;
+  const modelSettingsDirty = kvDirty || ctxDirty || specDirty;
   const [customPresets, setCustomPresets] = useState<Preset[]>(() =>
     loadSavedCustomPresets(),
   );
@@ -580,6 +590,32 @@ export function ChatSettingsPanel({
                     </SelectContent>
                   </Select>
                 </div>
+                {!currentModelIsVision && (
+                  <div className="flex items-center justify-between gap-3">
+                    <div className="min-w-0">
+                      <div className="text-xs font-medium">
+                        Speculative Decoding
+                      </div>
+                      <div className="text-[11px] text-muted-foreground">
+                        Speed up generation with no VRAM cost.
+                      </div>
+                    </div>
+                    <Select
+                      value={speculativeType ?? "off"}
+                      onValueChange={(v) => {
+                        setSpeculativeType(v === "off" ? null : v);
+                      }}
+                    >
+                      <SelectTrigger className="h-7 w-[120px] text-xs">
+                        <SelectValue />
+                      </SelectTrigger>
+                      <SelectContent>
+                        <SelectItem value="ngram-mod">On</SelectItem>
+                        <SelectItem value="off">Off</SelectItem>
+                      </SelectContent>
+                    </Select>
+                  </div>
+                )}
                 {modelSettingsDirty && (
                   <div className="flex flex-wrap gap-1.5 pt-1">
                     <button
@@ -594,6 +630,7 @@ export function ChatSettingsPanel({
                       onClick={() => {
                         setCustomContextLength(null);
                         setKvCacheDtype(loadedKvCacheDtype);
+                        setSpeculativeType(loadedSpeculativeType);
                       }}
                       className="rounded-md border px-2.5 py-1 text-[11px] font-medium text-muted-foreground transition-colors hover:bg-accent"
                     >

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

@@ -250,6 +250,7 @@ export function useChatModelRuntime() {
         const ggufNativeContextLength = statusRes.is_gguf
           ? (statusRes.native_context_length ?? null)
           : null;
+        const currentSpecType = statusRes.speculative_type ?? null;
         useChatRuntimeStore.setState({
           supportsReasoning,
           reasoningAlwaysOn,
@@ -257,6 +258,8 @@ export function useChatModelRuntime() {
           ggufContextLength: currentGgufContextLength,
           ggufMaxContextLength,
           ggufNativeContextLength,
+          speculativeType: currentSpecType,
+          loadedSpeculativeType: currentSpecType,
         });
 
         // Set reasoning default for Qwen3.5 small models
@@ -393,7 +396,7 @@ export function useChatModelRuntime() {
               previousWasUnloaded = true;
             }
 
-            const { chatTemplateOverride, kvCacheDtype, customContextLength, ggufContextLength } = useChatRuntimeStore.getState();
+            const { chatTemplateOverride, kvCacheDtype, customContextLength, ggufContextLength, speculativeType } = useChatRuntimeStore.getState();
             // GGUF: use custom context length, or 0 = model's native context
             // Non-GGUF: use the Max Seq Length slider value
             const effectiveMaxSeqLength = customContextLength != null
@@ -409,6 +412,7 @@ export function useChatModelRuntime() {
               trust_remote_code: paramsBeforeLoad.trustRemoteCode ?? false,
               chat_template_override: chatTemplateOverride,
               cache_type_kv: kvCacheDtype,
+              speculative_type: speculativeType,
             });
 
             // If cancelled while loading, don't update UI to show
@@ -431,6 +435,7 @@ export function useChatModelRuntime() {
               }
             }
             const loadedKv = loadResponse.cache_type_kv ?? null;
+            const loadedSpec = loadResponse.speculative_type ?? null;
             const nativeCtx = loadResponse.is_gguf
               ? (loadResponse.context_length ?? 131072)
               : null;
@@ -457,6 +462,8 @@ export function useChatModelRuntime() {
               codeToolsEnabled: loadResponse.supports_tools ?? false,
               kvCacheDtype: loadedKv,
               loadedKvCacheDtype: loadedKv,
+              speculativeType: loadedSpec,
+              loadedSpeculativeType: loadedSpec,
               customContextLength: keepCustomCtx,
               defaultChatTemplate: loadResponse.chat_template ?? null,
               chatTemplateOverride: null,

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

@@ -679,9 +679,33 @@ function ThreadNewChatSwitch({
   const isLoading = useAuiState(({ threads }) => threads.isLoading);
 
   useEffect(() => {
-    if (!isLoading) {
-      aui.threads().switchToNewThread();
+    if (isLoading) {
+      return;
     }
+
+    let cancelled = false;
+    // Clear immediately so the adapter never picks up a stale thread ID
+    // from a previous chat while we initialize the new one.
+    useChatRuntimeStore.getState().setActiveThreadId(null);
+
+    void (async () => {
+      try {
+        aui.threads().switchToNewThread();
+        const { remoteId } = await aui.threadListItem().initialize();
+        if (!cancelled) {
+          useChatRuntimeStore.getState().setActiveThreadId(remoteId);
+        }
+      } catch (error) {
+        if (!cancelled) {
+          useChatRuntimeStore.getState().setActiveThreadId(null);
+        }
+        console.error("Failed to initialize new chat thread", error);
+      }
+    })();
+
+    return () => {
+      cancelled = true;
+    };
   }, [aui, isLoading, nonce]);
 
   return null;
@@ -730,7 +754,7 @@ export function ChatRuntimeProvider({
 
   return (
     <AssistantRuntimeProvider runtime={runtime} aui={aui}>
-      <ActiveThreadSync enabled={modelType === "base" && !pairId} />
+      <ActiveThreadSync enabled={modelType === "base" && !pairId && !newThreadNonce} />
       {initialThreadId && <ThreadAutoSwitch threadId={initialThreadId} />}
       {!initialThreadId && newThreadNonce && (
         <ThreadNewChatSwitch nonce={newThreadNonce} />

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

@@ -165,6 +165,8 @@ type ChatRuntimeStore = {
   toolCallTimeout: number;
   kvCacheDtype: string | null;
   loadedKvCacheDtype: string | null;
+  speculativeType: string | null;
+  loadedSpeculativeType: string | null;
   customContextLength: number | null;
   defaultChatTemplate: string | null;
   chatTemplateOverride: string | null;
@@ -198,6 +200,7 @@ type ChatRuntimeStore = {
   setMaxToolCallsPerMessage: (value: number) => void;
   setToolCallTimeout: (value: number) => void;
   setKvCacheDtype: (dtype: string | null) => void;
+  setSpeculativeType: (type: string | null) => void;
   setCustomContextLength: (v: number | null) => void;
   setChatTemplateOverride: (template: string | null) => void;
   setPendingAudio: (base64: string, name: string) => void;
@@ -230,6 +233,8 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
   toolCallTimeout: loadInt(TOOL_CALL_TIMEOUT_KEY, 5),
   kvCacheDtype: null,
   loadedKvCacheDtype: null,
+  speculativeType: "ngram-mod",
+  loadedSpeculativeType: null,
   customContextLength: null,
   defaultChatTemplate: null,
   chatTemplateOverride: null,
@@ -302,6 +307,8 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
       toolStatus: null,
       kvCacheDtype: null,
       loadedKvCacheDtype: null,
+      speculativeType: "ngram-mod",
+      loadedSpeculativeType: null,
       customContextLength: null,
       defaultChatTemplate: null,
       chatTemplateOverride: null,
@@ -327,6 +334,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
       return { toolCallTimeout };
     }),
   setKvCacheDtype: (kvCacheDtype) => set({ kvCacheDtype }),
+  setSpeculativeType: (speculativeType) => set({ speculativeType }),
   setCustomContextLength: (customContextLength) => set({ customContextLength }),
   setChatTemplateOverride: (chatTemplateOverride) => set({ chatTemplateOverride }),
   setPendingAudio: (base64, name) =>

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

@@ -41,6 +41,7 @@ export interface LoadModelRequest {
   trust_remote_code?: boolean;
   chat_template_override?: string | null;
   cache_type_kv?: string | null;
+  speculative_type?: string | null;
 }
 
 export interface ValidateModelResponse {
@@ -93,6 +94,7 @@ export interface LoadModelResponse {
   supports_tools?: boolean;
   cache_type_kv?: string | null;
   chat_template?: string | null;
+  speculative_type?: string | null;
 }
 
 export interface UnloadModelRequest {
@@ -123,6 +125,7 @@ export interface InferenceStatusResponse {
   context_length?: number | null;
   max_context_length?: number | null;
   native_context_length?: number | null;
+  speculative_type?: string | null;
 }
 
 export interface AudioGenerationResponse {

studio/frontend/src/features/recipe-studio/components/executions/execution-data-tab.tsx

@@ -12,6 +12,7 @@ import {
   DropdownMenuLabel,
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
+import { Spinner } from "@/components/ui/spinner";
 import { cn } from "@/lib/utils";
 import { isExecutionInProgress } from "../../executions/execution-helpers";
 import type { RecipeExecutionRecord } from "../../execution-types";
@@ -121,7 +122,21 @@ export function ExecutionDataTab({
         </div>
       </div>
       {execution.dataset.length === 0 ? (
-        <p className="text-xs text-muted-foreground">No rows returned.</p>
+        isExecutionInProgress(execution.status) ? (
+          <div className="flex flex-col items-center justify-center gap-3 py-12 text-center">
+            <Spinner className="size-5" />
+            <div className="space-y-1">
+              <p className="text-sm font-medium text-muted-foreground">
+                Generating data…
+              </p>
+              <p className="text-xs text-muted-foreground">
+                Check the Overview tab for live terminal logs.
+              </p>
+            </div>
+          </div>
+        ) : (
+          <p className="text-xs text-muted-foreground">No rows returned.</p>
+        )
       ) : tableColumns.length === 0 ? (
         <p className="text-xs text-muted-foreground">
           All columns hidden. Use Columns to show at least one.

studio/frontend/src/features/recipe-studio/components/executions/execution-overview-tab.tsx

@@ -116,10 +116,12 @@ export function ExecutionOverviewTab({
                 />
               </div>
               <div className="space-y-1.5 text-xs">
-                <p className="flex items-center justify-between gap-3">
-                  <span className="text-muted-foreground">LLM columns</span>
-                  <span className="font-semibold">{formatMetricValue(llmColumnCount)}</span>
-                </p>
+                {llmColumnCount > 0 && (
+                  <p className="flex items-center justify-between gap-3">
+                    <span className="text-muted-foreground">LLM columns</span>
+                    <span className="font-semibold">{formatMetricValue(llmColumnCount)}</span>
+                  </p>
+                )}
                 <p className="flex items-center justify-between gap-3">
                   <span className="text-muted-foreground">Null rate</span>
                   <span className="font-semibold">{nullRate?.toFixed(1) ?? "--"}%</span>
@@ -164,40 +166,42 @@ export function ExecutionOverviewTab({
               </div>
             </div>
           </div>
-          <div className="rounded-xl border border-border/60 bg-card/55 p-3">
-            <div className="mb-2 flex items-center justify-between">
-              <p className="text-xs text-muted-foreground">Model usage</p>
-              <HugeiconsIcon icon={Flag02Icon} className="size-4 text-muted-foreground" />
-            </div>
-            {modelUsageRows.length === 0 ? (
-              <p className="text-xs text-muted-foreground">No model usage yet.</p>
-            ) : (
-              <div className="overflow-hidden rounded-lg border border-border/60 bg-card/50">
-                <Table>
-                  <TableHeader>
-                    <TableRow>
-                      <TableHead>Model</TableHead>
-                      <TableHead className="text-right">Input</TableHead>
-                      <TableHead className="text-right">Output</TableHead>
-                    </TableRow>
-                  </TableHeader>
-                  <TableBody>
-                    {modelUsageRows.map((usage) => (
-                      <TableRow key={usage.model}>
-                        <TableCell className="max-w-[320px] truncate">{usage.model}</TableCell>
-                        <TableCell className="text-right">
-                          {formatMetricValue(usage.input)}
-                        </TableCell>
-                        <TableCell className="text-right">
-                          {formatMetricValue(usage.output)}
-                        </TableCell>
-                      </TableRow>
-                    ))}
-                  </TableBody>
-                </Table>
+          {(llmColumnCount > 0 || modelUsageRows.length > 0) && (
+            <div className="rounded-xl border border-border/60 bg-card/55 p-3">
+              <div className="mb-2 flex items-center justify-between">
+                <p className="text-xs text-muted-foreground">Model usage</p>
+                <HugeiconsIcon icon={Flag02Icon} className="size-4 text-muted-foreground" />
               </div>
-            )}
-          </div>
+              {modelUsageRows.length === 0 ? (
+                <p className="text-xs text-muted-foreground">No model usage yet.</p>
+              ) : (
+                <div className="overflow-hidden rounded-lg border border-border/60 bg-card/50">
+                  <Table>
+                    <TableHeader>
+                      <TableRow>
+                        <TableHead>Model</TableHead>
+                        <TableHead className="text-right">Input</TableHead>
+                        <TableHead className="text-right">Output</TableHead>
+                      </TableRow>
+                    </TableHeader>
+                    <TableBody>
+                      {modelUsageRows.map((usage) => (
+                        <TableRow key={usage.model}>
+                          <TableCell className="max-w-[320px] truncate">{usage.model}</TableCell>
+                          <TableCell className="text-right">
+                            {formatMetricValue(usage.input)}
+                          </TableCell>
+                          <TableCell className="text-right">
+                            {formatMetricValue(usage.output)}
+                          </TableCell>
+                        </TableRow>
+                      ))}
+                    </TableBody>
+                  </Table>
+                </div>
+              )}
+            </div>
+          )}
         </div>
       )}
       <div className="overflow-hidden rounded-xl corner-squircle border">

studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx

@@ -59,7 +59,7 @@ export function ExecutionsView({
     typeof value === "number" && Number.isFinite(value)
       ? `${value.toLocaleString()} s`
       : "--";
-  const [detailTab, setDetailTab] = useState("overview");
+  const [detailTab, setDetailTab] = useState("data");
   const [hiddenDatasetColumnsByExecution, setHiddenDatasetColumnsByExecution] = useState<
     Record<string, string[]>
   >({});
@@ -341,12 +341,16 @@ export function ExecutionsView({
   }, [selectedExecution]);
 
   useEffect(() => {
-    if (!terminalRef.current) {
+    setDetailTab("data");
+  }, [selectedExecution?.id]);
+
+  useEffect(() => {
+    if (detailTab !== "overview" || !terminalRef.current) {
       return;
     }
     shouldStickTerminalToBottomRef.current = true;
     terminalRef.current.scrollTop = terminalRef.current.scrollHeight;
-  }, [selectedExecution?.id]);
+  }, [detailTab, selectedExecution?.id]);
 
   useEffect(() => {
     if (!terminalRef.current) {
@@ -440,9 +444,9 @@ export function ExecutionsView({
             <Tabs value={detailTab} onValueChange={setDetailTab}>
               <div className="flex items-center justify-between gap-2">
                 <TabsList className="border border-border/60 bg-card/40">
+                  <TabsTrigger value="data">Data</TabsTrigger>
                   <TabsTrigger value="overview">Overview</TabsTrigger>
                   <TabsTrigger value="columns">Columns</TabsTrigger>
-                  <TabsTrigger value="data">Data</TabsTrigger>
                   <TabsTrigger value="raw">Raw</TabsTrigger>
                 </TabsList>
                 <div className="flex items-center gap-2">

unsloth/models/_utils.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "2026.4.1"
+__version__ = "2026.4.2"
 
 __all__ = [
     "SUPPORTS_BFLOAT16",

unsloth/models/loader.py

@@ -114,6 +114,17 @@ FORCE_FLOAT32 = [
 
 global DISABLE_COMPILE_MODEL_NAMES
 # Must be alphabetically sorted for each entry
+
+
+def _strip_unsloth_bnb_4bit_suffix(model_name: str) -> str:
+    """Remove Unsloth 4bit suffixes without lowercasing (HF cache dirs are case-sensitive)."""
+    s = model_name
+    for suffix in ("-unsloth-bnb-4bit", "-bnb-4bit"):
+        if len(s) >= len(suffix) and s.lower().endswith(suffix.lower()):
+            s = s[: -len(suffix)]
+    return s
+
+
 DISABLE_COMPILE_MODEL_NAMES = [
     "aya_vision",
     "modernbert",
@@ -404,8 +415,7 @@ class FastLanguageModel(FastLlamaModel):
         if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
             ("-unsloth-bnb-4bit", "-bnb-4bit")
         ):
-            model_name = model_name.lower().removesuffix("-unsloth-bnb-4bit")
-            model_name = model_name.lower().removesuffix("-bnb-4bit")
+            model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
         # Change -BF16 to all False for 4bit, 8bit etc
         if model_name.lower().endswith("-bf16"):
             load_in_4bit = False
@@ -551,8 +561,7 @@ class FastLanguageModel(FastLlamaModel):
             if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
                 ("-unsloth-bnb-4bit", "-bnb-4bit")
             ):
-                model_name = model_name.lower().removesuffix("-unsloth-bnb-4bit")
-                model_name = model_name.lower().removesuffix("-bnb-4bit")
+                model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
             # Change -BF16 to all False for 4bit, 8bit etc
             if model_name.lower().endswith("-bf16"):
                 load_in_4bit = False
@@ -1019,8 +1028,7 @@ class FastModel(FastBaseModel):
         if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
             ("-unsloth-bnb-4bit", "-bnb-4bit")
         ):
-            model_name = model_name.lower().removesuffix("-unsloth-bnb-4bit")
-            model_name = model_name.lower().removesuffix("-bnb-4bit")
+            model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
         # Change -BF16 to all False for 4bit, 8bit etc
         if model_name.lower().endswith("-bf16"):
             load_in_4bit = False
@@ -1320,8 +1328,7 @@ class FastModel(FastBaseModel):
             if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
                 ("-unsloth-bnb-4bit", "-bnb-4bit")
             ):
-                model_name = model_name.lower().removesuffix("-unsloth-bnb-4bit")
-                model_name = model_name.lower().removesuffix("-bnb-4bit")
+                model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
             # Change -BF16 to all False for 4bit, 8bit etc
             if model_name.lower().endswith("-bf16"):
                 load_in_4bit = False

unsloth/models/loader_utils.py

@@ -162,7 +162,7 @@ def __get_model_name(
         # Support returning original full -bnb-4bit name if specified specifically
         # since we'll map it to the dynamic version instead
         if lower_model_name.endswith("-bnb-4bit"):
-            return lower_model_name
+            return model_name
 
         new_model_name = FLOAT_TO_INT_MAPPER[lower_model_name]
         # logger.warning_once(