Merge remote-tracking branch 'origin/main' into fold-integration

# Conflicts:
#	scripts/scan_packages_baseline.json

.github/workflows/consolidated-tests-ci.yml

@@ -364,6 +364,7 @@ jobs:
             tests/utils/test_attention_masks.py \
             tests/utils/test_trunc_normal_patch.py \
             tests/python/test_fast_language_model_text_only.py \
+            tests/test_prefetch_snapshot_scope.py \
             --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap'
           # The deselected test monkeypatches flash_attn_varlen_func, which is
           # only bound on the module when `flash_attn` is importable. flash_attn

.github/workflows/lockfile-audit.yml

@@ -60,11 +60,11 @@ jobs:
     runs-on: ubuntu-latest
     timeout-minutes: 5
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0  # v7.0.0
         with:
           persist-credentials: false
 
-      - uses: actions/setup-python@v5
+      - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1  # v6.3.0
         with:
           python-version: '3.12'
 

.github/workflows/studio-windows-inference-smoke.yml

@@ -1610,6 +1610,13 @@ jobs:
       - name: Install Pester v5
         shell: pwsh
         run: |
+          # PSGallery is intermittently absent from the repository list on GitHub's Windows
+          # runners, which makes `Set-PSRepository PSGallery` fail with "No repository with the
+          # name 'PSGallery' was found." Re-register the default gallery first so the policy
+          # change and module install below always have a repository to target.
+          if (-not (Get-PSRepository -Name PSGallery -ErrorAction SilentlyContinue)) {
+            Register-PSRepository -Default -ErrorAction SilentlyContinue
+          }
           Set-PSRepository PSGallery -InstallationPolicy Trusted
           Install-Module Pester -MinimumVersion 5.5.0 -Force -SkipPublisherCheck -Scope CurrentUser
           Import-Module Pester -MinimumVersion 5.5.0

.gitignore

@@ -11,6 +11,8 @@ outputs/
 exports/
 /datasets/
 studio/backend/assets/datasets/
+# Generated async worker / reviewer transcripts (never part of the product).
+studio/backend/async_task_outputs/
 unsloth_training_checkpoints/
 *.gguf
 *.safetensors

scripts/scan_packages_baseline.json

@@ -66,6 +66,14 @@
    "evidence": "Obfusc: L87:     __import__(name)\nExec: L735:         exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")",
    "evidence_hash": "3cb7d8247dea7dd3d7b21ededc0181c58c50099aeb73c9138a286f3d1ad92d4f"
   },
+  {
+   "package": "cffi",
+   "file": "cffi/_cffi_gen_src.py",
+   "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+   "severity": "HIGH",
+   "evidence": "Obfusc: L52:     compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53:     exec(compiled, globs, globs)",
+   "evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9"
+  },
   {
    "package": "cffi",
    "file": "cffi/setuptools_ext.py",
@@ -159,16 +167,16 @@
    "file": "fastapi/routing.py",
    "check": "C2 polling/beaconing loop detected",
    "severity": "CRITICAL",
-   "evidence": "L586:                                 while True: sha256:bef9ea429314fad39e063895a37dc5cfe9b04561f3d1acbb3c99abb4e92e6cfe",
-   "evidence_hash": "b15773e1bc249713156a349278ea60f7c0e3dd7d537affe929ab51089e1942bb"
+   "evidence": "L586:                                 while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45",
+   "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5"
   },
   {
    "package": "fastapi",
    "file": "fastapi/routing.py",
    "check": "C2 polling/beaconing loop detected",
    "severity": "CRITICAL",
-   "evidence": "L586:                                 while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45",
-   "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5"
+   "evidence": "L586:                                 while True: sha256:bef9ea429314fad39e063895a37dc5cfe9b04561f3d1acbb3c99abb4e92e6cfe",
+   "evidence_hash": "b15773e1bc249713156a349278ea60f7c0e3dd7d537affe929ab51089e1942bb"
   },
   {
    "package": "fastmcp-slim",
@@ -186,6 +194,14 @@
    "evidence": "FS: L624:     history.replaceState(null, \"\", url); sha256:fd8dbfa8af4dea2ce43f4d441f3f81239de341b76a2eb0a33c446f6757ce5f43\nNetwork: L1291:     with httpx.Client(timeout=30.0) as client: | L1305:     with httpx.Client(timeout=30.0) as client: | L1335:     with httpx.Client(timeout=30.0) as client: | L1537:         client = httpx.AsyncClient(\nL1538:             timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1539:         ) | L1701:     async with httpx.AsyncClient(trust_env=False) as client: | L1769:                     with socket.socket(family, socket.SOCK_STREAM) as s:",
    "evidence_hash": "6ada4a9111213bdee5ea24c70a72ec4acdc8ffe0de4a01fd9835bc261ccab8f8"
   },
+  {
+   "package": "fastmcp-slim",
+   "file": "fastmcp/cli/apps_dev.py",
+   "check": "Enumerates filesystem AND makes network calls",
+   "severity": "CRITICAL",
+   "evidence": "FS: L637:     history.replaceState(null, \"\", url); sha256:17068ba5bfed62c3a3007ec8bf3e0ea41ef6529b9e6112064d9afb3be9231436\nNetwork: L1304:     with httpx.Client(timeout=30.0) as client: | L1318:     with httpx.Client(timeout=30.0) as client: | L1348:     with httpx.Client(timeout=30.0) as client: | L1549:         client = httpx.AsyncClient(\nL1550:             timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551:         ) | L1713:     async with httpx.AsyncClient(trust_env=False) as client: | L1781:                     with socket.socket(family, socket.SOCK_STREAM) as s:",
+   "evidence_hash": "e5325edfada6499540e6f0c24a0868979d275522e2b6a180aa9b5dd3280681b4"
+  },
   {
    "package": "fastmcp-slim",
    "file": "fastmcp/server/auth/providers/jwt.py",
@@ -255,16 +271,16 @@
    "file": "huggingface_hub/hf_api.py",
    "check": "C2 polling/beaconing loop detected",
    "severity": "CRITICAL",
-   "evidence": "L4600:         while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272",
-   "evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602"
+   "evidence": "L4613:         while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6",
+   "evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4"
   },
   {
    "package": "huggingface-hub",
    "file": "huggingface_hub/hf_api.py",
    "check": "C2 polling/beaconing loop detected",
    "severity": "CRITICAL",
-   "evidence": "L4613:         while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6",
-   "evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4"
+   "evidence": "L4600:         while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272",
+   "evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602"
   },
   {
    "package": "huggingface-hub",

scripts/verify_import_hoist.py

@@ -564,6 +564,12 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
     for n, tids in b["module_import_targets"].items():
         if tids & after_used:
             continue  # resolved -> fine
+        # `from __future__ import ...` is a compiler directive, not a runtime
+        # binding: the name (`annotations`, ...) is never loaded, so it can never
+        # "resolve" to a use. Skip it so a legitimately-added future import
+        # (e.g. `annotations` for lazy PEP 604 `X | None` on py3.9) is not flagged.
+        if all(t.startswith("from:__future__:") for t in tids):
+            continue
         newly_added = bool(tids - before_module_targets)
         was_used_before = bool(tids & before_used)
         if newly_added or was_used_before:
@@ -588,9 +594,23 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
     #    package object and only *add* submodule attributes (e.g. adding
     #    `import urllib.error` next to `import urllib.request`). Nothing the name
     #    resolved to before is lost, so no reference is re-pointed -- skip it.
+    #
+    #    A deliberate *relocation* is also benign and must not block: when a name
+    #    keeps its spelling but its import source is moved A -> B in THIS diff (the
+    #    old `from A import x` is removed at module level and a new `from B import x`
+    #    is added), the swap is intentional, not a silent re-point to a pre-existing
+    #    different object. This mirrors the relocation tolerance already applied to
+    #    TARGET-MISSING. The dangerous case -- the name now resolving to a target
+    #    that already existed before (shadow/clash) -- is NOT exempted.
+    removed_module_targets = before_module_targets - after_module_targets
     for key, tafter in b["target_by_use"].items():
         tbefore = a["target_by_use"].get(key)
         if tbefore and tbefore != tafter and (tbefore - tafter):
+            lost = tbefore - tafter
+            gained = tafter - tbefore
+            relocated = lost <= removed_module_targets and gained <= added_module_targets
+            if relocated:
+                continue
             findings.append(
                 (
                     "BLOCKER",

studio/backend/core/inference/__init__.py

@@ -7,13 +7,16 @@ Inference submodule - backend for model loading and generation.
 The default get_inference_backend() returns an InferenceOrchestrator that
 delegates to a subprocess. The original InferenceBackend runs inside the
 subprocess and can be imported directly from .inference when needed.
+
+Public names are resolved lazily (PEP 562): importing this package -- or a
+dependency-light leaf like ``core.inference.chat_eos`` -- must NOT eagerly pull
+the orchestrator / llama_cpp import chain (httpx, subprocess plumbing, the ML
+backend and its Studio dependencies). Those load only when a public name is
+actually accessed, so standalone helpers stay unit-testable without the full
+inference stack.
 """
 
-from .orchestrator import InferenceOrchestrator, get_inference_backend
-from .llama_cpp import LlamaCppBackend
-
-# Expose InferenceOrchestrator as InferenceBackend for backward compat.
-InferenceBackend = InferenceOrchestrator
+from typing import TYPE_CHECKING
 
 __all__ = [
     "InferenceBackend",
@@ -21,3 +24,33 @@ __all__ = [
     "get_inference_backend",
     "LlamaCppBackend",
 ]
+
+# name -> (submodule, attribute); InferenceBackend aliases InferenceOrchestrator.
+_LAZY_ATTRS = {
+    "InferenceOrchestrator": ("orchestrator", "InferenceOrchestrator"),
+    "InferenceBackend": ("orchestrator", "InferenceOrchestrator"),
+    "get_inference_backend": ("orchestrator", "get_inference_backend"),
+    "LlamaCppBackend": ("llama_cpp", "LlamaCppBackend"),
+}
+
+
+def __getattr__(name):
+    try:
+        submodule, attr = _LAZY_ATTRS[name]
+    except KeyError:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
+    from importlib import import_module
+
+    value = getattr(import_module(f"{__name__}.{submodule}"), attr)
+    globals()[name] = value  # cache so later access skips __getattr__
+    return value
+
+
+def __dir__():
+    return sorted(set(globals()) | set(__all__))
+
+
+if TYPE_CHECKING:  # keep static analysers / IDEs aware of the lazy names
+    from .llama_cpp import LlamaCppBackend
+    from .orchestrator import InferenceOrchestrator, get_inference_backend
+    InferenceBackend = InferenceOrchestrator

studio/backend/core/inference/chat_eos.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Resolve a chat model's assistant-turn-end stop tokens.
+
+Some checkpoints set eos_token_id to a bare document terminator (Qwen3.5 ships
+config eos ``<|endoftext|>`` though chat turns end with ``<|im_end|>``, and its
+small chat variants ship no generation_config), so generation runs past the turn
+and loops -- re-emitting tool calls or hallucinating ``<|im_start|>`` turns.
+
+Turn-end markers are derived from the tokenizer's ``chat_template`` (the tokens it
+actually uses to end a turn), not raw vocab membership: a base/coder model can
+carry ChatML control tokens in a shared vocab without using them, and a loader
+may have synced ``eos_token`` to the document terminator. Dependency-light (no
+torch / unsloth) so it is unit-testable without the full inference stack.
+"""
+
+from typing import Optional
+
+# Canonical assistant-turn-end markers per chat family.
+_CHAT_TURN_END_TOKENS = (
+    "<|im_end|>",  # ChatML: Qwen, Yi
+    "<|eot_id|>",  # Llama 3.x
+    "<|eom_id|>",  # Llama 3.x tool turns
+    "<end_of_turn>",  # Gemma
+    "<turn|>",  # Gemma-4
+    "<|end|>",  # Phi
+    "<|end_of_turn|>",  # OpenChat / Starling (barred, distinct from Gemma's)
+)
+# harmony/gpt-oss uses <|end|> as a channel delimiter, not the turn end, and has
+# its own streamer, so its eos is left untouched.
+_HARMONY_MARKERS = ("<|channel|>", "<|constrain|>")
+
+
+def _eos_id_set(eos_token_id) -> set:
+    if isinstance(eos_token_id, (list, tuple)):
+        return {int(t) for t in eos_token_id if t is not None}
+    if eos_token_id is not None:
+        return {int(eos_token_id)}
+    return set()
+
+
+def _collect_template_text(chat_template) -> str:
+    """Flatten a tokenizer ``chat_template`` into one scannable string.
+
+    Usually the template is a single jinja string, but multi-variant models
+    (e.g. Hermes-3: a ``default`` plus a ``tool_use`` template) expose it as a
+    ``{name: template}`` dict -- or, as stored in tokenizer_config.json, a list
+    of ``{"name": ..., "template": ...}`` dicts. Scanning only the ``str`` case
+    would skip turn-end detection for those valid models, so gather every string
+    leaf (variant names are harmless: they never contain the markers).
+    """
+    if isinstance(chat_template, str):
+        return chat_template
+    if isinstance(chat_template, dict):
+        values = chat_template.values()
+    elif isinstance(chat_template, (list, tuple)):
+        values = chat_template
+    else:
+        return ""
+    parts = [_collect_template_text(v) for v in values]
+    return "\n".join(p for p in parts if p)
+
+
+def resolve_chat_turn_end_eos_ids_using(template_tokenizer, id_tokenizer) -> list:
+    """eos of ``id_tokenizer`` plus any canonical turn-end marker the
+    ``template_tokenizer``'s chat_template uses, resolved to ids on ``id_tokenizer`` --
+    the tokenizer generation actually uses.
+
+    Pass the same tokenizer for both at load time. After a mapped ``get_chat_template``
+    pass the MAPPED tokenizer as ``template_tokenizer`` (it carries the effective
+    template) and the ORIGINAL generation tokenizer as ``id_tokenizer``: a mapped
+    template registered ``map_eos_token=True`` can hand back a tokenizer whose vocab
+    folds the turn-end token onto the doc-eos id, and generate_stream re-reads the
+    original tokenizer, so resolving ids on the mapped tokenizer would store the wrong
+    (doc-eos) id and let generation run past the real turn marker."""
+    ids = _eos_id_set(getattr(id_tokenizer, "eos_token_id", None))
+    template = _collect_template_text(getattr(template_tokenizer, "chat_template", None))
+    if not template or any(h in template for h in _HARMONY_MARKERS):
+        return sorted(ids)
+    unk = getattr(id_tokenizer, "unk_token_id", None)
+    for marker in _CHAT_TURN_END_TOKENS:
+        if marker in template:
+            try:
+                tid = id_tokenizer.convert_tokens_to_ids(marker)
+            except Exception:
+                tid = None
+            if tid is not None and tid != unk and int(tid) >= 0:
+                ids.add(int(tid))
+    return sorted(ids)
+
+
+def resolve_chat_turn_end_eos_ids(tokenizer) -> list:
+    """tokenizer.eos plus any canonical turn-end marker the model's chat_template
+    actually uses. Cheap (convert_tokens_to_ids per marker, no get_vocab); intended
+    to be resolved once at load. Returns eos unchanged for harmony templates."""
+    return resolve_chat_turn_end_eos_ids_using(tokenizer, tokenizer)
+
+
+def chat_eos_repair(current_eos, turn_end_ids) -> Optional[list]:
+    """Merged eos_token_id list, or None if ``current_eos`` already covers every
+    resolved turn-end id. Used to repair a model's generation_config at load so
+    every ``.generate()`` path (vision, tool loops) stops at the turn boundary."""
+    if not turn_end_ids:
+        return None
+    current_set = _eos_id_set(current_eos)
+    if set(turn_end_ids) <= current_set:
+        return None
+    return sorted(current_set | set(turn_end_ids))

studio/backend/core/inference/chat_template_helpers.py

@@ -3,12 +3,60 @@
 
 """
 Dependency-light wrapper around tokenizer.apply_chat_template with a kwarg
-fallback for templates that reject reasoning/tools args.
+fallback for templates that reject reasoning/tools args, plus the shared
+native-chat-template fallback used by the transformers and MLX backends.
 """
 
+import copy
+import json
+import logging
 from typing import Optional
 
 
+logger = logging.getLogger(__name__)
+
+
+def _normalize_tool_call_arguments(messages: list) -> list:
+    """Coerce each assistant ``tool_calls[].function.arguments`` from a JSON
+    string to a dict.
+
+    The OpenAI wire format carries ``arguments`` as a JSON string, but some chat
+    templates (e.g. the stricter Qwen tool templates shipped with mlx-community
+    checkpoints) iterate ``arguments.items()`` and raise
+    ``TypeError: Can only get item pairs from a mapping.`` on the string form
+    when a prior tool call is re-rendered on the next turn. A dict works on both
+    strict and lenient templates, so parse the string; leave non-JSON or non-dict
+    values untouched. Returns the original list unchanged when nothing needed
+    coercing (no copy)."""
+    mutated = False
+    out: list = []
+    for msg in messages:
+        tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None
+        if not tool_calls:
+            out.append(msg)
+            continue
+        new_calls = []
+        msg_changed = False
+        for call in tool_calls:
+            fn = call.get("function") if isinstance(call, dict) else None
+            args = fn.get("arguments") if isinstance(fn, dict) else None
+            if isinstance(args, str):
+                try:
+                    parsed = json.loads(args)
+                except (ValueError, TypeError):
+                    parsed = None
+                if isinstance(parsed, dict):
+                    call = {**call, "function": {**fn, "arguments": parsed}}
+                    msg_changed = True
+            new_calls.append(call)
+        if msg_changed:
+            out.append({**msg, "tool_calls": new_calls})
+            mutated = True
+        else:
+            out.append(msg)
+    return out if mutated else messages
+
+
 def apply_chat_template_for_generation(
     tokenizer,
     messages: list,
@@ -38,21 +86,209 @@ def apply_chat_template_for_generation(
         attempts.append(dict(reasoning_kwargs))
     attempts.append({})
 
-    last_exc: Optional[Exception] = None
-    for kwargs in attempts:
+    def _render(msgs: list) -> str:
+        last_exc: Optional[Exception] = None
+        for kwargs in attempts:
+            try:
+                return tokenizer.apply_chat_template(
+                    msgs,
+                    tokenize = False,
+                    add_generation_prompt = True,
+                    **kwargs,
+                )
+            except TypeError as e:
+                last_exc = e
+                continue
+            except Exception as e:
+                last_exc = e
+                break
+        if last_exc is not None:
+            raise last_exc
+        raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")
+
+    try:
+        return _render(messages)
+    except Exception:
+        # Strict tool templates reject the JSON-string ``arguments`` form via
+        # TypeError or a broad Jinja raise_exception, so retry with dicts coerced.
+        # Original messages render first, so working templates stay byte-identical.
+        normalized = _normalize_tool_call_arguments(messages)
+        if normalized is messages:
+            raise
+        return _render(normalized)
+
+
+def render_native_template(
+    *,
+    model_info: dict,
+    active_model_name: Optional[str],
+    messages: list,
+    tools: list,
+    enable_thinking: Optional[bool] = None,
+    reasoning_effort: Optional[str] = None,
+    preserve_thinking: Optional[bool] = None,
+    apply_fn = None,
+    hf_token: Optional[str] = None,
+) -> Optional[str]:
+    """Render ``messages`` + ``tools`` with the model's NATIVE chat template.
+
+    Some Unsloth override templates (e.g. ``mistral``, ``gemma-4``) do not emit
+    the ``tools`` schema, so a tool-calling turn silently stops advertising tools.
+    The native template ships in the model repo and carries the family's
+    tool-calling syntax. It is loaded straight from the repo (bypassing any
+    override on the live tokenizer) and cached on ``model_info``. Returns the
+    rendered prompt only if the native template actually emits the tools (render
+    differs with vs without tools); otherwise ``None``.
+
+    ``hf_token`` is the token the model was loaded with -- passed to the repo load
+    so a gated/private model's native template can still be fetched (otherwise the
+    fallback fails silently and keeps the override prompt that dropped tools).
+
+    ``trust_remote_code`` is sourced from ``model_info`` (the value the model was
+    actually loaded with) rather than a call-site argument, so the native-template
+    reload uses exactly the consent already granted at load. A custom-code tokenizer
+    repo raises in ``AutoTokenizer.from_pretrained`` unless ``trust_remote_code`` is
+    passed, so without this the fallback fails silently and keeps the tool-dropping
+    prompt for a model the user already consented to run remote code for. For a LoRA
+    adapter the reload targets the base model, whose remote code was gated and loaded
+    under the same stored flag, so re-passing it executes no unconsented code.
+    """
+    # ``apply_fn`` lets a backend inject its own render; defaults to the module helper.
+    if apply_fn is None:
+        apply_fn = apply_chat_template_for_generation
+    native_tpl = model_info.get("native_chat_template")
+    if native_tpl is None:
+        # A LoRA adapter's native template lives on the base model, not the adapter id.
+        template_source = model_info.get("base_model") or active_model_name
+        # Re-use the load-time trust_remote_code so a custom-code tokenizer repo can
+        # instantiate its class (the stored flag already covers template_source).
+        trust_remote_code = bool(model_info.get("trust_remote_code", False))
         try:
-            return tokenizer.apply_chat_template(
-                messages,
-                tokenize = False,
-                add_generation_prompt = True,
-                **kwargs,
+            from transformers import AutoTokenizer
+            nt = AutoTokenizer.from_pretrained(
+                template_source,
+                token = hf_token if hf_token and hf_token.strip() else None,
+                trust_remote_code = trust_remote_code,
             )
-        except TypeError as e:
-            last_exc = e
-            continue
-        except Exception as e:
-            last_exc = e
-            break
-    if last_exc is not None:
-        raise last_exc
-    raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")
+            native_tpl = nt.chat_template or False
+        except Exception as exc:
+            logger.warning(
+                "Could not load native chat template for '%s': %s",
+                template_source,
+                exc,
+            )
+            # A failed fetch is not "no template": leave the sentinel unset so the next
+            # call retries (caching False would pin the tool-dropping override).
+            return None
+        model_info["native_chat_template"] = native_tpl
+    if not native_tpl:
+        return None
+
+    tokenizer = model_info.get("tokenizer") or model_info.get("processor")
+    if tokenizer is None:
+        return None
+    tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
+    # Render on a shallow copy: mutating the shared tokenizer.chat_template (outside the
+    # generation lock) races concurrent requests.
+    try:
+        render_tokenizer = copy.copy(tokenizer)
+        render_tokenizer.chat_template = native_tpl
+    except Exception as exc:
+        logger.warning(
+            "Could not clone tokenizer for native-template render of '%s': %s",
+            active_model_name,
+            exc,
+        )
+        return None
+    try:
+        with_tools = apply_fn(
+            render_tokenizer,
+            messages,
+            tools = tools,
+            enable_thinking = enable_thinking,
+            reasoning_effort = reasoning_effort,
+            preserve_thinking = preserve_thinking,
+        )
+        no_tools = apply_fn(
+            render_tokenizer,
+            messages,
+            tools = None,
+            enable_thinking = enable_thinking,
+            reasoning_effort = reasoning_effort,
+            preserve_thinking = preserve_thinking,
+        )
+    except Exception as exc:
+        logger.warning(
+            "Native-template tool render failed for '%s': %s",
+            active_model_name,
+            exc,
+        )
+        return None
+    return with_tools if with_tools != no_tools else None
+
+
+def render_with_native_template_fallback(
+    *,
+    formatted_prompt: str,
+    tokenizer,
+    model_info: dict,
+    active_model_name: Optional[str],
+    messages: list,
+    tools: Optional[list],
+    enable_thinking: Optional[bool] = None,
+    reasoning_effort: Optional[str] = None,
+    preserve_thinking: Optional[bool] = None,
+    apply_fn = None,
+    hf_token: Optional[str] = None,
+) -> str:
+    """Return ``formatted_prompt``, swapping in a native-template render when an
+    override template dropped the ``tools`` schema.
+
+    If ``tools`` were requested but the live render is identical with and without
+    them (detected by comparison, robust against tool names in the system prompt),
+    re-render with the model's native template. Shared by the transformers and MLX
+    backends so both advertise tools consistently. ``hf_token`` is forwarded so a
+    gated/private model's native template can still be fetched."""
+    if not tools:
+        return formatted_prompt
+    if apply_fn is None:
+        apply_fn = apply_chat_template_for_generation
+    # Probe whether the live template dropped the schema. A tools-requiring template
+    # can raise here; on any error keep the valid tools prompt rather than lose it.
+    try:
+        probe_no_tools = apply_fn(
+            tokenizer,
+            messages,
+            tools = None,
+            enable_thinking = enable_thinking,
+            reasoning_effort = reasoning_effort,
+            preserve_thinking = preserve_thinking,
+        )
+    except Exception as exc:
+        logger.warning(
+            "No-tools probe failed for '%s'; keeping the existing tools prompt: %s",
+            active_model_name,
+            exc,
+        )
+        return formatted_prompt
+    if formatted_prompt != probe_no_tools:
+        return formatted_prompt  # template already emits the tools schema
+    native_prompt = render_native_template(
+        model_info = model_info,
+        active_model_name = active_model_name,
+        messages = messages,
+        tools = tools,
+        enable_thinking = enable_thinking,
+        reasoning_effort = reasoning_effort,
+        preserve_thinking = preserve_thinking,
+        apply_fn = apply_fn,
+        hf_token = hf_token,
+    )
+    if native_prompt:
+        logger.info(
+            "Override template for '%s' dropped tool schemas; using the model's "
+            "native template for this tool-calling turn.",
+            active_model_name,
+        )
+        return native_prompt
+    return formatted_prompt

studio/backend/core/inference/inference.py

@@ -27,6 +27,10 @@ from utils.hardware import (
 from core.inference.audio_codecs import AudioCodecManager
 from core.inference.runtime_context import runtime_context_length
 from core.inference.message_content import content_to_text
+from core.inference.chat_eos import (
+    chat_eos_repair,
+    resolve_chat_turn_end_eos_ids_using,
+)
 from io import StringIO
 import structlog
 from loggers import get_logger
@@ -210,6 +214,50 @@ class InferenceBackend:
         # API uses -1 to disable top-k; transformers uses 0.
         return 0 if top_k < 0 else top_k
 
+    def _resolve_chat_eos(self, model_name: str) -> None:
+        """Resolve this chat model's assistant-turn-end stop tokens once at load,
+        cache them in model_info, and repair generation_config so every
+        ``.generate()`` path stops at the turn boundary.
+
+        Some checkpoints (e.g. Qwen3.5 / Qwen3.6 small chat models) end turns with
+        ``<|im_end|>`` but ship ``config.eos_token_id = <|endoftext|>`` and no
+        ``generation_config.json``, so paths that read ``generation_config`` (the
+        vision path, tool loops) run past the turn and loop. Turn-end markers are
+        derived from the chat_template (see chat_eos.resolve_chat_turn_end_eos_ids),
+        so base/coder models and harmony templates are left untouched.
+        """
+        info = self.models.get(model_name) or {}
+        model = info.get("model")
+        container = info.get("tokenizer")
+        tokenizer = getattr(container, "tokenizer", container)  # unwrap processors
+        if model is None or tokenizer is None:
+            return
+        # Vision models carry the chat_template on the processor, not the inner
+        # tokenizer. Read markers from whichever has one, but resolve ids on the
+        # generation tokenizer, else the vision path misses the turn-end token.
+        template_source = container if getattr(container, "chat_template", None) else tokenizer
+        try:
+            turn_end_ids = resolve_chat_turn_end_eos_ids_using(template_source, tokenizer)
+        except Exception as e:  # never block a load on eos resolution
+            logger.warning("Chat turn-end eos resolution failed for %s: %s", model_name, e)
+            return
+        info["chat_turn_end_eos_ids"] = turn_end_ids
+
+        gen = getattr(model, "generation_config", None)
+        if gen is None:
+            return
+        repaired = chat_eos_repair(gen.eos_token_id, turn_end_ids)
+        if repaired is None:
+            return
+        previous = gen.eos_token_id
+        gen.eos_token_id = repaired
+        logger.info(
+            "Repaired generation_config.eos_token_id for %s: %s -> %s",
+            model_name,
+            previous,
+            repaired,
+        )
+
     def load_model(
         self,
         config: ModelConfig,
@@ -221,6 +269,9 @@ class InferenceBackend:
         gpu_ids: Optional[list[int]] = None,
     ) -> bool:
         """Load any model: base, LoRA adapter, text, or vision."""
+        # Keep the token so the native-template fallback can fetch a
+        # gated model's repo template later during generation.
+        self._hf_token = hf_token
         # GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it.
         if max_seq_length <= 0:
             max_seq_length = 2048
@@ -231,6 +282,8 @@ class InferenceBackend:
             # Already loaded?
             if model_name in self.models and self.models[model_name].get("model"):
                 logger.info(f"Model {model_name} already loaded")
+                if hf_token:
+                    self.models[model_name]["hf_token"] = hf_token
                 self.active_model_name = model_name
                 return True
 
@@ -246,6 +299,14 @@ class InferenceBackend:
             )
 
             self.models[model_name] = {
+                # Per-model token: the native-template fallback must use the
+                # token this model was loaded with, not whichever loaded last.
+                "hf_token": hf_token,
+                # Per-model consent: the native-template reload must re-use the
+                # exact trust_remote_code this model (and a LoRA's base) was loaded
+                # with, so a custom-code tokenizer repo can be re-fetched without
+                # executing any code the user did not already consent to.
+                "trust_remote_code": trust_remote_code,
                 "is_vision": config.is_vision,
                 "is_lora": config.is_lora,
                 "is_audio": config.is_audio,
@@ -496,6 +557,7 @@ class InferenceBackend:
                 max_seq_length,
             )
 
+            self._resolve_chat_eos(model_name)
             self._load_chat_template_info(model_name)
 
             self.active_model_name = model_name
@@ -946,6 +1008,22 @@ class InferenceBackend:
                     tokenizer,
                     chat_template = template_name,
                 )
+                # The mapper installs the effective template only now, at generate
+                # time, so re-resolve and UNION into the load-time cache (never
+                # overwrite). get_chat_template can return a remapped tokenizer
+                # (turn-end folded onto doc-eos) while generate_stream reads the
+                # original, so take marker strings from the mapped template but
+                # resolve their ids on the original.
+                try:
+                    _gen_tok = model_info.get("tokenizer") or tokenizer
+                    refreshed = resolve_chat_turn_end_eos_ids_using(
+                        getattr(tokenizer, "tokenizer", tokenizer),
+                        getattr(_gen_tok, "tokenizer", _gen_tok),
+                    )
+                    existing = model_info.get("chat_turn_end_eos_ids") or []
+                    model_info["chat_turn_end_eos_ids"] = sorted(set(existing) | set(refreshed))
+                except Exception as e:
+                    logger.warning(f"Could not refresh chat turn-end eos after template: {e}")
             else:
                 logger.info(
                     f"No registered Unsloth template for {self.active_model_name}, using tokenizer default"
@@ -975,6 +1053,27 @@ class InferenceBackend:
                 reasoning_effort = reasoning_effort,
                 preserve_thinking = preserve_thinking,
             )
+
+            # If tools were requested but the (possibly overridden) template ignored
+            # them, fall back to the model's native template (shared with MLX).
+            from core.inference.chat_template_helpers import (
+                render_with_native_template_fallback,
+            )
+
+            formatted_prompt = render_with_native_template_fallback(
+                formatted_prompt = formatted_prompt,
+                tokenizer = tokenizer,
+                model_info = model_info,
+                active_model_name = self.active_model_name,
+                messages = template_messages,
+                tools = tools,
+                enable_thinking = enable_thinking,
+                reasoning_effort = reasoning_effort,
+                preserve_thinking = preserve_thinking,
+                apply_fn = self._apply_chat_template_for_generation,
+                hf_token = model_info.get("hf_token"),
+            )
+
             logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...")
         except Exception as e:
             logger.error(f"Error applying chat template: {e}")
@@ -1382,7 +1481,8 @@ class InferenceBackend:
                 min_p = min_p,
                 repetition_penalty = repetition_penalty,
                 do_sample = temperature > 0,
-                eos_token_id = tokenizer.eos_token_id,
+                # Resolved once at load (chat_template-derived turn-end tokens).
+                eos_token_id = model_info.get("chat_turn_end_eos_ids") or tokenizer.eos_token_id,
                 pad_token_id = tokenizer.eos_token_id
                 if tokenizer.pad_token_id is None
                 else tokenizer.pad_token_id,

studio/backend/core/inference/llama_cpp.py

@@ -38,9 +38,25 @@ from core.inference.llama_server_args import (
     strip_shadowing_flags,
     strip_split_mode_only,
 )
-from core.tool_healing import (
+
+# Share strip / signal constants with the multi-format parser so BUFFERING also
+# catches Llama-3 / Mistral / Gemma 4 (legacy helper only knew <tool_call> / <function=).
+from core.inference.tool_call_parser import (
+    _GEMMA_BARE_TC_PREFIX_RE,
+    _GEMMA_BARE_TC_RE,
     _TOOL_ALL_PATS,
-    strip_tool_call_markup,
+    _balanced_brace_end,
+    _strip_function_xml_calls,
+    _strip_gemma_wrapperless_calls,
+    _strip_glm_calls,
+    _strip_mistral_closed_calls,
+    TOOL_XML_SIGNALS as _SHARED_TOOL_XML_SIGNALS,
+    RAG_MAX_SEARCHES_PER_TURN,
+    RAG_SEARCH_CAP_NUDGE,
+    parse_tool_calls_from_text as _shared_parse_tool_calls_from_text,
+    strip_leading_bare_json_call,
+    strip_llama3_leading_sentinels,
+    strip_tool_markup as _shared_strip_tool_markup,
 )
 from utils.native_path_leases import child_env_without_native_path_secret
 from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
@@ -48,12 +64,6 @@ from utils.subprocess_compat import (
     windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
 )
 from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs
-from core.inference.tool_call_parser import (
-    RAG_MAX_SEARCHES_PER_TURN,
-    RAG_SEARCH_CAP_NUDGE,
-    TOOL_XML_SIGNALS,
-    parse_tool_calls_from_text as _shared_parse_tool_calls_from_text,
-)
 from core.inference.tool_loop_controller import (
     ToolLoopController,
     tool_event_provenance,
@@ -220,7 +230,7 @@ _INTENT_SIGNAL = re.compile(
     r"\b(?:now i|next i)\b"
     r")"
 )
-_MAX_REPROMPTS = 1
+_MAX_REPROMPTS = 3
 
 # Default max_tokens to the effective context when known. The floor is high
 # enough for reasoning-heavy GGUFs and max_tokens-omitting API clients.
@@ -232,6 +242,10 @@ _DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0  # 20 min
 _PROVISIONAL_ARGS_MIN_CHARS = 256
 _DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0  # 2 min
 _REPROMPT_MAX_CHARS = 2000
+# Cap tool calls from a single TEXTUAL-fallback turn (mirrors the safetensors
+# loop). Structured delta.tool_calls are grammar-bounded by llama-server; text
+# parsed from content is not, so one runaway turn could fan out unbounded.
+_MAX_TOOL_CALLS_PER_TURN = 8
 _FORCED_REPEAT_PLAN_SIGNAL = re.compile(
     r"\b(?:i\s+will|i'll|let\s+me|going\s+to|need\s+to|call|use|run|search|fetch|render)\b",
     re.I,
@@ -545,6 +559,13 @@ _TOOL_TEMPLATE_MARKERS = (
     "'role' == 'tool'",
     'message.role == "tool"',
     "message.role == 'tool'",
+    # DeepSeek: no top-level ``{% if tools %}`` block; it gates emission on
+    # ``message['role'] == 'tool'`` plus ``message['tool_calls'] is defined``.
+    "message['role'] == 'tool'",
+    'message["role"] == "tool"',
+    "message['tool_calls']",
+    'message["tool_calls"]',
+    "tool_calls is defined",
 )
 
 
@@ -7881,12 +7902,17 @@ class LlamaCppBackend:
     # ── Message building (OpenAI format) ──────────────────────────
 
     @staticmethod
-    def _parse_tool_calls_from_text(content: str, *, allow_incomplete: bool = True) -> list[dict]:
-        """Thin wrapper around the shared parser in tool_call_parser
-        so safetensors and llama_cpp pick up the same fixes."""
+    def _parse_tool_calls_from_text(
+        content: str,
+        *,
+        allow_incomplete: bool = True,
+        enabled_tool_names: Optional[set] = None,
+    ) -> list[dict]:
+        """Wrapper around the shared parser; ``enabled_tool_names`` gates the markerless bare-JSON form."""
         return _shared_parse_tool_calls_from_text(
             content,
             allow_incomplete = allow_incomplete,
+            enabled_tool_names = enabled_tool_names,
         )
 
     @staticmethod
@@ -8398,6 +8424,10 @@ class LlamaCppBackend:
                 "duration_ms": round((time.monotonic() - started_at) * 1000.0),
             }
 
+        # Enabled-name gate for the markerless Gemma strip (disabled/example
+        # names stay visible). Set per iteration; None = pre-loop name-agnostic.
+        _enabled_tool_names = None
+
         def _strip_tool_markup(
             text: str,
             *,
@@ -8406,11 +8436,22 @@ class LlamaCppBackend:
         ) -> str:
             if not (auto_heal_tool_calls or force):
                 return text
-            return strip_tool_call_markup(text, final = final)
+            return _shared_strip_tool_markup(
+                text, final = final, enabled_tool_names = _enabled_tool_names
+            )
 
         def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str:
             if not (auto_heal_tool_calls or force):
                 return text
+            # Shared parser patterns (not the legacy tool_healing set) so textual
+            # Mistral/python_tag calls entering DRAINING never leak. Balanced strips
+            # first (nested JSON removed whole); no final trim so length compares hold.
+            text = _strip_mistral_closed_calls(text)
+            text = _strip_gemma_wrapperless_calls(text, _enabled_tool_names)
+            # Parser-accurate scans close at each call's REAL terminator before
+            # the regex arms: literal markup inside a value is data.
+            text = _strip_function_xml_calls(text, final = True)
+            text = _strip_glm_calls(text, final = True)
             for pat in _TOOL_ALL_PATS:
                 text = pat.sub("", text)
             return text
@@ -8456,6 +8497,13 @@ class LlamaCppBackend:
                 cumulative_display += "<think>" + reasoning_accum + "</think>"
             cumulative_display += content_buffer
 
+        def _looks_like_enabled_bare_json(text: str, enabled_tool_names: set) -> bool:
+            """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False."""
+            probe = strip_llama3_leading_sentinels(text.lstrip())
+            if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)):
+                return False
+            return strip_leading_bare_json_call(probe, enabled_tool_names) != probe
+
         tool_controller = ToolLoopController(
             tools = tools,
             auto_heal_tool_calls = auto_heal_tool_calls,
@@ -8469,6 +8517,8 @@ class LlamaCppBackend:
             )
 
         _MAX_BUFFER_CHARS = 32
+        # Hold a leading ``{`` well past the 32-char XML cap until it balances (mirrors safetensors).
+        _MAX_BARE_JSON_BUFFER = 16384
         _append_budget_exhausted_nudge = True
         # RAG: cap knowledge-base searches per assistant turn. The controller is
         # tool-agnostic, so this gate stays in the loop.
@@ -8481,6 +8531,9 @@ class LlamaCppBackend:
         # "Hello!" won't match. Pattern compiled at module level
         # (_INTENT_SIGNAL).
         _reprompt_count = 0
+        # Gates ``max_tool_iterations`` on real tool turns (not the enlarged range) so reserved
+        # re-prompt slots don't extend the budget. Mirrors the safetensors guard.
+        _tool_iters_done = 0
         _forced_tool_call_pending = False
 
         # Reserve extra iterations for re-prompts so they don't consume the
@@ -8489,12 +8542,21 @@ class LlamaCppBackend:
         for iteration in range(max_tool_iterations + _extra):
             if cancel_event is not None and cancel_event.is_set():
                 return
+            # Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget.
+            _turn_executed_real_tool = False
 
             active_tools = tool_controller.active_tools()
             if not active_tools:
                 _append_budget_exhausted_nudge = False
                 break
-            _tool_xml_signals = TOOL_XML_SIGNALS
+            # Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call.
+            _enabled_tool_names = {
+                (tool.get("function") or {}).get("name")
+                for tool in active_tools
+                if (tool.get("function") or {}).get("name")
+            }
+            # Shared signal tuple so GGUF BUFFERING wakes on every format the parser knows (like safetensors).
+            _tool_xml_signals = _SHARED_TOOL_XML_SIGNALS
 
             # Build payload -- stream: True so we detect tool signals
             # in the first 1-2 chunks without a non-streaming penalty.
@@ -8777,7 +8839,48 @@ class LlamaCppBackend:
                                                 is_prefix = True
                                                 break
 
-                                        if is_match:
+                                        # Signal-less call shapes (mirror the safetensors
+                                        # loop): Llama-3.2 bare {"name":..} and Gemma
+                                        # call:NAME{...} would otherwise stream raw.
+                                        _hold_buffer = False
+                                        # Whole buffer is the call (no visible prefix) -- drain silently.
+                                        _drain_silently = False
+                                        if not is_match and not is_prefix:
+                                            _bare = strip_llama3_leading_sentinels(stripped_buf)
+                                            if _bare.startswith("{"):
+                                                if _balanced_brace_end(_bare, 0) is None:
+                                                    if len(stripped_buf) < _MAX_BARE_JSON_BUFFER:
+                                                        _hold_buffer = True
+                                                    elif _looks_like_enabled_bare_json(
+                                                        _bare, _enabled_tool_names
+                                                    ):
+                                                        # Oversized still-open enabled call: drain
+                                                        # rather than leak; a giant ordinary JSON
+                                                        # answer still streams.
+                                                        _drain_silently = True
+                                                elif self._parse_tool_calls_from_text(
+                                                    content_buffer,
+                                                    allow_incomplete = auto_heal_tool_calls,
+                                                    enabled_tool_names = _enabled_tool_names,
+                                                ):
+                                                    _drain_silently = True
+                                            elif (
+                                                "call:".startswith(stripped_buf)
+                                                or _GEMMA_BARE_TC_PREFIX_RE.match(stripped_buf)
+                                                is not None
+                                                or _GEMMA_BARE_TC_RE.match(stripped_buf) is not None
+                                            ):
+                                                # Whitespace-tolerant like the parser.
+                                                if _GEMMA_BARE_TC_RE.match(stripped_buf):
+                                                    _drain_silently = True
+                                                elif len(stripped_buf) < _MAX_BUFFER_CHARS:
+                                                    _hold_buffer = True
+
+                                        if _drain_silently:
+                                            # No visible prefix -- the buffered text IS
+                                            # the call; drain without yielding it.
+                                            detect_state = _S_DRAINING
+                                        elif is_match:
                                             # Tool signal -- flush any visible
                                             # prefix before DRAINING so the
                                             # route sends it before tool_start.
@@ -8794,7 +8897,9 @@ class LlamaCppBackend:
                                                         "text": cleaned,
                                                     }
                                             detect_state = _S_DRAINING
-                                        elif is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS:
+                                        elif _hold_buffer or (
+                                            is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS
+                                        ):
                                             pass  # keep buffering
                                         else:
                                             # Not a tool -- flush buffer
@@ -8821,8 +8926,17 @@ class LlamaCppBackend:
                 # ── Resolve BUFFERING at stream end ──
                 if detect_state == _S_BUFFERING:
                     stripped_buf = content_buffer.lstrip()
+                    # A held bare-JSON fragment has no XML signal; route it to DRAINING (the signal-only
+                    # gate below would flush the raw JSON to the user).
+                    _bare_eos = strip_llama3_leading_sentinels(stripped_buf)
+                    # Gate on enabled names so an ordinary JSON answer isn't routed to DRAINING and dropped.
+                    _is_bare_tc = bool(active_tools) and _looks_like_enabled_bare_json(
+                        _bare_eos, _enabled_tool_names
+                    )
                     if stripped_buf and any(s in stripped_buf for s in _tool_xml_signals):
                         detect_state = _S_DRAINING
+                    elif _is_bare_tc:
+                        detect_state = _S_DRAINING
                     elif content_accum or reasoning_accum:
                         detect_state = _S_STREAMING
                         if content_buffer:
@@ -8848,20 +8962,26 @@ class LlamaCppBackend:
                                     "text": cumulative_display,
                                 }
                     else:
+                        # Held buffer was no tool signal and no enabled bare-JSON call: a leading ``{`` is an
+                        # ordinary JSON answer and must be shown; any other partial-markup prefix is dropped.
+                        _held = strip_llama3_leading_sentinels(content_buffer.lstrip())
+                        if _held.startswith("{") and not _suppress_visible_output:
+                            yield {"type": "content", "text": _held}
                         return
 
                 # ── STREAMING path: no tool call ──
                 if detect_state == _S_STREAMING:
-                    # Safety net: check for XML tool signals in content. The
+                    # Safety net: re-parse the full content for tool calls. The
                     # route layer resets prev_text on tool_start, so post-tool
                     # synthesis streams correctly even if content was emitted
                     # before the tool XML.
-                    _safety_tc = None
-                    if any(s in content_accum for s in _tool_xml_signals):
-                        _safety_tc = self._parse_tool_calls_from_text(
-                            content_accum,
-                            allow_incomplete = auto_heal_tool_calls,
-                        )
+                    # Unconditional (not gated on _tool_xml_signals): bare-JSON and Gemma wrapper-less
+                    # calls carry no XML signal, so a signal gate would let them slip past.
+                    _safety_tc = self._parse_tool_calls_from_text(
+                        content_accum,
+                        allow_incomplete = auto_heal_tool_calls,
+                        enabled_tool_names = _enabled_tool_names,
+                    )
                     if not _safety_tc:
                         # ── Re-prompt on plan-without-action ──
                         # If the model described its intent (forward-looking
@@ -8978,10 +9098,13 @@ class LlamaCppBackend:
                             for i in sorted(tool_calls_acc)
                             if (tool_calls_acc[i].get("function", {}).get("name", "").strip())
                         ] or None
-                    if not tool_calls and any(s in content_accum for s in _tool_xml_signals):
+                    if not tool_calls:
+                        # Unconditional re-parse: we only reach DRAINING when the buffer looked like a
+                        # call, and bare-JSON / Gemma wrapper-less calls carry no XML signal to gate on.
                         tool_calls = self._parse_tool_calls_from_text(
                             content_accum,
                             allow_incomplete = auto_heal_tool_calls,
+                            enabled_tool_names = _enabled_tool_names,
                         )
                     if tool_calls and not has_structured_tc:
                         content_text = _strip_tool_markup(
@@ -8989,6 +9112,11 @@ class LlamaCppBackend:
                             final = True,
                             force = True,
                         )
+                        # ``_strip_tool_markup`` only knows XML; also drop a leading bare-JSON call so the
+                        # executed call isn't replayed as text or next-turn history.
+                        content_text = strip_leading_bare_json_call(
+                            content_text, _enabled_tool_names
+                        )
                     if tool_calls:
                         logger.info(
                             f"Parsed {len(tool_calls)} tool call(s) from "
@@ -9002,6 +9130,13 @@ class LlamaCppBackend:
                         if content_accum:
                             # Strip leaked tool-call XML before yielding.
                             content_accum = _strip_tool_markup(content_accum, final = True)
+                        # A truncated bare-JSON call has no XML markup to strip and didn't parse. With
+                        # Auto-Heal on, drop a leading ENABLED-tool fragment (ordinary JSON answers untouched);
+                        # off keeps it visible per the strict contract.
+                        if content_accum and active_tools and auto_heal_tool_calls:
+                            content_accum = strip_leading_bare_json_call(
+                                content_accum, _enabled_tool_names
+                            )
                         if content_accum:
                             yield {"type": "content", "text": content_accum}
                         _meta = _build_metadata_event(
@@ -9019,6 +9154,29 @@ class LlamaCppBackend:
                 _accumulated_predicted_ms += _it.get("predicted_ms", 0)
                 _accumulated_predicted_n += _it.get("predicted_n", 0)
 
+                # Collapse exact-duplicate calls and cap the count for the TEXTUAL
+                # fallback (mirrors the safetensors loop; see _MAX_TOOL_CALLS_PER_TURN).
+                if tool_calls and not has_structured_tc and len(tool_calls) > 1:
+                    _seen_keys: set = set()
+                    _deduped: list = []
+                    for _tc in tool_calls:
+                        _fn = _tc.get("function", {}) or {}
+                        _key = (_fn.get("name", ""), str(_fn.get("arguments", "")))
+                        if _key in _seen_keys:
+                            continue
+                        _seen_keys.add(_key)
+                        _deduped.append(_tc)
+                        if len(_deduped) >= _MAX_TOOL_CALLS_PER_TURN:
+                            break
+                    if len(_deduped) != len(tool_calls):
+                        logger.info(
+                            "GGUF textual fallback: collapsed %d repeated tool call(s) "
+                            "in one turn to %d",
+                            len(tool_calls),
+                            len(_deduped),
+                        )
+                    tool_calls = _deduped
+
                 # disable_parallel_tool_use: execute only the first tool call
                 # this turn. Truncate before building assistant_msg so the
                 # conversation stays consistent and extra calls are never executed.
@@ -9144,6 +9302,8 @@ class LlamaCppBackend:
                             _kb_search_count += 1
                     completion = tool_controller.record_result(decision, result)
                     resolved_provisional_tool_call_ids.add(decision.tool_call_id)
+                    # A tool ran this turn, so it counts against the caller's budget.
+                    _turn_executed_real_tool = True
                     yield completion.tool_end_event()
                     conversation.append(completion.tool_message())
 
@@ -9167,6 +9327,12 @@ class LlamaCppBackend:
                 if tool_controller.force_final_answer or not tool_controller.active_tools():
                     _append_budget_exhausted_nudge = False
                     break
+                # Count only real tool turns against the cap so reserved re-prompt slots can't become
+                # extra tool rounds; a no-op correction turn doesn't consume budget (GGUF parity).
+                if _turn_executed_real_tool:
+                    _tool_iters_done += 1
+                    if _tool_iters_done >= max_tool_iterations:
+                        break
                 continue
 
             except httpx.ConnectError:

studio/backend/core/inference/mlx_inference.py

@@ -104,6 +104,9 @@ class MLXInferenceBackend:
     ) -> bool:
         import mlx.core as mx
 
+        # Keep the token so the native-template fallback can fetch a
+        # gated model's repo template later during generation.
+        self._hf_token = hf_token
         model_name = config.identifier if hasattr(config, "identifier") else str(config)
         is_vision = getattr(config, "is_vision", False)
 
@@ -168,11 +171,20 @@ class MLXInferenceBackend:
 
         self.active_model_name = model_name
         self.models[model_name] = {
+            # Per-model token for the native-template fallback (matches transformers).
+            "hf_token": hf_token,
+            # Per-model consent for the native-template reload: re-use the exact
+            # trust_remote_code this model was loaded with (matches transformers).
+            "trust_remote_code": trust_remote_code,
             "model": self._model,
             "tokenizer": self._tokenizer,
             "processor": self._processor,
             "is_vision": is_vision,
             "is_lora": getattr(config, "is_lora", False),
+            # For a LoRA adapter the native chat template lives on the base model.
+            "base_model": getattr(config, "base_model", None)
+            if getattr(config, "is_lora", False)
+            else None,
             "is_audio": False,
             "audio_type": None,
             "has_audio_input": False,
@@ -355,6 +367,7 @@ class MLXInferenceBackend:
 
         from core.inference.chat_template_helpers import (
             apply_chat_template_for_generation,
+            render_with_native_template_fallback,
         )
 
         prompt = apply_chat_template_for_generation(
@@ -368,6 +381,25 @@ class MLXInferenceBackend:
         if prompt is None:
             raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible")
 
+        # Same parity fix as the transformers backend: if the template dropped the
+        # requested tools, fall back to the native template so MLX text models keep
+        # advertising them. ``self._tokenizer`` is this entry's model_info tokenizer,
+        # so probe and native render share a renderer. (The VLM path renders via the
+        # processor for image tokens and is intentionally not wired here.)
+        model_info = self.models.get(self.active_model_name, {})
+        prompt = render_with_native_template_fallback(
+            formatted_prompt = prompt,
+            tokenizer = self._tokenizer,
+            model_info = model_info,
+            active_model_name = self.active_model_name,
+            messages = messages,
+            tools = tools,
+            enable_thinking = enable_thinking,
+            reasoning_effort = reasoning_effort,
+            preserve_thinking = preserve_thinking,
+            hf_token = model_info.get("hf_token"),
+        )
+
         sampler = make_sampler(
             temp = temperature,
             top_p = top_p,

studio/backend/core/inference/passthrough_healing.py

@@ -29,10 +29,26 @@ import os
 from collections.abc import Mapping
 from typing import Any, Optional
 
-from core.inference.tool_call_parser import TOOL_XML_SIGNALS, has_tool_signal
 from core.inference.tool_loop_controller import coerce_tool_arguments
 from core.tool_healing import parse_tool_calls_from_text
 
+# Signals limited to the formats parse_tool_calls_from_text (core.tool_healing)
+# actually promotes. The parser module's broader signal list also covers Llama
+# <|python_tag|> and Mistral [TOOL_CALLS] for the streaming DRAIN buffers whose
+# full parser handles them; buffering those here would hold a streamed
+# client-tool call until finalization and then flush it as prose (this healer
+# cannot promote them), so the passthrough keeps its own aligned list.
+_HEAL_SIGNALS = (
+    "<tool_call>",
+    "<|tool_call>",
+    "<function=",
+)
+
+
+def _has_heal_signal(text: str) -> bool:
+    return any(s in text for s in _HEAL_SIGNALS)
+
+
 # Read once at import (same convention as the other UNSLOTH_* switches).
 _HEALING_DISABLED = os.environ.get("UNSLOTH_DISABLE_TOOL_CALL_HEALING", "0") == "1"
 # Nudging is OPT-IN: per-request nudge_tool_calls=true, or flip the process
@@ -44,7 +60,7 @@ def nudge_enabled(request_flag: Optional[bool]) -> bool:
     return _NUDGE_DEFAULT if request_flag is None else bool(request_flag)
 
 
-_MAX_SIGNAL_LEN = max(len(s) for s in TOOL_XML_SIGNALS)
+_MAX_SIGNAL_LEN = max(len(s) for s in _HEAL_SIGNALS)
 # A suspected-but-unclosed tool block larger than this is declared a false
 # alarm and flushed, bounding memory on a model rambling XML-lookalike text.
 _MAX_HOLD_CHARS = 64 * 1024
@@ -198,7 +214,7 @@ def heal_openai_message_events(
     if not isinstance(msg, dict) or msg.get("tool_calls"):
         return None
     content = msg.get("content")
-    if not isinstance(content, str) or not has_tool_signal(content):
+    if not isinstance(content, str) or not _has_heal_signal(content):
         return None
     parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True)
     tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None
@@ -248,7 +264,7 @@ def heal_openai_message(
 
 def _earliest_signal(buffer: str) -> int:
     best = -1
-    for signal in TOOL_XML_SIGNALS:
+    for signal in _HEAL_SIGNALS:
         index = buffer.find(signal)
         if index >= 0 and (best < 0 or index < best):
             best = index
@@ -275,7 +291,7 @@ def _partial_signal_suffix(buffer: str) -> int:
     """Length of the longest buffer suffix that is a proper prefix of a signal."""
     for length in range(min(len(buffer), _MAX_SIGNAL_LEN - 1), 0, -1):
         tail = buffer[-length:]
-        if any(signal.startswith(tail) for signal in TOOL_XML_SIGNALS):
+        if any(signal.startswith(tail) for signal in _HEAL_SIGNALS):
             return length
     return 0
 
@@ -508,7 +524,7 @@ def nudge_should_retry(
     if not message or message.get("tool_calls"):
         return False
     text = message.get("content")
-    if not isinstance(text, str) or not has_tool_signal(text):
+    if not isinstance(text, str) or not _has_heal_signal(text):
         return False
     return not _heal_would_promote(text, allowed_tools, tools)
 

studio/backend/core/inference/safetensors_agentic.py

@@ -21,12 +21,22 @@ from typing import Callable, Generator, Optional
 from loggers import get_logger
 
 from core.inference.tool_call_parser import (
+    _GEMMA_BARE_TC_PREFIX_RE,
+    _GEMMA_BARE_TC_RE,
     _TOOL_ALL_PATS,
+    _balanced_brace_end,
+    _strip_function_xml_calls,
+    _strip_gemma_wrapperless_calls,
+    _strip_glm_calls,
+    _strip_mistral_closed_calls,
+    _strip_mistral_reasoning,
     BUDGET_EXHAUSTED_NUDGE,
     RAG_MAX_SEARCHES_PER_TURN,
     RAG_SEARCH_CAP_NUDGE,
     TOOL_XML_SIGNALS,
     parse_tool_calls_from_text,
+    strip_leading_bare_json_call,
+    strip_llama3_leading_sentinels,
     strip_tool_markup,
 )
 from core.inference.tool_loop_controller import (
@@ -50,16 +60,63 @@ logger = get_logger(__name__)
 # Buffer cap while disambiguating a possible tool-call prefix.
 _MAX_BUFFER_CHARS = 32
 
+# Memory bound for holding a leading bare-JSON object whose top-level "{" never balances.
+_MAX_BARE_JSON_BUFFER = 16384
+
+# Forward-looking intent ("I'll", "First,", "Step 1:") = planning, not answering; nudge a call.
+# Negative lookahead drops negated forms ("I will not") so a refusal doesn't trigger it. Mirrors GGUF.
+_INTENT_SIGNAL = re.compile(
+    r"(?i)("
+    r"\b(i['’](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)"
+    r"|\b(?:first\b|step \d+:?|here['’]?s (?:my |the |a )?(?:plan|approach))"
+    r"|\b(?:now i|next i)\b"
+    r")"
+)
+_MAX_REPROMPTS = 3
+_REPROMPT_MAX_CHARS = 2000
+# Templated so the nudge names the caller's enabled tools, not a hardcoded set. Mirrors GGUF tool_hint.
+_REPROMPT_INSTRUCTION_TEMPLATE = (
+    "STOP. Do NOT write code or explain. You MUST call a tool NOW. Call {tool_hint} immediately."
+)
+
+# No grammar constraint here (unlike llama-server's lazy grammar): collapse
+# exact-duplicate calls and cap the count so a runaway turn cannot fan out.
+_MAX_TOOL_CALLS_PER_TURN = 8
+
+
+def _active_tool_names(active_tools: list[dict]) -> list[str]:
+    names = [
+        (tool.get("function") or {}).get("name")
+        for tool in active_tools
+        if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
+    ]
+    return [name for name in names if name]
+
 
 def strip_tool_markup_streaming(
     text: str,
     *,
     auto_heal_tool_calls: bool = True,
     tool_protocol_active: bool = False,
+    enabled_tool_names: Optional[set] = None,
 ) -> str:
-    """Strip open-ended tool XML from display text without trimming whitespace."""
+    """Strip open-ended tool XML from display text without trimming whitespace.
+    ``enabled_tool_names`` gates the markerless Gemma ``call:NAME{...}`` strip so a
+    disabled/example name in prose is kept (mirrors the parser gate)."""
     if not (auto_heal_tool_calls or tool_protocol_active):
         return text
+    # Mirror the final strip's scan order so streaming and final display agree:
+    # balanced strips first (nested JSON removed whole), then the guarded
+    # function-XML/GLM scans that close at each call's REAL terminator, so literal
+    # markup inside argument values is data and trailing prose survives. No final
+    # trim so streaming length comparisons hold. Leading Magistral [THINK]...[/THINK]
+    # is dropped (bracket form, not the reasoning channel's <think>); an unclosed
+    # [THINK] holds until [/THINK] so the cleaned text stays monotonic.
+    text = _strip_mistral_reasoning(text)
+    text = _strip_mistral_closed_calls(text)
+    text = _strip_gemma_wrapperless_calls(text, enabled_tool_names)
+    text = _strip_function_xml_calls(text, final = True)
+    text = _strip_glm_calls(text, final = True)
     for pat in _TOOL_ALL_PATS:
         text = pat.sub("", text)
     return text
@@ -70,10 +127,11 @@ def _strip_tool_markup_final(
     *,
     auto_heal_tool_calls: bool,
     tool_protocol_active: bool = False,
+    enabled_tool_names: Optional[set] = None,
 ) -> str:
     if not (auto_heal_tool_calls or tool_protocol_active):
         return text
-    return strip_tool_markup(text, final = True)
+    return strip_tool_markup(text, final = True, enabled_tool_names = enabled_tool_names)
 
 
 def _status_for_tool(tool_name: str, arguments: dict) -> str:
@@ -81,6 +139,14 @@ def _status_for_tool(tool_name: str, arguments: dict) -> str:
     return status_for_tool(tool_name, arguments)
 
 
+def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) -> bool:
+    """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False."""
+    probe = strip_llama3_leading_sentinels(text.lstrip())
+    if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)):
+        return False
+    return strip_leading_bare_json_call(probe, enabled_tool_names) != probe
+
+
 _FUNCTION_SIGNAL_RE = re.compile(r"<function=([\w-]+)>")
 _TOOL_CALL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w-]+)"')
 
@@ -198,6 +264,11 @@ def run_safetensors_tool_loop(
     kb_search_count = 0
     final_attempt_done = False
     next_call_id = 0
+    reprompt_count = 0
+    # Real tool-call turns completed. Only turns that actually executed a tool count
+    # against ``max_tool_iterations``; a duplicate/disabled no-op correction turn (and a
+    # plan-without-action re-prompt) must not consume budget, matching the GGUF loop.
+    _executed_tool_iters = 0
 
     def _tool_succeeded(tool_name: str) -> bool:
         key_prefix = f"{tool_name}:"
@@ -215,9 +286,13 @@ def run_safetensors_tool_loop(
     _state_streaming = 1
     _state_draining = 2
 
-    for iteration in range(max_tool_iterations + 1):
+    # Reserve re-prompt slots so they don't eat the caller's tool budget.
+    _extra_iters = _MAX_REPROMPTS if max_tool_iterations > 0 else 0
+    for iteration in range(max_tool_iterations + _extra_iters + 1):
         if cancel_event is not None and cancel_event.is_set():
             return
+        # Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget.
+        _turn_executed_real_tool = False
 
         if final_attempt_done:
             active_tools: list[dict] = []
@@ -229,6 +304,8 @@ def run_safetensors_tool_loop(
 
         tool_protocol_active = not final_attempt_done and (unrestricted_tools or bool(active_tools))
         tool_xml_signals = TOOL_XML_SIGNALS if tool_protocol_active else ()
+        # Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call.
+        _enabled_tool_names = None if unrestricted_tools else set(_active_tool_names(active_tools))
 
         detect_state = _state_buffering
         content_buffer = ""
@@ -315,6 +392,7 @@ def run_safetensors_tool_loop(
                         before_tool,
                         auto_heal_tool_calls = auto_heal_tool_calls,
                         tool_protocol_active = tool_protocol_active,
+                        enabled_tool_names = _enabled_tool_names,
                     )
                     if len(cleaned_before) > len(last_emitted):
                         last_emitted = cleaned_before
@@ -345,6 +423,7 @@ def run_safetensors_tool_loop(
                     cumulative_display,
                     auto_heal_tool_calls = auto_heal_tool_calls,
                     tool_protocol_active = tool_protocol_active,
+                    enabled_tool_names = _enabled_tool_names,
                 )
                 if len(cleaned) > len(last_emitted):
                     last_emitted = cleaned
@@ -367,6 +446,64 @@ def run_safetensors_tool_loop(
                     is_prefix = True
                     break
 
+            # Llama-3.2 ``custom_tools`` emits a bare ``{"name":..,"parameters":..}`` with no XML
+            # signal. Hold a leading ``{`` (after any sentinel) until it closes: drain if it parses
+            # as a call, else stream as content. Non-call text is always recovered downstream.
+            bare_probe = strip_llama3_leading_sentinels(stripped)
+            if (
+                not is_match
+                and not is_prefix
+                and tool_protocol_active
+                and bare_probe.startswith("{")
+            ):
+                if _balanced_brace_end(bare_probe, 0) is None:
+                    if len(stripped) < _MAX_BARE_JSON_BUFFER:
+                        continue  # object still open -- keep buffering
+                    elif _looks_like_enabled_bare_json(bare_probe, _enabled_tool_names):
+                        # Oversized still-open ENABLED-tool call: stop holding (memory bound) but
+                        # DRAIN instead of leaking the raw prefix; a giant ordinary JSON answer still streams.
+                        detect_state = _state_draining
+                        continue
+                elif parse_tool_calls_from_text(
+                    content_buffer,
+                    id_offset = next_call_id,
+                    allow_incomplete = auto_heal_tool_calls,
+                    enabled_tool_names = _enabled_tool_names,
+                ):
+                    # Closed object that parses as a bare-JSON call -- drain silently.
+                    detect_state = _state_draining
+                    continue
+                # Closed non-call object (or oversized non-call) -- stream as text.
+
+            # Gemma wrapper-less ``call:NAME{...}`` has no tool_xml_signals entry:
+            # buffer it here or it streams raw until the end-of-turn safety net.
+            # ``(?<!\w)`` keeps "recall:" out; the prefix regex is whitespace-tolerant.
+            if (
+                not is_match
+                and not is_prefix
+                and tool_protocol_active
+                and (
+                    "call:".startswith(stripped)
+                    or _GEMMA_BARE_TC_PREFIX_RE.match(stripped) is not None
+                    or _GEMMA_BARE_TC_RE.match(stripped) is not None
+                )
+            ):
+                if _GEMMA_BARE_TC_RE.match(stripped):
+                    detect_state = _state_draining
+                    continue
+                # A ``call:`` / ``call:partial_name`` prefix with no ``{`` yet: keep
+                # buffering the variable-length name instead of leaking ``call:longname``.
+                # Names can exceed 32 chars (OpenAI 64, MCP longer), so a fixed cap would
+                # flush real calls raw. The prefix regex self-terminates on ordinary prose
+                # and the ``{`` drains above; bound generously like the bare-JSON path.
+                if _GEMMA_BARE_TC_PREFIX_RE.match(stripped) is not None:
+                    if len(stripped) < _MAX_BARE_JSON_BUFFER:
+                        continue
+                    detect_state = _state_draining
+                    continue
+                if len(stripped) < _MAX_BUFFER_CHARS:
+                    continue  # bare "call:" prefix still forming
+
             if is_match:
                 # Tool signal -- flush any visible prefix before DRAINING
                 # so the route sends it before tool_start.
@@ -375,6 +512,7 @@ def run_safetensors_tool_loop(
                     cumulative_display,
                     auto_heal_tool_calls = auto_heal_tool_calls,
                     tool_protocol_active = tool_protocol_active,
+                    enabled_tool_names = _enabled_tool_names,
                 )
                 if len(cleaned) > len(last_emitted):
                     last_emitted = cleaned
@@ -407,6 +545,7 @@ def run_safetensors_tool_loop(
                     cumulative_display,
                     auto_heal_tool_calls = auto_heal_tool_calls,
                     tool_protocol_active = tool_protocol_active,
+                    enabled_tool_names = _enabled_tool_names,
                 )
                 if len(cleaned) > len(last_emitted):
                     last_emitted = cleaned
@@ -419,44 +558,76 @@ def run_safetensors_tool_loop(
         if detect_state == _state_buffering:
             # Buffer never resolved -- tool XML or plain content?
             stripped = content_buffer.lstrip()
+            _bare_eos = strip_llama3_leading_sentinels(stripped)
             if (
                 stripped
                 and tool_protocol_active
                 and any(sig in stripped for sig in tool_xml_signals)
             ):
                 detect_state = _state_draining
+            elif tool_protocol_active and _looks_like_enabled_bare_json(
+                _bare_eos, _enabled_tool_names
+            ):
+                # A held bare-JSON ENABLED-tool fragment has no XML signal; DRAIN it (an ordinary
+                # JSON answer falls through to the else and streams as content, GGUF parity).
+                detect_state = _state_draining
             else:
+                # Drain and fall through to STREAMING so the intent re-prompt + safety-net parser
+                # still fire on short emissions like "Let me search." that never exit BUFFERING.
                 if content_buffer:
                     cumulative_display += content_buffer
-                    yield {
-                        "type": "content",
-                        "text": _strip_tool_markup_final(
-                            cumulative_display,
-                            auto_heal_tool_calls = auto_heal_tool_calls,
-                            tool_protocol_active = False,
-                        ),
-                    }
-                yield {"type": "status", "text": ""}
-                return
+                    cleaned = strip_tool_markup(
+                        cumulative_display, final = True, enabled_tool_names = _enabled_tool_names
+                    )
+                    if len(cleaned) > len(last_emitted):
+                        last_emitted = cleaned
+                        yield {"type": "content", "text": cleaned}
+                detect_state = _state_streaming
 
         if detect_state == _state_streaming:
-            # No tool detected mid-stream -- check for late tool XML.
-            safety_tc = None
-            saw_tool_signal = tool_protocol_active and any(
-                sig in content_accum for sig in tool_xml_signals
+            # Run the parser even with no XML signal (the Llama-3.2 bare-JSON form carries none); it's
+            # strict so plain answers stay untouched. Mirrors GGUF.
+            safety_tc = parse_tool_calls_from_text(
+                content_accum,
+                id_offset = next_call_id,
+                allow_incomplete = auto_heal_tool_calls,
+                enabled_tool_names = _enabled_tool_names,
             )
-            if saw_tool_signal:
-                safety_tc = parse_tool_calls_from_text(
-                    content_accum,
-                    id_offset = next_call_id,
-                    allow_incomplete = auto_heal_tool_calls,
-                )
             if not safety_tc:
-                # Final answer: if a literal tool marker in prose was stripped
-                # during streaming but did not parse as a real call, restore the
-                # raw cumulative text for core callers. Route-level cleanup can
-                # still apply the Auto-Heal display policy.
-                if saw_tool_signal and content_accum:
+                # Re-prompt only when the model planned without acting (intent
+                # signal); "4" / "Hello!" never trigger. Mirrors GGUF.
+                _stripped = content_accum.strip()
+                if (
+                    tools
+                    and auto_heal_tool_calls
+                    and reprompt_count < _MAX_REPROMPTS
+                    and 0 < len(_stripped) < _REPROMPT_MAX_CHARS
+                    and _INTENT_SIGNAL.search(_stripped)
+                    and not final_attempt_done
+                ):
+                    reprompt_count += 1
+                    logger.info(
+                        "Safetensors re-prompt %d/%d: model planned without "
+                        "calling tools (%d chars)",
+                        reprompt_count,
+                        _MAX_REPROMPTS,
+                        len(_stripped),
+                    )
+                    tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool"
+                    conversation.append({"role": "assistant", "content": _stripped})
+                    conversation.append(
+                        {
+                            "role": "user",
+                            "content": _REPROMPT_INSTRUCTION_TEMPLATE.format(tool_hint = tool_hint),
+                        }
+                    )
+                    yield {"type": "status", "text": ""}
+                    continue
+
+                # Final answer. If a literal tool marker in prose was buffered but
+                # never parsed as a call, restore the raw text so the prose surfaces
+                # in full; route-level cleanup still applies the Auto-Heal policy.
+                if content_accum and any(sig in content_accum for sig in tool_xml_signals):
                     yield {"type": "content", "text": content_accum}
                 yield {"type": "status", "text": ""}
                 return
@@ -465,6 +636,7 @@ def run_safetensors_tool_loop(
                 content_accum,
                 auto_heal_tool_calls = auto_heal_tool_calls,
                 tool_protocol_active = True,
+                enabled_tool_names = _enabled_tool_names,
             )
             logger.info(
                 "Safetensors safety net: parsed %d tool call(s) from streamed content",
@@ -476,20 +648,25 @@ def run_safetensors_tool_loop(
                 content_accum,
                 id_offset = next_call_id,
                 allow_incomplete = auto_heal_tool_calls,
+                enabled_tool_names = _enabled_tool_names,
             )
             if not tool_calls:
                 # Parser found nothing. Auto-Heal-enabled display cleanup
                 # strips unparseable tool XML; disabled Auto-Heal preserves
                 # the raw text so literal/malformed markup stays visible.
                 if content_accum:
-                    yield {
-                        "type": "content",
-                        "text": _strip_tool_markup_final(
-                            content_accum,
-                            auto_heal_tool_calls = auto_heal_tool_calls,
-                            tool_protocol_active = False,
-                        ),
-                    }
+                    _drain_text = _strip_tool_markup_final(
+                        content_accum,
+                        auto_heal_tool_calls = auto_heal_tool_calls,
+                        tool_protocol_active = False,
+                        enabled_tool_names = _enabled_tool_names,
+                    )
+                    # Drained bare-JSON call that didn't parse: with Auto-Heal on, drop the fragment
+                    # (plain JSON answers are left untouched); off keeps it visible per the strict contract.
+                    if tool_protocol_active and auto_heal_tool_calls:
+                        _drain_text = strip_leading_bare_json_call(_drain_text, _enabled_tool_names)
+                    if _drain_text:
+                        yield {"type": "content", "text": _drain_text}
                 if provisional_render_html_started and not provisional_resolved:
                     provisional_resolved = True
                     yield {
@@ -505,10 +682,14 @@ def run_safetensors_tool_loop(
                 content_accum,
                 auto_heal_tool_calls = auto_heal_tool_calls,
                 tool_protocol_active = True,
+                enabled_tool_names = _enabled_tool_names,
             )
 
         if tool_calls:
             next_call_id += len(tool_calls)
+            # Strip a leading bare-JSON call from the kept content so it isn't replayed as text or
+            # next-turn history (``_strip_tool_markup_final`` only knows XML). No-op for plain JSON answers.
+            content_text = strip_leading_bare_json_call(content_text, _enabled_tool_names)
 
         if final_attempt_done:
             # Final-answer turn re-called a tool -- stop the loop.
@@ -517,6 +698,27 @@ def run_safetensors_tool_loop(
             yield {"type": "status", "text": ""}
             return
 
+        # Collapse exact-duplicate calls and cap the count (runaway-turn guard).
+        if tool_calls:
+            seen_keys: set = set()
+            deduped: list = []
+            for _tc in tool_calls:
+                _fn = _tc.get("function", {}) or {}
+                _key = (_fn.get("name", ""), str(_fn.get("arguments", "")))
+                if _key in seen_keys:
+                    continue
+                seen_keys.add(_key)
+                deduped.append(_tc)
+                if len(deduped) >= _MAX_TOOL_CALLS_PER_TURN:
+                    break
+            if len(deduped) != len(tool_calls):
+                logger.info(
+                    "Safetensors: collapsed %d repeated tool call(s) in one turn to %d",
+                    len(tool_calls),
+                    len(deduped),
+                )
+            tool_calls = deduped
+
         assistant_msg: dict = {"role": "assistant", "content": content_text}
         assistant_appended = False
 
@@ -634,6 +836,8 @@ def run_safetensors_tool_loop(
             completion = tool_controller.record_result(decision, result)
             if provisional_match:
                 provisional_resolved = True
+            # A tool ran this turn, so it counts against the caller's budget.
+            _turn_executed_real_tool = True
             yield completion.tool_end_event()
             conversation.append(completion.tool_message())
 
@@ -646,7 +850,11 @@ def run_safetensors_tool_loop(
         if not unrestricted_tools and not tool_controller.active_tools():
             final_attempt_done = True
             continue
-        if iteration + 1 >= max_tool_iterations and not final_attempt_done:
+        # Count only turns that executed a tool against the cap; a no-op correction turn doesn't
+        # consume budget so the model gets its nudge and another tool-enabled turn (GGUF parity).
+        if _turn_executed_real_tool:
+            _executed_tool_iters += 1
+        if _executed_tool_iters >= max_tool_iterations and not final_attempt_done:
             # Budget exhausted; nudge a final plain answer.
             final_attempt_done = True
             conversation.append({"role": "user", "content": BUDGET_EXHAUSTED_NUDGE})

studio/backend/core/tool_healing.py

@@ -10,40 +10,65 @@ orchestrator, structlog, httpx, or the rest of the studio backend.
 import json
 import re
 
-# Pre-compiled patterns for tool XML stripping. The hyphen in the name
-# char-class lets dashed MCP tool/parameter names (mcp__srv__list-issues,
-# issue-number) parse alongside the built-ins.
+# Strip patterns. The name-class hyphen matches dashed MCP names. Closed pairs
+# strip first so a closed call goes as a unit before any to-EOF sweep reaches
+# nested markup; only the final list adds the .*$ EOF sweeps.
+_TC_JSON_CLOSED_PAT = re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL)
+_TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?<tool_call\|>", re.DOTALL)
+_TC_FUNC_CLOSED_PAT = re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL)
+_TC_GEMMA_END_PAT = re.compile(r"<tool_call\|>")
 _TOOL_CLOSED_PATS = [
-    re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
-    re.compile(r"<\|tool_call>.*?<tool_call\|>", re.DOTALL),
-    re.compile(r"<tool_call\|>"),
-    re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL),
+    _TC_JSON_CLOSED_PAT,
+    _TC_GEMMA_CLOSED_PAT,
+    _TC_FUNC_CLOSED_PAT,
+    _TC_GEMMA_END_PAT,
 ]
 _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
-    re.compile(r"<tool_call>.*$", re.DOTALL),
     re.compile(r"<\|tool_call>.*$", re.DOTALL),
+    re.compile(r"<tool_call>.*$", re.DOTALL),
     re.compile(r"<function=[\w-]+>.*$", re.DOTALL),
 ]
+# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in
+# their argument data cannot make the helper truncate the block and its tail.
+_TOOL_CLOSED_BLOCK_PATS = [_TC_JSON_CLOSED_PAT, _TC_FUNC_CLOSED_PAT]
+# A lazy closed-pair pattern whose close token is absent rescans to EOF from
+# every opener (quadratic, re-run per streamed token); skip that doomed pass.
+_PAT_REQUIRED_TOKEN = {
+    _TC_JSON_CLOSED_PAT: "</tool_call>",
+    _TC_GEMMA_CLOSED_PAT: "<tool_call|>",
+    _TC_FUNC_CLOSED_PAT: "</function>",
+}
+
+
+def strip_tool_patterns(text: str, patterns) -> str:
+    """Apply ``patterns`` in order, skipping closed-pair passes with no close token."""
+    for pat in patterns:
+        token = _PAT_REQUIRED_TOKEN.get(pat)
+        if token is not None and token not in text:
+            continue
+        text = pat.sub("", text)
+    return text
+
 
 # Pre-compiled patterns for tool-call XML parsing.
 _TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
-_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>call:([\w-]+)\s*\{")
+# Name class allows dots/hyphens for dotted Gemma names; whitespace-tolerant around
+# ``call`` / ``:`` since drift emits ``call: name{`` and ``call : name{``.
+_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w.\-]+)\s*\{")
 _TC_FUNC_START_RE = re.compile(r"<function=([\w-]+)>\s*")
 _TC_END_TAG_RE = re.compile(r"</tool_call>")
 _TC_GEMMA_END_TAG_RE = re.compile(r"<tool_call\|>")
 _TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
-_TC_PARAM_START_RE = re.compile(r"<parameter=([\w-]+)>\s*")
+# Horizontal whitespace only so the newline + value indentation survive (_trim_param_value trims one newline).
+_TC_PARAM_START_RE = re.compile(r"<parameter=([\w-]+)>[^\S\n]*")
 _TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
 _GEMMA_QUOTE = '<|"|>'
 _PARAM_CLOSE_TAG = "</parameter>"
 _FUNC_CLOSE_TAG = "</function>"
-# A bare (unquoted) Gemma value ends at `}` or at a comma that begins the next
-# `key:` pair. A comma NOT followed by a key token is part of the value (e.g.
-# `location:New York, NY`), so it must not terminate the value. The key token
-# must be identifier-shaped (start with a letter or underscore); a comma
-# followed by digits-then-colon is value text such as a timestamp or ratio
-# (`meet at 10:00, 11:00 tomorrow`), not a new key.
-_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w-]*\s*:")
+# A bare (unquoted) Gemma value ends at `}` or at a comma beginning the next
+# identifier-shaped `key:` pair; a comma before a non-key (`New York, NY`,
+# `10:00, 11:00`) stays in the value. Dots let a dotted key end the value.
+_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:")
 
 
 def _balanced_brace_end(
@@ -139,14 +164,8 @@ def _split_top_level_commas(src: str) -> list:
 
 
 def _quote_gemma_array_elements(body: str) -> str:
-    """Normalise the elements of a Gemma array value so json.loads succeeds.
-
-    Gemma may emit ``labels:[bug,ui]`` without per-element quotes, or arrays of
-    objects (``items:[{path:a}]``) whose keys/values also lack quotes; left
-    as-is json.loads fails and the whole call is dropped. Bare string elements
-    are quoted, object and nested-array elements are normalised recursively, and
-    quoted strings (already normalised from ``<|"|>``), numbers, and JSON
-    literals are preserved."""
+    """Normalise a Gemma array value (``labels:[bug,ui]``) so json.loads succeeds:
+    quote bare strings, recurse into objects/arrays, keep quoted/JSON literals."""
     out: list[str] = []
     for element in _split_top_level_commas(body):
         stripped = element.strip()
@@ -154,11 +173,9 @@ def _quote_gemma_array_elements(body: str) -> str:
             out.append(element)
             continue
         if stripped[0] == "{":
-            # Object element: quote its keys/bare values like a top-level object.
             out.append(_quote_gemma_object_keys(stripped))
             continue
         if stripped[0] == "[":
-            # Nested array: normalise its elements too.
             inner_end = _balanced_bracket_end(stripped, 0)
             if inner_end == len(stripped) - 1:
                 out.append("[" + _quote_gemma_array_elements(stripped[1:inner_end]) + "]")
@@ -223,7 +240,9 @@ def _quote_gemma_object_keys(src: str) -> str:
         while i < len(src) and src[i].isspace():
             i += 1
         key_name_start = i
-        while i < len(src) and (src[i].isalnum() or src[i] in "_-"):
+        # Dots match the parser's key/name charset: Gemma emits dotted argument keys
+        # (user.name:...) for namespaced schemas.
+        while i < len(src) and (src[i].isalnum() or src[i] in "_-."):
             i += 1
         key_name = src[key_name_start:i]
         colon_pos = i
@@ -235,15 +254,12 @@ def _quote_gemma_object_keys(src: str) -> str:
             parts.append(src[i:colon_pos])
             parts.append(":")
             i = colon_pos + 1
-            # Gemma may emit bare string values ({unit:celsius}); quote them so
-            # json.loads succeeds. JSON scalars/objects/arrays/quoted stay as-is.
+            # Quote bare string values ({unit:celsius}); JSON stays as-is.
             ws = i
             while i < len(src) and src[i].isspace():
                 i += 1
             parts.append(src[ws:i])
             if i < len(src) and src[i] == "[":
-                # Array value: quote bare string elements (e.g. labels:[bug,ui])
-                # so json.loads succeeds instead of dropping the call.
                 arr_end = _balanced_bracket_end(src, i)
                 if arr_end < 0:
                     parts.append(src[i:])
@@ -253,9 +269,7 @@ def _quote_gemma_object_keys(src: str) -> str:
                     i = arr_end + 1
             elif i < len(src) and src[i] not in '"{':
                 v_start = i
-                # Consume the bare value up to `}` or a comma that starts the
-                # next key:value pair; a comma inside the value (e.g.
-                # `New York, NY`) does not terminate it.
+                # Bare value: up to `}` or a comma that starts the next key:pair.
                 while i < len(src):
                     if src[i] == "}":
                         break
@@ -267,7 +281,8 @@ def _quote_gemma_object_keys(src: str) -> str:
                     json.loads(raw.strip())
                     parts.append(raw)
                 except (json.JSONDecodeError, ValueError):
-                    parts.append(json.dumps(raw.strip()) if raw.strip() else raw)
+                    # Quote bare value; empty ({k:}) becomes "" so json.loads sees {"k":""} not invalid {"k":}.
+                    parts.append(json.dumps(raw.strip()))
         else:
             parts.append(src[key_start:i])
     return "".join(parts)
@@ -291,9 +306,97 @@ def _inside_open_parameter(content: str, pos: int) -> bool:
         last_param_start = match.start()
     if last_param_start < 0:
         return False
-    last_param_close = content.rfind(_PARAM_CLOSE_TAG, 0, pos)
-    last_func_close = content.rfind(_FUNC_CLOSE_TAG, 0, pos)
-    return last_param_start > max(last_param_close, last_func_close)
+    # The parameter's OWN close tag decides: if it closes after ``pos`` the position is
+    # argument data (even across literal function closes); an unclosed one falls back to func close.
+    own_close = content.find(_PARAM_CLOSE_TAG, last_param_start)
+    if own_close >= 0:
+        return own_close > pos
+    func_close = content.find(_FUNC_CLOSE_TAG, last_param_start)
+    return func_close < 0 or pos < func_close
+
+
+def _func_close_index(content: str, body_start: int, body: str) -> int:
+    """Index in ``body`` of the first ``</function>`` that is not argument
+    data (not inside an open parameter value); -1 when every close is data.
+    Taking the LAST close swallowed prose between the real close and a
+    literal ``</function>`` mentioned later in the answer."""
+    idx = body.find(_FUNC_CLOSE_TAG)
+    while idx >= 0:
+        if not _inside_open_parameter(content, body_start + idx):
+            return idx
+        idx = body.find(_FUNC_CLOSE_TAG, idx + 1)
+    return -1
+
+
+def _trim_param_value(val: str) -> str:
+    """Trim only the wrapping newline (not str.strip) so code/diff argument indentation survives."""
+    if val.startswith("\n"):
+        val = val[1:]
+    if val.endswith("\n"):
+        val = val[:-1]
+    return val
+
+
+def _marker_coverage(content: str, markers) -> list[tuple[int, int]]:
+    """Coverage ``[start, end]`` per marker, used to skip markers that are another
+    call's data. Closes pair to markers via a per-format stack so an inner close
+    is not mistaken for the outer's. Unbalanced braces cover to EOF; balanced with
+    a paired close cover through it (markers before the close are data); balanced
+    without one cover only the braces, so a later sibling is still recovered."""
+    n = len(content)
+    brace_regions = [(s, be) for (s, be, _k, _m) in markers if be >= 0]
+    events = []  # (position, order) with order 0 = braces-done, 1 = close marker
+    for idx, (_start, brace_end, _kind, _m) in enumerate(markers):
+        if brace_end >= 0:
+            events.append((brace_end, 0, _kind, idx))
+    for kind, close_re in (("json", _TC_END_TAG_RE), ("gemma", _TC_GEMMA_END_TAG_RE)):
+        for cm in close_re.finditer(content):
+            # A close inside another call's balanced braces is quoted data; it
+            # must not pop an earlier close-less marker and swallow a sibling.
+            if any(s < cm.start() < be for s, be in brace_regions):
+                continue
+            events.append((cm.start(), 1, kind, cm.end()))
+    events.sort(key = lambda e: (e[0], e[1]))
+    waiting = {"json": [], "gemma": []}
+    close_end_for: dict[int, int] = {}
+    for _pos, order, kind, payload in events:
+        if order == 0:
+            waiting[kind].append(payload)  # marker index, now awaiting its close
+        elif waiting[kind]:
+            close_end_for[waiting[kind].pop()] = payload  # innermost open marker closes here
+    coverage = []
+    for idx, (start, brace_end, _kind, _m) in enumerate(markers):
+        if brace_end < 0:
+            coverage.append((start, n))
+        elif idx in close_end_for:
+            coverage.append((start, close_end_for[idx]))
+        else:
+            coverage.append((start, brace_end))
+    return coverage
+
+
+def _build_markers(content: str):
+    """JSON/Gemma tool markers as ``(start, brace_end, kind, match)`` in document
+    order; ``brace_end < 0`` marks an unbalanced (to-EOF) open."""
+    markers = []
+    for start_re, gemma, kind in (
+        (_TC_JSON_START_RE, False, "json"),
+        (_TC_GEMMA_START_RE, True, "gemma"),
+    ):
+        for m in start_re.finditer(content):
+            if _inside_open_parameter(content, m.start()):
+                continue
+            brace_end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = gemma)
+            markers.append((m.start(), brace_end, kind, m))
+    markers.sort(key = lambda c: c[0])
+    return markers
+
+
+def marker_coverage(content: str) -> list[tuple[int, int]]:
+    """Coverage spans of JSON/Gemma tool markers so other parsers can treat markup
+    inside a marker's coverage (even a marker that failed to parse) as that call's
+    data rather than a sibling call."""
+    return _marker_coverage(content, _build_markers(content))
 
 
 def parse_tool_calls_from_text(
@@ -317,47 +420,40 @@ def parse_tool_calls_from_text(
     """
     tool_calls: list[dict] = []
     call_spans: list[tuple] = []
-    # Collect every supported call format with spans, then emit in document
-    # order. A marker inside another call's argument string is data, not a
-    # separate executable call.
-    parsed_items = []  # (start, span_end, name, arguments)
-    candidates = []  # (start, brace_end, kind, match)
-    for m in _TC_JSON_START_RE.finditer(content):
-        if _inside_open_parameter(content, m.start()):
-            continue
-        end = _balanced_brace_end(content, m.end() - 1)
-        if end >= 0:
-            candidates.append((m.start(), end, "json", m))
-    for m in _TC_GEMMA_START_RE.finditer(content):
-        if _inside_open_parameter(content, m.start()):
-            continue
-        end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = True)
-        if end >= 0:
-            candidates.append((m.start(), end, "gemma", m))
-    candidates.sort(key = lambda c: c[0])
-
-    candidate_spans = [(s, e) for s, e, _kind, _m in candidates]
-    for idx, (start, end, kind, m) in enumerate(candidates):
-        if any(s <= start and end <= e for j, (s, e) in enumerate(candidate_spans) if j != idx):
+    # Collect JSON/Gemma markers; _marker_coverage decides nesting. A marker inside
+    # another call's coverage, or an open <parameter=> value, is data not executed.
+    markers = _build_markers(content)
+    coverage = _marker_coverage(content, markers)
+    parsed_items = []  # (start, span_end, name, arguments) in document order
+    for idx, (start, brace_end, kind, m) in enumerate(markers):
+        # A marker starting inside another's coverage is that call's data. The
+        # end is exclusive so a marker at a close's end is an adjacent sibling.
+        if any(s <= start < e for j, (s, e) in enumerate(coverage) if j != idx):
             continue
+        if brace_end < 0:
+            continue  # unclosed: not parseable; the fallback still excludes its XML
         if not allow_incomplete:
-            tail = content[end + 1 :].lstrip()
+            tail = content[brace_end + 1 :].lstrip()
             close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE
             if close_re.match(tail) is None:
                 continue
         try:
             if kind == "json":
-                obj = json.loads(content[m.end() - 1 : end + 1])
+                obj = json.loads(content[m.end() - 1 : brace_end + 1])
                 name = obj.get("name", "")
-                arguments = obj.get("arguments", {})
+                # Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside a Hermes <tool_call>).
+                arguments = obj.get("arguments")
+                if arguments is None:
+                    arguments = obj.get("parameters", {})
                 if isinstance(arguments, dict):
                     arguments = json.dumps(arguments)
             else:
                 name = m.group(1)
-                arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : end]))
+                arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end]))
         except (json.JSONDecodeError, ValueError):
             continue
-        span_end = end + 1
+        # Span reaches through the close tag when present, else just the braces.
+        span_end = brace_end + 1
         close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE
         ws = len(content[span_end:]) - len(content[span_end:].lstrip())
         close_m = close_re.match(content, span_end + ws)
@@ -365,11 +461,15 @@ def parse_tool_calls_from_text(
             span_end = close_m.end()
         parsed_items.append((start, span_end, name, arguments))
 
+    # Function-XML calls promote in document order alongside marker calls (the
+    # #6801 contract). A <function=> inside any marker's coverage is excluded --
+    # even if that marker failed to parse -- so nested XML cannot escape; one
+    # after a balanced close-less marker is a sibling, not swallowed to EOF.
     func_starts = [
         fm
         for fm in _TC_FUNC_START_RE.finditer(content)
         if not _inside_open_parameter(content, fm.start())
-        and not any(s <= fm.start() <= e for s, e in candidate_spans)
+        and not any(s <= fm.start() < e for s, e in coverage)
     ]
     for idx, fm in enumerate(func_starts):
         func_name = fm.group(1)
@@ -382,7 +482,7 @@ def parse_tool_calls_from_text(
             body_end = len(content)
         body_end = min(body_end, next_func)
         body = content[body_start:body_end]
-        close_idx = body.rfind(_FUNC_CLOSE_TAG)
+        close_idx = _func_close_index(content, body_start, body)
         if close_idx >= 0:
             span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG)
             body = body[:close_idx]
@@ -404,7 +504,7 @@ def parse_tool_calls_from_text(
                 val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
             else:
                 val = _TC_PARAM_CLOSE_RE.sub("", val)
-            arguments[pm.group(1)] = val.strip()
+            arguments[pm.group(1)] = _trim_param_value(val)
         else:
             valid_params = True
             for pidx, pm in enumerate(param_starts):
@@ -422,7 +522,7 @@ def parse_tool_calls_from_text(
                     val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
                 else:
                     val = _TC_PARAM_CLOSE_RE.sub("", val)
-                arguments[param_name] = val.strip()
+                arguments[param_name] = _trim_param_value(val)
             if not valid_params:
                 continue
 
@@ -444,11 +544,107 @@ def parse_tool_calls_from_text(
             }
         )
         call_spans.append((start, span_end))
+
     if with_spans:
         return tool_calls, call_spans
     return tool_calls
 
 
+def _strip_gemma_native_spans(text: str, *, final: bool) -> str:
+    """Remove complete Gemma-native spans, brace/quote-balanced so a literal
+    ``<tool_call|>`` in a quoted argument cannot truncate the span. An incomplete
+    span is dropped to EOF when ``final``, else kept (still streaming)."""
+    out: list[str] = []
+    cursor = 0
+    for match in _TC_GEMMA_START_RE.finditer(text):
+        start = match.start()
+        if start < cursor:
+            continue
+        brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True)
+        if brace_end < 0:
+            # Unbalanced: nothing completes from here on. Drop the rest if final,
+            # else keep it; stop either way (rescanning would be quadratic).
+            if final:
+                out.append(text[cursor:start])
+                cursor = len(text)
+            break
+        # Junk between } and <tool_call|> is malformed-call markup: strip through
+        # the close, keep text after it. No close anywhere means stop (linear).
+        close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1)
+        if close is None:
+            if final:
+                out.append(text[cursor:start])
+                cursor = len(text)
+            break
+        out.append(text[cursor:start])
+        cursor = close.end()
+    out.append(text[cursor:])
+    return "".join(out)
+
+
+def _gemma_span_ranges(text: str) -> list:
+    """``(start, end)`` of each complete Gemma-native span; same walk as
+    ``_strip_gemma_native_spans`` without stripping."""
+    ranges: list[tuple] = []
+    cursor = 0
+    for match in _TC_GEMMA_START_RE.finditer(text):
+        start = match.start()
+        if start < cursor:
+            continue
+        brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True)
+        if brace_end < 0:
+            break
+        close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1)
+        if close is None:
+            break
+        ranges.append((start, close.end()))
+        cursor = close.end()
+    return ranges
+
+
+def _strip_closed_blocks_outside_gemma(text: str) -> str:
+    """Closed JSON/function pre-pass that skips matches starting inside a complete
+    Gemma span: deleting across the span boundary would mangle the Gemma close and
+    truncate the tail. A skipped match resumes at the covering span's end, so a
+    real function-XML call after the span is still stripped."""
+    ranges = _gemma_span_ranges(text)
+    if not ranges:
+        return strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS)
+    for pat in _TOOL_CLOSED_BLOCK_PATS:
+        token = _PAT_REQUIRED_TOKEN.get(pat)
+        if token is not None and token not in text:
+            continue
+        out: list[str] = []
+        pos = 0
+        while True:
+            m = pat.search(text, pos)
+            if m is None:
+                out.append(text[pos:])
+                break
+            covering = next((r for r in ranges if r[0] <= m.start() < r[1]), None)
+            if covering is not None:
+                out.append(text[pos : covering[1]])
+                pos = covering[1]
+                continue
+            out.append(text[pos : m.start()])
+            pos = m.end()
+        new_text = "".join(out)
+        if new_text != text:
+            text = new_text
+            ranges = _gemma_span_ranges(text)
+    return text
+
+
+def strip_tool_markup_final(text: str) -> str:
+    """Final display strip, shared with the streaming wrappers so all paths order
+    the passes identically: Gemma-aware closed JSON/function blocks first, then
+    well-formed Gemma spans (quote-aware), then the regex sweeps mop up malformed
+    spans and drop any unclosed remainder to EOF. Whitespace is kept."""
+    text = _strip_closed_blocks_outside_gemma(text)
+    text = _strip_gemma_native_spans(text, final = True)
+    return strip_tool_patterns(text, _TOOL_ALL_PATS)
+
+
 def strip_tool_call_markup(text: str, *, final: bool = False) -> str:
     """Strip tool-call XML markup from text.
 
@@ -456,7 +652,9 @@ def strip_tool_call_markup(text: str, *, final: bool = False) -> str:
     When ``final`` is True, trailing incomplete tool-call blocks are removed
     too, and the result is stripped of surrounding whitespace.
     """
-    patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
-    for pat in patterns:
-        text = pat.sub("", text)
-    return text.strip() if final else text
+    if final:
+        return strip_tool_markup_final(text).strip()
+    # Non-final: same ordering as the final path, but incomplete blocks are kept.
+    text = _strip_closed_blocks_outside_gemma(text)
+    text = _strip_gemma_native_spans(text, final = False)
+    return strip_tool_patterns(text, _TOOL_CLOSED_PATS)

studio/backend/routes/inference.py

@@ -603,6 +603,17 @@ def _chat_content_chunk(completion_id, created, model_name, text) -> str:
     )
 
 
+def _chat_reasoning_chunk(completion_id, created, model_name, text) -> str:
+    """Like ``_chat_content_chunk`` but on ``reasoning_content`` (renders the UI thinking block)."""
+    return _chat_chunk_sse(
+        completion_id,
+        created,
+        model_name,
+        delta = ChoiceDelta(reasoning_content = text),
+        finish_reason = None,
+    )
+
+
 def _chat_final_chunk(completion_id, created, model_name, finish_reason) -> str:
     """Terminal stop chunk (empty delta) carrying the finish reason."""
     return _chat_chunk_sse(
@@ -841,17 +852,14 @@ class _SameTaskStreamingResponse(StreamingResponse):
         **kwargs,
     ) -> None:
         super().__init__(*args, **kwargs)
-        # Async callable invoked when the client disconnects before the body
-        # iterator is ever advanced. A generator that never started cannot run
-        # its own try/finally, so a stream that acquires resources before its
-        # first yield (the passthrough opens an upstream httpx stream eagerly)
-        # passes this to release them.
+        # Released when the client disconnects before the body iterator starts:
+        # its try/finally never runs, so a stream that opens resources before the
+        # first yield (the passthrough's upstream httpx stream) passes this.
         self._unstarted_cleanup = unstarted_cleanup
 
     async def __call__(self, scope, receive, send) -> None:
-        # Track whether the body iterator was ever advanced: send() only emits a
-        # body message after the generator yields its first chunk, so a failure
-        # before then means it never entered its try/finally.
+        # send() emits a body message only after the first chunk, so no body
+        # message means the generator never entered its try/finally.
         body_started = False
 
         async def _tracking_send(message) -> None:
@@ -862,15 +870,11 @@ class _SameTaskStreamingResponse(StreamingResponse):
 
         try:
             await self.stream_response(_tracking_send)
-        except OSError:
-            # Client disconnected mid-send.
+        except OSError:  # client disconnected mid-send
             if body_started:
-                # The generator produced at least one chunk and is suspended in
-                # its try/finally. Throw CancelledError into it (not aclose's
-                # GeneratorExit) so its `except asyncio.CancelledError` handler
-                # runs and finishes any api_monitor entry; GeneratorExit would
-                # skip it and only run `finally`. Fall back to aclose() without
-                # athrow.
+                # Generator is suspended in its try/finally: throw CancelledError
+                # (not aclose's GeneratorExit) so its handler finishes the
+                # api_monitor entry. Fall back to aclose() without athrow.
                 athrow = getattr(self.body_iterator, "athrow", None)
                 if athrow is not None:
                     try:
@@ -882,16 +886,16 @@ class _SameTaskStreamingResponse(StreamingResponse):
                     if aclose is not None:
                         await aclose()
             else:
-                # http.response.start failed before the body iterator advanced,
-                # so its try/finally never armed and aclose()/athrow() are no-ops
-                # on an unstarted generator. Release any resources acquired
-                # before the first yield via the explicit cleanup hook.
+                # Generator never started; aclose()/athrow() are no-ops on it, so
+                # release eager resources via the hook. getattr guards a response
+                # built through __new__ without __init__ (tests, pickling).
                 aclose = getattr(self.body_iterator, "aclose", None)
                 if aclose is not None:
                     await aclose()
-                if self._unstarted_cleanup is not None:
+                cleanup = getattr(self, "_unstarted_cleanup", None)
+                if cleanup is not None:
                     try:
-                        await self._unstarted_cleanup()
+                        await cleanup()
                     except Exception:
                         pass
             raise ClientDisconnect()
@@ -899,6 +903,16 @@ class _SameTaskStreamingResponse(StreamingResponse):
             await self.background()
 
 
+def _tracked_cancel_unstarted_cleanup(tracker):
+    """unstarted_cleanup that exits ``tracker`` on a pre-start disconnect, when
+    the generator's finally (which normally exits it) never runs."""
+
+    async def _cleanup() -> None:
+        tracker.__exit__(None, None, None)
+
+    return _cleanup
+
+
 async def _aclose_stream_resources(
     *,
     watchers = (),
@@ -1148,6 +1162,13 @@ from core.inference.key_exchange import decrypt_api_key
 from core.inference.model_ids import public_model_id
 from core.inference.api_monitor import api_monitor
 from core.inference.llama_http import nonstreaming_client
+from core.inference.tool_call_parser import (
+    _strip_function_xml_calls,
+    _strip_gemma_wrapperless_calls,
+    _strip_glm_calls,
+    _strip_mistral_closed_calls,
+)
+from core.inference.tool_call_parser import TOOL_XML_SIGNALS as _PARSER_TOOL_SIGNALS
 from core.inference.passthrough_healing import (
     StreamToolCallHealer,
     heal_gate,
@@ -1306,6 +1327,11 @@ async def artifact_preview_frame(allow_network: bool = False):
     )
 
 
+# Whitespace/escape-tolerant bare-JSON tool-template detector (matches pretty-printed and
+# JSON-escaped ``{"name":`` plus the ``"function"`` alias), mirroring the parser's tolerance.
+_BARE_JSON_NAME_MARKER_RE = _re.compile(r'\{\s*\\?"(?:name|function)\\?"\s*:')
+
+
 def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict:
     """Classify reasoning/tool capabilities via the GGUF classifier so flags
     match across backends. gpt-oss is overridden: Harmony routes reasoning and
@@ -1316,17 +1342,22 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict:
         model_identifier = model_id,
         log_source = "safetensors",
     )
-    # Our safetensors loop only parses <tool_call>{json}</tool_call>,
-    # <function=name>...</function>, and Gemma native <|tool_call>...<tool_call|>.
-    # Llama uses <|python_tag|>, Mistral uses [TOOL_CALLS]; advertising tools for
-    # those enables a pill the parser can't honour. GGUF is unaffected --
-    # llama-server normalises every format into structured deltas.
+    # Markers any supported parser recognises (template advertises tools but
+    # uses none -> drop the pill). Reuse the parser's own signal list so this
+    # gate never drifts (a hand-maintained copy lost the DeepSeek variants);
+    # ``<arg_key>`` is GLM's unique signal, absent from the shared set. The
+    # bare-JSON ``{"name":`` form is matched below with the whitespace/escape-
+    # tolerant ``_BARE_JSON_NAME_MARKER_RE`` so pretty-printed or escaped
+    # templates are not mis-classified as tool-less.
+    _PARSER_MARKERS = (
+        *_PARSER_TOOL_SIGNALS,
+        "<arg_key>",
+    )
     if (
         flags.get("supports_tools")
         and chat_template
-        and "<tool_call>" not in chat_template
-        and "<function=" not in chat_template
-        and "<|tool_call>" not in chat_template
+        and not any(m in chat_template for m in _PARSER_MARKERS)
+        and not _BARE_JSON_NAME_MARKER_RE.search(chat_template)
     ):
         logger.info(
             "safetensors: template advertises tools but uses an "
@@ -1347,6 +1378,39 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict:
     return flags
 
 
+def _sf_reasoning_prefill_mode(
+    features: dict,
+    enable_thinking: Optional[bool],
+    template: Optional[str] = None,
+    reasoning_effort: Optional[str] = None,
+) -> bool:
+    """Whether this request begins INSIDE an unclosed ``<think>`` (Qwen3/Qwen3.5/GLM prefill it).
+
+    Gated on the STANDARD ``<think>``/``</think>`` markers: a bespoke reasoning channel (e.g. gemma)
+    never emits ``</think>``, so prefilled mode would swallow the whole answer -- excluded, as are
+    gpt-oss and thinking-disabled requests. ``enable_thinking=None`` defaults ON, so plain requests prefill.
+    """
+    if features.get("reasoning_style") not in ("enable_thinking", "enable_thinking_effort"):
+        return False
+    tpl = template or ""
+    if "</think>" not in tpl and "<think>" not in tpl:
+        return False
+    if features.get("reasoning_always_on"):
+        return True
+    if not features.get("supports_reasoning"):
+        return False
+    if enable_thinking is False:
+        return False
+    # A reasoning_effort="none" request disables thinking for enable_thinking_effort
+    # (GLM-5.2) models the same way enable_thinking=False does (see
+    # ``_request_reasoning_kwargs``). Without this, the model emits no ``</think>`` and
+    # a plain answer is swallowed whole into reasoning_content, leaving the visible
+    # response empty.
+    if features.get("reasoning_style") == "enable_thinking_effort" and reasoning_effort == "none":
+        return False
+    return True
+
+
 def _effective_enable_tools(payload) -> Optional[bool]:
     """Resolve `payload.enable_tools` against the process-level tool policy.
 
@@ -1617,30 +1681,83 @@ def _apply_rag_nudge(nudge: str, tools: list[dict], *, rag_scope) -> str:
     return nudge + " " + _RAG_GROUNDING_NUDGE
 
 
-# Strip tool-call XML the speculative buffer in core/inference/llama_cpp.py
-# split across the visible/DRAIN boundary. Four leak shapes:
-#   1. well-formed `<tool_call>...</tool_call>` / `<function=...>...</function>`
-#   2. orphan opening to EOF (close was DRAINED)
-#   3. bare orphan close (open was DRAINED)
-#   4. tail-only `</parameter>` (outer close truncated by EOS); anchored to
-#      `\Z` so mid-text `<parameter>` in user code samples survives.
+# Strip leaked tool-call markup: every shared-parser format plus the four leak
+# shapes llama_cpp.py's speculative buffer splits across the visible/DRAIN
+# boundary. Mistral [TOOL_CALLS] uses the parser's balanced-brace helper (a
+# non-greedy regex would truncate nested JSON); the DeepSeek opener alternation
+# is the parser's own, so a signal we parse is never left un-stripped.
+from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC
+
 _TOOL_XML_RE = _re.compile(
-    # Hyphen in the name char-class matches MCP tool names with dashes
-    # (mcp__srv__list-issues) that would otherwise leak past this strip.
-    r"<(?:tool_call|function=[\w-]+)>.*?(?:</(?:tool_call|function)>|\Z)"
+    # Arm order/notes: the closed ``<function=...>`` arm runs first and extends
+    # to the call's REAL close so a literal ``</function>`` in a value does not
+    # leak the tail; the combined arm still catches ``<tool_call>`` and orphan
+    # tails. The python_tag arm bounds only on REAL Llama control sentinels
+    # (stopping at any ``<|`` truncated on literal ``<|x|>`` tokens in values).
+    # The last arms cover DeepSeek envelopes (all opener variants), Kimi section
+    # blocks, and bare Kimi calls. Name class ``[\w.\-]`` mirrors the parser.
+    # Those three arms carry a call-shaped lookahead (matching the parser's
+    # ``_TOOL_ALL_PATS``): a prose answer that merely mentions a marker
+    # (``See <|tool_call_begin|> in the docs``) is only stripped when a real
+    # call actually follows the marker, or the marker is a bare fragment at EOF.
+    r'<function(?:=[\w.\-]+|\s+name="[\w.\-]+")>(?:(?!<function(?:=[\w.\-]+|\s+name="[\w.\-]+")>).)*</function>'
+    r'|<(?:tool_call|function(?:=[\w.\-]+|\s+name="[\w.\-]+"))>.*?(?:</(?:tool_call|function)>|\Z)'
     r"|<\|tool_call>.*?(?:<tool_call\|>|\Z)"
     r"|</(?:tool_call|function)>"
     r"|<tool_call\|>"
-    r"|</parameter>\s*\Z",
+    r"|<\|python_tag\|>(?:[^<]|<(?!\|(?:eot_id|eom_id|python_tag|start_header_id|end_header_id|begin_of_text|finetune_right_pad_id)\|))*"
+    r"|"
+    + _DS_OPEN_SRC
+    + r"(?=\s*(?:<｜tool▁call▁begin｜>|function)|\s*$).*?(?:<｜tool▁calls▁end｜>|\Z)"
+    r"|<\|tool_calls_section_begin\|>(?=\s*<\|tool_call_begin\|>|\s*$).*?(?:<\|tool_calls_section_end\|>|\Z)"
+    r"|<\|tool_call_begin\|>(?=\s*[A-Za-z_][\w.\-]*:\d|\s*$).*?(?:<\|tool_call_end\|>|\Z)"
+    # ``</param>`` is the attribute-form alias of ``</parameter>`` (the parser accepts
+    # both); strip a tail-only orphan close of either spelling.
+    r"|</(?:parameter|param)>\s*\Z",
     _re.DOTALL,
 )
 
 
-def _strip_tool_xml_for_display(text: str, *, auto_heal_tool_calls: bool) -> str:
-    """Apply route-level XML leak cleanup only when Auto-Heal is enabled."""
+def _gemma_strip_gate(tools) -> set:
+    """Enabled tool NAMES gating the wrapper-less Gemma strip (mirrors the
+    parser/loop gate: only an enabled ``call:foo{...}`` is a call). With NO tools
+    enabled this returns an EMPTY set, not ``None``: every ``call:NAME{...}`` is
+    then prose, and ``None`` would strip-all and delete a legitimate answer."""
+    names = {
+        (t.get("function") or {}).get("name")
+        for t in (tools or [])
+        if isinstance(t, dict) and isinstance(t.get("function"), dict)
+    }
+    names.discard(None)
+    return names
+
+
+def _strip_tool_xml(text: str, enabled_tool_names: Optional[set] = None) -> str:
+    """Combine the parser's scan-based strips (Mistral balanced-brace, gated
+    Gemma wrapper-less, GLM real-close, guarded function-XML) with
+    ``_TOOL_XML_RE`` -- the scan strips close at each call's REAL terminator so
+    literal markup inside argument values is data, not a leaked tail.
+    ``enabled_tool_names`` gates the Gemma strip; ``None`` strips every closed call."""
+    cleaned = _strip_glm_calls(
+        _strip_gemma_wrapperless_calls(_strip_mistral_closed_calls(text), enabled_tool_names),
+        final = True,
+    )
+    cleaned = _strip_function_xml_calls(cleaned, final = True)
+    return _TOOL_XML_RE.sub("", cleaned)
+
+
+def _strip_tool_xml_for_display(
+    text: str,
+    *,
+    auto_heal_tool_calls: bool,
+    enabled_tool_names: Optional[set] = None,
+) -> str:
+    """Route-level leak cleanup (Auto-Heal only). Delegates to ``_strip_tool_xml``
+    so the Mistral balanced-brace pass runs too (``_TOOL_XML_RE`` alone has no
+    ``[TOOL_CALLS]`` arm). ``enabled_tool_names`` gates the Gemma strip."""
     if not auto_heal_tool_calls:
         return text
-    return _TOOL_XML_RE.sub("", text)
+    return _strip_tool_xml(text, enabled_tool_names)
 
 
 logger = get_logger(__name__)
@@ -4046,12 +4163,9 @@ async def generate_stream(
             _DONE = object()
             while True:
                 if cancel_event.is_set():
-                    # The disconnect watcher set cancel_event between chunks.
-                    # Reset the backend here: closing the Python generator does
-                    # not signal a subprocess backend, so without this it keeps
-                    # decoding after the client is gone. The finally's reset is
-                    # guarded on cancel_event being unset, so it will not run
-                    # again for this path.
+                    # Watcher set cancel_event between chunks. Reset here: closing
+                    # the generator does not signal a subprocess backend, so it would
+                    # keep decoding. The finally's reset is guarded, so no double-run.
                     backend.reset_generation_state()
                     break
                 chunk = await asyncio.to_thread(next, gen, _DONE)
@@ -5661,6 +5775,7 @@ async def openai_chat_completions(
 
                 return _SameTaskStreamingResponse(
                     audio_input_stream(),
+                    unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker),
                     media_type = "text/event-stream",
                     headers = {
                         "Cache-Control": "no-cache",
@@ -5936,6 +6051,7 @@ async def openai_chat_completions(
                     _msg["content"] = _strip_tool_xml_for_display(
                         _msg["content"],
                         auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
+                        enabled_tool_names = _gemma_strip_gate(tools_to_use),
                     ).strip()
 
             def gguf_generate_with_tools():
@@ -6069,6 +6185,7 @@ async def openai_chat_completions(
                         clean_cumulative = _strip_tool_xml_for_display(
                             raw_cumulative,
                             auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
+                            enabled_tool_names = _gemma_strip_gate(tools_to_use),
                         )
                         new_text = clean_cumulative[len(prev_text) :]
                         prev_text = clean_cumulative
@@ -6140,6 +6257,7 @@ async def openai_chat_completions(
             if payload.stream:
                 return _SameTaskStreamingResponse(
                     gguf_tool_stream(),
+                    unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker),
                     media_type = "text/event-stream",
                     headers = {
                         "Cache-Control": "no-cache",
@@ -6174,6 +6292,7 @@ async def openai_chat_completions(
                             full_text = _strip_tool_xml_for_display(
                                 event.get("text", ""),
                                 auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
+                                enabled_tool_names = _gemma_strip_gate(tools_to_use),
                             )
                     return full_text, usage, finish
                 finally:
@@ -6396,6 +6515,7 @@ async def openai_chat_completions(
 
             return _SameTaskStreamingResponse(
                 gguf_stream_chunks(),
+                unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker),
                 media_type = "text/event-stream",
                 headers = {
                     "Cache-Control": "no-cache",
@@ -6545,6 +6665,22 @@ async def openai_chat_completions(
     _sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template")
     _sf_features = _detect_safetensors_features(backend, _sf_tpl)
 
+    # Split prefilled-``<think>`` output into reasoning_content deltas (GGUF parity) so the UI
+    # renders the thinking block for safetensors AND MLX.
+    _sf_parse_think = bool(
+        _sf_features.get("supports_reasoning") or _sf_features.get("reasoning_always_on")
+    )
+    # Prefilled-open only for prefill styles with thinking on this request; gpt-oss excluded.
+    _sf_reasoning_prefilled = _sf_reasoning_prefill_mode(
+        _sf_features, payload.enable_thinking, _sf_tpl, payload.reasoning_effort
+    )
+
+    def _new_sf_reasoning_extractor():
+        return _ResponsesReasoningExtractor(
+            parse_think_markers = _sf_parse_think,
+            reasoning_prefilled = _sf_reasoning_prefilled,
+        )
+
     cancel_event = threading.Event()
     completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
     created = int(time.time())
@@ -6631,6 +6767,7 @@ async def openai_chat_completions(
                         "content": _strip_tool_xml_for_display(
                             _msg["content"],
                             auto_heal_tool_calls = _sf_auto_heal_tool_calls,
+                            enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use),
                         ).strip(),
                     }
                 )
@@ -6686,6 +6823,19 @@ async def openai_chat_completions(
 
                 gen = sf_generate_with_tools()
                 prev_text = ""
+                reasoning_extractor = _new_sf_reasoning_extractor()
+
+                def _sf_flush_reasoning():
+                    # Drain the extractor at a turn boundary / stream end (GGUF parity); only visible text reaches the monitor.
+                    fr, fv = reasoning_extractor.finish()
+                    out = []
+                    if fr:
+                        out.append(_chat_reasoning_chunk(completion_id, created, model_name, fr))
+                    if fv:
+                        api_monitor.append_reply(monitor_id, fv)
+                        out.append(_chat_content_chunk(completion_id, created, model_name, fv))
+                    return out
+
                 while True:
                     if cancel_event.is_set():
                         backend.reset_generation_state()
@@ -6702,7 +6852,11 @@ async def openai_chat_completions(
 
                     if event["type"] == "status":
                         if not event["text"]:
+                            # Iteration boundary: flush reasoning, then start a fresh extractor for the next turn.
+                            for _c in _sf_flush_reasoning():
+                                yield _c
                             prev_text = ""
+                            reasoning_extractor = _new_sf_reasoning_extractor()
                         status_data = json.dumps(
                             {
                                 "type": "tool_status",
@@ -6714,7 +6868,11 @@ async def openai_chat_completions(
 
                     if event["type"] in ("tool_start", "tool_end"):
                         if event["type"] == "tool_start":
+                            # Flush reasoning before the tool_start line so the thinking block closes ahead of the tool card.
+                            for _c in _sf_flush_reasoning():
+                                yield _c
                             prev_text = ""
+                            reasoning_extractor = _new_sf_reasoning_extractor()
                         yield f"data: {json.dumps(event)}\n\n"
                         continue
 
@@ -6723,14 +6881,24 @@ async def openai_chat_completions(
                     clean_cumulative = _strip_tool_xml_for_display(
                         raw_cumulative,
                         auto_heal_tool_calls = _sf_auto_heal_tool_calls,
+                        enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use),
                     )
                     new_text = clean_cumulative[len(prev_text) :]
                     prev_text = clean_cumulative
                     if not new_text:
                         continue
-                    api_monitor.append_reply(monitor_id, new_text)
-                    yield _chat_content_chunk(completion_id, created, model_name, new_text)
+                    # Split reasoning vs visible; only visible reaches the monitor.
+                    reasoning_delta, visible_delta = reasoning_extractor.feed(new_text)
+                    if reasoning_delta:
+                        yield _chat_reasoning_chunk(
+                            completion_id, created, model_name, reasoning_delta
+                        )
+                    if visible_delta:
+                        api_monitor.append_reply(monitor_id, visible_delta)
+                        yield _chat_content_chunk(completion_id, created, model_name, visible_delta)
 
+                for _c in _sf_flush_reasoning():
+                    yield _c
                 yield _chat_final_chunk(completion_id, created, model_name, "stop")
                 # Usage chunk from the last turn, same shape as the
                 # GGUF tool loop's metadata. Request-scoped holder, so
@@ -6783,6 +6951,7 @@ async def openai_chat_completions(
         if payload.stream:
             return _SameTaskStreamingResponse(
                 sf_tool_stream(),
+                unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_sf_tracker),
                 media_type = "text/event-stream",
                 headers = {
                     "Cache-Control": "no-cache",
@@ -6804,22 +6973,32 @@ async def openai_chat_completions(
                         full_text = _strip_tool_xml_for_display(
                             event.get("text", ""),
                             auto_heal_tool_calls = _sf_auto_heal_tool_calls,
+                            enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use),
                         )
                 return full_text
 
             content_text = await asyncio.to_thread(_drain_to_text)
-            api_monitor.set_reply(monitor_id, content_text)
+            # Split prefilled <think> reasoning out of the visible answer (GGUF parity); monitor gets visible text only.
+            _reasoning_text, _visible_text = _extract_responses_reasoning(
+                content_text,
+                parse_think_markers = _sf_parse_think,
+                reasoning_prefilled = _sf_reasoning_prefilled,
+            )
+            api_monitor.set_reply(monitor_id, _visible_text)
             _stats = _sf_stats_holder.get("stats")
             if _stats:
                 _monitor_usage(monitor_id, _stats.get("usage"))
             api_monitor.finish(monitor_id, "cancelled" if cancel_event.is_set() else "completed")
+            _sf_msg_kwargs = {"content": _visible_text}
+            if _reasoning_text:
+                _sf_msg_kwargs["reasoning_content"] = _reasoning_text
             response = ChatCompletion(
                 id = completion_id,
                 created = created,
                 model = model_name,
                 choices = [
                     CompletionChoice(
-                        message = CompletionMessage(content = content_text),
+                        message = CompletionMessage(**_sf_msg_kwargs),
                         finish_reason = "stop",
                     )
                 ],
@@ -6898,6 +7077,8 @@ async def openai_chat_completions(
                 yield _chat_role_chunk(completion_id, created, model_name)
 
                 prev_text = ""
+                # Split prefilled <think> into reasoning_content deltas (GGUF parity). Single turn (no per-turn reset); also serves MLX.
+                reasoning_extractor = _new_sf_reasoning_extractor()
                 # Run the sync generator in a thread pool to avoid blocking the
                 # event loop. Critical for compare mode: two SSE requests arrive
                 # concurrently but the orchestrator serializes them via
@@ -6926,9 +7107,21 @@ async def openai_chat_completions(
                     prev_text = cumulative
                     if not new_text:
                         continue
-                    api_monitor.append_reply(monitor_id, new_text)
-                    yield _chat_content_chunk(completion_id, created, model_name, new_text)
+                    reasoning_delta, visible_delta = reasoning_extractor.feed(new_text)
+                    if reasoning_delta:
+                        yield _chat_reasoning_chunk(
+                            completion_id, created, model_name, reasoning_delta
+                        )
+                    if visible_delta:
+                        api_monitor.append_reply(monitor_id, visible_delta)
+                        yield _chat_content_chunk(completion_id, created, model_name, visible_delta)
 
+                final_reasoning, final_visible = reasoning_extractor.finish()
+                if final_reasoning:
+                    yield _chat_reasoning_chunk(completion_id, created, model_name, final_reasoning)
+                if final_visible:
+                    api_monitor.append_reply(monitor_id, final_visible)
+                    yield _chat_content_chunk(completion_id, created, model_name, final_visible)
                 yield _chat_final_chunk(completion_id, created, model_name, "stop")
                 # Usage chunk (choices=[], usage set), same shape as the
                 # GGUF path so the speed popover works for MLX too.
@@ -6975,6 +7168,7 @@ async def openai_chat_completions(
 
         return _SameTaskStreamingResponse(
             stream_chunks(),
+            unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker),
             media_type = "text/event-stream",
             headers = {
                 "Cache-Control": "no-cache",
@@ -6990,18 +7184,27 @@ async def openai_chat_completions(
             for token in generate():
                 full_text = token
 
+            # Split prefilled <think> reasoning from the visible answer (GGUF parity); also covers MLX.
+            _reasoning_text, _visible_text = _extract_responses_reasoning(
+                full_text,
+                parse_think_markers = _sf_parse_think,
+                reasoning_prefilled = _sf_reasoning_prefilled,
+            )
+            _plain_msg_kwargs = {"content": _visible_text}
+            if _reasoning_text:
+                _plain_msg_kwargs["reasoning_content"] = _reasoning_text
             response = ChatCompletion(
                 id = completion_id,
                 created = created,
                 model = model_name,
                 choices = [
                     CompletionChoice(
-                        message = CompletionMessage(content = full_text),
+                        message = CompletionMessage(**_plain_msg_kwargs),
                         finish_reason = "stop",
                     )
                 ],
             )
-            api_monitor.set_reply(monitor_id, full_text)
+            api_monitor.set_reply(monitor_id, _visible_text)
             _stats = stats_holder.get("stats")
             if _stats:
                 _monitor_usage(monitor_id, _stats.get("usage"))
@@ -7824,10 +8027,18 @@ def _responses_marker_holdback(text: str, markers: tuple[str, ...]) -> int:
 class _ResponsesReasoningExtractor:
     """Split local <think> markup into Responses reasoning and visible text."""
 
-    def __init__(self, *, parse_think_markers: bool = False) -> None:
+    def __init__(
+        self,
+        *,
+        parse_think_markers: bool = False,
+        reasoning_prefilled: bool = False,
+    ) -> None:
         self._buffer = ""
-        self._in_reasoning = False
-        self._parse_think_markers = parse_think_markers
+        # ``reasoning_prefilled``: output begins INSIDE an unclosed ``<think>`` (Qwen3/GLM prefill),
+        # so start in reasoning to capture leading text until the first ``</think>``. Callers default False.
+        self._in_reasoning = reasoning_prefilled
+        # Splitting requires marker parsing; a prefilled open implies it.
+        self._parse_think_markers = parse_think_markers or reasoning_prefilled
 
     def feed(
         self,
@@ -7850,14 +8061,21 @@ class _ResponsesReasoningExtractor:
             if self._in_reasoning:
                 close_idx = self._buffer.find(_RESPONSES_THINK_CLOSE)
                 if close_idx != -1:
-                    reasoning_parts.append(self._buffer[:close_idx])
+                    reasoning_parts.append(
+                        self._buffer[:close_idx].replace(_RESPONSES_THINK_OPEN, "")
+                    )
                     self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :]
                     self._in_reasoning = False
                     continue
-                keep = _responses_marker_holdback(self._buffer, (_RESPONSES_THINK_CLOSE,))
+                # Hold back a trailing partial of EITHER marker: the close (clean chunk-boundary split)
+                # and a stray open (so a re-emitted ``<think>`` isn't leaked into the reasoning drawer).
+                keep = _responses_marker_holdback(
+                    self._buffer, (_RESPONSES_THINK_CLOSE, _RESPONSES_THINK_OPEN)
+                )
                 if keep == len(self._buffer):
                     break
-                reasoning_parts.append(self._buffer[:-keep] if keep else self._buffer)
+                emit = self._buffer[:-keep] if keep else self._buffer
+                reasoning_parts.append(emit.replace(_RESPONSES_THINK_OPEN, ""))
                 self._buffer = self._buffer[-keep:] if keep else ""
                 break
 
@@ -7894,7 +8112,7 @@ class _ResponsesReasoningExtractor:
             return "", remaining
         if self._in_reasoning:
             self._in_reasoning = False
-            return remaining, ""
+            return remaining.replace(_RESPONSES_THINK_OPEN, ""), ""
         return "", remaining.replace(_RESPONSES_THINK_CLOSE, "")
 
 
@@ -7903,8 +8121,12 @@ def _extract_responses_reasoning(
     reasoning_content: Any = None,
     *,
     parse_think_markers: bool = False,
+    reasoning_prefilled: bool = False,
 ) -> tuple[str, str]:
-    extractor = _ResponsesReasoningExtractor(parse_think_markers = parse_think_markers)
+    extractor = _ResponsesReasoningExtractor(
+        parse_think_markers = parse_think_markers,
+        reasoning_prefilled = reasoning_prefilled,
+    )
     reasoning, visible = extractor.feed(text, reasoning_content)
     final_reasoning, final_visible = extractor.finish()
     return reasoning + final_reasoning, visible + final_visible
@@ -9734,7 +9956,9 @@ async def anthropic_messages(
         # Strip stale tool-call XML from conversation
         for _msg in openai_messages:
             if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str):
-                _msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip()
+                _msg["content"] = _strip_tool_xml(
+                    _msg["content"], _gemma_strip_gate(openai_tools)
+                ).strip()
 
         def _run_tool_gen():
             return llama_backend.generate_chat_completion_with_tools(
@@ -9779,6 +10003,7 @@ async def anthropic_messages(
                 message_id,
                 model_name,
                 disable_parallel_tool_use = _disable_parallel,
+                openai_tools = openai_tools,
             )
         )
 
@@ -9857,11 +10082,8 @@ async def _anthropic_tool_stream(
         drop_until_tool_end = False
 
         gen = run_gen()
-        # Concurrent disconnect watcher: the loop only polls is_disconnected()
-        # between events, so a client disconnect during a long prefill or
-        # generation step would otherwise hold the decode slot until the next
-        # event or a failed send. The watcher sets cancel_event so the backend
-        # stops promptly.
+        # Watcher to cancel on disconnect: the in-loop poll fires only between
+        # events, so a mid-prefill disconnect would otherwise hold the decode slot.
         disconnect_watcher = asyncio.create_task(
             _await_disconnect_then_cancel(request, cancel_event)
         )
@@ -9888,7 +10110,7 @@ async def _anthropic_tool_stream(
                 # content event that was purely tool XML doesn't count as text.
                 if etype == "content":
                     event = dict(event)
-                    event["text"] = _TOOL_XML_RE.sub("", event["text"])
+                    event["text"] = _strip_tool_xml(event["text"], _gemma_strip_gate(openai_tools))
                 # disable_parallel_tool_use: keep only the first tool_use block,
                 # dropping every later tool_start and its paired tool_end (robust
                 # to empty tool-call ids — tracked by state, not id matching).
@@ -9953,11 +10175,8 @@ async def _anthropic_plain_stream(
         captured_finish_reason = None
 
         gen = run_gen()
-        # Concurrent disconnect watcher: the loop only polls is_disconnected()
-        # between chunks, so a client disconnect during a long prefill or
-        # generation step would otherwise hold the decode slot until the next
-        # chunk or a failed send. The watcher sets cancel_event so the backend
-        # stops promptly.
+        # Watcher to cancel on disconnect: the in-loop poll fires only between
+        # chunks, so a mid-prefill disconnect would otherwise hold the decode slot.
         disconnect_watcher = asyncio.create_task(
             _await_disconnect_then_cancel(request, cancel_event)
         )
@@ -10046,6 +10265,7 @@ async def _anthropic_tool_non_streaming(
     message_id,
     model_name,
     disable_parallel_tool_use = False,
+    openai_tools = None,
 ):
     """Non-streaming response for the tool-calling path.
 
@@ -10074,7 +10294,7 @@ async def _anthropic_tool_non_streaming(
         etype = event.get("type", "")
         if etype == "content":
             # Strip leaked tool-call XML
-            clean = _TOOL_XML_RE.sub("", event["text"])
+            clean = _strip_tool_xml(event["text"], _gemma_strip_gate(openai_tools))
             new = clean[len(prev_text) :]
             prev_text = clean
             if new:
@@ -10544,9 +10764,13 @@ async def _anthropic_passthrough_non_streaming(
         text = message.get("content") or ""
         if text:
             # Keep unpromoted bytes when healing is active; legacy stripping is
-            # only for opted-out or no-client-tool requests.
+            # only for opted-out or no-client-tool requests. Use the full
+            # _strip_tool_xml pass so Mistral [TOOL_CALLS] and guarded
+            # function-XML leaks are cleaned too, not just _TOOL_XML_RE forms,
+            # with the Gemma display gate so a disabled/example call:NAME{...}
+            # in prose survives.
             if not healing_active:
-                text = _TOOL_XML_RE.sub("", text)
+                text = _strip_tool_xml(text, _gemma_strip_gate(openai_tools))
             text = text.strip()
             if text:
                 content_blocks.append(AnthropicResponseTextBlock(text = text))
@@ -10909,6 +11133,10 @@ async def _openai_passthrough_stream(
     response ``id``, ``finish_reason`` (including ``"tool_calls"``),
     ``delta.tool_calls``, and any client-requested trailing ``usage`` chunk so
     the client sees a standard OpenAI response.
+
+    Reasoning/tool-call splitting is delegated to llama-server (``--jinja
+    --reasoning-format auto``), so ``delta.content`` carries no raw markup and is
+    deliberately not re-parsed locally, unlike the ``/completion`` paths.
     """
     target_url = f"{llama_backend.base_url}/v1/chat/completions"
     body = _build_openai_passthrough_body(
@@ -11325,11 +11553,9 @@ async def _openai_passthrough_stream(
                                 delta = choice.get("delta")
                                 if isinstance(delta, dict) and delta.get("tool_calls"):
                                     saw_tool_call_delta = True
-                        # Detect an upstream error chunk independently of API
-                        # monitoring: when monitor_id is None (skip_api_monitor),
-                        # _monitor_openai_sse_line returns before inspecting the
-                        # error, so without this the synthetic-finish guard would
-                        # emit a successful finish_reason after a failed stream.
+                        # Detect an error chunk independently of API monitoring
+                        # (skip_api_monitor returns early), else the synthetic
+                        # finish would fire after a failed stream.
                         if _monitor_openai_error_message(chunk_data):
                             saw_stream_error = True
                     # With healing active, a content-bearing line may be replaced by

studio/backend/tests/test_chat_eos_template_refresh.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Mapper models whose own tokenizer ships no chat_template have their turn-end
+eos resolved at LOAD from an empty template (document eos only). The effective
+template is installed later, at generate time, via get_chat_template, so the
+turn-end-eos cache must be refreshed then; otherwise generate_stream runs past
+the ChatML <|im_end|> boundary and loops (the exact bug this PR fixes).
+"""
+
+import sys
+from pathlib import Path
+
+import pytest
+
+_BACKEND = Path(__file__).resolve().parent.parent
+if str(_BACKEND) not in sys.path:
+    sys.path.insert(0, str(_BACKEND))
+
+# These tests construct InferenceBackend, pulling the full stack. CI may lack
+# unsloth/unsloth_zoo (ImportError) or have a broken CUDA/bitsandbytes setup
+# (RuntimeError); skip at module level so collection is not aborted (exit 2).
+try:
+    from core.inference import inference as inf_mod  # noqa: E402
+    from core.inference.inference import InferenceBackend  # noqa: E402
+except (ImportError, RuntimeError) as exc:  # pragma: no cover - env-dependent
+    pytest.skip(
+        f"full inference backend unavailable ({type(exc).__name__}: {exc})",
+        allow_module_level = True,
+    )
+
+_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}"
+_GEMMA = "{% for m in messages %}<start_of_turn>{{m.role}}\n{{m.content}}<end_of_turn>{% endfor %}"
+
+
+class _FakeTokenizer:
+    def __init__(
+        self,
+        eos_id,
+        chat_template = "",
+        token_ids = None,
+    ):
+        self.eos_token_id = eos_id
+        self.chat_template = chat_template
+        self.pad_token_id = eos_id
+        self.unk_token_id = None
+        self._ids = dict(token_ids or {})
+
+    def convert_tokens_to_ids(self, tok):
+        return self._ids.get(tok)
+
+
+def test_turn_end_eos_refreshed_after_generate_time_template(monkeypatch):
+    import utils.datasets as ds
+
+    backend = InferenceBackend.__new__(InferenceBackend)
+    backend.active_model_name = "unsloth/qwen2.5-0.5b"
+
+    # No chat_template at load, so the cache stored only the document eos, though
+    # <|im_end|> is atomic in the vocab (unused until the mapper installs a template).
+    bare_tok = _FakeTokenizer(151643, chat_template = "", token_ids = {"<|im_end|>": 151645})
+    model_info = {
+        "tokenizer": bare_tok,
+        "is_vision": False,
+        "chat_turn_end_eos_ids": [151643],
+    }
+    backend.models = {backend.active_model_name: model_info}
+
+    # The mapper installs a ChatML template (turns end with <|im_end|>) at generate time.
+    templated_tok = _FakeTokenizer(151643, chat_template = _CHATML, token_ids = {"<|im_end|>": 151645})
+    monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: templated_tok)
+    monkeypatch.setattr(
+        ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "qwen-2.5"}, raising = False
+    )
+
+    # Stub the tail so the generator runs through the refresh without a real model.
+    monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False)
+    monkeypatch.setattr(
+        backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False
+    )
+    monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False)
+
+    list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}]))
+
+    # After the template is applied the cache must include the ChatML turn-end id.
+    assert model_info["chat_turn_end_eos_ids"] == [151643, 151645]
+
+
+def test_turn_end_eos_refresh_preserves_load_time_ids_on_destructive_swap(monkeypatch):
+    # Regression: get_chat_template can return a remapped tokenizer (Gemma: <end_of_turn>
+    # folded onto the eos id) while generate_stream re-reads the original. Resolving on
+    # the swap yields a narrower set, so the refresh must UNION, never overwrite.
+    import utils.datasets as ds
+
+    backend = InferenceBackend.__new__(InferenceBackend)
+    backend.active_model_name = "unsloth/gemma-2b-it"
+
+    # Original tokenizer (used by generate_stream): <end_of_turn>=107 distinct from
+    # eos=1, so the load-time cache resolved to [1, 107].
+    orig_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"<end_of_turn>": 107})
+    model_info = {
+        "tokenizer": orig_tok,
+        "is_vision": False,
+        "chat_turn_end_eos_ids": [1, 107],
+    }
+    backend.models = {backend.active_model_name: model_info}
+
+    # Destructively-swapped tokenizer: <end_of_turn> now maps onto eos id 1, so
+    # resolving on it yields only [1] (drops 107).
+    swapped_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"<end_of_turn>": 1})
+    monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: swapped_tok)
+    monkeypatch.setattr(
+        ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "gemma-3"}, raising = False
+    )
+
+    monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False)
+    monkeypatch.setattr(
+        backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False
+    )
+    monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False)
+
+    list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}]))
+
+    # The load-time <end_of_turn>=107 must survive: overwriting with the swapped
+    # [1] would regress and loop past the turn.
+    assert model_info["chat_turn_end_eos_ids"] == [1, 107]
+
+
+def test_turn_end_eos_refresh_resolves_marker_id_on_original_not_remapped(monkeypatch):
+    # Yi-style map_eos_token=True: the original carries <|im_end|> at its own id, but
+    # get_chat_template folds it onto the doc-eos id. generate_stream uses the original,
+    # so read marker strings from the mapped template but ids from the original.
+    import utils.datasets as ds
+
+    backend = InferenceBackend.__new__(InferenceBackend)
+    backend.active_model_name = "01-ai/yi-6b"
+
+    # Original: no template of its own, doc eos = 2, <|im_end|> atomic = 7.
+    orig_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7})
+    model_info = {
+        "tokenizer": orig_tok,
+        "is_vision": False,
+        "chat_turn_end_eos_ids": [2],
+    }
+    backend.models = {backend.active_model_name: model_info}
+
+    # Remapped tokenizer: ChatML template, but <|im_end|> folded onto doc-eos id 2.
+    remapped_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2})
+    monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: remapped_tok)
+    monkeypatch.setattr(
+        ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "chatml"}, raising = False
+    )
+
+    monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False)
+    monkeypatch.setattr(
+        backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False
+    )
+    monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False)
+
+    list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}]))
+
+    # The real <|im_end|>=7 (original vocab) must be recovered, not the remapped 2.
+    assert model_info["chat_turn_end_eos_ids"] == [2, 7]
+
+
+class _FakeProcessor:
+    """A ProcessorMixin-like container: carries the chat_template itself and
+    wraps the real text tokenizer as ``.tokenizer`` (the vision layout)."""
+
+    def __init__(self, chat_template, tokenizer):
+        self.chat_template = chat_template
+        self.tokenizer = tokenizer
+
+
+def test_resolve_chat_eos_reads_vision_processor_template():
+    # Vision model: the chat_template lives on the processor while the inner tokenizer
+    # ships none. _resolve_chat_eos must read the marker from the processor but resolve
+    # its id on the inner tokenizer, and repair generation_config.
+    from types import SimpleNamespace
+
+    inner_tok = _FakeTokenizer(1, chat_template = "", token_ids = {"<end_of_turn>": 107})
+    processor = _FakeProcessor(_GEMMA, inner_tok)
+    model = SimpleNamespace(generation_config = SimpleNamespace(eos_token_id = 1))
+
+    backend = InferenceBackend.__new__(InferenceBackend)
+    backend.active_model_name = "unsloth/gemma-3-4b-it"
+    model_info = {"model": model, "tokenizer": processor, "processor": processor, "is_vision": True}
+    backend.models = {backend.active_model_name: model_info}
+
+    backend._resolve_chat_eos(backend.active_model_name)
+
+    assert model_info["chat_turn_end_eos_ids"] == [1, 107]
+    # generation_config repaired so the vision .generate() path stops at the turn.
+    assert model.generation_config.eos_token_id == [1, 107]

studio/backend/tests/test_chat_template_tool_arguments.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""apply_chat_template_for_generation must coerce assistant tool_call arguments
+from the OpenAI JSON-string form to a dict before rendering. Strict tool
+templates (e.g. mlx-community Qwen3.5 checkpoints) iterate arguments.items() and
+raise "Can only get item pairs from a mapping." on the string form when a prior
+tool call is re-rendered on the next turn (MLX + transformers paths).
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+
+_BACKEND = Path(__file__).resolve().parent.parent
+if str(_BACKEND) not in sys.path:
+    sys.path.insert(0, str(_BACKEND))
+
+from core.inference.chat_template_helpers import (  # noqa: E402
+    _normalize_tool_call_arguments,
+    apply_chat_template_for_generation,
+)
+
+
+def _conv(arguments):
+    return [
+        {"role": "user", "content": "weather?"},
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {
+                    "type": "function",
+                    "id": "c1",
+                    "function": {"name": "web_search", "arguments": arguments},
+                }
+            ],
+        },
+        {"role": "tool", "name": "web_search", "content": "21C sunny"},
+    ]
+
+
+class _StrictTemplateTokenizer:
+    """Mimics a strict Qwen tool template: rejects string tool_call arguments."""
+
+    def apply_chat_template(
+        self,
+        messages,
+        *,
+        tokenize = False,
+        add_generation_prompt = True,
+        **kw,
+    ):
+        for msg in messages:
+            for call in msg.get("tool_calls", []) or []:
+                args = call.get("function", {}).get("arguments")
+                if isinstance(args, str):
+                    raise TypeError("Can only get item pairs from a mapping.")
+        return "RENDERED"
+
+
+def test_string_arguments_are_parsed_to_dict():
+    out = _normalize_tool_call_arguments(_conv('{"query": "sweden"}'))
+    args = out[1]["tool_calls"][0]["function"]["arguments"]
+    assert args == {"query": "sweden"}
+
+
+def test_dict_arguments_untouched_and_no_copy():
+    conv = _conv({"query": "sweden"})
+    assert _normalize_tool_call_arguments(conv) is conv
+
+
+def test_non_json_string_left_as_is():
+    out = _normalize_tool_call_arguments(_conv("not json"))
+    assert out[1]["tool_calls"][0]["function"]["arguments"] == "not json"
+
+
+def test_render_succeeds_on_strict_template_with_string_arguments():
+    # Regression: strict template + string args used to raise.
+    result = apply_chat_template_for_generation(_StrictTemplateTokenizer(), _conv('{"query": "x"}'))
+    assert result == "RENDERED"
+
+
+class _RecordingTokenizer:
+    """Lenient template: renders whatever arguments it is given (string or dict)."""
+
+    def __init__(self):
+        self.seen_arguments = None
+
+    def apply_chat_template(
+        self,
+        messages,
+        *,
+        tokenize = False,
+        add_generation_prompt = True,
+        **kw,
+    ):
+        for msg in messages:
+            for call in msg.get("tool_calls", []) or []:
+                self.seen_arguments = call.get("function", {}).get("arguments")
+        return "RENDERED"
+
+
+def test_lenient_template_receives_original_string_untouched():
+    # Lenient template must see the exact original string, not a coerced dict.
+    tok = _RecordingTokenizer()
+    apply_chat_template_for_generation(tok, _conv('{"query": "x"}'))
+    assert tok.seen_arguments == '{"query": "x"}'
+
+
+def test_messages_without_tool_calls_pass_through_unchanged():
+    conv = [{"role": "user", "content": "hi"}]
+    assert _normalize_tool_call_arguments(conv) is conv
+
+
+class _RaiseExceptionTemplateTokenizer:
+    """Mimics the bundled gemma-4.jinja: rejects string tool_call arguments via
+    ``raise_exception(...)``, which surfaces as a Jinja error, NOT a TypeError."""
+
+    def apply_chat_template(
+        self,
+        messages,
+        *,
+        tokenize = False,
+        add_generation_prompt = True,
+        **kw,
+    ):
+        for msg in messages:
+            for call in msg.get("tool_calls", []) or []:
+                args = call.get("function", {}).get("arguments")
+                if isinstance(args, str):
+                    raise ValueError(
+                        "chat_template: tool_calls[].function.arguments must be a "
+                        "JSON object (mapping), not a string."
+                    )
+        return "RENDERED"
+
+
+def test_render_succeeds_on_raise_exception_template_with_string_arguments():
+    # Regression: gemma-4.jinja rejects string args via a non-TypeError; retry must still coerce.
+    result = apply_chat_template_for_generation(
+        _RaiseExceptionTemplateTokenizer(), _conv('{"query": "x"}')
+    )
+    assert result == "RENDERED"
+
+
+def test_unrelated_template_error_still_propagates_with_dict_args():
+    # Failure unrelated to string args (dict args, nothing to coerce) must propagate.
+    class _AlwaysRaises:
+        def apply_chat_template(self, messages, **kw):
+            raise ValueError("template is broken")
+
+    with pytest.raises(ValueError, match = "broken"):
+        apply_chat_template_for_generation(_AlwaysRaises(), _conv({"query": "x"}))

studio/backend/tests/test_chat_turn_end_eos.py

@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""chat_eos: resolve assistant-turn-end stop tokens from the chat_template and
+repair generation_config so a chat model whose eos is a bare document terminator
+(Qwen3.5: config eos <|endoftext|>, turns end with <|im_end|>) stops at the turn
+boundary instead of running past it and looping. Dependency-light: imported here
+without the full inference stack.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+_BACKEND = Path(__file__).resolve().parent.parent
+if str(_BACKEND) not in sys.path:
+    sys.path.insert(0, str(_BACKEND))
+
+from core.inference.chat_eos import (  # noqa: E402
+    chat_eos_repair,
+    resolve_chat_turn_end_eos_ids,
+    resolve_chat_turn_end_eos_ids_using,
+)
+
+
+class _FakeTokenizer:
+    def __init__(
+        self,
+        eos_id,
+        chat_template = "",
+        token_ids = None,
+        unk_token_id = None,
+    ):
+        self.eos_token_id = eos_id
+        self.chat_template = chat_template
+        self.unk_token_id = unk_token_id
+        self._ids = dict(token_ids or {})
+
+    def convert_tokens_to_ids(self, tok):
+        return self._ids.get(tok, self.unk_token_id)
+
+
+# ---- resolve_chat_turn_end_eos_ids ---------------------------------------
+
+_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}"
+
+
+def test_qwen35_adds_im_end_from_template():
+    # eos synced to <|endoftext|> (248044); template uses <|im_end|> (248046).
+    tok = _FakeTokenizer(248044, chat_template = _CHATML, token_ids = {"<|im_end|>": 248046})
+    assert resolve_chat_turn_end_eos_ids(tok) == [248044, 248046]
+
+
+def test_marker_in_vocab_but_not_in_template_is_ignored():
+    # Base/coder model: <|im_end|> is in the vocab but the template does not use
+    # it, so it must not become a stop token.
+    tok = _FakeTokenizer(248044, chat_template = "{{ messages }}", token_ids = {"<|im_end|>": 248046})
+    assert resolve_chat_turn_end_eos_ids(tok) == [248044]
+
+
+def test_harmony_template_is_left_untouched():
+    # gpt-oss/harmony: <|end|> is a channel delimiter, not the turn end.
+    harmony = "<|start|>assistant<|channel|>analysis<|message|>...<|end|>"
+    tok = _FakeTokenizer(200002, chat_template = harmony, token_ids = {"<|end|>": 200007})
+    assert resolve_chat_turn_end_eos_ids(tok) == [200002]
+
+
+def test_llama3_eot_id_from_template():
+    tok = _FakeTokenizer(128001, chat_template = "...<|eot_id|>...", token_ids = {"<|eot_id|>": 128009})
+    assert resolve_chat_turn_end_eos_ids(tok) == [128001, 128009]
+
+
+def test_gemma4_turn_marker_from_template():
+    # Gemma-4 ends turns with <turn|> while keeping a document eos, so <turn|> must
+    # be added as a stop token.
+    tok = _FakeTokenizer(
+        1, chat_template = "...<start_of_turn>...<turn|>...", token_ids = {"<turn|>": 106}
+    )
+    assert resolve_chat_turn_end_eos_ids(tok) == [1, 106]
+
+
+def test_resolve_using_reads_markers_from_template_but_ids_from_generation_tokenizer():
+    # map_eos_token=True: the mapped template remaps <|im_end|> onto the doc-eos id,
+    # but the original keeps it atomic. Reading marker STRINGS from the template but
+    # IDS on the original recovers the real turn-end id (7), not the doc-eos id (2).
+    template_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2})
+    id_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7})
+    assert resolve_chat_turn_end_eos_ids_using(template_tok, id_tok) == [2, 7]
+    # Same tokenizer for both reproduces the plain resolve (load-time behaviour).
+    assert resolve_chat_turn_end_eos_ids_using(template_tok, template_tok) == [2]
+
+
+def test_list_eos_preserved():
+    tok = _FakeTokenizer([1, 2], chat_template = _CHATML, token_ids = {"<|im_end|>": 2})
+    assert resolve_chat_turn_end_eos_ids(tok) == [1, 2]
+
+
+def test_missing_marker_maps_to_unk_and_is_skipped():
+    tok = _FakeTokenizer(7, chat_template = _CHATML, token_ids = {}, unk_token_id = 0)
+    assert resolve_chat_turn_end_eos_ids(tok) == [7]
+
+
+def test_starling_barred_end_of_turn_from_template():
+    # OpenChat/Starling end turns with the BARRED <|end_of_turn|> (distinct from
+    # Gemma's <end_of_turn>). eos synced to </s>=2, turn marker at 32000.
+    starling = "GPT4 Correct Assistant: hi<|end_of_turn|>"
+    tok = _FakeTokenizer(2, chat_template = starling, token_ids = {"<|end_of_turn|>": 32000})
+    assert resolve_chat_turn_end_eos_ids(tok) == [2, 32000]
+
+
+def test_dict_chat_template_scans_all_variants():
+    # Hermes-3 style: chat_template is a {name: template} dict. Detection must scan
+    # every variant, not bail because the container is not a plain str.
+    tmpl = {"default": "{{ messages }}", "tool_use": _CHATML}
+    tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5})
+    assert resolve_chat_turn_end_eos_ids(tok) == [2, 5]
+
+
+def test_list_of_dicts_chat_template_scans_all_variants():
+    # tokenizer_config.json stores multi-templates as a list of {name, template}.
+    tmpl = [{"name": "default", "template": _CHATML}]
+    tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5})
+    assert resolve_chat_turn_end_eos_ids(tok) == [2, 5]
+
+
+def test_dict_harmony_template_left_untouched():
+    # A multi-variant container whose variant is harmony must still be left alone.
+    tmpl = {"default": "<|start|>assistant<|channel|>analysis<|message|>...<|end|>"}
+    tok = _FakeTokenizer(200002, chat_template = tmpl, token_ids = {"<|end|>": 200007})
+    assert resolve_chat_turn_end_eos_ids(tok) == [200002]
+
+
+# ---- chat_eos_repair ------------------------------------------------------
+
+
+def test_repair_adds_missing_turn_end():
+    assert chat_eos_repair(248044, [248044, 248046]) == [248044, 248046]
+
+
+def test_repair_from_missing_generation_config_eos():
+    assert chat_eos_repair(None, [248046]) == [248046]
+
+
+def test_repair_noop_when_already_covered():
+    assert chat_eos_repair([248046, 248044], [248046]) is None
+
+
+def test_repair_noop_when_no_turn_end_ids():
+    assert chat_eos_repair(248044, []) is None

studio/backend/tests/test_gemma_tool_parse_edge_cases.py

@@ -1,15 +1,8 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""Edge cases in Gemma-native tool-call parsing.
-
-Covers two failure modes:
-  1. A bare (unquoted) string argument that contains a comma, e.g.
-     ``location:New York, NY`` -- the comma must not be treated as the next
-     key boundary, or the whole call is dropped.
-  2. A tool-call marker that appears INSIDE another call's argument string is
-     data, not a real call, so it must not be promoted to a second tool call.
-"""
+"""Gemma-native tool-call parsing edge cases: commas inside bare string values,
+and markers inside another call's argument data staying data."""
 
 from __future__ import annotations
 
@@ -21,7 +14,11 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
     sys.path.insert(0, _BACKEND_DIR)
 
-from core.inference.tool_call_parser import parse_tool_calls_from_text
+from core.inference.tool_call_parser import (
+    _gemma_parse_value,
+    parse_tool_calls_from_text,
+)
+from core.tool_healing import strip_tool_call_markup
 
 
 def _args(call: dict) -> dict:
@@ -40,14 +37,22 @@ def test_bare_string_argument_with_comma_is_kept():
 def test_normal_multi_key_arguments_still_split():
     calls = parse_tool_calls_from_text('<|tool_call>call:f{a:1,b:hello,c:"x,y"}<tool_call|>')
     assert len(calls) == 1, calls
-    # Numbers stay numeric, bare strings get quoted, an explicit quoted comma
-    # stays inside its value.
     assert _args(calls[0]) == {"a": 1, "b": "hello", "c": "x,y"}
 
 
+def test_empty_bare_value_becomes_empty_string_not_dropped():
+    # An empty bare value (``{query:}``) must serialise as ``""`` (``{"query":}`` is invalid JSON and dropped the call).
+    calls = parse_tool_calls_from_text("<|tool_call>call:search{query:,unit:celsius}<tool_call|>")
+    assert len(calls) == 1, calls
+    assert _args(calls[0]) == {"query": "", "unit": "celsius"}
+
+    only = parse_tool_calls_from_text("<|tool_call>call:get{q:}<tool_call|>")
+    assert len(only) == 1, only
+    assert _args(only[0]) == {"q": ""}
+
+
 def test_bare_value_with_timestamps_after_comma_is_kept():
-    # A comma followed by digits-then-colon (a timestamp/ratio) is value text,
-    # not a new key, so the whole query must be preserved as one argument.
+    # A comma before digits-then-colon (timestamp/ratio) is value text, not a key.
     calls = parse_tool_calls_from_text(
         "<|tool_call>call:remind{query:meet at 10:00, 11:00 tomorrow,priority:high}<tool_call|>"
     )
@@ -55,9 +60,16 @@ def test_bare_value_with_timestamps_after_comma_is_kept():
     assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow", "priority": "high"}
 
 
+def test_wrapperless_bare_value_with_timestamps_after_comma_is_kept():
+    # The wrapper-less Gemma form (no <|tool_call> markers) goes through the
+    # _gemma_parse_stripped_body scanner and its _GEMMA_KEY_RE.
+    calls = parse_tool_calls_from_text("call:web_search{query:meet at 10:00, 11:00 tomorrow}")
+    assert len(calls) == 1, calls
+    assert calls[0]["function"]["name"] == "web_search"
+    assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow"}
+
+
 def test_marker_inside_json_argument_is_not_a_second_call():
-    # A python call whose `code` argument contains a Gemma marker string. The
-    # marker is data and must not execute as a second `terminal` call.
     content = (
         '<tool_call>{"name":"python","arguments":{"code":'
         '"x = 1  # <|tool_call>call:terminal{command:ls}<tool_call|>"}}</tool_call>'
@@ -75,8 +87,6 @@ def test_two_separate_gemma_calls_both_parse():
 
 
 def test_mixed_format_calls_preserve_document_order():
-    # A Gemma-native call precedes a JSON-format call in the text; tools execute
-    # in returned order, so `create` must come before `read`.
     content = (
         "<|tool_call>call:create{path:a}<tool_call|> then "
         '<tool_call>{"name":"read","arguments":{"path":"a"}}</tool_call>'
@@ -86,8 +96,6 @@ def test_mixed_format_calls_preserve_document_order():
 
 
 def test_json_marker_inside_gemma_argument_is_not_a_second_call():
-    # The reverse of the JSON-outer case: a JSON-style marker inside a Gemma
-    # call's quoted argument is code text, not a second `terminal` call.
     content = (
         '<|tool_call>call:python{code:<|"|>'
         'print(<tool_call>{"name":"terminal","arguments":{"command":"ls"}}</tool_call>)'
@@ -98,18 +106,14 @@ def test_json_marker_inside_gemma_argument_is_not_a_second_call():
 
 
 def test_nested_gemma_marker_in_unquoted_arg_does_not_run_inner_call():
-    # An UNQUOTED Gemma value containing a literal marker: the outer object fails
-    # to normalize (the inner braces/marker break the JSON), but the inner marker
-    # is nested in the outer candidate span, so it must not be promoted to a
-    # standalone `terminal` call. The safe outcome is no executed tool call.
+    # An UNQUOTED Gemma value containing a literal marker: the marker is nested in the outer
+    # candidate span, so it must not be promoted to a standalone `terminal` call (no tool call).
     content = "<|tool_call>call:python{code:<|tool_call>call:terminal{command:ls}<tool_call|>}<tool_call|>"
     calls = parse_tool_calls_from_text(content)
     assert "terminal" not in [c["function"]["name"] for c in calls], calls
 
 
 def test_bare_string_array_argument_is_quoted():
-    # Gemma may emit an array of bare strings without per-element quotes; they
-    # must be quoted so the call is not dropped.
     calls = parse_tool_calls_from_text("<|tool_call>call:label{labels:[bug,ui]}<tool_call|>")
     assert len(calls) == 1, calls
     assert _args(calls[0]) == {"labels": ["bug", "ui"]}
@@ -123,8 +127,6 @@ def test_array_keeps_numbers_and_quoted_elements():
 
 
 def test_array_of_objects_is_normalised():
-    # Arrays of objects are a common tool-schema shape; their (unquoted) keys and
-    # bare values must be normalised too, not left verbatim, or the call drops.
     calls = parse_tool_calls_from_text(
         "<|tool_call>call:batch{items:[{path:a,mode:r},{path:b,mode:w}]}<tool_call|>"
     )
@@ -138,9 +140,6 @@ def test_nested_array_elements_are_normalised():
 
 
 def test_gemma_marker_inside_xml_parameter_is_not_a_second_call():
-    # An XML-style <function=...> call whose <parameter=code> value contains a
-    # Gemma marker: the marker is the parameter's data, not a separate terminal
-    # call, so only the python call must be returned.
     content = (
         "<tool_call><function=python><parameter=code>"
         "x = 1  # <|tool_call>call:terminal{command:ls}<tool_call|>"
@@ -159,3 +158,240 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call():
     )
     calls = parse_tool_calls_from_text(content)
     assert [c["function"]["name"] for c in calls] == ["python"], calls
+
+
+def test_wrapperless_nested_object_argument_is_parsed():
+    # skip_special_tokens stream: wrapper and <|"|> markers stripped, so a nested object arrives bare.
+    calls = parse_tool_calls_from_text("call:f{loc:{city:NYC},n:3}")
+    assert len(calls) == 1
+    assert _args(calls[0]) == {"loc": {"city": "NYC"}, "n": 3}
+
+
+def test_wrapperless_array_argument_is_parsed():
+    calls = parse_tool_calls_from_text("call:label{labels:[bug,ui],n:2}")
+    assert len(calls) == 1
+    assert _args(calls[0]) == {"labels": ["bug", "ui"], "n": 2}
+
+
+def test_wrapperless_deeply_nested_object_and_array_are_preserved():
+    # The single-pass parser must keep multi-level nesting (objects inside
+    # objects, arrays inside arrays) intact, not flatten or drop it.
+    calls = parse_tool_calls_from_text(
+        "call:f{loc:{city:NYC,geo:{lat:1,lng:2}},tags:[a,b,[c,d]],n:3}"
+    )
+    assert len(calls) == 1
+    assert _args(calls[0]) == {
+        "loc": {"city": "NYC", "geo": {"lat": 1, "lng": 2}},
+        "tags": ["a", "b", ["c", "d"]],
+        "n": 3,
+    }
+
+
+def test_gemma_parse_array_advances_on_stray_brace():
+    # Regression: a stray '}' / ']' / ',' where an array element is expected must
+    # not stall _gemma_parse_value at the same index (it looped forever before).
+    from core.inference.tool_call_parser import _gemma_parse_array
+
+    items, end, closed = _gemma_parse_array("[a,}]", 0)
+    assert end == 5 and closed is True  # consumed through the closing ']'
+    assert items[0] == "a"
+
+
+def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping():
+    # Parse keeps the quoted close marker as data; strip removes the whole span.
+    text = '<|tool_call>call:python{code:<|"|>print("<tool_call|>")<|"|>}<tool_call|>'
+    calls = parse_tool_calls_from_text(text)
+    assert len(calls) == 1, calls
+    assert _args(calls[0]) == {"code": 'print("<tool_call|>")'}
+    assert strip_tool_call_markup("before " + text + " after") == "before  after"
+    assert strip_tool_call_markup("before " + text + " after", final = True) == "before  after"
+
+
+def test_nested_xml_in_malformed_gemma_call_does_not_execute():
+    # The failed Gemma candidate's span still covers its nested <function=>.
+    text = (
+        "<|tool_call>call:outer{code:<function=terminal><parameter=command>id"
+        "</parameter></function></tool_call>, broken:{x}}<tool_call|>"
+    )
+    for allow_incomplete in (True, False):
+        calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
+        assert "terminal" not in [c["function"]["name"] for c in calls], calls
+
+
+def test_unbalanced_gemma_call_with_xml_does_not_execute():
+    # Unclosed braces cover to EOF, so the trailing <function=> is excluded.
+    text = (
+        "<|tool_call>call:outer{code:<function=terminal>"
+        "<parameter=command>id</parameter></function>"
+    )
+    for allow_incomplete in (True, False):
+        calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
+        assert "terminal" not in [c["function"]["name"] for c in calls], calls
+
+
+def test_standalone_function_xml_still_parses():
+    text = "<function=terminal><parameter=command>id</parameter></function>"
+    calls = parse_tool_calls_from_text(text)
+    assert [c["function"]["name"] for c in calls] == ["terminal"], calls
+
+
+def test_xml_between_braces_and_close_marker_does_not_execute():
+    # Coverage runs to the close marker, so <function=> in the gap is data.
+    text = (
+        "<|tool_call>call:outer{broken:{x}}<function=terminal>"
+        "<parameter=command>id</parameter></function><tool_call|>"
+    )
+    for allow_incomplete in (True, False):
+        calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
+        assert "terminal" not in [c["function"]["name"] for c in calls], calls
+
+
+def test_balanced_inner_call_inside_unclosed_outer_does_not_execute():
+    text = "<|tool_call>call:outer{code:<|tool_call>call:terminal{command:id}<tool_call|>"
+    for allow_incomplete in (True, False):
+        calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
+        assert "terminal" not in [c["function"]["name"] for c in calls], calls
+
+
+def test_strip_preserves_text_after_malformed_gemma_close():
+    # Junk before the close is a malformed span: strip through it, keep the tail.
+    text = "pre <|tool_call>call:t{a:1} note <tool_call|> post"
+    assert strip_tool_call_markup(text) == "pre  post"
+    assert strip_tool_call_markup(text, final = True) == "pre  post"
+
+
+def test_malformed_closed_gemma_span_is_stripped():
+    assert (
+        strip_tool_call_markup('before <|tool_call>{"name":"x"}<tool_call|> after')
+        == "before  after"
+    )
+
+
+def test_valid_call_after_missing_close_is_recovered():
+    # A close-less call covers only its braces, so the later call is recovered.
+    text = "<|tool_call>call:a{x:1} <|tool_call>call:b{y:2}<tool_call|>"
+    names_inc = [
+        c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = True)
+    ]
+    assert "b" in names_inc, names_inc
+    names_strict = [
+        c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False)
+    ]
+    assert names_strict == ["b"], names_strict
+
+
+def test_strip_non_final_keeps_incomplete_gemma_block():
+    text = "before <|tool_call>call:t{"
+    assert strip_tool_call_markup(text) == text
+    assert strip_tool_call_markup(text, final = True) == "before"
+
+
+def test_json_call_between_gemma_braces_and_close_does_not_execute():
+    # A JSON call between the outer's braces and its close is covered data.
+    text = (
+        "<|tool_call>call:outer{broken:{x}}"
+        '<tool_call>{"name":"terminal","arguments":{"command":"id"}}</tool_call>'
+        "<tool_call|>"
+    )
+    for allow_incomplete in (True, False):
+        calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
+        assert "terminal" not in [c["function"]["name"] for c in calls], calls
+
+
+def test_gemma_call_between_gemma_braces_and_close_does_not_execute():
+    # Same escape with a Gemma-native inner marker.
+    text = "<|tool_call>call:outer{broken:{x}}<|tool_call>call:terminal{command:id}<tool_call|><tool_call|>"
+    for allow_incomplete in (True, False):
+        calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
+        assert "terminal" not in [c["function"]["name"] for c in calls], calls
+
+
+def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener():
+    # The to-EOF Gemma sweep must not eat visible text after </function>.
+    text = (
+        'before <function=python><parameter=code>print("<|tool_call>")</parameter></function> after'
+    )
+    assert strip_tool_call_markup(text, final = True) == "before  after"
+    assert strip_tool_call_markup(text) == "before  after"
+
+
+def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener():
+    # A call-form Gemma opener quoted in a closed block must not truncate it.
+    xml = "<function=python><parameter=code><|tool_call>call:t{</parameter></function>"
+    json_block = (
+        '<tool_call>{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}</tool_call>'
+    )
+    for block in (xml, json_block):
+        text = "before " + block + " after"
+        assert strip_tool_call_markup(text, final = True) == "before  after", block
+        assert strip_tool_call_markup(text) == "before  after", block
+
+
+def test_function_sibling_after_close_less_gemma_marker_is_recovered():
+    # The close-less marker covers only its braces; the XML sibling is recovered.
+    text = (
+        "<|tool_call>call:bad{broken:{x}} "
+        "<function=terminal><parameter=command>id</parameter></function>"
+    )
+    for allow_incomplete in (True, False):
+        calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete)
+        assert [c["function"]["name"] for c in calls] == ["terminal"], calls
+
+
+def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered():
+    # A close token quoted in the later call must not extend the earlier
+    # close-less marker's coverage over that call.
+    gemma = '<|tool_call>call:a{x:1} <|tool_call>call:b{note:<|"|></tool_call><|"|>}<tool_call|>'
+    names = [
+        c["function"]["name"] for c in parse_tool_calls_from_text(gemma, allow_incomplete = False)
+    ]
+    assert names == ["b"], names
+    json_text = (
+        '<tool_call>{"name":"a","arguments":{}} '
+        '<tool_call>{"name":"b","arguments":{"x":"</tool_call>"}}</tool_call>'
+    )
+    names_j = [
+        c["function"]["name"] for c in parse_tool_calls_from_text(json_text, allow_incomplete = False)
+    ]
+    assert "b" in names_j, names_j
+
+
+def test_gemma_parse_value_always_advances_on_stray_delimiter():
+    # A stray delimiter (`,`, `}`, `]`) at the primitive position must still advance the
+    # index by at least one, or a caller looping on it spins forever at 100% CPU (DoS).
+    for delim in (",", "}", "]"):
+        text = delim + "rest"
+        value, nxt, _explicit = _gemma_parse_value(text, 0)
+        assert nxt > 0, (delim, value, nxt)
+
+
+def test_malformed_gemma_array_does_not_hang():
+    # ``[},]`` puts a stray ``}`` at the primitive position inside a list body.
+    # On the buggy parser this hangs the server; guard with a wall-clock timeout
+    # so the regression fails loudly instead of blocking CI forever.
+    import threading
+
+    result: dict = {}
+
+    def _run():
+        result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:[},]}<tool_call|>")
+
+    t = threading.Thread(target = _run, daemon = True)
+    t.start()
+    t.join(timeout = 10.0)
+    assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed array input"
+
+
+def test_malformed_gemma_mapping_value_does_not_hang():
+    # A stray ``}`` where a mapping value is expected must also terminate.
+    import threading
+
+    result: dict = {}
+
+    def _run():
+        result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:}},b:1}<tool_call|>")
+
+    t = threading.Thread(target = _run, daemon = True)
+    t.start()
+    t.join(timeout = 10.0)
+    assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed mapping input"

studio/backend/tests/test_hf_xet_fallback.py

@@ -1,18 +1,16 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""Unit tests for utils.hf_xet_fallback: the no-progress watchdog, the Xet->HTTP
-transport policy, and the HF_HUB_DISABLE_XET precondition the fallback rests on.
-CPU-only, no network, no real subprocess (the per-attempt download seam is
-monkeypatched).
+"""Tests for the Studio shim over the shared unsloth_zoo Xet -> HTTP fallback.
+
+The transport-policy matrix is tested once in unsloth_zoo; here we assert only the
+Studio seam: re-exporting the shared API and injecting the marker-aware
+prepare_cache_for_transport on the HTTP retry. CPU-only, no network, no real subprocess.
 """
 
 from __future__ import annotations
 
-import subprocess
 import sys
-import threading
-import time
 import types as _types
 from pathlib import Path
 
@@ -22,9 +20,8 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
     sys.path.insert(0, _BACKEND_DIR)
 
-# Stub heavy/unavailable deps before importing the module under test. Use the
-# real structlog when present; a bare stub left in sys.modules would break later
-# modules that log at import time.
+# Stub heavy/unavailable deps before importing the module under test. Use real structlog when present;
+# a bare stub would break later modules that log at import time.
 _loggers_stub = _types.ModuleType("loggers")
 _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
 sys.modules.setdefault("loggers", _loggers_stub)
@@ -34,171 +31,59 @@ except ImportError:
     sys.modules["structlog"] = _types.ModuleType("structlog")
 
 import huggingface_hub
-from huggingface_hub import constants as hf_constants
+
+try:
+    import unsloth_zoo.hf_xet_fallback as _shared_mod
+    shared = _shared_mod
+except Exception:  # noqa: BLE001 - still collect degraded-path tests when unsloth_zoo is unavailable
+    shared = None
 
 import utils.hf_xet_fallback as xf
 
 
-# --------------------------------------------------------------------------- #
-# Watchdog: fires only on a constant-size .incomplete, sparse-aware byte total.
-# --------------------------------------------------------------------------- #
-REPO = "ztest/xet-watchdog"
-
-
-@pytest.fixture
-def hf_cache(tmp_path, monkeypatch):
-    monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
-    return tmp_path
-
-
-def _blobs_dir(root: Path, repo_id: str = REPO) -> Path:
-    d = root / f"models--{repo_id.replace('/', '--')}" / "blobs"
-    d.mkdir(parents = True, exist_ok = True)
-    return d
-
-
-def _wait(
-    predicate,
-    timeout: float = 2.0,
-    step: float = 0.02,
-) -> bool:
-    deadline = time.monotonic() + timeout
-    while time.monotonic() < deadline:
-        if predicate():
-            return True
-        time.sleep(step)
-    return predicate()
-
-
-def test_constant_incomplete_fires_stall(hf_cache):
-    blobs = _blobs_dir(hf_cache)
-    (blobs / "deadbeef.incomplete").write_bytes(b"\0" * 1024)  # never grows
-
-    calls: list[str] = []
-    stop = xf.start_watchdog(
-        repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3
-    )
-    try:
-        assert _wait(
-            lambda: len(calls) >= 1, timeout = 3.0
-        ), "watchdog never fired on a constant-size .incomplete"
-    finally:
-        stop.set()
-    assert "stalled" in calls[0].lower()
-
-
-def test_growing_incomplete_never_stalls(hf_cache):
-    blobs = _blobs_dir(hf_cache)
-    part = blobs / "growing.incomplete"
-    part.write_bytes(b"\0" * 1024)
-
-    grow_stop = threading.Event()
-
-    def _grow():
-        size = 1024
-        while not grow_stop.wait(0.05):
-            size += 4096
-            part.write_bytes(b"\0" * size)
-
-    grower = threading.Thread(target = _grow, daemon = True)
-    grower.start()
-
-    calls: list[str] = []
-    stop = xf.start_watchdog(
-        repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3
-    )
-    try:
-        time.sleep(1.0)  # well past stall_timeout, but bytes keep growing
-        assert calls == [], "watchdog fired despite continuous progress"
-    finally:
-        stop.set()
-        grow_stop.set()
-
-
-def test_no_incomplete_never_stalls(hf_cache):
-    blobs = _blobs_dir(hf_cache)
-    (blobs / "finalized_blob").write_bytes(b"\0" * 4096)  # no .incomplete
-
-    calls: list[str] = []
-    stop = xf.start_watchdog(
-        repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3
-    )
-    try:
-        time.sleep(0.8)
-        assert calls == [], "watchdog fired with no active .incomplete"
-    finally:
-        stop.set()
-
-
-def test_stall_fires_at_most_once(hf_cache):
-    blobs = _blobs_dir(hf_cache)
-    (blobs / "frozen.incomplete").write_bytes(b"\0" * 2048)
-
-    calls: list[str] = []
-    stop = xf.start_watchdog(
-        repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.2
-    )
-    try:
-        assert _wait(lambda: len(calls) >= 1, timeout = 3.0)
-        time.sleep(0.6)  # keep ticking; must not fire again
-        assert len(calls) == 1, f"on_stall fired {len(calls)} times, expected exactly 1"
-    finally:
-        stop.set()
-
-
-def test_get_state_empty_cache(hf_cache):
-    assert xf.get_hf_download_state([REPO]) == (0, False)
-
-
-def test_get_state_absent_cache_root(tmp_path, monkeypatch):
-    monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path / "no-such-cache"))
-    assert xf.get_hf_download_state([REPO]) == (0, False)
-
-
-def test_get_state_skips_local_paths(hf_cache):
-    # Filesystem paths are not HF repo IDs and must be ignored without error.
-    assert xf.get_hf_download_state(["/abs/path", "./rel", "~user", "c:\\x"]) == (0, False)
-
-
-def test_get_state_sparse_aware(hf_cache):
-    blobs = _blobs_dir(hf_cache)
-    sparse = blobs / "sparse.incomplete"
-    with open(sparse, "wb") as f:
-        f.truncate(64 * 1024 * 1024)  # large apparent size, few allocated blocks
-    st = sparse.stat()
-    if getattr(st, "st_blocks", 0) == 0:
-        pytest.skip("filesystem does not report st_blocks; sparse accounting unavailable")
-    total, has_incomplete = xf.get_hf_download_state([REPO])
-    assert has_incomplete is True
-    assert total < st.st_size, "sparse partial counted at apparent size, not allocated blocks"
-
-
-# --------------------------------------------------------------------------- #
-# Transport policy: cached short-circuit, cancel, error propagation, and the
-# single Xet->HTTP fallback. _run_download_attempt is faked, so no real spawn.
-# --------------------------------------------------------------------------- #
 DL_REPO, FILE = "ztest/xet-dl", "model-Q4_K_XL.gguf"
 
 
-@pytest.fixture(autouse = True)
-def _no_real_cache_hit(monkeypatch):
-    """Default: the cached probe misses; tests override it to force a hit."""
+def _requires_shared():
+    if shared is None:
+        pytest.skip("unsloth_zoo.hf_xet_fallback is not installed in this environment")
+
+
+def test_shim_reexports_shared_api():
+    _requires_shared()
+    assert xf.DownloadStallError is shared.DownloadStallError
+    for name in (
+        "start_watchdog",
+        "get_hf_download_state",
+        "child_should_disable_xet",
+        "hf_hub_download_with_xet_fallback",
+        "snapshot_download_with_xet_fallback",
+    ):
+        assert hasattr(xf, name), f"shim missing {name}"
+
+
+def test_child_should_disable_xet_truth_table():
+    assert xf.child_should_disable_xet({"disable_xet": True}) is True
+    assert xf.child_should_disable_xet({"disable_xet": False}) is False
+    assert xf.child_should_disable_xet({}) is False
+
+
+def test_shim_injects_studio_prepare_on_http_retry(monkeypatch):
+    """A Xet stall retries over HTTP and the shim runs Studio's marker-aware
+    ``prepare_cache_for_transport(..., 'http')`` before the retry."""
+    _requires_shared()
+    for var in ("UNSLOTH_DISABLE_XET", "UNSLOTH_STABLE_DOWNLOADS", "HF_HUB_DISABLE_XET"):
+        monkeypatch.delenv(var, raising = False)
     monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None)
 
+    seen_disable_xet = []
 
-class _FakeAttempt:
-    """Records calls to the download seam and returns scripted results."""
-
-    def __init__(self, results):
-        self._results = list(results)
-        self.calls = []
-
-    def __call__(
-        self,
+    def fake_attempt(
         repo_id,
-        filename,
-        token,
         *,
+        kind,
+        params,
+        token,
         repo_type,
         disable_xet,
         cancel_event,
@@ -208,146 +93,243 @@ class _FakeAttempt:
         on_status,
         force_download = False,
     ):
-        self.calls.append(
-            _types.SimpleNamespace(
-                repo_id = repo_id,
-                filename = filename,
-                disable_xet = disable_xet,
-                repo_type = repo_type,
-            )
-        )
-        return self._results[len(self.calls) - 1]
+        seen_disable_xet.append(disable_xet)
+        return ("ok", "/cache/model.gguf") if disable_xet else ("stall", None)
 
+    monkeypatch.setattr(shared, "_run_download_attempt", fake_attempt)
 
-def _install(monkeypatch, results):
-    fake = _FakeAttempt(results)
-    monkeypatch.setattr(xf, "_run_download_attempt", fake)
-    return fake
-
-
-def test_cached_file_short_circuits(monkeypatch, tmp_path):
-    cached = tmp_path / "cached.gguf"
-    cached.write_bytes(b"\0" * 8)
-    monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: str(cached))
-    fake = _install(monkeypatch, [])  # must not be called
-
-    out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
-    assert out == str(cached)
-    assert fake.calls == [], "spawned a download for an already-cached file"
-
-
-def test_cancel_before_start_raises_no_attempt(monkeypatch):
-    fake = _install(monkeypatch, [])
-    ev = threading.Event()
-    ev.set()
-    with pytest.raises(RuntimeError, match = "Cancelled"):
-        xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None, cancel_event = ev)
-    assert fake.calls == []
-
-
-def test_nonstall_error_propagates_without_fallback(monkeypatch):
-    fake = _install(monkeypatch, [("error", "RepositoryNotFoundError: 404 not found")])
-    with pytest.raises(RuntimeError, match = "RepositoryNotFoundError"):
-        xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
-    assert len(fake.calls) == 1, "deterministic error must not trigger an HTTP fallback"
-    assert fake.calls[0].disable_xet is False
-
-
-def test_immediate_success_uses_xet_only(monkeypatch):
-    prepared = []
-    monkeypatch.setattr(
-        "hub.utils.download_registry.prepare_cache_for_transport",
-        lambda *a, **k: prepared.append(a),
-    )
-    fake = _install(monkeypatch, [("ok", "/cache/model.gguf")])
-    out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
-    assert out == "/cache/model.gguf"
-    assert len(fake.calls) == 1 and fake.calls[0].disable_xet is False
-    assert prepared == [], "no cache prep should run when Xet succeeds first try"
-
-
-def test_stall_then_http_fallback_succeeds(monkeypatch):
     prepared = []
     monkeypatch.setattr(
         "hub.utils.download_registry.prepare_cache_for_transport",
         lambda repo_type, repo_id, mode, *a, **k: prepared.append((repo_type, repo_id, mode)),
     )
-    fake = _install(monkeypatch, [("stall", None), ("ok", "/cache/model.gguf")])
 
     out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
     assert out == "/cache/model.gguf"
-    assert len(fake.calls) == 2
-    assert fake.calls[0].disable_xet is False  # Xet first
-    assert fake.calls[1].disable_xet is True  # HTTP fallback
-    assert prepared == [("model", DL_REPO, "http")], "must prep cache for HTTP before the retry"
+    assert seen_disable_xet == [False, True]  # Xet first, then HTTP
+    assert prepared == [("model", DL_REPO, "http")], "shim must run Studio's marker-aware prep"
 
 
-def test_second_stall_raises_download_stall_error(monkeypatch):
-    monkeypatch.setattr(
-        "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None
-    )
-    fake = _install(monkeypatch, [("stall", None), ("stall", None)])
-    with pytest.raises(xf.DownloadStallError):
-        xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
-    assert len(fake.calls) == 2
+def test_shim_snapshot_injects_studio_prepare(monkeypatch):
+    """The snapshot wrapper forwards Studio's marker-aware prep, like the file wrapper."""
+    captured = {}
+
+    def fake_snapshot(repo_id, **kwargs):
+        captured["repo_id"] = repo_id
+        captured["prepare_for_http_fn"] = kwargs.get("prepare_for_http_fn")
+        return "/tmp/snap-dir"
+
+    monkeypatch.setattr(xf, "_shared_snapshot_download_with_xet_fallback", fake_snapshot)
+    out = xf.snapshot_download_with_xet_fallback("org/model")
+    assert out == "/tmp/snap-dir"
+    assert captured["repo_id"] == "org/model"
+    assert captured["prepare_for_http_fn"] is xf._studio_prepare_for_http
 
 
-def test_cancelled_midattempt_raises_no_fallback(monkeypatch):
-    fake = _install(monkeypatch, [("cancelled", None)])
-    with pytest.raises(RuntimeError, match = "Cancelled"):
-        xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
-    assert len(fake.calls) == 1
+def test_degrades_gracefully_without_shared_helper(monkeypatch):
+    """On an older unsloth_zoo lacking the shared helper, the shim still imports (Studio
+    boots) and exposes stub API doing plain HF downloads with the watchdog disabled."""
+    import importlib
+
+    class _BlockShared:
+        def find_spec(
+            self,
+            name,
+            path = None,
+            target = None,
+        ):
+            if name == "unsloth_zoo.hf_xet_fallback":
+                raise ModuleNotFoundError(f"No module named '{name}'", name = name)
+            return None
+
+    finder = _BlockShared()
+    saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None)
+    saved_shim = sys.modules.pop("utils.hf_xet_fallback", None)
+    sys.meta_path.insert(0, finder)
+    try:
+        degraded = importlib.import_module("utils.hf_xet_fallback")
+
+        # Boots without raising and mirrors the shared API surface.
+        assert issubclass(degraded.DownloadStallError, RuntimeError)
+        assert degraded.child_should_disable_xet({"disable_xet": True}) is True
+        assert degraded.get_hf_download_state(["x"]) is None  # unmeasurable
+        event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None)
+        assert hasattr(event, "set") and not event.is_set()  # never fires
+
+        # Degraded mode still emits heartbeats so the inactivity deadline is not tripped.
+        import time as _time
+
+        beats = []
+        hb_stop = degraded.start_watchdog(
+            repo_ids = ["x"],
+            on_stall = lambda m: None,
+            on_heartbeat = beats.append,
+            interval = 0.02,
+        )
+        try:
+            deadline = _time.monotonic() + 2.0
+            while not beats and _time.monotonic() < deadline:
+                _time.sleep(0.02)
+            assert beats, "degraded watchdog emitted no heartbeat"
+        finally:
+            hb_stop.set()
+
+        # Downloads fall back to plain huggingface_hub (no watchdog, no crash).
+        called = {}
+
+        def _fake_snapshot(repo_id, **kwargs):
+            called["repo_id"] = repo_id
+            return "/snap-dir"
+
+        monkeypatch.setattr(huggingface_hub, "snapshot_download", _fake_snapshot)
+        assert degraded.snapshot_download_with_xet_fallback("org/model") == "/snap-dir"
+        assert called["repo_id"] == "org/model"
+
+        # Cancellation still holds: an already-set cancel_event aborts before the HF download.
+        import threading as _threading
+
+        cancelled = _threading.Event()
+        cancelled.set()
+        called.clear()
+        with pytest.raises(RuntimeError, match = "Cancelled"):
+            degraded.snapshot_download_with_xet_fallback("org/model", cancel_event = cancelled)
+        assert "repo_id" not in called, "degraded download ran despite cancellation"
+    finally:
+        sys.meta_path.remove(finder)
+        sys.modules.pop("utils.hf_xet_fallback", None)
+        if saved_shared is not None:
+            sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared
+        if saved_shim is not None:
+            sys.modules["utils.hf_xet_fallback"] = saved_shim
 
 
-def test_per_file_independent_fallback(monkeypatch):
-    """A stalled shard falls back; a sibling shard that succeeds does not."""
-    monkeypatch.setattr(
-        "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None
-    )
-    fake = _install(monkeypatch, [("ok", "/a"), ("stall", None), ("ok", "/b")])
-    assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardA.gguf", None) == "/a"
-    assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardB.gguf", None) == "/b"
-    assert [c.disable_xet for c in fake.calls] == [False, False, True]
+def test_degrades_when_unsloth_zoo_entirely_absent():
+    """When unsloth_zoo is absent entirely, the import raises
+    ModuleNotFoundError(name='unsloth_zoo') (top-level package). Guard that the shim still
+    degrades and does not re-raise, breaking every Studio import that pulls it in."""
+    import importlib
+
+    class _BlockZoo:
+        def find_spec(
+            self,
+            name,
+            path = None,
+            target = None,
+        ):
+            # Whole package absent, so ModuleNotFoundError.name is the top-level 'unsloth_zoo'.
+            if name == "unsloth_zoo" or name.startswith("unsloth_zoo."):
+                raise ModuleNotFoundError("No module named 'unsloth_zoo'", name = "unsloth_zoo")
+            return None
+
+    finder = _BlockZoo()
+    saved = {
+        k: v
+        for k, v in list(sys.modules.items())
+        if k == "unsloth_zoo" or k.startswith("unsloth_zoo.")
+    }
+    for k in saved:
+        del sys.modules[k]
+    saved_shim = sys.modules.pop("utils.hf_xet_fallback", None)
+    sys.meta_path.insert(0, finder)
+    try:
+        degraded = importlib.import_module("utils.hf_xet_fallback")
+        # Boots without raising and exposes the stub API.
+        assert issubclass(degraded.DownloadStallError, RuntimeError)
+        assert degraded.get_hf_download_state(["x"]) is None
+        event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None)
+        assert hasattr(event, "set") and not event.is_set()
+    finally:
+        sys.meta_path.remove(finder)
+        sys.modules.pop("utils.hf_xet_fallback", None)
+        sys.modules.update(saved)
+        if saved_shim is not None:
+            sys.modules["utils.hf_xet_fallback"] = saved_shim
 
 
-# --------------------------------------------------------------------------- #
-# Precondition: HF_HUB_DISABLE_XET is read at import time, so assert its effect
-# in a FRESH interpreter (huggingface/huggingface_hub#3266 once ignored it).
-# --------------------------------------------------------------------------- #
-def _safe_path() -> str:
+def test_degrades_when_shared_helper_import_raises_importerror():
+    """unsloth_zoo can be installed yet fail to import when torch is missing (llama.cpp/GGUF-only
+    Studio), raising ImportError not ModuleNotFoundError. The shim must degrade for that too."""
+    import importlib
+
+    class _BlockWithImportError:
+        def find_spec(
+            self,
+            name,
+            path = None,
+            target = None,
+        ):
+            if name == "unsloth_zoo.hf_xet_fallback":
+                # Mirror a torch-less install: a plain ImportError with no .name.
+                raise ImportError("Unsloth: Pytorch is not installed.")
+            return None
+
+    finder = _BlockWithImportError()
+    saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None)
+    saved_zoo = sys.modules.pop("unsloth_zoo", None)
+    saved_shim = sys.modules.pop("utils.hf_xet_fallback", None)
+    sys.meta_path.insert(0, finder)
+    try:
+        degraded = importlib.import_module("utils.hf_xet_fallback")
+        assert issubclass(degraded.DownloadStallError, RuntimeError)
+        assert degraded.get_hf_download_state(["x"]) is None
+        event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None)
+        assert hasattr(event, "set") and not event.is_set()
+    finally:
+        sys.meta_path.remove(finder)
+        sys.modules.pop("utils.hf_xet_fallback", None)
+        if saved_shared is not None:
+            sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared
+        if saved_zoo is not None:
+            sys.modules["unsloth_zoo"] = saved_zoo
+        if saved_shim is not None:
+            sys.modules["utils.hf_xet_fallback"] = saved_shim
+
+
+def test_retries_under_light_gpu_init_when_import_fails(monkeypatch):
+    """GPU detection in unsloth_zoo's __init__ raises NotImplementedError on a GPU-less host. The shim
+    retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails."""
+    import importlib
     import os
-    return os.environ.get("PATH", "")
 
+    monkeypatch.delenv("UNSLOTH_ZOO_DISABLE_GPU_INIT", raising = False)
+    seen_env = []
 
-def test_disable_xet_constant_set_in_fresh_interpreter():
-    code = (
-        "from huggingface_hub import constants as c; "
-        "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is True else 17)"
-    )
-    proc = subprocess.run(
-        [sys.executable, "-c", code],
-        env = {"HF_HUB_DISABLE_XET": "1", "PATH": _safe_path()},
-        capture_output = True,
-        text = True,
-    )
-    assert proc.returncode == 0, (
-        f"HF_HUB_DISABLE_XET=1 did not set constants.HF_HUB_DISABLE_XET=True "
-        f"(rc={proc.returncode}): {proc.stderr}"
-    )
+    class _GpuGatedBlocker:
+        def find_spec(
+            self,
+            name,
+            path = None,
+            target = None,
+        ):
+            # Crash is in unsloth_zoo's __init__, so intercept "unsloth_zoo" itself (the parent).
+            if name == "unsloth_zoo":
+                # Record the env each attempt sees; raise the no-GPU error both times so the shim
+                # degrades.
+                seen_env.append(os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT"))
+                raise NotImplementedError("Unsloth cannot find any torch accelerator")
+            return None
 
-
-def test_default_leaves_xet_enabled():
-    code = (
-        "from huggingface_hub import constants as c; "
-        "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is False else 17)"
-    )
-    proc = subprocess.run(
-        [sys.executable, "-c", code],
-        env = {"PATH": _safe_path()},  # no HF_HUB_DISABLE_XET
-        capture_output = True,
-        text = True,
-    )
-    assert proc.returncode == 0, (
-        f"without the env var, constants.HF_HUB_DISABLE_XET was not False "
-        f"(rc={proc.returncode}): {proc.stderr}"
-    )
+    finder = _GpuGatedBlocker()
+    saved = {
+        k: v
+        for k, v in list(sys.modules.items())
+        if k == "unsloth_zoo" or k.startswith("unsloth_zoo.")
+    }
+    for k in saved:
+        del sys.modules[k]
+    saved_shim = sys.modules.pop("utils.hf_xet_fallback", None)
+    sys.meta_path.insert(0, finder)
+    try:
+        degraded = importlib.import_module("utils.hf_xet_fallback")
+        # First attempt without the light env, then a retry with it set.
+        assert seen_env == [None, "1"], seen_env
+        # Both attempts raised -> Studio still boots in degraded mode.
+        assert issubclass(degraded.DownloadStallError, RuntimeError)
+        # The env override must not leak past the import.
+        assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None
+    finally:
+        sys.meta_path.remove(finder)
+        sys.modules.pop("utils.hf_xet_fallback", None)
+        sys.modules.update(saved)
+        if saved_shim is not None:
+            sys.modules["utils.hf_xet_fallback"] = saved_shim

studio/backend/tests/test_llama_cpp_tool_loop.py

@@ -20,7 +20,11 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
     sys.path.insert(0, _BACKEND_DIR)
 
-from core.inference.llama_cpp import _PROVISIONAL_ARGS_MIN_CHARS, LlamaCppBackend
+from core.inference.llama_cpp import (
+    _MAX_REPROMPTS,
+    _PROVISIONAL_ARGS_MIN_CHARS,
+    LlamaCppBackend,
+)
 from state import tool_approvals
 from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision
 
@@ -1036,9 +1040,11 @@ def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch):
 def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch):
     """No-tool re-prompt attempts should not concatenate into the UI."""
 
-    streams = [
-        [_sse({"content": "I will use render_html now."}), _done()],
-        [_sse({"content": "Understood. I will use render_html now."}), _done()],
+    # One initial response plus one stream per re-prompt; derive the count from the shared cap.
+    streams = [[_sse({"content": "I will use render_html now."}), _done()]]
+    streams += [
+        [_sse({"content": "Understood. I will use render_html now."}), _done()]
+        for _ in range(_MAX_REPROMPTS)
     ]
     payloads: list[dict] = []
     backend = _make_backend(monkeypatch, streams, payloads)
@@ -1073,7 +1079,7 @@ def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch):
 
     content_texts = [event.get("text", "") for event in events if event.get("type") == "content"]
     assert content_texts == ["I will use render_html now."]
-    assert len(payloads) == 2
+    assert len(payloads) == _MAX_REPROMPTS + 1
 
 
 def test_forced_reprompt_plain_final_answer_is_visible(monkeypatch):
@@ -1200,6 +1206,66 @@ def test_auto_heal_disabled_parses_well_formed_xml_when_tools_enabled(monkeypatc
     )
 
 
+def test_textual_mistral_marker_not_leaked_when_inline_with_preface(monkeypatch):
+    # Textual Mistral ``[TOOL_CALLS]`` inline with visible preface: the DRAINING flush must use the
+    # shared parser patterns (which know ``[TOOL_CALLS]``); the legacy set leaked the marker to clients.
+    streams = [
+        [_sse({"content": 'Let me search. [TOOL_CALLS]web_search{"query":"cats"}'}), _done()],
+        [_sse({"content": "done"}), _done()],
+    ]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, streams, payloads)
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return "result"
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "search"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 1,
+        )
+    )
+
+    assert calls == [("web_search", {"query": "cats"})]
+    content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
+    assert all("[TOOL_CALLS]" not in t for t in content_texts), content_texts
+    assert any("Let me search." in t for t in content_texts)
+
+
+def test_textual_llama_python_tag_marker_not_leaked(monkeypatch):
+    # Same leak class for the Llama-3 built-in ``<|python_tag|>NAME.call(...)`` form.
+    streams = [
+        [_sse({"content": '<|python_tag|>web_search.call(query="cats")'}), _done()],
+        [_sse({"content": "done"}), _done()],
+    ]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, streams, payloads)
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return "result"
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "search"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 1,
+        )
+    )
+
+    assert calls == [("web_search", {"query": "cats"})]
+    content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
+    assert all("<|python_tag|>" not in t for t in content_texts), content_texts
+
+
 def test_reprompted_tool_call_still_streams_final_answer(monkeypatch):
     """Suppression ends once a forced re-prompt actually calls a tool."""
 
@@ -1738,6 +1804,255 @@ def test_empty_tool_call_id_does_not_emit_provisional_card(monkeypatch):
     assert calls == [("python", {"code": big_code})]
 
 
+def _streamed_content(text: str, frag: int = 4) -> list[str]:
+    """Stream content token-by-token like llama-server; ``frag`` sets the chunk size."""
+    chunks = [_sse({"content": text[i : i + frag]}) for i in range(0, len(text), frag)]
+    chunks.append(_done())
+    return chunks
+
+
+def test_bare_json_tool_call_streamed_is_not_leaked_and_executes(monkeypatch):
+    """A wrapper-less bare-JSON call must be held while incomplete, drained silently, and executed with nothing leaking."""
+
+    bare_call = '{"name": "web_search", "parameters": {"query": "weather in Sydney"}}'
+    first_stream = _streamed_content(bare_call)
+    final_stream = [_sse({"content": "It is sunny in Sydney."}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return "Weather: sunny, 22C."
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "weather in Sydney?"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 1,
+        )
+    )
+
+    # The tool ran with the parsed arguments.
+    assert calls == [("web_search", {"query": "weather in Sydney"})]
+    assert any(
+        event.get("type") == "tool_end" and event.get("tool_name") == "web_search"
+        for event in events
+    )
+
+    # The bare JSON never leaked to the user-visible stream.
+    content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
+    assert all('"name"' not in t for t in content_texts), content_texts
+    assert all("web_search" not in t for t in content_texts), content_texts
+    # The post-tool synthesis is still streamed.
+    assert any("sunny in Sydney" in t for t in content_texts), content_texts
+
+
+def test_ordinary_json_with_name_key_is_shown_not_treated_as_tool_call(monkeypatch):
+    """Markerless JSON with a non-enabled name is the answer, not a phantom call."""
+
+    answer = '{"name": "Alice", "parameters": {"age": 30}}'
+    first_stream = _streamed_content(answer)
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [first_stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+    monkeypatch.setattr(
+        "core.inference.tools.execute_tool",
+        lambda n, a, **_k: (calls.append((n, a)) or "x"),
+    )
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "give me a person record"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 1,
+        )
+    )
+
+    assert calls == [], calls
+    content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
+    assert any("Alice" in t for t in content_texts), content_texts
+
+
+def test_incomplete_bare_json_truncation_is_not_leaked(monkeypatch):
+    """If generation is cut off mid bare-JSON object (no closing brace), the held
+    fragment must be stripped at stream end rather than dumped to the user."""
+
+    truncated = '{"name": "web_search", "parameters": {"query": "weather in S'
+    stream = _streamed_content(truncated)
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [stream], payloads)
+
+    monkeypatch.setattr(
+        "core.inference.tools.execute_tool",
+        lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("no complete call")),
+    )
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "weather?"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 1,
+        )
+    )
+
+    content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
+    assert all('{"name"' not in t for t in content_texts), content_texts
+
+
+def test_gguf_truncated_ordinary_json_with_name_key_is_shown_not_suppressed(monkeypatch):
+    """A truncated markerless object whose "name" is NOT an enabled tool (a person
+    record cut off mid-stream, ``{"name":"Alice","age":``) must still be shown. The
+    end-of-stream ``_is_bare_tc`` heuristic routed any ``{...,"name",...}`` fragment
+    to DRAINING (dropped); it is now gated on the enabled tool names so only a real
+    truncated tool call is suppressed, ordinary JSON streams through."""
+
+    truncated = '{"name": "Alice", "age": 30, "bio": "loves '
+    stream = _streamed_content(truncated)
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+    monkeypatch.setattr(
+        "core.inference.tools.execute_tool",
+        lambda n, a, **_k: (calls.append((n, a)) or "x"),
+    )
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "start a person record"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 1,
+        )
+    )
+
+    assert calls == [], calls
+    content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
+    assert any("Alice" in t for t in content_texts), content_texts
+
+
+def test_gguf_truncated_disabled_name_json_is_preserved_when_tools_active(monkeypatch):
+    """A truncated JSON answer with a non-enabled name must still be shown (resolvers are gated on enabled names)."""
+
+    truncated = '{"name": "Alice", "parameters": {"age": 30'
+    stream = _streamed_content(truncated)
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+    monkeypatch.setattr(
+        "core.inference.tools.execute_tool",
+        lambda n, a, **_k: (calls.append((n, a)) or "x"),
+    )
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "give json"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 1,
+        )
+    )
+
+    assert calls == [], calls
+    content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
+    assert any("Alice" in t for t in content_texts), content_texts
+
+
+def test_gguf_truncated_enabled_name_json_is_still_suppressed(monkeypatch):
+    """Counterpart guard: a truncated ENABLED-tool bare call (``web_search``) cut off
+    mid-JSON still must NOT leak -- the gate only spares disabled / non-tool names."""
+
+    truncated = '{"name": "web_search", "parameters": {"query": "weather in S'
+    stream = _streamed_content(truncated)
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [stream], payloads)
+
+    monkeypatch.setattr(
+        "core.inference.tools.execute_tool",
+        lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("no complete call")),
+    )
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "weather?"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 1,
+        )
+    )
+
+    content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
+    assert all("web_search" not in t for t in content_texts), content_texts
+    assert all('{"name"' not in t for t in content_texts), content_texts
+
+
+def test_gguf_oversized_disabled_name_json_is_preserved(monkeypatch):
+    """An oversized still-open JSON answer with a non-enabled name streams as content, not a phantom drain."""
+
+    cap = 16384
+    big = "A" * (cap + 5000)
+    answer = '{"name":"Alice","parameters":{"bio":"' + big  # never closes
+    first_stream = [_sse({"content": answer[i : i + 2000]}) for i in range(0, len(answer), 2000)]
+    first_stream.append(_done())
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [first_stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+    monkeypatch.setattr(
+        "core.inference.tools.execute_tool",
+        lambda n, a, **_k: (calls.append((n, a)) or "x"),
+    )
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "long json"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 1,
+        )
+    )
+
+    assert calls == [], calls
+    content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
+    assert any("Alice" in t for t in content_texts), content_texts[:1]
+
+
+def test_gemma_wrapperless_call_streamed_is_not_leaked_and_executes(monkeypatch):
+    """Gemma 4 GGUF (skip_special_tokens) streams a wrapper-less ``call:NAME{..}``
+    with no XML signal. Like bare JSON, the BUFFERING scan must recognise it via
+    _GEMMA_BARE_TC_RE, drain it silently, and execute the tool -- never leaking
+    the ``call:`` markup to the user-visible stream."""
+
+    gemma_call = 'call:web_search{query:"weather in Sydney"}'
+    first_stream = _streamed_content(gemma_call)
+    final_stream = [_sse({"content": "It is sunny in Sydney."}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return "Weather: sunny, 22C."
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "weather in Sydney?"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 1,
+        )
+    )
+
+    assert calls == [("web_search", {"query": "weather in Sydney"})]
+    content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
+    assert all("call:" not in t for t in content_texts), content_texts
+    assert any("sunny in Sydney" in t for t in content_texts), content_texts
+
+
 def _usage_done(usage: dict, finish_reason: str = "stop") -> str:
     """A terminal SSE chunk carrying llama-server's ``usage`` block, the way the
     real server reports it on the final chunk of a completion."""
@@ -1813,3 +2128,191 @@ def test_metadata_event_omits_prompt_tokens_details_when_absent(monkeypatch):
     metadata = [e for e in events if e.get("type") == "metadata"]
     assert metadata, "expected a metadata event"
     assert "prompt_tokens_details" not in metadata[-1]["usage"]
+
+
+def test_gguf_oversized_bare_json_not_leaked_and_executes(monkeypatch):
+    """An oversized bare-JSON call drains rather than streams, and still executes via the safety net."""
+
+    cap = 16384
+    big = "A" * (cap + 5000)
+    full = '{"name":"python","parameters":{"code":"' + big + '"}}'
+    first_stream = [_sse({"content": full[i : i + 2000]}) for i in range(0, len(full), 2000)]
+    first_stream.append(_done())
+    final_stream = [_sse({"content": "done"}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+    monkeypatch.setattr(
+        "core.inference.tools.execute_tool",
+        lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"),
+    )
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "run"}],
+            tools = [{"type": "function", "function": {"name": "python"}}],
+            max_tool_iterations = 1,
+        )
+    )
+
+    content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
+    assert not any(t.lstrip().startswith('{"name') for t in content_texts), content_texts[:1]
+    assert calls and calls[0][0] == "python"
+    assert len(calls[0][1].get("code", "")) > cap
+
+
+def test_gguf_bare_json_call_not_replayed_in_next_turn_content(monkeypatch):
+    """After a bare-JSON call executes, the kept assistant message must not carry the raw call as content."""
+
+    import copy
+
+    first_stream = [
+        _sse({"content": '{"name":"web_search","parameters":{"query":"cats"}}'}),
+        _done(),
+    ]
+    final_stream = [_sse({"content": "Found."}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", lambda *_a, **_k: "RESULT")
+
+    list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "cats"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 2,
+        )
+    )
+
+    assert len(payloads) >= 2
+    asst = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"]
+    assert asst and not any('"name"' in (m.get("content") or "") for m in asst), asst
+
+
+def test_gguf_textual_fallback_caps_distinct_tool_calls_per_turn(monkeypatch):
+    """A single textual-fallback turn that parses many DISTINCT tool calls must be
+    capped at _MAX_TOOL_CALLS_PER_TURN (structured delta.tool_calls are grammar
+    bounded by llama-server; text parsed from content is not). Mirrors the
+    safetensors loop so one runaway turn cannot fan out into dozens of executions."""
+    from core.inference.llama_cpp import _MAX_TOOL_CALLS_PER_TURN
+
+    n = _MAX_TOOL_CALLS_PER_TURN + 4
+    blocks = "".join(
+        '<tool_call>{"name":"t%d","arguments":{"i":%d}}</tool_call>' % (i, i) for i in range(n)
+    )
+    first_stream = [_sse({"content": blocks}), _done()]
+    final_stream = [_sse({"content": "done"}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+    monkeypatch.setattr(
+        "core.inference.tools.execute_tool",
+        lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"),
+    )
+
+    list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "go"}],
+            tools = [{"type": "function", "function": {"name": f"t{i}"}} for i in range(n)],
+            max_tool_iterations = 1,
+        )
+    )
+
+    assert len(calls) == _MAX_TOOL_CALLS_PER_TURN, [c[0] for c in calls]
+    # The cap keeps the first calls in order (no reordering / drop of leading ones).
+    assert [c[0] for c in calls] == [f"t{i}" for i in range(_MAX_TOOL_CALLS_PER_TURN)]
+
+
+def test_gguf_textual_fallback_collapses_duplicate_tool_calls(monkeypatch):
+    """Exact-duplicate textual calls in one turn collapse to a single execution."""
+    blocks = '<tool_call>{"name":"web_search","arguments":{"query":"cats"}}</tool_call>' * 5
+    first_stream = [_sse({"content": blocks}), _done()]
+    final_stream = [_sse({"content": "done"}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+    monkeypatch.setattr(
+        "core.inference.tools.execute_tool",
+        lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"),
+    )
+
+    list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "cats"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 1,
+        )
+    )
+
+    assert len(calls) == 1, [c[0] for c in calls]
+
+
+def test_gguf_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(monkeypatch):
+    """Auto-Heal OFF keeps a truncated enabled-name fragment visible; ON suppresses it (strip gated on auto_heal_tool_calls)."""
+
+    trunc = '{"name":"web_search","parameters":{"query":"weather'
+
+    def _run(auto_heal):
+        stream = [_sse({"content": trunc}), _done()]
+        backend = _make_backend(monkeypatch, [stream], [])
+        calls: list[tuple[str, dict]] = []
+        monkeypatch.setattr(
+            "core.inference.tools.execute_tool",
+            lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
+        )
+        events = list(
+            backend.generate_chat_completion_with_tools(
+                messages = [{"role": "user", "content": "x"}],
+                tools = [{"type": "function", "function": {"name": "web_search"}}],
+                max_tool_iterations = 1,
+                auto_heal_tool_calls = auto_heal,
+            )
+        )
+        contents = "".join(e.get("text", "") for e in events if e.get("type") == "content")
+        return calls, contents
+
+    calls_off, contents_off = _run(False)
+    assert calls_off == [], calls_off
+    assert "web_search" in contents_off, contents_off
+
+    calls_on, contents_on = _run(True)
+    assert calls_on == [], calls_on
+    assert "web_search" not in contents_on, contents_on
+
+
+def test_gguf_valid_tool_calls_respect_max_tool_iterations(monkeypatch):
+    """Re-prompt slots must not extend the tool budget: stop after ``max_tool_iterations`` executed rounds."""
+    # More tool-call streams than the budget: if re-prompt slots leaked into the budget (the bug) the
+    # loop would run 2+3=5 rounds; honouring it stops after 2, then a tool-less final-answer pass.
+    streams = [
+        _structured_tool_call("web_search", {"query": f"q{i}"}, f"call_{i}") for i in range(6)
+    ]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, streams, payloads)
+
+    calls: list[tuple[str, dict]] = []
+    monkeypatch.setattr(
+        "core.inference.tools.execute_tool",
+        lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
+    )
+
+    list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "search repeatedly"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 2,
+        )
+    )
+
+    # Exactly two executed tool rounds, then one final-answer pass.
+    assert len(calls) == 2, calls
+    assert len(payloads) == 3, len(payloads)
+    # The final pass is the budget-exhausted nudge and carries no tools.
+    assert _tool_names(payloads[2]) == [], _tool_names(payloads[2])
+    assert any(
+        m.get("role") == "user" and "used all available tool calls" in m.get("content", "")
+        for m in payloads[2]["messages"]
+    ), payloads[2]["messages"]

studio/backend/tests/test_mcp_servers.py

@@ -587,10 +587,12 @@ def test_tool_xml_strip_handles_hyphenated_function_names():
     import re as _re
     from pathlib import Path
 
+    from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC
+
     src = (Path(__file__).resolve().parent.parent / "routes/inference.py").read_text()
     m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", src, _re.DOTALL)
     assert m, "could not extract _TOOL_XML_RE"
-    ns: dict = {"_re": _re}
+    ns: dict = {"_re": _re, "_DS_OPEN_SRC": _DS_OPEN_SRC}
     exec(f"_TOOL_XML_RE = _re.compile({m.group(1)})", ns)
     rx = ns["_TOOL_XML_RE"]
     stripped = rx.sub(

studio/backend/tests/test_mlx_inference_backend.py

@@ -100,6 +100,32 @@ def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch):
     ]
     assert backend._is_vlm is False
     assert isinstance(backend._tokenizer, _DummyTokenizer)
+    # Non-LoRA text model: no base_model on the record.
+    assert backend.models["fake/text"]["base_model"] is None
+
+
+def test_mlx_text_lora_record_keeps_base_model_for_native_template(monkeypatch):
+    # A LoRA adapter's own tokenizer often ships no chat template; the native tool-calling template
+    # lives on the base model.
+    _install_fake_mlx(monkeypatch)
+    calls = []
+    _install_fake_fast_mlx(monkeypatch, calls)
+
+    from core.inference.mlx_inference import MLXInferenceBackend
+
+    backend = MLXInferenceBackend()
+    config = SimpleNamespace(
+        identifier = "fake/text-adapter",
+        is_vision = False,
+        is_lora = True,
+        base_model = "fake/text-base",
+    )
+
+    assert backend.load_model(config, max_seq_length = 4096, hf_token = "hf-token")
+
+    record = backend.models["fake/text-adapter"]
+    assert record["is_lora"] is True
+    assert record["base_model"] == "fake/text-base"
 
 
 def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewrite(
@@ -188,12 +214,12 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
     _install_fake_mlx(monkeypatch)
     from core.inference.mlx_inference import MLXInferenceBackend
 
-    captured = {}
+    # The text path renders once with tools, then the native-template fallback makes a second no-
+    # tools probe call (tools=None) to detect whether the template dropped the schema.
+    captured_calls = []
 
     def _fake_apply(tokenizer, messages, **kwargs):
-        captured["tokenizer"] = tokenizer
-        captured["messages"] = messages
-        captured["kwargs"] = kwargs
+        captured_calls.append({"tokenizer": tokenizer, "messages": messages, "kwargs": kwargs})
         return "<rendered prompt>"
 
     monkeypatch.setattr(
@@ -248,8 +274,15 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
         )
     )
     assert out == ["hi"]
-    # The toggled kwargs must reach the chat-template helper.
-    assert captured["kwargs"]["tools"] == [{"function": {"name": "web_search"}}]
-    assert captured["kwargs"]["enable_thinking"] is True
-    assert captured["kwargs"]["reasoning_effort"] == "medium"
-    assert captured["kwargs"]["preserve_thinking"] is True
+    # The toggled kwargs must reach the chat-template helper on the real render
+    # (one of the calls carries the tools; the fallback probe passes tools=None).
+    tool_renders = [
+        c
+        for c in captured_calls
+        if c["kwargs"].get("tools") == [{"function": {"name": "web_search"}}]
+    ]
+    assert tool_renders, captured_calls
+    render = tool_renders[0]
+    assert render["kwargs"]["enable_thinking"] is True
+    assert render["kwargs"]["reasoning_effort"] == "medium"
+    assert render["kwargs"]["preserve_thinking"] is True

studio/backend/tests/test_model_update_robustness.py

@@ -5,8 +5,8 @@
 Covers:
   * GGUF variant listing computes update_available from the already-fetched
     sibling metadata instead of a second Hub call.
-  * hf_hub_download_with_xet_fallback(force_download=True) bypasses the
-    try_to_load_from_cache cache-first early-return.
+  * hf_hub_download_with_xet_fallback forwards force_download through the shim to the
+    shared unsloth_zoo helper (which owns the cache-first early-return and its bypass).
 
 The cache "Update" action now runs through the download manager as a normal
 managed download (so it shows in the Downloads panel with progress + cancel),
@@ -341,44 +341,26 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path):
 # ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ───
 
 
-def test_force_download_bypasses_cache_first_early_return(monkeypatch):
-    """force_download=True skips the try_to_load_from_cache early-return and
-    proceeds to the real download path; force_download=False returns the cached
-    path without ever attempting a download (X2/F2)."""
-    import huggingface_hub as hf
+def test_force_download_is_forwarded_through_the_shim(monkeypatch):
+    """The shim's contract is to forward force_download unchanged to the shared helper (which owns the
+    cache-first early-return and bypass). Verify both False and True reach it (X2/F2)."""
     import utils.hf_xet_fallback as X
 
-    cached_path = "/cache/blob/cached.gguf"
+    seen = []
 
-    # Pretend the blob IS cached on disk (try_to_load_from_cache is imported
-    # inside the function from huggingface_hub, and os.path.exists must agree).
-    monkeypatch.setattr(hf, "try_to_load_from_cache", lambda *a, **k: cached_path, raising = False)
-    monkeypatch.setattr(X.os.path, "exists", lambda p: True, raising = False)
+    def fake_shared(repo_id, filename, token, **kwargs):
+        seen.append(kwargs.get("force_download"))
+        return "/downloaded/path"
 
-    attempts = []
+    monkeypatch.setattr(X, "_shared_hf_hub_download_with_xet_fallback", fake_shared, raising = True)
 
-    def fake_attempt(repo_id, filename, token, **kwargs):
-        attempts.append(
-            {"repo_id": repo_id, "filename": filename, "force": kwargs.get("force_download")}
-        )
-        return ("ok", "/freshly/downloaded/path")
-
-    monkeypatch.setattr(X, "_run_download_attempt", fake_attempt, raising = True)
-
-    # force_download=False: cache-first early-return, no download attempt.
-    out = X.hf_hub_download_with_xet_fallback(
+    X.hf_hub_download_with_xet_fallback(
         "unsloth/repo", "model.gguf", token = None, force_download = False
     )
-    assert out == cached_path
-    assert attempts == []  # never reached the real download
-
-    # force_download=True: bypass the early-return, run the real download.
-    out2 = X.hf_hub_download_with_xet_fallback(
+    X.hf_hub_download_with_xet_fallback(
         "unsloth/repo", "model.gguf", token = None, force_download = True
     )
-    assert out2 == "/freshly/downloaded/path"
-    assert len(attempts) == 1
-    assert attempts[0]["force"] is True
+    assert seen == [False, True]  # the shim forwards force_download to the shared helper unchanged
 
 
 # ── multi-revision GGUF blob comparison and update reclaim ──

studio/backend/tests/test_native_template_trust_remote_code.py

@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Regression tests for trust_remote_code in the native-template fallback.
+
+``render_native_template`` re-fetches a model's native chat template from its
+repo when an Unsloth override template (mistral, gemma-4) dropped the tools
+schema. For a model loaded with ``trust_remote_code=True`` whose tokenizer repo
+carries custom code, the secondary ``AutoTokenizer.from_pretrained`` must re-use
+that same consent or transformers raises (it requires ``trust_remote_code`` to
+instantiate a custom tokenizer class), the ``except`` swallows it, and the
+request silently keeps the tool-dropping prompt even though the user already
+consented to remote code for the model load.
+
+These tests pin that the stored ``trust_remote_code`` is threaded to the reload,
+that the reload is skipped (returns ``None`` without executing code) when no
+consent is stored, and that both backend ``model_info`` dicts persist the flag at
+load time so the read lands on a value ``load_model`` actually set.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+
+import pytest
+
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+# ``chat_template_helpers`` is dependency-light (copy / logging / typing, with the
+# transformers import deferred inside the function). Load it directly so the test
+# runs without importing the heavy ``core.inference`` package (unsloth / torch).
+_HELPERS_PATH = Path(_BACKEND_DIR) / "core" / "inference" / "chat_template_helpers.py"
+_spec = importlib.util.spec_from_file_location("_native_tpl_trc_test", _HELPERS_PATH)
+chat_template_helpers = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(chat_template_helpers)
+
+render_native_template = chat_template_helpers.render_native_template
+
+
+# A native template that emits a tools section only when tools are provided, so the
+# with-tools vs no-tools render differs and ``render_native_template`` accepts it.
+_NATIVE_TEMPLATE = (
+    "{% for m in messages %}{{ m['role'] }}: {{ m['content'] }}\n{% endfor %}"
+    "{% if tools %}[AVAILABLE_TOOLS]{{ tools }}[/AVAILABLE_TOOLS]\n{% endif %}"
+    "{% if add_generation_prompt %}assistant:{% endif %}"
+)
+
+_MESSAGES = [{"role": "user", "content": "what is the weather"}]
+_TOOLS = [{"type": "function", "function": {"name": "get_weather"}}]
+
+
+class _JinjaTokenizer:
+    """Minimal tokenizer whose ``apply_chat_template`` renders ``self.chat_template``.
+
+    Stands in for the live model tokenizer that ``render_native_template`` shallow-
+    copies and re-points at the native template before rendering.
+    """
+
+    def __init__(self, chat_template):
+        self.chat_template = chat_template
+
+    def apply_chat_template(
+        self,
+        messages,
+        tokenize = False,
+        add_generation_prompt = True,
+        tools = None,
+        **kwargs,
+    ):
+        from jinja2 import BaseLoader, Environment
+        env = Environment(loader = BaseLoader())
+        return env.from_string(self.chat_template).render(
+            messages = messages,
+            tools = tools,
+            add_generation_prompt = add_generation_prompt,
+        )
+
+
+def _install_custom_code_tokenizer(monkeypatch):
+    """Patch ``AutoTokenizer.from_pretrained`` to mimic a custom-code repo: raise
+    unless ``trust_remote_code`` is truthy, else return a tokenizer carrying the
+    native template. Records the ``trust_remote_code`` it was called with."""
+    pytest.importorskip("jinja2")
+    from transformers import AutoTokenizer
+
+    calls = {}
+
+    def fake_from_pretrained(
+        model_id,
+        *args,
+        trust_remote_code = False,
+        token = None,
+        **kwargs,
+    ):
+        calls["trust_remote_code"] = trust_remote_code
+        calls["model_id"] = model_id
+        calls["token"] = token
+        if not trust_remote_code:
+            # Mirrors transformers.dynamic_module_utils.resolve_trust_remote_code:
+            # has_remote_code and not has_local_code and not trust_remote_code -> ValueError.
+            raise ValueError(
+                f"The repository {model_id} contains custom code which must be executed "
+                "to correctly load the model. Please pass the argument "
+                "`trust_remote_code=True` to allow custom code to be run."
+            )
+        return _JinjaTokenizer(_NATIVE_TEMPLATE)
+
+    monkeypatch.setattr(AutoTokenizer, "from_pretrained", staticmethod(fake_from_pretrained))
+    return calls
+
+
+def _model_info(trust_remote_code):
+    return {
+        "native_chat_template": None,  # force the repo reload path
+        "base_model": None,  # non-LoRA: template_source == active_model_name
+        "trust_remote_code": trust_remote_code,
+        # Live tokenizer that gets shallow-copied + re-pointed at the native template.
+        "tokenizer": _JinjaTokenizer("OVERRIDE-THAT-DROPS-TOOLS"),
+    }
+
+
+def test_native_reload_passes_stored_trust_remote_code(monkeypatch):
+    """With ``trust_remote_code`` stored on ``model_info`` the custom-code reload
+    succeeds and the tools-advertising native prompt is returned. This FAILS before
+    the fix (reload omits the flag, raises, is swallowed, returns None)."""
+    calls = _install_custom_code_tokenizer(monkeypatch)
+    model_info = _model_info(trust_remote_code = True)
+
+    out = render_native_template(
+        model_info = model_info,
+        active_model_name = "acme/custom-tokenizer-model",
+        messages = _MESSAGES,
+        tools = _TOOLS,
+    )
+
+    assert out is not None, "native fallback should render the tools prompt with consent"
+    assert "[AVAILABLE_TOOLS]" in out
+    assert "get_weather" in out
+    assert calls["trust_remote_code"] is True  # the stored consent was threaded through
+    # A successful fetch is cached so the next tool turn skips the reload.
+    assert model_info["native_chat_template"] == _NATIVE_TEMPLATE
+
+
+def test_native_reload_without_consent_returns_none(monkeypatch):
+    """Without stored consent the custom-code reload raises, is swallowed, and
+    ``render_native_template`` returns None (no unconsented code execution). Proves
+    the stored flag -- not a hard-coded True -- drives the reload."""
+    calls = _install_custom_code_tokenizer(monkeypatch)
+    model_info = _model_info(trust_remote_code = False)
+
+    out = render_native_template(
+        model_info = model_info,
+        active_model_name = "acme/custom-tokenizer-model",
+        messages = _MESSAGES,
+        tools = _TOOLS,
+    )
+
+    assert out is None
+    assert calls["trust_remote_code"] is False
+    # A failed fetch must not be cached as "no template" (would pin the tool drop).
+    assert model_info["native_chat_template"] is None
+
+
+def test_backend_model_info_persists_trust_remote_code():
+    """Both backends must store ``trust_remote_code`` on their per-model info dict so
+    ``render_native_template`` can source the consent value. Guards against the read
+    landing on a key ``load_model`` never sets (which would silently no-op the fix)."""
+    inf = (Path(_BACKEND_DIR) / "core" / "inference" / "inference.py").read_text()
+    mlx = (Path(_BACKEND_DIR) / "core" / "inference" / "mlx_inference.py").read_text()
+    assert '"trust_remote_code": trust_remote_code,' in inf
+    assert '"trust_remote_code": trust_remote_code,' in mlx

studio/backend/tests/test_responses_tool_passthrough.py

@@ -59,6 +59,7 @@ from models.inference import (
     ResponsesUsage,
 )
 from routes.inference import (
+    _ResponsesReasoningExtractor,
     _SameTaskStreamingResponse,
     _build_chat_request,
     _chat_tool_calls_to_responses_output,
@@ -795,6 +796,7 @@ class TestResponsesNonStreamingAdapter:
 
     def test_monitor_records_translated_visible_text(self, monkeypatch):
         import routes.inference as inf_mod
+        import routes.inference as inf_mod
 
         async def fake_chat_completions(chat_req, request):
             assert request.state.skip_api_monitor is True
@@ -1988,6 +1990,126 @@ class TestTranslatedMessagesValidate:
             ChatMessage(**m.model_dump(exclude_none = True))
 
 
+# reasoning_prefilled mode: Qwen3/GLM enable_thinking templates prefill an unclosed <think>, so
+# generation begins inside the think block and emits only the closing </think>; the extractor starts in reasoning.
+class TestReasoningPrefilledExtractor:
+    def test_prefilled_single_feed_splits_lone_close(self):
+        # T1: reasoning...</think>answer with a prefilled (unseen) open tag.
+        reasoning, visible = _extract_responses_reasoning(
+            "plan</think>answer",
+            parse_think_markers = True,
+            reasoning_prefilled = True,
+        )
+        assert reasoning == "plan"
+        assert visible == "answer"
+
+    def test_prefilled_never_closed_is_all_reasoning(self):
+        # T2: truncated mid-thought (no </think>) -> all reasoning (GGUF parity).
+        reasoning, visible = _extract_responses_reasoning(
+            "still thinking with no close",
+            parse_think_markers = True,
+            reasoning_prefilled = True,
+        )
+        assert reasoning == "still thinking with no close"
+        assert visible == ""
+
+    def test_prefilled_close_split_across_feeds(self):
+        # T3: </think> straddles two feed() calls; holdback resolves it.
+        ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True)
+        r1, v1 = ex.feed("plan</th")
+        r2, v2 = ex.feed("ink>ans")
+        fr, fv = ex.finish()
+        assert (r1 + r2 + fr) == "plan"
+        assert (v1 + v2 + fv) == "ans"
+
+    def test_prefilled_close_split_one_char_per_feed(self):
+        # T4: every char in its own feed still splits correctly.
+        ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True)
+        reasoning, visible = "", ""
+        for ch in "plan</think>x":
+            r, v = ex.feed(ch)
+            reasoning += r
+            visible += v
+        fr, fv = ex.finish()
+        assert (reasoning + fr) == "plan"
+        assert (visible + fv) == "x"
+
+    def test_prefilled_empty_generation(self):
+        # T5: nothing generated.
+        reasoning, visible = _extract_responses_reasoning(
+            "",
+            parse_think_markers = True,
+            reasoning_prefilled = True,
+        )
+        assert reasoning == ""
+        assert visible == ""
+
+    def test_prefilled_whitespace_after_close_is_visible(self):
+        # T6: Qwen commonly emits </think>\n\n before the answer.
+        reasoning, visible = _extract_responses_reasoning(
+            "plan</think>\n\nanswer",
+            parse_think_markers = True,
+            reasoning_prefilled = True,
+        )
+        assert reasoning == "plan"
+        assert visible == "\n\nanswer"
+
+    def test_prefilled_stray_open_tag_is_suppressed(self):
+        # T7: a re-emitted literal <think> inside prefilled reasoning is dropped,
+        # not leaked into the drawer (covers enable_thinking_effort full-tag output).
+        reasoning, visible = _extract_responses_reasoning(
+            "a<think>b</think>c",
+            parse_think_markers = True,
+            reasoning_prefilled = True,
+        )
+        assert reasoning == "ab"
+        assert visible == "c"
+        assert "<think>" not in reasoning
+
+    def test_prefilled_close_at_start_empty_reasoning(self):
+        # T8: model closed immediately (empty reasoning) then answered.
+        reasoning, visible = _extract_responses_reasoning(
+            "</think>hi",
+            parse_think_markers = True,
+            reasoning_prefilled = True,
+        )
+        assert reasoning == ""
+        assert visible == "hi"
+
+    def test_not_prefilled_lone_close_preserves_current_behavior(self):
+        # T9: GGUF-parity guard -- WITHOUT prefilled, a lone </think> keeps the
+        # pre-fix behavior (reasoning stays visible, tag dropped). Ensures GGUF and
+        # every existing caller are byte-identical.
+        reasoning, visible = _extract_responses_reasoning(
+            "reasoning</think>ans",
+            parse_think_markers = True,
+            reasoning_prefilled = False,
+        )
+        assert reasoning == ""
+        assert visible == "reasoningans"
+
+    def test_not_prefilled_full_pair_still_splits(self):
+        # T10: normal explicit <think>..</think> (GGUF / Harmony) unchanged.
+        reasoning, visible = _extract_responses_reasoning(
+            "<think>r</think>v",
+            parse_think_markers = True,
+            reasoning_prefilled = False,
+        )
+        assert reasoning == "r"
+        assert visible == "v"
+
+    def test_prefilled_ignored_when_markers_not_parsed(self):
+        # T11: a non-reasoning model (parse_think_markers False) still passes text
+        # straight through even if reasoning_prefilled were mistakenly set False.
+        reasoning, visible = _extract_responses_reasoning(
+            "just an answer",
+            parse_think_markers = False,
+            reasoning_prefilled = False,
+        )
+        assert reasoning == ""
+        assert visible == "just an answer"
+
+
 # =====================================================================
 # Streaming passthrough healing — text-form calls promoted in order
 # =====================================================================

studio/backend/tests/test_safetensors_capability_advertise.py

@@ -11,6 +11,8 @@ from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
+import pytest
+
 _backend_root = Path(__file__).resolve().parent.parent
 if str(_backend_root) not in sys.path:
     sys.path.insert(0, str(_backend_root))
@@ -127,9 +129,8 @@ def test_detect_safetensors_features_gptoss_disables_tools():
     assert flags["supports_tools"] is False
 
 
-# Llama-3 / Mistral advertise tools but emit <|python_tag|> / [TOOL_CALLS],
-# which our parser can't read. The route helper must not flip supports_tools=True
-# for them, else the UI enables a pill the agentic loop can't honour.
+# Llama-3 / Mistral / Gemma 4 tool-call formats are now parser-supported, so supports_tools=True
+# must hold for all of them; only templates matching none of the five known markers are suppressed.
 
 LLAMA3_TEMPLATE = """
 {%- if tools %}
@@ -161,27 +162,186 @@ MISTRAL_TEMPLATE = """
 {%- endfor %}
 """
 
+GEMMA4_TEMPLATE = """
+{%- if tools %}
+  {{- 'Tools available. Emit calls as ' }}
+  {{- '<|tool_call>call:NAME{key:<|"|>val<|"|>}<tool_call|>' }}
+  {%- for tool in tools %}
+    {{- tool | tojson }}
+  {%- endfor %}
+{%- endif %}
+"""
 
-def test_detect_safetensors_features_llama3_template_suppresses_tools():
-    """Llama-3 emits <|python_tag|>; safetensors loop cannot parse it."""
+
+def test_detect_safetensors_features_llama3_template_keeps_tools_on():
+    """Llama-3 emits <|python_tag|>; parser now supports it."""
     from routes.inference import _detect_safetensors_features
 
     backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
     flags = _detect_safetensors_features(backend, LLAMA3_TEMPLATE)
-    assert flags["supports_tools"] is False
+    assert flags["supports_tools"] is True
 
 
-def test_detect_safetensors_features_mistral_template_suppresses_tools():
-    """Mistral emits [TOOL_CALLS]; safetensors loop cannot parse it."""
+def test_detect_safetensors_features_mistral_template_keeps_tools_on():
+    """Mistral emits [TOOL_CALLS]; parser now supports it."""
     from routes.inference import _detect_safetensors_features
 
     backend = SimpleNamespace(active_model_name = "unsloth/mistral-7b-instruct-v0.3")
     flags = _detect_safetensors_features(backend, MISTRAL_TEMPLATE)
+    assert flags["supports_tools"] is True
+
+
+def test_detect_safetensors_features_gemma4_template_keeps_tools_on():
+    """Gemma 4 emits <|tool_call>; parser now supports it."""
+    from routes.inference import _detect_safetensors_features
+
+    backend = SimpleNamespace(active_model_name = "unsloth/gemma-4-E2B-it-UD-MLX-4bit")
+    flags = _detect_safetensors_features(backend, GEMMA4_TEMPLATE)
+    assert flags["supports_tools"] is True
+
+
+# DeepSeek V3 / V3.1 / R1 emit ``<｜tool▁calls▁begin｜>...`` blocks.
+# Note the full-width pipe (U+FF5C) and lower-1/8-block (U+2581).
+DEEPSEEK_TEMPLATE = """
+{%- if tools %}
+  {%- for tool in tools %}
+    {{- tool | tojson }}
+  {%- endfor %}
+{%- endif %}
+{%- for message in messages %}
+  {%- if message.role == 'assistant' and message.tool_calls %}
+    {%- for tc in message.tool_calls %}
+      {{- '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tc.function.name +
+          '<｜tool▁sep｜>' + tc.function.arguments + '<｜tool▁call▁end｜>' }}
+    {%- endfor %}
+  {%- endif %}
+{%- endfor %}
+"""
+
+
+def test_detect_safetensors_features_deepseek_template_keeps_tools_on():
+    """DeepSeek emits ``<｜tool▁calls▁begin｜>...``; parser now supports it."""
+    from routes.inference import _detect_safetensors_features
+
+    backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1")
+    flags = _detect_safetensors_features(backend, DEEPSEEK_TEMPLATE)
+    assert flags["supports_tools"] is True
+
+
+# GLM 4.5 / 4.6 / 4.7 emit ``<tool_call>NAME\n<arg_key>...<arg_value>...
+GLM_TEMPLATE = """
+{%- if tools %}
+  For each function call, output the function name and arguments within
+  the following XML format:
+  <tool_call>{function-name}
+  <arg_key>{arg-key}</arg_key>
+  <arg_value>{arg-value}</arg_value>
+  </tool_call>
+  {%- for tool in tools %}
+    {{- tool | tojson }}
+  {%- endfor %}
+{%- endif %}
+"""
+
+
+def test_detect_safetensors_features_glm_template_keeps_tools_on():
+    """GLM 4.x emits ``<tool_call>NAME\\n<arg_key>...``; parser handles it."""
+    from routes.inference import _detect_safetensors_features
+
+    backend = SimpleNamespace(active_model_name = "unsloth/GLM-4.6")
+    flags = _detect_safetensors_features(backend, GLM_TEMPLATE)
+    assert flags["supports_tools"] is True
+
+
+# Kimi K2 / Moonshot uses ``<|tool_calls_section_begin|>...`` blocks
+# with ``functions.NAME:IDX`` as the per-call id.
+KIMI_TEMPLATE = """
+{%- if tools %}
+  <|im_system|>tool_declare<|im_middle|>{{ tools | tojson }}<|im_end|>
+{%- endif %}
+{%- for message in messages %}
+  {%- if message.role == 'assistant' and message.tool_calls %}
+    <|tool_calls_section_begin|>
+    {%- for tc in message.tool_calls %}
+      <|tool_call_begin|>{{ tc.id }}<|tool_call_argument_begin|>{{ tc.function.arguments | tojson }}<|tool_call_end|>
+    {%- endfor %}
+    <|tool_calls_section_end|>
+  {%- endif %}
+{%- endfor %}
+"""
+
+
+def test_detect_safetensors_features_kimi_template_keeps_tools_on():
+    """Kimi K2 emits ``<|tool_calls_section_begin|>...``; parser handles it."""
+    from routes.inference import _detect_safetensors_features
+
+    backend = SimpleNamespace(active_model_name = "unsloth/Kimi-K2-Instruct")
+    flags = _detect_safetensors_features(backend, KIMI_TEMPLATE)
+    assert flags["supports_tools"] is True
+
+
+LLAMA3_2_BARE_JSON_TEMPLATE = """
+{%- if tools %}
+  {{- 'Given the following functions, respond with JSON for a function call.' }}
+  {{- 'Respond in the format {"name": function name, "parameters": dictionary}.' }}
+  {%- for tool in tools %}
+    {{- tool | tojson }}
+  {%- endfor %}
+{%- endif %}
+{%- for message in messages %}
+  {%- if 'tool_calls' in message %}
+    {{- '{"name": "' + message.tool_calls[0].function.name + '", '}}
+    {{- '"parameters": ' + (message.tool_calls[0].function.arguments | tojson) + '}' }}
+  {%- endif %}
+{%- endfor %}
+"""
+
+
+def test_detect_safetensors_features_llama3_2_bare_json_keeps_tools_on():
+    """Llama-3.2 bare JSON is supported, so the pill stays enabled."""
+    from routes.inference import _detect_safetensors_features
+
+    backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
+    flags = _detect_safetensors_features(backend, LLAMA3_2_BARE_JSON_TEMPLATE)
+    assert flags["supports_tools"] is True
+
+
+MINICPM5_ATTRIBUTE_TEMPLATE = """
+{%- if tools %}
+  {{- 'Available tools. Emit calls as ' }}
+  {{- '<function name="NAME"><parameter name="key">value</parameter></function>' }}
+  {%- for tool in tools %}
+    {{- tool | tojson }}
+  {%- endfor %}
+{%- endif %}
+"""
+
+
+def test_detect_safetensors_features_attribute_function_form_keeps_tools_on():
+    """The attribute form ``<function name="...">`` must be whitelisted or the pill is wrongly suppressed."""
+    from routes.inference import _detect_safetensors_features
+
+    backend = SimpleNamespace(active_model_name = "openbmb/MiniCPM-5")
+    flags = _detect_safetensors_features(backend, MINICPM5_ATTRIBUTE_TEMPLATE)
+    assert flags["supports_tools"] is True
+
+
+def test_detect_safetensors_features_unknown_format_suppresses_tools():
+    """Tools advertised with no known marker must be suppressed."""
+    from routes.inference import _detect_safetensors_features
+
+    tpl = (
+        "{%- if tools %}<|im_start|>system\n"
+        "Emit tool calls as JSON-RPC notifications inside the response."
+        "<|im_end|>{%- endif %}"
+    )
+    backend = SimpleNamespace(active_model_name = "custom/unknown-tool-format")
+    flags = _detect_safetensors_features(backend, tpl)
     assert flags["supports_tools"] is False
 
 
 def test_detect_safetensors_features_qwen_tool_call_keeps_tools_on():
-    """Sanity check: gate only suppresses non-Qwen formats."""
+    """Sanity check: Qwen <tool_call> marker still flips supports_tools."""
     from routes.inference import _detect_safetensors_features
 
     backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B")
@@ -454,3 +614,161 @@ def test_route_layer_emits_supports_tools_true_for_qwen3_safetensors():
     assert flags["supports_tools"] is True
     assert flags["supports_reasoning"] is True
     assert flags["supports_preserve_thinking"] is True
+
+
+@pytest.mark.parametrize(
+    "opener",
+    [
+        "<｜tool▁calls▁begin｜>",  # canonical
+        "<｜tool_calls_begin｜>",  # ASCII underscores
+        "<｜tool▁calls｜>",  # short form
+        "<｜tool calls begin｜>",  # spaces
+        "<｜tool\\_calls\\_begin｜>",  # escaped underscores
+    ],
+)
+def test_detect_safetensors_features_deepseek_opener_variants_keep_tools_on(opener):
+    # Every DeepSeek opener the parser accepts must keep supports_tools on; the route gate derives
+    # its markers from the parser's TOOL_XML_SIGNALS so it can no longer drift behind the parser ...
+    from routes.inference import _detect_safetensors_features
+
+    tpl = (
+        "{%- if tools %}tools{%- endif %}"
+        + opener
+        + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_time{}"
+        "<｜tool▁call▁end｜><｜tool▁calls▁end｜>"
+    )
+    backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1")
+    flags = _detect_safetensors_features(backend, tpl)
+    assert flags["supports_tools"] is True
+
+
+# Templates that advertise tools ({%- if tools %}) and prompt the bare-JSON
+# call form, but whose ``{"name":`` example is pretty-printed or JSON-escaped.
+_WHITESPACE_BARE_JSON_TEMPLATE = (
+    "{%- if tools %}\n"
+    "To call a tool, output JSON of the form:\n"
+    '{ "name" : "function_name", "parameters": { } }\n'
+    "{%- endif %}\n"
+    "{{ messages }}"
+)
+_ESCAPED_BARE_JSON_TEMPLATE = (
+    "{%- if tools %}\n"
+    'Respond with {\\"name\\": \\"fn\\", \\"parameters\\": {}}\n'
+    "{%- endif %}\n"
+    "{{ messages }}"
+)
+_TOOLS_ADVERTISED_NO_PARSEABLE_FORM = (
+    "{%- if tools %}\nYou may use the available tools.\n{%- endif %}\n{{ messages }}"
+)
+
+
+def test_detect_safetensors_features_keeps_tools_for_pretty_printed_bare_json():
+    # A pretty-printed bare-JSON example (``{ "name" :``) must keep supports_tools since the parser
+    # accepts that whitespace via raw_decode.
+    from routes.inference import _detect_safetensors_features
+
+    backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
+    flags = _detect_safetensors_features(backend, _WHITESPACE_BARE_JSON_TEMPLATE)
+    assert flags["supports_tools"] is True
+
+
+def test_detect_safetensors_features_keeps_tools_for_escaped_bare_json():
+    from routes.inference import _detect_safetensors_features
+
+    backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
+    flags = _detect_safetensors_features(backend, _ESCAPED_BARE_JSON_TEMPLATE)
+    assert flags["supports_tools"] is True
+
+
+def test_detect_safetensors_features_drops_tools_when_no_parseable_form():
+    # Negative control: tools advertised but no parser-recognised emission form at
+    # all -> the pill is still dropped (the gate is not now matching everything).
+    from routes.inference import _detect_safetensors_features
+
+    backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
+    flags = _detect_safetensors_features(backend, _TOOLS_ADVERTISED_NO_PARSEABLE_FORM)
+    assert flags["supports_tools"] is False
+
+
+def test_detect_safetensors_features_keeps_tools_for_function_alias_bare_json():
+    # A template documenting the parser-supported {"function":...} bare-JSON alias
+    # must keep supports_tools, mirroring the {"name":...} form.
+    from routes.inference import _detect_safetensors_features
+
+    tpl = (
+        "{%- if tools %}\n"
+        'Respond with {"function": "fn", "parameters": {}}\n'
+        "{%- endif %}\n"
+        "{{ messages }}"
+    )
+    backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct")
+    flags = _detect_safetensors_features(backend, tpl)
+    assert flags["supports_tools"] is True
+
+
+# _sf_reasoning_prefill_mode gates the prefilled-<think> extractor so safetensors/MLX reach
+# GGUF reasoning-block parity for enable_thinking models.
+class TestSafetensorsReasoningPrefillGate:
+    # A minimal Qwen3-style template with the standard <think>/</think> markers.
+    _QWEN_TPL = "{% if enable_thinking %}<think>{% endif %}...</think>..."
+    # gemma-style bespoke reasoning channel -- no standard markers.
+    _GEMMA_TPL = "{% if enable_thinking %}<|think|>{% endif %}<|channel>thought<channel|>"
+
+    def _features(self, **over):
+        base = {
+            "supports_reasoning": True,
+            "reasoning_always_on": False,
+            "reasoning_style": "enable_thinking",
+        }
+        base.update(over)
+        return base
+
+    def test_g1_enable_thinking_true(self):
+        # G1: Qwen3.5 template + explicit enable_thinking=True -> prefilled.
+        from routes.inference import _sf_reasoning_prefill_mode
+        assert _sf_reasoning_prefill_mode(self._features(), True, self._QWEN_TPL) is True
+
+    def test_g2_enable_thinking_none_defaults_on(self):
+        # G2: default request (None) -> prefilled (Qwen3/GLM templates default on).
+        from routes.inference import _sf_reasoning_prefill_mode
+        assert _sf_reasoning_prefill_mode(self._features(), None, self._QWEN_TPL) is True
+
+    def test_g3_enable_thinking_false(self):
+        # G3: thinking explicitly off -> not prefilled.
+        from routes.inference import _sf_reasoning_prefill_mode
+        assert _sf_reasoning_prefill_mode(self._features(), False, self._QWEN_TPL) is False
+
+    def test_g4_gpt_oss_reasoning_effort_excluded(self):
+        # G4: gpt-oss uses explicit tags via HarmonyTextStreamer -> normal mode.
+        from routes.inference import _sf_reasoning_prefill_mode
+        feats = self._features(reasoning_style = "reasoning_effort")
+        assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False
+
+    def test_g5_enable_thinking_effort_included(self):
+        # G5: GLM-style enable_thinking_effort also prefills.
+        from routes.inference import _sf_reasoning_prefill_mode
+        feats = self._features(reasoning_style = "enable_thinking_effort")
+        assert _sf_reasoning_prefill_mode(feats, None, self._QWEN_TPL) is True
+
+    def test_g6_non_reasoning_model(self):
+        # G6: no reasoning capability -> never prefilled.
+        from routes.inference import _sf_reasoning_prefill_mode
+        feats = self._features(supports_reasoning = False, reasoning_style = None)
+        assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False
+
+    def test_g7_reasoning_always_on(self):
+        # G7: hardcoded-<think> template -> prefilled regardless of the flag.
+        from routes.inference import _sf_reasoning_prefill_mode
+        feats = self._features(reasoning_always_on = True)
+        assert _sf_reasoning_prefill_mode(feats, False, self._QWEN_TPL) is True
+
+    def test_g8_gemma_bespoke_channel_excluded(self):
+        # G8: gemma's <|think|>/<|channel> format has no </think> -> NOT prefilled
+        # (would otherwise swallow the whole answer as reasoning). Regression guard.
+        from routes.inference import _sf_reasoning_prefill_mode
+        assert _sf_reasoning_prefill_mode(self._features(), True, self._GEMMA_TPL) is False
+
+    def test_g9_missing_template_not_prefilled(self):
+        # G9: no template available -> conservative (not prefilled).
+        from routes.inference import _sf_reasoning_prefill_mode
+        assert _sf_reasoning_prefill_mode(self._features(), True, None) is False

studio/backend/tests/test_safetensors_reasoning_stream.py

@@ -0,0 +1,186 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Safetensors/MLX reasoning-block parity with GGUF.
+
+enable_thinking templates prefill an unclosed ``<think>``, so the stream must split the leading
+text into ``reasoning_content`` deltas (per turn, monitor gets visible text only). Replays a copy
+of ``sf_tool_stream``'s reasoning loop from routes/inference.py against synthetic events.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+from routes.inference import (
+    _ResponsesReasoningExtractor,
+    _sf_reasoning_prefill_mode,
+    _strip_tool_xml_for_display,
+)
+
+
+def _replay_sf_reasoning_stream(events: list[dict], *, prefilled: bool) -> dict:
+    """Mirror sf_tool_stream's reasoning loop: diff cumulative snapshots, reset (flushing) on turn end."""
+    prev_text = ""
+    extractor = _ResponsesReasoningExtractor(
+        parse_think_markers = True, reasoning_prefilled = prefilled
+    )
+    reasoning_deltas: list[str] = []
+    visible_deltas: list[str] = []
+    monitor: list[str] = []
+    tool_starts: list[dict] = []
+    order: list[str] = []  # sequence of ("reasoning"|"visible"|"tool_start") events
+
+    def _flush():
+        fr, fv = extractor.finish()
+        if fr:
+            reasoning_deltas.append(fr)
+            order.append("reasoning")
+        if fv:
+            visible_deltas.append(fv)
+            monitor.append(fv)
+            order.append("visible")
+
+    for event in events:
+        etype = event["type"]
+        if etype == "status":
+            if not event["text"]:
+                _flush()
+                prev_text = ""
+                extractor = _ResponsesReasoningExtractor(
+                    parse_think_markers = True, reasoning_prefilled = prefilled
+                )
+            continue
+        if etype in ("tool_start", "tool_end"):
+            if etype == "tool_start":
+                _flush()
+                prev_text = ""
+                extractor = _ResponsesReasoningExtractor(
+                    parse_think_markers = True, reasoning_prefilled = prefilled
+                )
+                tool_starts.append(event)
+                order.append("tool_start")
+            continue
+        clean = _strip_tool_xml_for_display(event.get("text", ""), auto_heal_tool_calls = True)
+        new_text = clean[len(prev_text) :]
+        prev_text = clean
+        if not new_text:
+            continue
+        r, v = extractor.feed(new_text)
+        if r:
+            reasoning_deltas.append(r)
+            order.append("reasoning")
+        if v:
+            visible_deltas.append(v)
+            monitor.append(v)
+            order.append("visible")
+    _flush()
+    return {
+        "reasoning": "".join(reasoning_deltas),
+        "visible": "".join(visible_deltas),
+        "monitor": "".join(monitor),
+        "tool_starts": tool_starts,
+        "order": order,
+    }
+
+
+def test_s1_plain_stream_splits_prefilled_reasoning():
+    # S1: plain/MLX single turn -> reasoning delta + visible delta; monitor visible-only.
+    events = [
+        {"type": "content", "text": "Let me compute 17*23"},
+        {"type": "content", "text": "Let me compute 17*23 = 391</think>The answer is 391."},
+    ]
+    out = _replay_sf_reasoning_stream(events, prefilled = True)
+    assert out["reasoning"] == "Let me compute 17*23 = 391"
+    assert out["visible"] == "The answer is 391."
+    assert out["monitor"] == "The answer is 391."
+    assert "<think>" not in out["reasoning"] and "</think>" not in out["visible"]
+
+
+def test_s2_reasoning_flushed_before_tool_start():
+    # S2: reasoning streamed as reasoning_content, then flushed BEFORE tool_start.
+    events = [
+        {"type": "content", "text": "I should search"},
+        {"type": "content", "text": "I should search Sydney weather</think>"},
+        {"type": "tool_start", "tool_name": "web_search", "tool_call_id": "c0"},
+        {"type": "tool_end", "tool_name": "web_search", "tool_call_id": "c0"},
+        {"type": "status", "text": ""},
+        {"type": "content", "text": "Found it</think>Sydney is 21C today."},
+    ]
+    out = _replay_sf_reasoning_stream(events, prefilled = True)
+    # Both turns' reasoning surfaced, answer only from turn 2.
+    assert "I should search Sydney weather" in out["reasoning"]
+    assert "Found it" in out["reasoning"]
+    assert out["visible"] == "Sydney is 21C today."
+    assert out["monitor"] == "Sydney is 21C today."
+    # Ordering: the pre-tool reasoning is emitted before the tool_start.
+    assert out["order"].index("reasoning") < out["order"].index("tool_start")
+
+
+def test_s3_extractor_resets_each_turn():
+    # S3: multi-turn -> the two turns' reasoning are distinct (fresh extractor each).
+    events = [
+        {"type": "content", "text": "turn1 thoughts</think>partial"},
+        {"type": "status", "text": ""},
+        {"type": "content", "text": "turn2 thoughts</think>final answer"},
+    ]
+    out = _replay_sf_reasoning_stream(events, prefilled = True)
+    assert out["reasoning"] == "turn1 thoughtsturn2 thoughts"
+    assert out["visible"] == "partialfinal answer"
+
+
+def test_s4_harmony_full_tags_normal_mode():
+    # S4: gpt-oss / explicit-tag models use normal mode (prefilled=False).
+    events = [{"type": "content", "text": "<think>reasoning here</think>visible answer"}]
+    out = _replay_sf_reasoning_stream(events, prefilled = False)
+    assert out["reasoning"] == "reasoning here"
+    assert out["visible"] == "visible answer"
+
+
+def test_s5_thinking_off_no_reasoning_deltas():
+    # S5: thinking disabled -> not prefilled, no </think>, all content is visible.
+    events = [{"type": "content", "text": "Just the plain answer, no thinking."}]
+    out = _replay_sf_reasoning_stream(events, prefilled = False)
+    assert out["reasoning"] == ""
+    assert out["visible"] == "Just the plain answer, no thinking."
+    assert out["monitor"] == "Just the plain answer, no thinking."
+
+
+_THINK_TPL = "...{% if enable_thinking %}<think>{% endif %}...</think>..."
+
+
+def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort():
+    # GLM-5.2-style enable_thinking_effort: a request with reasoning_effort="none" (and
+    # enable_thinking omitted) disables thinking exactly like enable_thinking=False, so
+    # prefilled mode must be OFF. Otherwise the model emits no </think> and a plain
+    # answer is swallowed whole into reasoning_content, leaving the visible response
+    # empty (the exact bug: prefilled=True below eats the whole answer).
+    feats = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True}
+    assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "none") is False
+    # Thinking on (effort level or default) still prefills.
+    assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "high") is True
+    assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, None) is True
+    # An explicit enable_thinking=False also disables (unchanged).
+    assert _sf_reasoning_prefill_mode(feats, False, _THINK_TPL, "high") is False
+    # reasoning_always_on wins regardless of reasoning_effort.
+    always = {**feats, "reasoning_always_on": True}
+    assert _sf_reasoning_prefill_mode(always, None, _THINK_TPL, "none") is True
+    # Plain enable_thinking models (Qwen) have no "none" sentinel; unaffected.
+    plain = {"reasoning_style": "enable_thinking", "supports_reasoning": True}
+    assert _sf_reasoning_prefill_mode(plain, None, _THINK_TPL, "none") is True
+
+    # End-to-end: with the corrected prefilled=False, a plain no-</think> answer is
+    # emitted as visible content rather than swallowed into the thinking drawer.
+    events = [{"type": "content", "text": "The capital of France is Paris."}]
+    out = _replay_sf_reasoning_stream(events, prefilled = False)
+    assert out["visible"] == "The capital of France is Paris."
+    assert out["reasoning"] == ""
+    # The buggy prefilled=True path is what swallowed the whole answer (guard the delta).
+    swallowed = _replay_sf_reasoning_stream(events, prefilled = True)
+    assert swallowed["visible"] == ""
+    assert swallowed["reasoning"] == "The capital of France is Paris."

studio/backend/tests/test_safetensors_toolcall_wiring.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Deterministic backend-wiring test for the safetensors / MLX tool-calling path.
+
+The parser and the cumulative-text state machine are already covered exhaustively by
+``test_safetensors_tool_loop.py`` with fake generators. What that suite does not touch is the
+*backend's own tool-injection seam*: both ``InferenceBackend`` (transformers) and
+``MLXInferenceBackend`` render the prompt through the shared
+``apply_chat_template_for_generation(..., tools=...)`` helper and stream cumulative text into the
+shared ``run_safetensors_tool_loop`` (see ``core/inference/inference.py`` and
+``core/inference/mlx_inference.py`` -- both call the same helper and the same loop, so a single CPU
+test of that seam covers the macOS MLX path too).
+
+This test drives that exact seam with deterministic fakes -- a fake tokenizer that records the
+``tools`` it is handed, a canned tool-call generation, and a stub executor -- and asserts the full
+agentic chain end to end:
+
+    tools injected into the template -> loop parses the call -> tool dispatched once ->
+    tool result fed back -> generation re-entered -> final answer streamed.
+
+It is the deterministic, download-free stand-in for the real-model MLX / GGUF browser tool-calling
+end-to-end: it imports no torch / unsloth / mlx, so it runs in the portable Backend CI alongside the
+tool-call parser tests. Follow-up to the parser test PRs (#5620 / #5704).
+"""
+
+from core.inference.chat_template_helpers import apply_chat_template_for_generation
+from core.inference.safetensors_agentic import run_safetensors_tool_loop
+
+TOOL_NAME = "get_weather"
+TOOL_ARGS = {"city": "Paris"}
+FAKE_TOOL = {
+    "type": "function",
+    "function": {
+        "name": TOOL_NAME,
+        "description": "Get the current weather for a city.",
+        "parameters": {
+            "type": "object",
+            "properties": {"city": {"type": "string"}},
+            "required": ["city"],
+        },
+    },
+}
+# Full parser matrix lives in test_safetensors_tool_loop.py.
+TOOL_CALL_TEXT = '<tool_call>{"name": "get_weather", "arguments": {"city": "Paris"}}</tool_call>'
+FINAL_ANSWER = "The weather in Paris is sunny and 22C."
+TOOL_RESULT = "Paris: sunny, 22C"
+
+
+class RecordingTokenizer:
+    """Fake tokenizer that records the ``tools`` handed to ``apply_chat_template``.
+
+    Modelled on ``TestChatTemplateHelper._Tok`` in ``test_safetensors_tool_loop.py``: it accepts the
+    real helper's kwargs and returns a canned prompt, so the test can assert the backend seam actually
+    forwarded the tool schema -- a silent drop on a chat-template fallback would leave ``tools_seen``
+    holding ``None``.
+    """
+
+    def __init__(self):
+        self.tools_seen: list = []
+        self.call_count = 0
+
+    def apply_chat_template(
+        self,
+        messages,
+        *,
+        tokenize = False,
+        add_generation_prompt = True,
+        **kwargs,
+    ):
+        self.call_count += 1
+        self.tools_seen.append(kwargs.get("tools"))
+        return "PROMPT"
+
+
+class StubExecutor:
+    """Stand-in for ``core.inference.tools.execute_tool``: records calls, returns a fixed result.
+
+    A fake tool name plus this stub means no real python / terminal / web / RAG side effect can run.
+    """
+
+    def __init__(self, result: str):
+        self.result = result
+        self.calls: list[tuple[str, dict]] = []
+
+    def __call__(
+        self,
+        name,
+        arguments,
+        *,
+        cancel_event = None,
+        timeout = None,
+        session_id = None,
+        rag_scope = None,
+        disable_sandbox = False,
+    ):
+        self.calls.append((name, arguments))
+        return self.result
+
+
+def _collect(generator, max_events = 200):
+    events = []
+    for ev in generator:
+        events.append(ev)
+        if len(events) >= max_events:
+            break
+    return events
+
+
+def _tool_names(tools):
+    return [(t.get("function") or {}).get("name") for t in (tools or [])]
+
+
+def test_backend_seam_injects_tools_and_drives_full_tool_loop():
+    """The shared backend seam forwards tools into the chat template, and the loop parses the call,
+    dispatches it once, feeds the result back, and re-enters generation for the final answer."""
+    tok = RecordingTokenizer()
+    executor = StubExecutor(TOOL_RESULT)
+    turns = iter([TOOL_CALL_TEXT, FINAL_ANSWER])
+    active_tools_seen: list = []
+    conversations_seen: list = []
+
+    def single_turn(conversation, *, active_tools = None):
+        # Mirror the real _single_turn: render via the shared helper, then yield cumulative snapshots.
+        active_tools_seen.append(active_tools)
+        conversations_seen.append([dict(m) for m in conversation])
+        apply_chat_template_for_generation(tok, conversation, tools = active_tools)
+        text = next(turns)
+        mid = len(text) // 2
+        acc = ""
+        for chunk in (text[:mid], text[mid:]):
+            acc += chunk
+            yield acc
+
+    events = _collect(
+        run_safetensors_tool_loop(
+            single_turn = single_turn,
+            messages = [{"role": "user", "content": "What is the weather in Paris?"}],
+            tools = [FAKE_TOOL],
+            execute_tool = executor,
+            max_tool_iterations = 3,
+        )
+    )
+
+    # 1. Helper forwarded the tool schema to the tokenizer (seam does not drop tools).
+    assert tok.tools_seen, "tokenizer.apply_chat_template was never called"
+    assert tok.tools_seen[0], "tool schema was dropped before reaching the tokenizer"
+    assert TOOL_NAME in _tool_names(tok.tools_seen[0])
+
+    # 2. Loop offered the tool to the first generation turn.
+    assert active_tools_seen and active_tools_seen[0] is not None
+    assert TOOL_NAME in _tool_names(active_tools_seen[0])
+
+    # 3 / 4 / 5. Exactly one tool_start, one dispatch with parsed args, one tool_end with the result.
+    tool_starts = [e for e in events if e["type"] == "tool_start"]
+    tool_ends = [e for e in events if e["type"] == "tool_end"]
+    assert len(tool_starts) == 1 and tool_starts[0]["tool_name"] == TOOL_NAME
+    assert executor.calls == [(TOOL_NAME, TOOL_ARGS)], executor.calls
+    assert len(tool_ends) == 1 and tool_ends[0]["result"] == TOOL_RESULT
+
+    # 6. Final answer streams after the tool result: loop appended it and re-entered generation.
+    contents = [e for e in events if e["type"] == "content"]
+    assert contents and FINAL_ANSWER in contents[-1]["text"]
+    last_tool_end_idx = max(i for i, e in enumerate(events) if e["type"] == "tool_end")
+    last_content_idx = max(i for i, e in enumerate(events) if e["type"] == "content")
+    assert last_content_idx > last_tool_end_idx, "final answer must stream after the tool result"
+
+    # 6b. Tool result fed back into the conversation before the final turn (6 alone misses this:
+    #     the fake generation ignores the conversation).
+    assert len(conversations_seen) >= 2, "loop did not re-enter generation after the tool call"
+    final_turn_convo = conversations_seen[1]
+    assert any(
+        TOOL_RESULT in str(m.get("content", "")) for m in final_turn_convo
+    ), "tool result was not fed back into the conversation before the final generation turn"
+
+    # 7. Guard: raw tool-call markup never leaked to the client as content.
+    for e in contents:
+        assert "<tool_call>" not in e["text"]
+        assert TOOL_NAME not in e["text"]

studio/backend/tests/test_tool_strip_guard.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+"""strip_tool_patterns must match the plain per-pattern loop while skipping the
+quadratic no-match rescan of a closed-pair sweep whose close token is absent."""
+
+import random
+import sys
+import time
+from pathlib import Path
+
+_BACKEND_ROOT = Path(__file__).resolve().parents[1]
+if str(_BACKEND_ROOT) not in sys.path:
+    sys.path.insert(0, str(_BACKEND_ROOT))
+
+from core.tool_healing import (
+    _TOOL_ALL_PATS,
+    _TOOL_CLOSED_PATS,
+    strip_tool_call_markup,
+    strip_tool_patterns,
+)
+
+
+def _naive(text, patterns):
+    for pat in patterns:
+        text = pat.sub("", text)
+    return text
+
+
+_TOKENS = [
+    "<tool_call>",
+    "</tool_call>",
+    "<|tool_call>",
+    "<tool_call|>",
+    "<function=x>",
+    "<function=mcp__s__a-b>",
+    "</function>",
+    "<parameter=p>",
+    "</parameter>",
+    "call:fn{",
+    "}",
+    "{",
+    '<|"|>',
+    "A",
+    " ",
+    "\n",
+    "id",
+    "x:1",
+    "</tool",
+    "call>",
+]
+
+
+def test_guard_matches_plain_loop_on_fuzz():
+    rng = random.Random(1234)
+    for patterns in (_TOOL_ALL_PATS, _TOOL_CLOSED_PATS):
+        for _ in range(20000):
+            s = "".join(rng.choice(_TOKENS) for _ in range(rng.randint(0, 10)))
+            assert strip_tool_patterns(s, patterns) == _naive(s, patterns), (s, patterns)
+
+
+def test_strip_markup_representative_cases_unchanged():
+    assert strip_tool_call_markup("a <tool_call>{}</tool_call> b") == "a  b"
+    assert strip_tool_call_markup("a <function=x><parameter=p>1</parameter></function> b") == "a  b"
+    # Non-final keeps an unclosed block; final strips it to EOF.
+    assert strip_tool_call_markup("a <tool_call>{partial") == "a <tool_call>{partial"
+    assert strip_tool_call_markup("a <tool_call>{partial", final = True) == "a"
+
+
+def test_no_quadratic_blowup_on_unclosed_markers():
+    # Unguarded, this took minutes.
+    big = "<tool_call>" * 20000 + "<function=x>" * 20000
+    t0 = time.perf_counter()
+    out = strip_tool_call_markup(big, final = True)
+    assert time.perf_counter() - t0 < 2.0
+    assert out == ""

studio/backend/tests/test_tool_xml_strip.py

@@ -24,18 +24,67 @@ import re as _re
 _src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text()
 _m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL)
 assert _m, "could not extract _TOOL_XML_RE source"
-_ns = {"_re": _re}
+# The lazy ``(.*?)\n\)`` could grab a shorter expression if an arm is ever wrapped;
+# pin the DeepSeek + bare-Kimi arms so a silent truncation fails loudly here.
+assert "_DS_OPEN_SRC" in _m.group(1) and "tool_call_begin" in _m.group(
+    1
+), "extracted _TOOL_XML_RE is missing expected arms (extraction truncated?)"
+# The regex reuses the parser's shared DeepSeek opener alternation; provide it so the extracted
+# ``_re.compile`` expression resolves the same source.
+from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC
+from core.inference.tool_call_parser import (
+    _strip_function_xml_calls,
+    _strip_gemma_wrapperless_calls,
+    _strip_glm_calls,
+    _strip_mistral_closed_calls,
+)
+
+from typing import Optional as _Optional
+
+_ns = {
+    "_re": _re,
+    "_DS_OPEN_SRC": _DS_OPEN_SRC,
+    "Optional": _Optional,
+    "_strip_mistral_closed_calls": _strip_mistral_closed_calls,
+    "_strip_gemma_wrapperless_calls": _strip_gemma_wrapperless_calls,
+    "_strip_glm_calls": _strip_glm_calls,
+    "_strip_function_xml_calls": _strip_function_xml_calls,
+}
 exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns)
 _TOOL_XML_RE = _ns["_TOOL_XML_RE"]
+
+# Signatures may span multiple lines and now carry the enabled_tool_names gate; match
+# the whole (possibly multi-line) signature up to ``-> str:`` then the indented body.
+_xml_helper = _re.search(
+    r"def _strip_tool_xml\((?:.|\n)*?\) -> str:\n(?:    .+\n)+",
+    _src,
+)
+assert _xml_helper, "could not extract _strip_tool_xml source"
+assert "_strip_mistral_closed_calls" in _xml_helper.group(
+    0
+), "extracted _strip_tool_xml no longer runs the Mistral balanced strip"
+exec(_xml_helper.group(0), _ns)
+_strip_tool_xml = _ns["_strip_tool_xml"]
+
 _helper = _re.search(
-    r"def _strip_tool_xml_for_display\(text: str, \*, auto_heal_tool_calls: bool\) -> str:\n"
-    r"(?:    .+\n)+",
+    r"def _strip_tool_xml_for_display\((?:.|\n)*?\) -> str:\n(?:    .+\n)+",
     _src,
 )
 assert _helper, "could not extract _strip_tool_xml_for_display source"
+# After the V1 fix the display helper delegates to _strip_tool_xml; confirm the
+# extracted body actually reached that call rather than truncating early.
+assert "_strip_tool_xml(" in _helper.group(0), "display helper no longer delegates"
 exec(_helper.group(0), _ns)
 _strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"]
 
+_gate_src = _re.search(
+    r"def _gemma_strip_gate\((?:.|\n)*?\) -> set:\n(?:    .+\n)+",
+    _src,
+)
+assert _gate_src, "could not extract _gemma_strip_gate source"
+exec(_gate_src.group(0), _ns)
+_gemma_strip_gate = _ns["_gemma_strip_gate"]
+
 
 # ── Well-formed pairs ─────────────────────────────────────────────
 
@@ -46,6 +95,16 @@ def test_route_display_strip_respects_disabled_auto_heal_contract():
     assert "<tool_call>" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
 
 
+def test_route_display_strip_removes_mistral_tool_calls_with_nested_json():
+    # _TOOL_XML_RE has no [TOOL_CALLS] arm, so the helper delegates to _strip_tool_xml for the Mistral
+    # balanced-brace strip (a non-greedy \{.*?\} would truncate nested JSON).
+    text = 'ok [TOOL_CALLS]web_search{"filters":{"date":"2024"},"query":"cats"} tail'
+    assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text
+    out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
+    assert "[TOOL_CALLS]" not in out and "web_search" not in out, out
+    assert out == "ok  tail"
+
+
 def test_strips_well_formed_tool_call():
     text = (
         "Let me search.\n"
@@ -73,6 +132,26 @@ def test_strips_function_only_well_formed():
     assert "Done." in cleaned
 
 
+def test_strips_function_attribute_form():
+    # Attribute form ``<function name="...">`` (MiniCPM-5 / MiniMax-M2) must strip from the route too
+    # (it previously leaked into the UI); a dotted/hyphenated name also strips.
+    text = (
+        'Sure.\n<function name="get_weather">\n'
+        "<parameter=city>\nSydney\n</parameter>\n</function>\nDone."
+    )
+    cleaned = _TOOL_XML_RE.sub("", text)
+    assert "<function name=" not in cleaned
+    assert "</function>" not in cleaned
+    assert "Sure." in cleaned and "Done." in cleaned
+
+    dotted = 'A <function name="srv.list-issues">x</function> B'
+    assert _TOOL_XML_RE.sub("", dotted) == "A  B"
+
+    # Auto-Heal-disabled display contract still preserves literal markup.
+    assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text
+    assert "<function name=" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
+
+
 # ── Orphan openings ───────────────────────────────────────────────
 
 
@@ -281,3 +360,241 @@ def test_no_catastrophic_backtracking_on_orphan_opening_spam():
     elapsed = time.perf_counter() - t0
     assert elapsed < 0.1, f"regex took {elapsed*1000:.0f}ms on 1000x orphan opens"
     assert "<tool_call>" not in cleaned
+
+
+# ── DeepSeek opener variants + bare Kimi (parse/strip symmetry) ──
+
+
+def test_strips_deepseek_space_opener_variant():
+    # The space-separated opener is parsed by the parser, so the display strip
+    # must remove it too (the shared opener alternation is reused here).
+    text = (
+        "pre <｜tool calls begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>"
+        '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post'
+    )
+    cleaned = _TOOL_XML_RE.sub("", text)
+    assert "tool" not in cleaned.replace("post", "").replace("pre", "")
+    assert cleaned == "pre  post"
+
+
+def test_strips_deepseek_escaped_underscore_opener_variant():
+    text = (
+        "pre <｜tool\\_calls\\_begin｜><｜tool▁call▁begin｜>get_y<｜tool▁sep｜>"
+        '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post'
+    )
+    cleaned = _TOOL_XML_RE.sub("", text)
+    assert cleaned == "pre  post"
+
+
+def test_strips_bare_kimi_call_without_section_wrapper():
+    # Kimi can emit a bare <|tool_call_begin|>...<|tool_call_end|> with no
+    # section wrapper; the parser accepts it, so the strip must cover it.
+    text = (
+        "pre <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>"
+        '{"a":1}<|tool_call_end|> post'
+    )
+    cleaned = _TOOL_XML_RE.sub("", text)
+    assert "tool_call_begin" not in cleaned
+    assert cleaned == "pre  post"
+
+
+@pytest.mark.parametrize(
+    "text",
+    [
+        # Prose that merely names a Kimi/DeepSeek marker (no real call follows) must
+        # survive: the call-shaped lookahead fires only on a real call or a bare EOF
+        # fragment, so an answer discussing the protocol is never truncated.
+        "See <|tool_call_begin|> in the docs. More prose after it.",
+        "The <|tool_calls_section_begin|> marker opens a batch. Read on.",
+        "DeepSeek uses <｜tool▁calls▁begin｜> to start a call block, then continues.",
+    ],
+)
+def test_deepseek_kimi_false_alarm_prose_is_kept(text):
+    # Regression for the route arm truncating a prose answer that references a marker
+    # without a following call (parser _TOOL_ALL_PATS already had this lookahead).
+    assert _TOOL_XML_RE.sub("", text) == text
+
+
+def test_deepseek_kimi_real_calls_still_strip_after_false_alarm_fix():
+    # The lookahead must not weaken real-call stripping: closed, truncated, and bare
+    # EOF-fragment forms all still get removed.
+    closed = (
+        "answer <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>"
+        '{"a":1}<|tool_call_end|> tail'
+    )
+    assert _TOOL_XML_RE.sub("", closed) == "answer  tail"
+    eof_fragment = "prefix <|tool_call_begin|>"
+    assert _TOOL_XML_RE.sub("", eof_fragment) == "prefix "
+    deepseek = (
+        "reply <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>"
+        '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜>'
+    )
+    assert _TOOL_XML_RE.sub("", deepseek) == "reply "
+
+
+# ── Llama-3 <|python_tag|> arm bounds on REAL sentinels only ──────
+
+
+# Llama-3 <|python_tag|> arm bounds on REAL sentinels only
+def test_python_tag_strip_consumes_literal_sentinel_in_arg():
+    # A <|python_tag|> tool call whose JSON argument carries a literal <|...|>
+    # token (here <|cite|>) must be stripped whole. The old `<(?!\|)` arm stopped
+    # at any `<|`, leaking the call tail (e.g. `<|cite|> here"}}`) into display.
+    text = '<|python_tag|>{"name": "send", "parameters": {"text": "use <|cite|> here"}}'
+    cleaned = _TOOL_XML_RE.sub("", text)
+    assert cleaned == "", f"python_tag call leaked at literal sentinel: {cleaned!r}"
+
+
+@pytest.mark.parametrize(
+    "sentinel",
+    [
+        "<|eot_id|>",
+        "<|eom_id|>",
+        "<|start_header_id|>",
+        "<|end_header_id|>",
+    ],
+)
+def test_python_tag_strip_stops_at_real_sentinel(sentinel):
+    # A genuine Llama control sentinel still bounds the strip so following
+    # assistant text is preserved (the arm must not swallow past it).
+    text = f'<|python_tag|>{{"name": "x", "parameters": {{}}}}{sentinel}visible answer'
+    cleaned = _TOOL_XML_RE.sub("", text)
+    assert (
+        cleaned == f"{sentinel}visible answer"
+    ), f"strip did not stop at real sentinel {sentinel!r}: {cleaned!r}"
+
+
+def test_python_tag_strip_restarts_on_second_python_tag():
+    # A second <|python_tag|> opens a new tool-call region, so the whole pair is
+    # stripped (the arm bounds the first, then the next match consumes the rest).
+    text = '<|python_tag|>{"name": "a"}<|python_tag|>{"name": "b"}'
+    cleaned = _TOOL_XML_RE.sub("", text)
+    assert cleaned == "", f"second python_tag region leaked: {cleaned!r}"
+
+
+def test_glm_call_with_literal_close_tag_in_arg_value_is_stripped_whole():
+    # GLM 4.x emits <tool_call>NAME<arg_key>k</arg_key><arg_value>v</arg_value> ...</tool_call>.
+    text = (
+        "<tool_call>web_search\n<arg_key>query</arg_key>\n"
+        "<arg_value>find </tool_call> here</arg_value>\n</tool_call> done"
+    )
+    out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
+    assert "</arg_value>" not in out
+    assert "<arg_key>" not in out
+    assert out.strip() == "done"
+
+
+def test_glm_normal_and_qwen_calls_still_stripped_by_route():
+    # Regression: a normal GLM call (no literal close tag) and a Qwen
+    # <tool_call>{json}</tool_call> are still stripped; trailing prose is kept.
+    glm = "<tool_call>get_time\n<arg_key>tz</arg_key>\n<arg_value>UTC</arg_value>\n</tool_call> ok"
+    assert _strip_tool_xml_for_display(glm, auto_heal_tool_calls = True).strip() == "ok"
+    qwen = '<tool_call>{"name":"web_search","arguments":{"q":"x"}}</tool_call> after'
+    assert _strip_tool_xml_for_display(qwen, auto_heal_tool_calls = True).strip() == "after"
+
+
+def test_route_strip_removes_param_alias_close_tag():
+    # The parser accepts the <param name="...">...</param> attribute-form alias of
+    # <parameter=...>; the route tail cleanup must strip an orphan </param> close too.
+    assert _strip_tool_xml_for_display("answer </param>", auto_heal_tool_calls = True) == "answer "
+    assert (
+        _strip_tool_xml_for_display("answer </parameter>", auto_heal_tool_calls = True) == "answer "
+    )
+
+
+def test_route_strip_uses_guarded_function_scan_for_literal_nested_markup():
+    # A literal <function=...></function> in a value must not truncate the strip: the route runs the
+    # parser's guarded function-XML scan before the regex, matching the core strip.
+    text = "<function=python><parameter=code><function=evil></function></parameter></function> tail"
+    assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "tail"
+
+
+def test_route_strip_gates_wrapperless_gemma_by_enabled_tools():
+    # The route strip must gate the markerless Gemma call:NAME{...} form on the enabled tool names,
+    # like the parser/loop, so a disabled/example name in prose is preserved in ...
+    prose = "To document syntax you write call:foo{query:example}. That shows the format."
+    assert "call:foo{query:example}" in _strip_tool_xml(prose, {"web_search"})
+    # An enabled name is still a real call and stripped.
+    assert "call:web_search" not in _strip_tool_xml(
+        "Answer. call:web_search{query:x}", {"web_search"}
+    )
+    # No gate (legacy) strips every closed call.
+    assert "call:foo" not in _strip_tool_xml(prose)
+
+
+def test_gemma_strip_gate_empty_tools_preserves_prose():
+    # With NO tools enabled the gate must return an EMPTY set (strip nothing), not None: None falls
+    # back to strip-all and deletes an answer that documents the call:NAME{...} syntax.
+    assert _gemma_strip_gate([]) == set()
+    assert _gemma_strip_gate(None) == set()
+    assert _gemma_strip_gate([{"function": {"name": "web_search"}}]) == {"web_search"}
+    prose = "To document syntax you write call:foo{query:example}. That shows the format."
+    assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate([]))
+    assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate(None))
+    # An enabled tool's real call is still stripped.
+    assert "call:web_search" not in _strip_tool_xml(
+        "Answer. call:web_search{query:x}",
+        _gemma_strip_gate([{"function": {"name": "web_search"}}]),
+    )
+
+
+def test_strip_keeps_prose_after_closed_function_call_with_literal_close():
+    # The call ends at its first non-data close: prose after it survives the
+    # strip even when it mentions a literal </function>.
+    from core.inference.tool_call_parser import strip_tool_markup
+    text = (
+        "<function=web_search><parameter=query>cats</parameter></function>"
+        " Done. The tag </function> closes a call."
+    )
+    assert strip_tool_markup(text, final = True) == "Done. The tag </function> closes a call."
+
+
+def test_final_strip_keeps_prose_mentioning_bare_markers():
+    # A false-alarm marker in a normal answer must not lose everything after
+    # it; only text that looks like that family's call start drops.
+    from core.inference.tool_call_parser import strip_tool_markup
+    for text in (
+        "See [TOOL_CALLS] docs for details. More prose after.",
+        "<|python_tag|> is the Llama marker. Explanation continues.",
+        "The <|tool_call> opener wraps Gemma calls.",
+    ):
+        assert strip_tool_markup(text, final = True) == text
+    # A bare marker at end-of-text is a fragment and still drops.
+    assert strip_tool_markup("Answer text [TOOL_CALLS]", final = True) == "Answer text"
+
+
+def test_final_strip_still_drops_truncated_marker_calls():
+    from core.inference.tool_call_parser import strip_tool_markup
+    for text in (
+        '[TOOL_CALLS][{"name":"web_search","argu',
+        '[TOOL_CALLS]web_search[ARGS]{"q":"x',
+        '<|python_tag|>{"name":"web_search","par',
+        '<|python_tag|>foo.call(items=["a',
+        "<|tool_call>call:web_search{query:tru",
+    ):
+        assert strip_tool_markup(text, final = True) == ""
+
+
+def test_chained_bare_json_strip_consumes_all_calls():
+    # The loops keep this text as next-turn history: a leftover executed call
+    # would be replayed alongside the structured tool_calls.
+    from core.inference.tool_call_parser import strip_leading_bare_json_call
+
+    enabled = {"web_search", "python"}
+    chained = (
+        '{"name":"web_search","parameters":{"q":"first"}};'
+        '{"name":"python","parameters":{"code":"x"}}'
+    )
+    assert strip_leading_bare_json_call(chained, enabled_tool_names = enabled) == ""
+    assert (
+        strip_leading_bare_json_call(chained + " trailing prose", enabled_tool_names = enabled)
+        == "trailing prose"
+    )
+    # The chain stops at a non-call answer object, which stays visible.
+    call_then_answer = (
+        '{"name":"web_search","parameters":{"q":"x"}};{"name":"web_search","result":"data"}'
+    )
+    assert (
+        strip_leading_bare_json_call(call_then_answer, enabled_tool_names = enabled)
+        == '{"name":"web_search","result":"data"}'
+    )

studio/backend/utils/hf_xet_fallback.py

@@ -1,341 +1,204 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""Xet-primary HF downloads with an automatic HTTP fallback on a no-progress stall.
+"""Studio shim over the shared ``unsloth_zoo.hf_xet_fallback`` Xet -> HTTP stall fallback.
 
-Xet (``hf_xet``) is the fast default but can hang with no progress and no
-exception, and a blocked native thread cannot be killed. Keep Xet primary; fall
-back to plain HTTP only when the parent observes a stall. ``HF_HUB_DISABLE_XET``
-is read at import time, so the fallback runs in a fresh ``spawn`` child (not a
-thread) that sets the env before importing ``huggingface_hub``. Cached files
-short-circuit with no child; deterministic errors (401/403/404/disk-full) and
-cancellation propagate without a fallback. Mirrors the safetensors inference
-recovery in core/inference/{orchestrator,worker}.py.
+Re-exports the shared API and injects Studio's marker-aware cache purge
+(``prepare_cache_for_transport``) so the download manager keeps its ``.transport``
+marker semantics on the HTTP retry.
 """
 
 from __future__ import annotations
 
-import multiprocessing as mp
-import os
-import queue
-import signal
-import sys
 import threading
-import time
 from typing import Any, Callable, Optional
 
-from loggers import get_logger
+_shared_import_error = None
+try:
+    import unsloth_zoo.hf_xet_fallback as _shared
+    _shared_available = True
+except Exception as _exc:  # noqa: BLE001 - any import failure must degrade, not crash
+    # unsloth_zoo's __init__ runs torch/GPU detection, which raises on a torch-less/GPU-less Studio
+    # host. The download helper needs none of it, so retry via the light UNSLOTH_ZOO_DISABLE_GPU_INIT
+    # path before giving up.
+    _shared_import_error = _exc
+    import os as _os
 
-logger = get_logger(__name__)
-
-_CTX = mp.get_context("spawn")
-
-# Defaults match the existing inference watchdog and hub shutdown deadline.
-DEFAULT_HEARTBEAT_INTERVAL = 30.0
-DEFAULT_STALL_TIMEOUT = 180.0
-DEFAULT_GRACE_PERIOD = 10.0
-_POLL_INTERVAL = 0.5
-
-
-class DownloadStallError(RuntimeError):
-    """Raised when no download progress is observed for too long.
-
-    Canonical home; orchestrator.py re-imports it so all paths share one type.
-    """
-
-
-def child_should_disable_xet(config: dict) -> bool:
-    """Single source of truth for the per-worker Xet env flip."""
-    return bool(config.get("disable_xet"))
-
-
-def get_hf_download_state(
-    repo_ids: Optional[list[str]] = None, *, repo_type: str = "model"
-) -> Optional[tuple[int, bool]]:
-    """Return ``(total_on_disk_bytes, has_incomplete)`` for the active HF cache.
-
-    Sparse-aware (st_blocks based) so a sparse Xet/``hf_transfer`` ``.incomplete``
-    is not mistaken for full-size progress. ``None`` means the state could not be
-    measured, so callers skip stall logic for that tick.
-    """
+    _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT")
+    _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1"
     try:
-        from hub.utils.hf_cache_state import (
-            blob_bytes_present,
-            has_active_incomplete_blobs,
-            hf_cache_root,
-            iter_active_repo_cache_dirs,
-        )
+        import unsloth_zoo.hf_xet_fallback as _shared
+        _shared_available = True
+        _shared_import_error = None
+    except Exception as _exc2:  # noqa: BLE001 - degrade so Studio still boots with plain HF downloads
+        _shared_import_error = _exc2
+        _shared_available = False
+    finally:
+        if _prev_gpu_init is None:
+            _os.environ.pop("UNSLOTH_ZOO_DISABLE_GPU_INIT", None)
+        else:
+            _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = _prev_gpu_init
 
-        if hf_cache_root() is None:
-            return (0, False)
+if _shared_available:
+    # Bind by assignment so each public name shares one module-level binding with the degraded branch.
+    DEFAULT_GRACE_PERIOD = _shared.DEFAULT_GRACE_PERIOD
+    DEFAULT_HEARTBEAT_INTERVAL = _shared.DEFAULT_HEARTBEAT_INTERVAL
+    DEFAULT_STALL_TIMEOUT = _shared.DEFAULT_STALL_TIMEOUT
+    DownloadStallError = _shared.DownloadStallError
+    child_should_disable_xet = _shared.child_should_disable_xet
+    get_hf_download_state = _shared.get_hf_download_state
+    start_watchdog = _shared.start_watchdog
+    _shared_hf_hub_download_with_xet_fallback = _shared.hf_hub_download_with_xet_fallback
+    _shared_snapshot_download_with_xet_fallback = _shared.snapshot_download_with_xet_fallback
+else:
+    # Degrade instead of crashing Studio: plain HF downloads, stall watchdog disabled. Thin stubs,
+    # not a second copy of the orchestration; recovery returns once unsloth_zoo is upgraded.
+    import logging as _logging
 
-        total = 0
-        has_incomplete = False
-        for repo_id in repo_ids or []:
-            # Skip local paths: HF IDs never start with / . ~ or contain "\".
-            if not repo_id or repo_id.startswith(("/", ".", "~")) or "\\" in repo_id:
-                continue
-            for entry in iter_active_repo_cache_dirs(repo_type, repo_id):
-                blobs_dir = entry / "blobs"
-                if not blobs_dir.is_dir():
-                    continue
-                for blob in blobs_dir.iterdir():
-                    try:
-                        if blob.is_file():
-                            total += blob_bytes_present(blob)
-                    except OSError:
-                        pass
-            if has_active_incomplete_blobs(repo_type, repo_id):
-                has_incomplete = True
-        return (total, has_incomplete)
-    except Exception as e:
-        logger.debug("Failed to determine HF download state: %s", e)
-        return None
+    _logging.getLogger(__name__).warning(
+        "unsloth_zoo.hf_xet_fallback unavailable (%s); the Xet stall watchdog is "
+        "disabled. Install/upgrade unsloth_zoo (and its torch dependency) to "
+        "re-enable automatic Xet -> HTTP download recovery.",
+        _shared_import_error,
+    )
 
+    DEFAULT_HEARTBEAT_INTERVAL = 30.0
+    DEFAULT_STALL_TIMEOUT = 180.0
+    DEFAULT_GRACE_PERIOD = 10.0
 
-def start_watchdog(
-    *,
-    repo_ids: list[str],
-    on_stall: Callable[[str], None],
-    repo_type: str = "model",
-    interval: float = DEFAULT_HEARTBEAT_INTERVAL,
-    stall_timeout: float = DEFAULT_STALL_TIMEOUT,
-    xet_disabled: bool = False,
-    on_heartbeat: Optional[Callable[[str], None]] = None,
-) -> threading.Event:
-    """Start a daemon thread that fires ``on_stall(message)`` exactly once iff a
-    ``*.incomplete`` is present AND the on-disk size is unchanged for
-    *stall_timeout* seconds. The timer resets while no ``*.incomplete`` exists, so
-    post-download init is never misread as a stall. Returns a stop event the
-    caller sets when the download phase ends.
-    """
-    stop = threading.Event()
-    transport = "https" if xet_disabled else "xet"
-    fired = False
+    class DownloadStallError(RuntimeError):
+        """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode."""
 
-    def _beat() -> None:
-        nonlocal fired
-        state = get_hf_download_state(repo_ids, repo_type = repo_type)
-        last_size = state[0] if state is not None else 0
-        last_change = time.monotonic()
+    def child_should_disable_xet(config: dict) -> bool:
+        return bool(config.get("disable_xet"))
 
-        while not stop.wait(interval):
-            state = get_hf_download_state(repo_ids, repo_type = repo_type)
-            now = time.monotonic()
+    def get_hf_download_state(*args: Any, **kwargs: Any) -> None:
+        return None  # unmeasurable -> the (absent) watchdog never fires
 
-            if state is None:
-                if on_heartbeat is not None:
+    def start_watchdog(
+        *,
+        on_heartbeat: "Optional[Callable[[str], None]]" = None,
+        interval: float = DEFAULT_HEARTBEAT_INTERVAL,
+        xet_disabled: bool = False,
+        **kwargs: Any,
+    ) -> "threading.Event":
+        # No stall detection, but keep emitting heartbeats so the orchestrator's inactivity deadline
+        # is not tripped during a long download.
+        stop = threading.Event()
+        if on_heartbeat is None:
+            return stop
+        transport = "https" if xet_disabled else "xet"
+
+        def _beat() -> None:
+            while not stop.wait(interval):
+                try:
                     on_heartbeat(f"Downloading ({transport} transport)...")
-                continue
+                except Exception:
+                    pass
 
-            current_size, has_incomplete = state
-            if current_size != last_size:
-                last_size = current_size
-                last_change = now
+        threading.Thread(
+            target = _beat,
+            daemon = True,
+            name = "hf-xet-degraded-heartbeat",
+        ).start()
+        return stop
 
-            # Reset unless .incomplete confirms an active download, so model init
-            # and lock waits are not counted as a stall.
-            if not has_incomplete:
-                last_change = now
-            elif now - last_change >= stall_timeout:
-                if not fired:
-                    fired = True
-                    on_stall(
-                        f"Download appears stalled ({transport} transport) "
-                        f"-- no progress for {int(now - last_change)}s"
-                    )
-                return
+    def _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool:
+        return cancel_event is not None and cancel_event.is_set()
 
-            if on_heartbeat is not None:
-                on_heartbeat(f"Downloading ({transport} transport)...")
+    def _shared_hf_hub_download_with_xet_fallback(
+        repo_id: str,
+        filename: str,
+        token: Optional[str],
+        *,
+        repo_type: str = "model",
+        revision: Optional[str] = None,
+        cache_dir: Optional[str] = None,
+        force_download: bool = False,
+        cancel_event: "Optional[threading.Event]" = None,
+        **_ignored: Any,
+    ) -> str:
+        # Keep the cancellation contract: do not start or return a download once cancelled.
+        if _degraded_cancelled(cancel_event):
+            raise RuntimeError("Cancelled")
 
-    threading.Thread(target = _beat, daemon = True, name = "hf-xet-watchdog").start()
-    return stop
-
-
-def _download_child_entry(
-    *,
-    repo_id: str,
-    filename: str,
-    token: Optional[str],
-    repo_type: str,
-    disable_xet: bool,
-    result_queue: Any,
-    force_download: bool = False,
-) -> None:
-    """Spawn-child entrypoint: download one file and report the result.
-
-    Top-level and picklable. Sets the Xet env BEFORE importing huggingface_hub,
-    forms its own process group so the parent can kill the whole transfer, and
-    never logs the token or signed URLs.
-    """
-    # Die with Studio on Linux (this mp child gets no parent-set preexec_fn).
-    try:
-        from utils.process_lifetime import bind_current_process_to_parent_lifetime
-        bind_current_process_to_parent_lifetime()
-    except Exception:
-        pass
-
-    if hasattr(os, "setsid"):
-        try:
-            os.setsid()
-        except OSError:
-            pass
-
-    if disable_xet:
-        os.environ["HF_HUB_DISABLE_XET"] = "1"
-        # Keep the HTTP writer sequential and resumable (hf_transfer leaves sparse
-        # partials a sequential resume cannot safely continue).
-        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
-    os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
-
-    # Test-only fault injection (never set in production): stall the Xet attempt
-    # so the watchdog + HTTP fallback can be exercised against a real repo.
-    if not disable_xet and os.environ.get("UNSLOTH_HF_XET_FORCE_STALL") == "1":
-        import time as _t
-        try:
-            from huggingface_hub.constants import HF_HUB_CACHE
-
-            blobs = os.path.join(HF_HUB_CACHE, "models--" + repo_id.replace("/", "--"), "blobs")
-            os.makedirs(blobs, exist_ok = True)
-            with open(os.path.join(blobs, "xet-force-stall.incomplete"), "wb") as fh:
-                fh.write(b"\0" * 4096)
-        except OSError:
-            pass
-        while True:
-            _t.sleep(3600)
-
-    try:
         from huggingface_hub import hf_hub_download
+
         path = hf_hub_download(
             repo_id = repo_id,
             filename = filename,
-            repo_type = repo_type,
             token = token,
+            repo_type = repo_type,
+            revision = revision,
+            cache_dir = cache_dir,
             force_download = force_download,
         )
-        result_queue.put({"ok": True, "path": path})
-    except BaseException as e:  # noqa: BLE001 - report every failure to the parent
-        error = f"{type(e).__name__}: {e}"
-        try:
-            from hub.utils.download_registry import scrub_secrets
-            error = scrub_secrets(error, hf_token = token)
-        except Exception:
-            pass
-        result_queue.put({"ok": False, "error": error})
+        if _degraded_cancelled(cancel_event):
+            raise RuntimeError("Cancelled")
+        return path
 
+    def _shared_snapshot_download_with_xet_fallback(
+        repo_id: str,
+        *,
+        revision: Optional[str] = None,
+        token: Optional[str] = None,
+        repo_type: str = "model",
+        cache_dir: Optional[str] = None,
+        allow_patterns: Optional[Any] = None,
+        ignore_patterns: Optional[Any] = None,
+        force_download: bool = False,
+        cancel_event: "Optional[threading.Event]" = None,
+        **_ignored: Any,
+    ) -> str:
+        if _degraded_cancelled(cancel_event):
+            raise RuntimeError("Cancelled")
 
-def _terminate_process_group(proc: "mp.process.BaseProcess", grace_period: float) -> None:
-    """Kill *proc* and its whole process group (Xet may spawn helper procs).
+        from huggingface_hub import snapshot_download
 
-    The child calls ``os.setsid()`` so its pgid equals its pid; signal via
-    ``os.killpg(pid, ...)`` -- NOT ``getpgid``, which before the child becomes a
-    group leader resolves to OUR group. SIGTERM, then SIGKILL after *grace_period*.
-    """
-    pid = proc.pid
-
-    def _signal_group(sig: int) -> None:
-        if pid is not None and hasattr(os, "killpg"):
-            try:
-                os.killpg(pid, sig)
-                return
-            except (ProcessLookupError, PermissionError, OSError):
-                pass
-        # Windows or pre-setsid: best effort on the single process.
-        try:
-            proc.terminate() if sig != getattr(signal, "SIGKILL", -9) else proc.kill()
-        except Exception:
-            pass
-
-    _signal_group(getattr(signal, "SIGTERM", signal.SIGINT))
-    proc.join(timeout = grace_period)
-    if proc.is_alive():
-        _signal_group(getattr(signal, "SIGKILL", signal.SIGTERM))
-        proc.join(timeout = 5.0)
-
-
-def _run_download_attempt(
-    repo_id: str,
-    filename: str,
-    token: Optional[str],
-    *,
-    repo_type: str,
-    disable_xet: bool,
-    cancel_event: Optional[threading.Event],
-    stall_timeout: float,
-    interval: float,
-    grace_period: float,
-    on_status: Optional[Callable[[str], None]],
-    force_download: bool = False,
-) -> tuple[str, Optional[str]]:
-    """Run one download in a spawn child supervised by the no-progress watchdog.
-
-    Returns ``("ok", path)``, ``("stall", None)``, ``("cancelled", None)``, or
-    ``("error", message)``. This is the seam tests monkeypatch to avoid spawning.
-    """
-    result_queue: Any = _CTX.Queue()
-    proc = _CTX.Process(
-        target = _download_child_entry,
-        kwargs = dict(
+        path = snapshot_download(
             repo_id = repo_id,
-            filename = filename,
-            token = token,
             repo_type = repo_type,
-            disable_xet = disable_xet,
-            result_queue = result_queue,
+            revision = revision,
+            token = token,
+            cache_dir = cache_dir,
+            allow_patterns = allow_patterns,
+            ignore_patterns = ignore_patterns,
             force_download = force_download,
-        ),
-        daemon = True,
-    )
-    proc.start()
-    from utils.process_lifetime import adopt_pid
-
-    adopt_pid(proc.pid)  # bind to parent lifetime (Windows job / sweep)
-
-    stalled = threading.Event()
-    stop_watchdog = start_watchdog(
-        repo_ids = [repo_id],
-        on_stall = lambda msg: stalled.set(),
-        repo_type = repo_type,
-        interval = interval,
-        stall_timeout = stall_timeout,
-        xet_disabled = disable_xet,
-        on_heartbeat = on_status,
-    )
-
-    result: Optional[dict] = None
-    try:
-        while proc.is_alive():
-            if cancel_event is not None and cancel_event.is_set():
-                _terminate_process_group(proc, grace_period)
-                return ("cancelled", None)
-            if stalled.is_set():
-                _terminate_process_group(proc, grace_period)
-                return ("stall", None)
-            try:
-                result = result_queue.get(timeout = _POLL_INTERVAL)
-                break
-            except queue.Empty:
-                continue
-        else:
-            # Process exited; drain any result it enqueued.
-            try:
-                result = result_queue.get_nowait()
-            except queue.Empty:
-                result = None
-    finally:
-        stop_watchdog.set()
-        proc.join(timeout = grace_period)
-
-    if result is None:
-        return (
-            "error",
-            f"download process for '{repo_id}/{filename}' exited "
-            f"(code={proc.exitcode}) without a result",
         )
-    if result.get("ok"):
-        return ("ok", result["path"])
-    return ("error", result.get("error") or "unknown download error")
+        if _degraded_cancelled(cancel_event):
+            raise RuntimeError("Cancelled")
+        return path
+
+
+__all__ = [
+    "DEFAULT_GRACE_PERIOD",
+    "DEFAULT_HEARTBEAT_INTERVAL",
+    "DEFAULT_STALL_TIMEOUT",
+    "DownloadStallError",
+    "child_should_disable_xet",
+    "get_hf_download_state",
+    "start_watchdog",
+    "hf_hub_download_with_xet_fallback",
+    "snapshot_download_with_xet_fallback",
+]
+
+
+def _studio_prepare_for_http(repo_type: str, repo_id: str) -> None:
+    """Studio's marker-aware purge before an HTTP resume, keeping the download manager's ``.transport``
+    accounting consistent (vs unsloth_zoo's generic default). Guarded: a purge failure is logged,
+    not fatal to the retry."""
+    try:
+        from hub.utils.download_registry import prepare_cache_for_transport
+        prepare_cache_for_transport(repo_type, repo_id, "http")
+    except Exception as exc:
+        try:
+            from loggers import get_logger
+            get_logger(__name__).debug(
+                "Studio prepare_cache_for_transport failed for %s: %s", repo_id, exc
+            )
+        except ModuleNotFoundError as logger_exc:
+            if logger_exc.name != "loggers":
+                raise
 
 
 def hf_hub_download_with_xet_fallback(
@@ -345,83 +208,32 @@ def hf_hub_download_with_xet_fallback(
     *,
     cancel_event: Optional[threading.Event] = None,
     repo_type: str = "model",
+    revision: Optional[str] = None,
     stall_timeout: float = DEFAULT_STALL_TIMEOUT,
     interval: float = DEFAULT_HEARTBEAT_INTERVAL,
     grace_period: float = DEFAULT_GRACE_PERIOD,
     on_status: Optional[Callable[[str], None]] = None,
     force_download: bool = False,
 ) -> str:
-    """Download a single file with Xet primary and HTTP as a stall-only fallback.
+    """Single-file download via the shared fallback with Studio's marker-aware HTTP-retry prep.
+    ``force_download`` re-fetches a newer blob over a cached one (Studio's model-update path)."""
+    return _shared_hf_hub_download_with_xet_fallback(
+        repo_id,
+        filename,
+        token,
+        cancel_event = cancel_event,
+        repo_type = repo_type,
+        revision = revision,
+        stall_timeout = stall_timeout,
+        interval = interval,
+        grace_period = grace_period,
+        on_status = on_status,
+        force_download = force_download,
+        prepare_for_http_fn = _studio_prepare_for_http,
+    )
 
-    Returns the local cache path. Raises ``RuntimeError("Cancelled")`` if
-    *cancel_event* is set, re-raises a deterministic child error unchanged (no
-    fallback), and raises ``DownloadStallError`` only if BOTH transports stall.
 
-    When *force_download* is True the cache-first early-return is skipped and the
-    flag is threaded to ``hf_hub_download`` so a newer remote blob is re-fetched
-    even if an older blob is already cached.
-    """
-    # Finalized blob already cached: return it with no child and no network.
-    # Skipped when force_download is set so an update re-fetches a newer blob.
-    if not force_download:
-        try:
-            from huggingface_hub import try_to_load_from_cache
-            cached = try_to_load_from_cache(repo_id, filename, repo_type = repo_type)
-            if isinstance(cached, str) and os.path.exists(cached):
-                return cached
-        except Exception as e:
-            logger.debug("Cached probe failed for %s/%s: %s", repo_id, filename, e)
-
-    if cancel_event is not None and cancel_event.is_set():
-        raise RuntimeError("Cancelled")
-
-    disable_xet = False
-    for attempt in range(2):
-        if disable_xet:
-            # Purge a non-HTTP partial before resuming over HTTP: an HTTP resume
-            # over a sparse Xet/hf_transfer partial silently corrupts the blob.
-            try:
-                from hub.utils.download_registry import prepare_cache_for_transport
-                prepare_cache_for_transport(repo_type, repo_id, "http")
-            except Exception as e:
-                logger.debug("prepare_cache_for_transport failed for %s: %s", repo_id, e)
-
-        kind, payload = _run_download_attempt(
-            repo_id,
-            filename,
-            token,
-            repo_type = repo_type,
-            disable_xet = disable_xet,
-            cancel_event = cancel_event,
-            stall_timeout = stall_timeout,
-            interval = interval,
-            grace_period = grace_period,
-            on_status = on_status,
-            force_download = force_download,
-        )
-
-        if kind == "ok":
-            return payload  # type: ignore[return-value]
-        if kind == "cancelled":
-            raise RuntimeError("Cancelled")
-        if kind == "error":
-            # Deterministic failure: the other transport would fail identically.
-            raise RuntimeError(payload)
-        # kind == "stall"
-        if attempt == 0 and not disable_xet:
-            logger.warning(
-                "Download stalled for '%s/%s' -- retrying with HF_HUB_DISABLE_XET=1",
-                repo_id,
-                filename,
-            )
-            if on_status is not None:
-                on_status(f"{repo_id}/{filename}: Xet stalled, retrying over HTTP")
-            disable_xet = True
-            continue
-        raise DownloadStallError(
-            f"Download stalled for '{repo_id}/{filename}' even with "
-            f"HF_HUB_DISABLE_XET=1 -- check your network connection"
-        )
-
-    # Unreachable: the loop either returns or raises on each attempt.
-    raise DownloadStallError(f"Download failed for '{repo_id}/{filename}'")
+def snapshot_download_with_xet_fallback(repo_id: str, **kwargs: Any) -> str:
+    """Whole-repo download via the shared fallback with Studio's marker-aware HTTP-retry prep."""
+    kwargs.setdefault("prepare_for_http_fn", _studio_prepare_for_http)
+    return _shared_snapshot_download_with_xet_fallback(repo_id, **kwargs)

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

@@ -42,10 +42,8 @@ export function ModelUpdateAction({
 }: ModelUpdateActionProps) {
   const [open, setOpen] = useState(false);
 
-  // The update is a managed download (it surfaces in the global Downloads panel
-  // with progress + cancel). When this exact repo+variant finishes, refresh the
-  // caller so the "update available" cue clears once the new revision is on
-  // disk. A ref keeps the subscription stable across renders without resubscribing.
+  // Refresh the caller when this repo+variant's download finishes so the "update available" cue
+  // clears. A ref keeps the subscription stable across renders.
   const onUpdatedRef = useRef(onUpdated);
   onUpdatedRef.current = onUpdated;
   useEffect(() => {
@@ -60,9 +58,8 @@ export function ModelUpdateAction({
   }, [repoId, variant]);
 
   const handleConfirm = useCallback(() => {
-    // Start the background re-download and close the dialog immediately; the
-    // Downloads panel owns progress + cancel from here. Only a failure to START
-    // surfaces a toast — a failed download reports itself in the panel.
+    // Start the re-download and close the dialog; the Downloads panel owns progress + cancel.
+    // Only a failure to START toasts (a failed download shows in the panel).
     void Promise.resolve()
       .then(onConfirm)
       .catch((err) => {

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

@@ -1505,11 +1505,8 @@ export function HubModelPicker({
   catalog?: CatalogGroup[];
 }) {
   const gpu = useGpuInfo();
-  // The currently-loaded/running model id. We read params.checkpoint from the
-  // runtime store (backend-mirrored from /api/inference/status.active_model, see
-  // chat-runtime-store) rather than the dropdown `isSelected` highlight (which is
-  // just `value === repo_id` and can reflect a staged, not-yet-loaded pick). Used
-  // to disable the cached-row update action for the model that's live in memory.
+  // Live model id from the runtime store (backend-mirrored active_model), not the dropdown
+  // highlight which can be a staged pick. Disables the update action for it.
   const loadedModelId = useChatRuntimeStore((s) => s.params.checkpoint);
   // Last-loaded timestamps power the "Recent" sort (vs "Downloaded" = file date).
   const loadTimes = useModelLoadTimes(value);
@@ -1849,11 +1846,8 @@ export function HubModelPicker({
     refreshLocalModelsList();
   }, [hfToken, refreshLocalModelsList]);
 
-  // Updates run as MANAGED downloads (they show in the global Downloads panel
-  // with manifest-based progress + a working Cancel), instead of a blocking
-  // call. The worker re-resolves `main` and pulls only changed blobs, so the
-  // cached copy stays usable until the new revision lands. The row's
-  // ModelUpdateAction refreshes the list when this repo+variant completes.
+  // Updates run as managed downloads (Downloads panel: progress + Cancel), not a blocking
+  // call. The worker pulls only changed blobs, so the cached copy stays usable until done.
   const startManagedUpdate = useCallback((repoId: string, variant: string, expectedBytes: number) => {
     return downloadManager
       .requestStart({

studio/frontend/src/components/floating-monitor.tsx

@@ -27,9 +27,11 @@ function usageTextClass(percent: number): string {
   return "text-primary";
 }
 
-function formatGb(value: number): string {
+function formatGiB(value: number): string {
+  // RAM/VRAM come from the backend in binary units (bytes / 1024**3), matching
+  // nvidia-smi and PyTorch, so label the readout GiB rather than GB.
   const digits = value >= 10 ? 1 : 2;
-  return `${value.toFixed(digits)} GB`;
+  return `${value.toFixed(digits)} GiB`;
 }
 
 export function FloatingMonitor() {
@@ -116,7 +118,7 @@ export function FloatingMonitor() {
               </span>
             </div>
             <div className="text-xs text-muted-foreground font-mono tabular-nums">
-              {formatGb(ramUsed)} / {formatGb(ramTotal)}
+              {formatGiB(ramUsed)} / {formatGiB(ramTotal)}
             </div>
             <Progress
               value={ramPercent}
@@ -144,7 +146,7 @@ export function FloatingMonitor() {
                 </span>
               </div>
               <div className="text-xs text-muted-foreground font-mono tabular-nums">
-                {formatGb(vramUsed)} / {formatGb(vramTotal)}
+                {formatGiB(vramUsed)} / {formatGiB(vramTotal)}
               </div>
               <Progress
                 value={vramPercent}

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

@@ -1085,11 +1085,11 @@ export function ModelsPage() {
   const { vramInfo, minMemory } = useHubModelVram(selectedModel, gpu);
 
   const gpuLabel = gpu.available
-    ? `${Math.round(gpu.memoryTotalGb)} GB`
+    ? `${Math.round(gpu.memoryTotalGb)} GiB`
     : "Unavailable";
   const ramLabel =
     gpu.systemRamTotalGb > 0
-      ? `${Math.round(gpu.systemRamTotalGb)} GB`
+      ? `${Math.round(gpu.systemRamTotalGb)} GiB`
       : "Unavailable";
   const coreLabel =
     gpu.cpuCore > 0 && gpu.cpuThread > 0

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

@@ -125,7 +125,7 @@ export function SummaryStep() {
               <span className="text-xs text-muted-foreground">GPU</span>
               <div className="flex items-center gap-2">
                 <span className="text-sm font-medium">{hw.gpuName ?? "---"}</span>
-                <Badge variant="secondary">{hw.vramTotalGb != null ? `${hw.vramTotalGb} GB` : "---"}</Badge>
+                <Badge variant="secondary">{hw.vramTotalGb != null ? `${hw.vramTotalGb} GiB` : "---"}</Badge>
               </div>
             </div>
           </div>

studio/frontend/src/features/settings/tabs/about-tab.tsx

@@ -158,7 +158,7 @@ export function AboutTab() {
               <code className="font-mono text-xs text-muted-foreground">
                 {gpu.name ?? "—"}
                 {gpu.vramTotalGb != null
-                  ? ` · ${Math.round(gpu.vramTotalGb)} GB`
+                  ? ` · ${Math.round(gpu.vramTotalGb)} GiB`
                   : ""}
               </code>
             </SettingsRow>

studio/frontend/src/features/settings/tabs/resources-tab.tsx

@@ -47,6 +47,15 @@ function formatGb(value: number | null | undefined): string {
   return `${safe.toFixed(digits)} GB`;
 }
 
+// RAM/VRAM come from the backend in binary units (bytes / 1024**3), matching
+// nvidia-smi and PyTorch, so label those readouts GiB. Disk stays on formatGb
+// because the backend reports disk in decimal GB (bytes / 1e9).
+function formatGiB(value: number | null | undefined): string {
+  const safe = isFiniteNumber(value) ? Math.max(0, value) : 0;
+  const digits = safe >= 10 ? 1 : 2;
+  return `${safe.toFixed(digits)} GiB`;
+}
+
 function formatMb(value: number | null | undefined): string {
   const safe = isFiniteNumber(value) ? Math.max(0, value) : 0;
   return `${Math.round(safe).toLocaleString()} MB`;
@@ -300,9 +309,9 @@ export function ResourcesTab() {
           />
           <MetricTile
             label={t("settings.resources.liveMonitor.ram")}
-            value={`${formatGb(metrics.ramUsed)} / ${formatGb(metrics.ramTotal)}`}
+            value={`${formatGiB(metrics.ramUsed)} / ${formatGiB(metrics.ramTotal)}`}
             detail={t("settings.resources.liveMonitor.free", {
-              value: formatGb(systemInfo.memory?.available_gb),
+              value: formatGiB(systemInfo.memory?.available_gb),
             })}
             percent={systemInfo.memory?.percent_used ?? 0}
           />
@@ -318,13 +327,13 @@ export function ResourcesTab() {
             label={t("settings.resources.liveMonitor.vram")}
             value={
               hasGpu
-                ? `${formatGb(metrics.vramUsed)} / ${formatGb(metrics.vramTotal)}`
+                ? `${formatGiB(metrics.vramUsed)} / ${formatGiB(metrics.vramTotal)}`
                 : t("settings.resources.liveMonitor.noGpu")
             }
             detail={
               hasGpu
                 ? t("settings.resources.liveMonitor.free", {
-                    value: formatGb(metrics.vramFree),
+                    value: formatGiB(metrics.vramFree),
                   })
                 : backendLabel
             }
@@ -373,17 +382,17 @@ export function ResourcesTab() {
                 <div className="grid gap-1 text-xs text-muted-foreground sm:grid-cols-3 sm:gap-2">
                   <span className="min-w-0 truncate font-mono tabular-nums">
                     {t("settings.resources.gpu.used", {
-                      value: formatGb(used),
+                      value: formatGiB(used),
                     })}
                   </span>
                   <span className="min-w-0 truncate font-mono tabular-nums sm:text-center">
                     {t("settings.resources.gpu.free", {
-                      value: formatGb(free),
+                      value: formatGiB(free),
                     })}
                   </span>
                   <span className="min-w-0 truncate font-mono tabular-nums sm:text-right">
                     {t("settings.resources.gpu.total", {
-                      value: formatGb(total),
+                      value: formatGiB(total),
                     })}
                   </span>
                 </div>

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

@@ -411,7 +411,7 @@ function LiveGpuPanel({
                   value={index}
                   className="bg-popover text-popover-foreground dark:bg-zinc-900 dark:text-zinc-100"
                 >
-                  GPU {device.visible_ordinal ?? index} - {device.backend} ({device.vram_total_gb ? `${Math.round(device.vram_total_gb)}GB` : "N/A"})
+                  GPU {device.visible_ordinal ?? index} - {device.backend} ({device.vram_total_gb ? `${Math.round(device.vram_total_gb)}GiB` : "N/A"})
                 </option>
               ))}
             </select>
@@ -446,7 +446,7 @@ function LiveGpuPanel({
           icon={<HugeiconsIcon icon={RamMemoryIcon} className="size-3.5" />}
           value={
             currentGpu.vram_used_gb != null && currentGpu.vram_total_gb != null
-              ? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GB`
+              ? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GiB`
               : "--"
           }
           pct={currentGpu.vram_utilization_pct ?? 0}

tests/saving/test_quant_method_none_normalization.py

@@ -0,0 +1,81 @@
+"""CPU-only regression for the quant-method normalization loops in save.py.
+
+`unsloth_save_pretrained_gguf` and `save_to_gguf_generic` each normalize the
+`quantization_method` list, mapping a ``None`` element to ``"q8_0"``. The mapping
+used to call ``quant_method.lower()`` as the first statement of the loop, so a
+``None`` element (e.g. ``quantization_method=[None]`` or ``["q4_k_m", None]``)
+raised ``AttributeError: 'NoneType' object has no attribute 'lower'`` and the
+``elif quant_method is None`` branch was unreachable dead code.
+
+The loop is inline inside two heavy functions (importing unsloth needs
+unsloth_zoo / a GPU), so - like test_is_gpt_oss_detection.py - we extract just the
+loop source via ``ast`` and exec it against sample inputs. That exercises the real
+source: it fails on the old ordering and passes once ``None`` is handled first.
+"""
+
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+import pytest
+
+SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py"
+SAVE_SRC = SAVE_PY.read_text(encoding = "utf-8")
+SAVE_TREE = ast.parse(SAVE_SRC, filename = str(SAVE_PY))
+
+# The target functions and the list variable each one appends the normalized method to.
+TARGETS = (
+    ("unsloth_save_pretrained_gguf", "quantization_methods"),
+    ("save_to_gguf_generic", "new_quantization_methods"),
+)
+
+
+def _func(tree, name):
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef) and node.name == name:
+            return node
+    raise AssertionError(f"function {name!r} not found in {SAVE_PY.name}")
+
+
+def _quant_loop(func_name):
+    # The quant-normalization `for` loop iterates `quantization_method`; grab its source.
+    func = _func(SAVE_TREE, func_name)
+    for node in ast.walk(func):
+        if (
+            isinstance(node, ast.For)
+            and isinstance(node.iter, ast.Call)
+            and isinstance(node.iter.func, ast.Name)
+            and node.iter.func.id == "enumerate"
+            and isinstance(node.iter.args[0], ast.Name)
+            and node.iter.args[0].id == "quantization_method"
+        ):
+            return node
+    raise AssertionError(f"quant-normalization loop not found in {func_name}")
+
+
+def _run_loop(func_name, out_var, quantization_method):
+    # exec just the extracted loop against a given input, returning the appended methods.
+    loop_src = ast.get_source_segment(SAVE_SRC, _quant_loop(func_name))
+    namespace = {out_var: [], "quantization_method": quantization_method}
+    exec(loop_src, {"__builtins__": __builtins__}, namespace)
+    return namespace[out_var]
+
+
+@pytest.mark.parametrize("func_name, out_var", TARGETS)
+def test_none_element_maps_to_q8_0(func_name, out_var):
+    # A bare None inside the list must map to q8_0, not raise AttributeError.
+    assert _run_loop(func_name, out_var, [None]) == ["q8_0"]
+
+
+@pytest.mark.parametrize("func_name, out_var", TARGETS)
+def test_none_mixed_with_strings(func_name, out_var):
+    # None resolves to q8_0 while sibling string methods are still normalized (lowercased).
+    assert _run_loop(func_name, out_var, ["Q4_K_M", None]) == ["q4_k_m", "q8_0"]
+
+
+@pytest.mark.parametrize("func_name, out_var", TARGETS)
+def test_string_methods_unchanged(func_name, out_var):
+    # The fix must not alter behavior for the ordinary string inputs.
+    methods = ["not_quantized", "fast_quantized", "quantized", "Q8_0"]
+    assert _run_loop(func_name, out_var, methods) == ["f16", "q8_0", "q4_k_m", "q8_0"]

tests/saving/test_torchao_remote_code_consent.py

@@ -0,0 +1,150 @@
+"""Regression tests for the export-time remote-code trust decision.
+
+FP8/FP4/INT quantization export re-reads the just-merged checkpoint. It used to enable
+trust_remote_code whenever the checkpoint's config carried an ``auto_map`` entry, so a model
+that loads fine with built-in classes (and therefore skips the load-time consent scan) could
+smuggle unvetted remote code that then runs at export. The export paths now derive
+trust_remote_code from ``_loaded_via_remote_code`` - the already approved load decision - instead.
+
+These run on CPU with no torch / unsloth import: they AST-extract the real helper from
+unsloth/save.py and exec it in isolation, plus assert the call sites dropped the auto_map trust.
+"""
+
+import ast
+from pathlib import Path
+
+_SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py"
+_SRC = _SAVE_PY.read_text(encoding = "utf-8")
+
+
+def _load_helper():
+    """Exec just `_loaded_via_remote_code` from save.py (no torch import) and return it."""
+    tree = ast.parse(_SRC)
+    fn = next(
+        n
+        for n in tree.body
+        if isinstance(n, ast.FunctionDef) and n.name == "_loaded_via_remote_code"
+    )
+    ns = {}
+    exec(compile(ast.Module(body = [fn], type_ignores = []), str(_SAVE_PY), "exec"), ns)
+    return ns["_loaded_via_remote_code"]
+
+
+_loaded_via_remote_code = _load_helper()
+
+
+def _obj(module_name, **attrs):
+    """A throwaway instance whose class __module__ is `module_name`, plus given attributes."""
+    cls = type("Fake", (), {})
+    cls.__module__ = module_name
+    inst = cls()
+    for k, v in attrs.items():
+        setattr(inst, k, v)
+    return inst
+
+
+def test_builtin_class_is_not_remote_code():
+    assert _loaded_via_remote_code(_obj("transformers.models.llama.modeling_llama")) is False
+
+
+def test_transformers_modules_class_is_remote_code():
+    assert _loaded_via_remote_code(_obj("transformers_modules.acme.modeling_x")) is True
+
+
+def test_none_is_not_remote_code():
+    assert _loaded_via_remote_code(None) is False
+
+
+def test_none_module_is_not_remote_code():
+    # A class whose __module__ is None must not raise AttributeError.
+    assert _loaded_via_remote_code(_obj(None)) is False
+
+
+def test_auto_map_in_config_alone_does_not_grant_trust():
+    # The core bypass: a built-in-loadable model whose config merely declares auto_map must NOT
+    # be treated as remote-code-loaded (that is exactly what enabled the consent-gate bypass).
+    cfg = type("Cfg", (), {"auto_map": {"AutoModelForCausalLM": "modeling_x.Model"}})()
+    assert (
+        _loaded_via_remote_code(_obj("transformers.models.llama.modeling_llama", config = cfg))
+        is False
+    )
+
+
+def test_peft_base_model_is_unwrapped():
+    base = _obj("transformers_modules.acme.modeling_x")
+    peft = _obj("peft.peft_model", get_base_model = lambda: base)
+    assert _loaded_via_remote_code(peft) is True
+
+
+def test_wrapper_model_attr_is_walked():
+    inner = _obj("transformers_modules.acme.modeling_x")
+    wrapper = _obj("peft.peft_model", model = inner)
+    assert _loaded_via_remote_code(wrapper) is True
+
+
+def test_wrapper_over_builtin_stays_false():
+    inner = _obj("transformers.models.llama.modeling_llama")
+    wrapper = _obj("peft.peft_model", model = inner)
+    assert _loaded_via_remote_code(wrapper) is False
+
+
+def test_processor_held_custom_tokenizer_is_detected():
+    # A built-in ProcessorMixin can hold an approved custom-code tokenizer; the walk must
+    # descend into processor components or the export reload loses that approved trust.
+    tok = _obj("transformers_modules.acme.tokenization_x")
+    proc = _obj("transformers.processing_utils", tokenizer = tok)
+    assert _loaded_via_remote_code(proc) is True
+
+
+def test_processor_held_custom_image_processor_is_detected():
+    ip = _obj("transformers_modules.acme.image_processing_x")
+    proc = _obj("transformers.processing_utils", image_processor = ip)
+    assert _loaded_via_remote_code(proc) is True
+
+
+def test_builtin_processor_with_builtin_components_stays_false():
+    proc = _obj(
+        "transformers.processing_utils",
+        tokenizer = _obj("transformers.tokenization_utils_fast"),
+        image_processor = _obj("transformers.image_processing_utils"),
+    )
+    assert _loaded_via_remote_code(proc) is False
+
+
+def test_cyclic_wrappers_terminate():
+    a = _obj("peft.peft_model")
+    b = _obj("peft.peft_model", model = a)
+    a.model = b
+    assert _loaded_via_remote_code(a) is False
+
+
+# -- call-site assertions: the auto_map-derived trust is gone from every export path -----------
+
+
+def test_torchao_export_derives_trust_from_load_decision():
+    assert "model_trust = _loaded_via_remote_code(model)" in _SRC
+    assert "tok_trust = _loaded_via_remote_code(tokenizer)" in _SRC
+    assert "trust_remote_code = model_trust" in _SRC
+    assert "trust_remote_code = tok_trust" in _SRC
+    # The staged-config auto_map scan that granted trust is removed.
+    assert 'if "auto_map" in json.load' not in _SRC
+
+
+def test_compressed_and_gguf_lora_paths_drop_auto_map_trust():
+    # No path derives a trust decision straight from config auto_map anymore, and no path
+    # collapses model and tokenizer trust into one flag.
+    assert 'bool(getattr(model.config, "auto_map", None))' not in _SRC
+    assert "_loaded_via_remote_code(model) or _loaded_via_remote_code(tokenizer)" not in _SRC
+    assert "if _loaded_via_remote_code(model):" in _SRC  # GGUF-LoRA converter flag
+
+
+def test_compressed_export_keeps_model_and_tokenizer_trust_separate():
+    # The subprocess gets one flag per component, so an approved custom tokenizer cannot
+    # enable an unapproved model's code during compressed quantization (or vice versa).
+    assert 'cmd.append("--trust-remote-code")' in _SRC
+    assert 'cmd.append("--trust-remote-code-tokenizer")' in _SRC
+    qsrc = (_SAVE_PY.parent / "_compressed_quantize.py").read_text(encoding = "utf-8")
+    assert 'ap.add_argument("--trust-remote-code-tokenizer", action = "store_true")' in qsrc
+    assert "trust_remote_code = args.trust_remote_code_tokenizer" in qsrc
+    # The model loads keep the model flag only.
+    assert "args.model, args.trust_remote_code)" in qsrc

tests/test_attn_impl_honor_explicit.py

@@ -0,0 +1,190 @@
+"""An explicit non-flash attention request must survive the flash disable path.
+
+When flash attention is disabled for a model, a caller who explicitly asked for
+"sdpa" or "flex_attention" should keep that choice instead of being downgraded
+to whatever the conservative supports_* fallback would pick.
+"""
+
+import pytest
+
+from unsloth.models._utils import (
+    _disable_flash_attention_if_needed,
+    resolve_attention_implementation,
+)
+
+
+def test_explicit_sdpa_is_honored_even_when_not_marked_supported():
+    config = {}
+    result = _disable_flash_attention_if_needed(
+        config,
+        attn_implementation = "sdpa",
+        supports_sdpa = False,  # conservative flag would have skipped sdpa
+        supports_flex_attention = False,
+        would_use_flash_attention = True,
+        disable_reason = "unit test forces flash disabled",
+    )
+    assert result == "sdpa"
+    assert config.get("_attn_implementation") == "sdpa"
+
+
+def test_explicit_flex_is_honored_when_supported():
+    config = {}
+    result = _disable_flash_attention_if_needed(
+        config,
+        attn_implementation = "flex_attention",
+        supports_sdpa = True,
+        supports_flex_attention = True,
+        would_use_flash_attention = True,
+        disable_reason = "unit test forces flash disabled",
+    )
+    assert result == "flex_attention"
+    assert config.get("_attn_implementation") == "flex_attention"
+
+
+def test_explicit_flex_falls_back_when_not_supported():
+    # flex_attention is False for known-broken/excluded configs (e.g. gpt_oss),
+    # so an explicit flex request must not select that backend - it falls back.
+    config = {}
+    result = _disable_flash_attention_if_needed(
+        config,
+        attn_implementation = "flex_attention",
+        supports_sdpa = True,
+        supports_flex_attention = False,
+        would_use_flash_attention = True,
+        disable_reason = "unit test forces flash disabled",
+    )
+    assert result == "sdpa"
+
+
+def test_synthesized_config_sdpa_is_not_treated_as_explicit():
+    # The language loader seeds the config with attn_implementation="sdpa"; when the
+    # caller passes nothing, that synthesized value must not override the flex fallback
+    # for a model that supports flex but not sdpa.
+    config = {"attn_implementation": "sdpa"}
+    result = _disable_flash_attention_if_needed(
+        config,
+        attn_implementation = None,
+        supports_sdpa = False,
+        supports_flex_attention = True,
+        would_use_flash_attention = False,
+        disable_reason = "unit test forces flash disabled",
+    )
+    assert result == "flex_attention"
+
+
+def test_no_disable_reason_returns_request_untouched():
+    result = _disable_flash_attention_if_needed(
+        {},
+        attn_implementation = "flash_attention_2",
+        disable_reason = None,
+    )
+    assert result == "flash_attention_2"
+
+
+def test_flash_request_still_falls_back_when_disabled():
+    config = {}
+    result = _disable_flash_attention_if_needed(
+        config,
+        attn_implementation = "flash_attention_2",
+        supports_sdpa = True,
+        would_use_flash_attention = True,
+        disable_reason = "unit test forces flash disabled",
+    )
+    assert result == "sdpa"
+
+
+def test_resolver_honors_explicit_sdpa_when_not_supported_and_flash_disabled():
+    # End-to-end through the public resolver: an explicit sdpa request with a
+    # flash-disabled config (oversized head dim) and supports_sdpa=False must not be
+    # rewritten to eager by the resolver's own not-supports_sdpa guard.
+    config = {"model_type": "test", "head_dim": 512}  # head_dim > 256 disables flash
+    result = resolve_attention_implementation(
+        model_class = None,
+        config = config,
+        requested_attn_implementation = "sdpa",
+        supports_sdpa = False,
+    )
+    assert result == "sdpa"
+    assert config.get("_attn_implementation") == "sdpa"
+
+
+def test_resolver_downgrades_non_explicit_sdpa_when_not_supported():
+    # No explicit request: the model resolution seeds sdpa/eager and the guard must
+    # still downgrade a synthesized sdpa to eager for a model that cannot run it.
+    config = {"model_type": "test", "attn_implementation": "sdpa"}
+    result = resolve_attention_implementation(
+        model_class = None,
+        config = config,
+        requested_attn_implementation = None,
+        supports_sdpa = False,
+    )
+    assert result == "eager"
+
+
+def test_resolver_downgrades_explicit_sdpa_for_sdpa_excluded_model():
+    # gpt_oss is in _SDPA_EXCLUDED_MODELS (sdpa is known-broken) and _FLASH_EXCLUDED_MODELS
+    # (flash disabled). Honoring an explicit sdpa request must not re-enable that broken
+    # backend: it downgrades to eager, mirroring how an explicit flex request falls back
+    # for _FLEX_EXCLUDED_MODELS. supports_sdpa=True proves the exclusion overrides even a
+    # model that otherwise advertises SDPA support.
+    config = {"model_type": "gpt_oss"}
+    result = resolve_attention_implementation(
+        model_class = None,
+        config = config,
+        requested_attn_implementation = "sdpa",
+        supports_sdpa = True,
+    )
+    assert result == "eager"
+    assert config.get("_attn_implementation") == "eager"
+
+
+@pytest.mark.parametrize("model_type", ["gemma3", "gemma3_text"])
+def test_resolver_downgrades_explicit_sdpa_for_disable_sdpa_model(model_type):
+    # gemma3 / gemma3_text are in DISABLE_SDPA_MODEL_NAMES: the loader forces
+    # supports_sdpa=False because their bundled SDPA modules are wrong. An explicit
+    # sdpa request with flash disabled must NOT re-enable that known-wrong path - it
+    # downgrades to eager, exactly like _SDPA_EXCLUDED_MODELS (gpt_oss). head_dim>256
+    # disables flash to mirror the real flash-disabled scenario.
+    config = {"model_type": model_type, "head_dim": 512}
+    result = resolve_attention_implementation(
+        model_class = None,
+        config = config,
+        requested_attn_implementation = "sdpa",
+        supports_sdpa = False,
+    )
+    assert result == "eager"
+    assert config.get("_attn_implementation") == "eager"
+
+
+def test_resolver_does_not_overmatch_gemma3n_for_explicit_sdpa():
+    # The "gemma3," trailing-comma guard must not match gemma3n: gemma3n is not in
+    # DISABLE_SDPA_MODEL_NAMES, so it stays a conservative (not known-wrong) model and an
+    # explicit sdpa request is still honored. Proves the substring match neither over- nor
+    # under-matches.
+    config = {"model_type": "gemma3n", "head_dim": 512}
+    result = resolve_attention_implementation(
+        model_class = None,
+        config = config,
+        requested_attn_implementation = "sdpa",
+        supports_sdpa = False,
+    )
+    assert result == "sdpa"
+    assert config.get("_attn_implementation") == "sdpa"
+
+
+def test_resolver_downgrades_synthesized_sdpa_for_disable_sdpa_model():
+    # A synthesized/default sdpa (requested is None; the value came from config) on a
+    # DISABLE_SDPA_MODEL_NAMES model must still downgrade to eager.
+    config = {"model_type": "gemma3", "attn_implementation": "sdpa"}
+    result = resolve_attention_implementation(
+        model_class = None,
+        config = config,
+        requested_attn_implementation = None,
+        supports_sdpa = False,
+    )
+    assert result == "eager"
+
+
+if __name__ == "__main__":
+    import sys
+    sys.exit(pytest.main([__file__, "-q"]))

tests/test_fp8_tiny_e8m0.py

@@ -0,0 +1,123 @@
+"""FP8 block-quant linear must handle tiny / non-tileable weights and e8m0 scales.
+
+Two things break the triton block path:
+  * a hidden dim not divisible by the activation block size (tiny test models),
+  * float8_e8m0fnu weight scales, which have no triton dtype mapping.
+The forward falls back to a torch-native blockwise dequant + bf16 matmul; this
+test checks that fallback runs finite forward + backward and matches a plain
+dequant reference.
+"""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason = "needs CUDA")
+
+
+def _reference(X, weight, scale, block):
+    # Expand the per-block scale to full weight shape and dequantize.
+    m, n = weight.shape
+    s = scale.to(torch.float32)
+    s = s.repeat_interleave(block[0], 0)[:m].repeat_interleave(block[1], 1)[:, :n]
+    W = (weight.to(torch.float32) * s).to(X.dtype)
+    return X @ W.T
+
+
+def test_tiny_non_tileable_forward_backward_matches_reference():
+    from unsloth.kernels.fp8 import FP8BlockQuantLinear
+
+    torch.manual_seed(0)
+    dev = "cuda"
+    block = [128, 128]
+    m, n = 8, 8  # non-tileable, in-dim % 128 != 0
+    weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16)  # (out=m, in=n)
+    scale = torch.rand(1, 1, device = dev, dtype = torch.float32) + 0.5
+    X = torch.randn(4, n, device = dev, dtype = torch.bfloat16, requires_grad = True)
+
+    out = FP8BlockQuantLinear.apply(X, weight, scale)
+    assert torch.isfinite(out).all(), "forward produced non-finite values"
+
+    ref = _reference(X.detach(), weight, scale, block)
+    torch.testing.assert_close(out, ref, atol = 5e-2, rtol = 5e-2)
+
+    out.sum().backward()
+    assert X.grad is not None and torch.isfinite(X.grad).all(), "backward non-finite"
+
+
+def test_e8m0_scale_is_upcast_and_runs():
+    from unsloth.kernels.fp8 import FP8BlockQuantLinear
+
+    if not hasattr(torch, "float8_e8m0fnu"):
+        pytest.skip("torch build lacks float8_e8m0fnu")
+
+    dev = "cuda"
+    m, n = 8, 8
+    weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16)
+    scale = (torch.rand(1, 1, device = dev) + 1.0).to(torch.float8_e8m0fnu)
+    X = torch.randn(4, n, device = dev, dtype = torch.bfloat16, requires_grad = True)
+
+    out = FP8BlockQuantLinear.apply(X, weight, scale)
+    assert torch.isfinite(out).all()
+    out.sum().backward()
+    assert torch.isfinite(X.grad).all()
+
+
+def test_rectangular_block_dequant_matches_reference():
+    # Rectangular blocks (block_size[0] != block_size[1]) that tile evenly used to
+    # route through the triton weight_dequant kernel, which uses a single BLOCK_SIZE
+    # for both axes and mis-indexes the column scale. Verify the torch expansion path
+    # now matches the reference for a 64x256 weight with block [64, 128] (scale 1x2).
+    from unsloth.kernels.fp8 import _blockwise_weight_dequant_any_shape
+
+    torch.manual_seed(0)
+    dev = "cuda"
+    block = [64, 128]
+    m, n = 64, 256  # evenly tiled: 64 % 64 == 0, 256 % 128 == 0
+    weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16)
+    # Distinct per-block column scales expose column mis-indexing.
+    scale = torch.tensor([[0.5, 3.0]], device = dev, dtype = torch.float32)
+
+    W_deq = _blockwise_weight_dequant_any_shape(weight, scale, block, torch.bfloat16)
+
+    s = scale.repeat_interleave(block[0], 0)[:m].repeat_interleave(block[1], 1)[:, :n]
+    ref = (weight.to(torch.float32) * s).to(torch.bfloat16)
+    torch.testing.assert_close(W_deq, ref, atol = 5e-3, rtol = 5e-3)
+
+
+def test_e8m0_scale_preserves_non_default_block_size_attr():
+    # An e8m0 scale carrying a non-default block_size attribute must keep it across
+    # the float32 upcast in forward; otherwise the lookup falls back to [128, 128]
+    # and a compatible layout is wrongly rejected as incompatible.
+    from unsloth.kernels.fp8 import FP8BlockQuantLinear
+
+    if not hasattr(torch, "float8_e8m0fnu"):
+        pytest.skip("torch build lacks float8_e8m0fnu")
+
+    torch.manual_seed(0)
+    dev = "cuda"
+    block = [64, 64]
+    # in-dim 96 is not divisible by block[1]=64 -> forward takes the torch dequant
+    # fallback (no fp8 matmul kernel). Scale shape (2, 2) validates for [64, 64] but
+    # not [128, 128] (which expects (1, 1)).
+    m, n = 128, 96
+    weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16)  # no block_size attr
+    scale_f = torch.rand(2, 2, device = dev) + 1.0
+    scale = scale_f.to(torch.float8_e8m0fnu)
+    scale.block_size = block  # attribute lives on the scale, not the weight
+    X = torch.randn(4, n, device = dev, dtype = torch.bfloat16, requires_grad = True)
+
+    # With [128, 128] this raises "not compatible with block size"; success proves
+    # the [64, 64] attribute survived the e8m0 -> float32 upcast.
+    out = FP8BlockQuantLinear.apply(X, weight, scale)
+    assert torch.isfinite(out).all()
+
+    ref = _reference(X.detach(), weight, scale.to(torch.float32), block)
+    torch.testing.assert_close(out, ref, atol = 5e-2, rtol = 5e-2)
+
+    out.sum().backward()
+    assert X.grad is not None and torch.isfinite(X.grad).all()
+
+
+if __name__ == "__main__":
+    import sys
+    sys.exit(pytest.main([__file__, "-q"]))

tests/test_moe_lora_targets.py

@@ -49,3 +49,190 @@ def test_explicit_dotted_module_target_does_not_discover_moe_parameters():
         )
         is None
     )
+
+
+@pytest.mark.parametrize(
+    "target_modules",
+    [
+        # Attention-only auto-regex lists every projection leaf (incl. gate/up/down)
+        # but its path segment is attention-only, so experts must NOT be targeted.
+        r"(?:\bmodel\.layers\.[\d]{1,}\.(?:self_attn|attention|attn|mixer)\.(?:q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj))",
+        ".*self_attn.*proj",
+        # An mlp path alternative with attention-only leaves is still attention-only.
+        r"model\.layers\.\d+\.(?:mlp|self_attn)\.(?:q_proj|k_proj|v_proj|o_proj)",
+    ],
+)
+def test_attention_only_regex_does_not_discover_moe_parameters(target_modules):
+    from unsloth.models._utils import get_moe_target_parameters
+    assert get_moe_target_parameters(_FakeMoeModel(), target_modules) is None
+
+
+def test_single_leaf_regex_targets_only_that_projection():
+    from unsloth.models._utils import get_moe_target_parameters
+    assert get_moe_target_parameters(_FakeMoeModel(), ".*experts.*down_proj") == [
+        "mlp.experts.down_proj",
+    ]
+    assert get_moe_target_parameters(_FakeMoeModel(), ".*mlp.*gate_proj") == [
+        "mlp.experts.gate_up_proj",
+    ]
+
+
+def test_auto_regex_mlp_tag_block_discovers_moe_on_fused_models():
+    # get_peft_regex on a fused-expert model lists only attention Linears as
+    # leaves; the mlp tag block is the remaining signal of MLP finetune intent.
+    from unsloth.models._utils import get_moe_target_parameters
+    both_auto = (
+        r"(?:\bmodel\.layers\.[\d]{1,}\."
+        r"(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense|mixer)\."
+        r"(?:(?:q_proj|k_proj|v_proj|o_proj)))"
+    )
+    assert get_moe_target_parameters(_FakeMoeModel(), both_auto) == [
+        "mlp.experts.gate_up_proj",
+        "mlp.experts.down_proj",
+    ]
+
+
+def test_explicit_attention_only_list_does_not_discover_moe_parameters():
+    # An explicit attention-only leaf list names no MLP projection, so experts
+    # must never be targeted. get_peft_model routes this ORIGINAL list (not the
+    # scoped regex) into detection precisely because family scoping makes
+    # get_peft_regex emit its full "mlp|feed_forward|ffn|dense" component block
+    # even for an attention-only request (see the regex below), which the
+    # string fallback cannot distinguish from the fused-expert auto regex.
+    from unsloth.models._utils import get_moe_target_parameters
+
+    attn_only_list = ["q_proj", "k_proj", "v_proj", "o_proj"]
+    assert get_moe_target_parameters(_FakeMoeModel(), attn_only_list) is None
+    assert get_moe_target_parameters(_FakeMoeModel(), tuple(attn_only_list)) is None
+
+    # The regex get_peft_regex emits for that same attention-only list under a
+    # vision-off family scope carries the mlp component block, so the string
+    # path would wrongly enable experts -- hence detection must use the list.
+    scoped_regex = (
+        r"(?:.*?(?:language|text).*?"
+        r"(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense|mixer).*?"
+        r"(?:q_proj|k_proj|v_proj|o_proj))"
+    )
+    assert get_moe_target_parameters(_FakeMoeModel(), scoped_regex) == [
+        "mlp.experts.gate_up_proj",
+        "mlp.experts.down_proj",
+    ]
+
+
+def test_frozen_mlp_full_list_does_not_discover_moe_parameters():
+    # Regression: an explicit list that names MLP leaves together with
+    # finetune_mlp_modules=False must NOT train experts. get_peft_regex scopes
+    # the MLP leaves out (its emitted regex carries no mlp tag block), so
+    # detection has to key on that SCOPED regex -- keying on the original list
+    # would let its gate/up/down leaves silently re-enable the frozen experts.
+    from unsloth.models._utils import (
+        _select_moe_detection_targets,
+        get_moe_target_parameters,
+    )
+
+    original_list = [
+        "q_proj",
+        "k_proj",
+        "v_proj",
+        "o_proj",
+        "gate_proj",
+        "up_proj",
+        "down_proj",
+    ]
+    # Representative of what get_peft_regex emits for that list under
+    # finetune_mlp_modules=False: attention-only path, no mlp component block.
+    scoped_regex = (
+        r"(?:.*?(?:language|text).*?"
+        r"(?:self_attn|attention|attn|mixer).*?"
+        r"(?:q_proj|k_proj|v_proj|o_proj))"
+    )
+    selected = _select_moe_detection_targets(
+        original_list,
+        scoped_regex,
+        finetune_mlp_modules = False,
+        finetune_language_layers = True,
+    )
+    assert selected is scoped_regex
+    assert get_moe_target_parameters(_FakeMoeModel(), selected) is None
+
+
+def test_frozen_language_full_list_does_not_discover_moe_parameters():
+    # Vision-only request (finetune_language_layers=False) with a full leaf list
+    # must not reach the language-model experts either.
+    from unsloth.models._utils import (
+        _select_moe_detection_targets,
+        get_moe_target_parameters,
+    )
+
+    original_list = ["q_proj", "gate_proj", "up_proj", "down_proj"]
+    scoped_regex = (
+        r"(?:.*?(?:vision|visual|image).*?"
+        r"(?:self_attn|attention|attn|mixer).*?"
+        r"(?:q_proj|k_proj|v_proj|o_proj))"
+    )
+    selected = _select_moe_detection_targets(
+        original_list,
+        scoped_regex,
+        finetune_mlp_modules = True,
+        finetune_language_layers = False,
+    )
+    assert selected is scoped_regex
+    assert get_moe_target_parameters(_FakeMoeModel(), selected) is None
+
+
+def test_in_scope_mlp_full_list_still_discovers_moe_parameters():
+    # With MLP and language both in scope, an explicit list that names MLP
+    # leaves SHOULD enable the experts (unchanged behavior): the original list
+    # is preferred and carries the gate/up/down intent.
+    from unsloth.models._utils import (
+        _select_moe_detection_targets,
+        get_moe_target_parameters,
+    )
+
+    original_list = [
+        "q_proj",
+        "k_proj",
+        "v_proj",
+        "o_proj",
+        "gate_proj",
+        "up_proj",
+        "down_proj",
+    ]
+    scoped_regex = r".*self_attn.*proj"  # unused: original list is preferred
+    selected = _select_moe_detection_targets(
+        original_list,
+        scoped_regex,
+        finetune_mlp_modules = True,
+        finetune_language_layers = True,
+    )
+    assert selected is original_list
+    assert get_moe_target_parameters(_FakeMoeModel(), selected) == [
+        "mlp.experts.gate_up_proj",
+        "mlp.experts.down_proj",
+    ]
+
+
+def test_attention_only_list_prefers_original_when_in_scope():
+    # The case the PR originally fixed: an attention-only list routed through
+    # get_peft_regex under a family scope (e.g. vision-off) still keeps experts
+    # off, because with MLP+language in scope detection uses the original
+    # attention-only list rather than the regex's spurious mlp component block.
+    from unsloth.models._utils import (
+        _select_moe_detection_targets,
+        get_moe_target_parameters,
+    )
+
+    attn_only_list = ["q_proj", "k_proj", "v_proj", "o_proj"]
+    scoped_regex = (  # carries the spurious mlp block get_peft_regex always adds
+        r"(?:.*?(?:language|text).*?"
+        r"(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense).*?"
+        r"(?:q_proj|k_proj|v_proj|o_proj))"
+    )
+    selected = _select_moe_detection_targets(
+        attn_only_list,
+        scoped_regex,
+        finetune_mlp_modules = True,
+        finetune_language_layers = True,
+    )
+    assert selected is attn_only_list
+    assert get_moe_target_parameters(_FakeMoeModel(), selected) is None

tests/test_prefetch_snapshot_scope.py

@@ -0,0 +1,916 @@
+# Unsloth Zoo - Utilities for Unsloth
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""Pure-CPU, no-network unit tests for prefetch snapshot scoping in unsloth/models/_utils.py.
+
+maybe_prefetch_hf_snapshot warms the HF cache before the in-process load. The warm must cover at
+least what the load reads (else the missing file falls to an unprotected in-process Xet fetch) but
+not pull weights the load never reads. These tests lock the allow/ignore patterns each mode hands
+snapshot_download_with_xet_fallback. The zoo downloader is monkeypatched to capture its kwargs.
+"""
+
+import fnmatch
+import sys
+import types
+
+import pytest
+
+from unsloth.models import _utils as U
+
+
+def _filter(names, allow_patterns, ignore_patterns):
+    """Mirror HF filter_repo_objects: keep on allow match (or None), drop on ignore match."""
+    kept = []
+    for name in names:
+        if allow_patterns is not None and not any(fnmatch.fnmatch(name, p) for p in allow_patterns):
+            continue
+        if ignore_patterns and any(fnmatch.fnmatch(name, p) for p in ignore_patterns):
+            continue
+        kept.append(name)
+    return kept
+
+
+@pytest.fixture
+def capture(monkeypatch):
+    """Run maybe_prefetch_hf_snapshot with a fake repo, capturing the patterns forwarded to a
+    fake injected zoo downloader (independent of the installed unsloth_zoo). Offline env cleared."""
+    monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
+    monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
+
+    state = {}
+
+    def fake_download(repo_id, **kw):
+        state["repo_id"] = repo_id
+        state["allow_patterns"] = kw.get("allow_patterns")
+        state["ignore_patterns"] = kw.get("ignore_patterns")
+        state["variant"] = kw.get("variant")
+        return "/tmp/fake-snapshot"
+
+    fake_module = types.ModuleType("unsloth_zoo.hf_xet_fallback")
+    fake_module.snapshot_download_with_xet_fallback = fake_download
+    fake_module.DownloadStallError = type("DownloadStallError", (RuntimeError,), {})
+    monkeypatch.setitem(sys.modules, "unsloth_zoo.hf_xet_fallback", fake_module)
+
+    # Neutralize the model_info network call by default; tests exercising format selection
+    # install their own.
+    import huggingface_hub
+
+    class _NoNetworkApi:
+        def model_info(self, *a, **k):
+            raise RuntimeError("no network in test")
+
+    monkeypatch.setattr(huggingface_hub, "HfApi", _NoNetworkApi)
+
+    def run(**call_kwargs):
+        state.clear()
+        ok = U.maybe_prefetch_hf_snapshot("some-org/some-repo", **call_kwargs)
+        return ok, state
+
+    return run
+
+
+# Representative repo listing: root weights + aux, subdir, adapter, checkpoint, merged weights.
+_SAMPLE_FILES = [
+    "config.json",
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "model-00001-of-00002.safetensors",
+    "model-00002-of-00002.safetensors",
+    "model.safetensors.index.json",
+    "pytorch_model.bin",
+    "fp16/model.safetensors",
+    "experimental/model-00001-of-00002.safetensors",
+    "checkpoint-500/model.safetensors",
+    "adapter_config.json",
+    "adapter_model.safetensors",
+]
+
+
+def test_weights_at_root_excludes_subdir_weights(capture):
+    """A root load ignores subdir weights (fp16/, experimental/, checkpoint-500/) but keeps root weights."""
+    ok, st = capture(weights_at_root = True, use_safetensors = True)
+    assert ok is True
+    assert st["allow_patterns"] is None
+    ig = st["ignore_patterns"]
+    assert "*/*.safetensors" in ig and "*/*.bin" in ig
+    kept = _filter(_SAMPLE_FILES, st["allow_patterns"], ig)
+    assert "model-00001-of-00002.safetensors" in kept
+    assert "model.safetensors.index.json" in kept
+    assert "config.json" in kept
+    assert "fp16/model.safetensors" not in kept
+    assert "experimental/model-00001-of-00002.safetensors" not in kept
+    assert "checkpoint-500/model.safetensors" not in kept
+
+
+def test_adapter_only_excludes_merged_weights(capture):
+    """An adapter warm keeps adapter files + root aux, not merged full-model weights."""
+    ok, st = capture(adapter_only = True)
+    assert ok is True
+    assert st["ignore_patterns"] is None
+    allow = st["allow_patterns"]
+    assert "adapter_config.json" in allow and "adapter_model*" in allow
+    kept = _filter(_SAMPLE_FILES, allow, st["ignore_patterns"])
+    assert "adapter_config.json" in kept
+    assert "adapter_model.safetensors" in kept
+    assert "config.json" in kept and "tokenizer.json" in kept
+    assert "model-00001-of-00002.safetensors" not in kept
+    assert "pytorch_model.bin" not in kept
+    assert "fp16/model.safetensors" not in kept
+
+
+def test_adapter_only_warms_sharded_adapter(capture):
+    """A sharded adapter is still covered by the adapter_model* glob."""
+    _, st = capture(adapter_only = True)
+    sharded = [
+        "adapter_config.json",
+        "adapter_model-00001-of-00002.safetensors",
+        "adapter_model-00002-of-00002.safetensors",
+        "adapter_model.safetensors.index.json",
+    ]
+    kept = _filter(sharded, st["allow_patterns"], st["ignore_patterns"])
+    assert set(kept) == set(sharded)
+
+
+def test_tokenizer_only_warms_only_aux_files(capture):
+    """A tokenizer-only repo warms tokenizer/config/vocab files, never weights."""
+    _, st = capture(tokenizer_only = True)
+    assert st["ignore_patterns"] is None
+    assert st["allow_patterns"] == list(U._ROOT_AUX_PREFETCH_PATTERNS)
+    kept = _filter(_SAMPLE_FILES, st["allow_patterns"], st["ignore_patterns"])
+    assert "tokenizer.json" in kept and "config.json" in kept
+    assert "model-00001-of-00002.safetensors" not in kept
+    assert "adapter_model.safetensors" not in kept
+
+
+def test_aux_warm_covers_arbitrary_remote_code_modules(capture):
+    """The aux warm must cover any *.py, since trust_remote_code auto_map names modules freely."""
+    _, st = capture(tokenizer_only = True)
+    allow = st["allow_patterns"]
+    assert "*.py" in allow
+    remote_code = [
+        "config.json",
+        "modeling.py",
+        "tokenization.py",
+        "my_custom_code.py",
+        "configuration_foo.py",
+    ]
+    kept = _filter(remote_code, allow, st["ignore_patterns"])
+    for name in ("modeling.py", "tokenization.py", "my_custom_code.py", "configuration_foo.py"):
+        assert name in kept, name
+
+
+def test_subfolder_warms_subfolder_plus_root_aux(capture):
+    """A subfolder load warms that subfolder's weights plus root aux; other subdirs/root weights skipped."""
+    _, st = capture(subfolder = "fp16")
+    allow = st["allow_patterns"]
+    assert "fp16/*" in allow
+    assert all(p in allow for p in U._ROOT_AUX_PREFETCH_PATTERNS)
+    kept = _filter(_SAMPLE_FILES, allow, st["ignore_patterns"])
+    assert "fp16/model.safetensors" in kept
+    assert "config.json" in kept
+    assert "experimental/model-00001-of-00002.safetensors" not in kept
+
+
+def test_subfolder_takes_precedence_over_weights_at_root(capture):
+    """When a subfolder is requested the subfolder branch wins over weights_at_root."""
+    _, st = capture(subfolder = "fp16", weights_at_root = True)
+    assert "fp16/*" in st["allow_patterns"]
+    kept = _filter(_SAMPLE_FILES, st["allow_patterns"], st["ignore_patterns"])
+    assert "fp16/model.safetensors" in kept
+
+
+def test_local_dir_is_not_warmed(capture, tmp_path):
+    """A local directory path skips the warm (returns False)."""
+    d = tmp_path / "local-model"
+    d.mkdir()
+    ok = U.maybe_prefetch_hf_snapshot(str(d), weights_at_root = True)
+    assert ok is False
+
+
+def _install_fake_model_info(monkeypatch, filenames):
+    """Make HfApi().model_info(...).siblings report filenames, with no network."""
+    import huggingface_hub
+
+    class _Sib:
+        def __init__(self, name):
+            self.rfilename = name
+
+    class _Info:
+        def __init__(self, names):
+            self.siblings = [_Sib(n) for n in names]
+
+    class _Api:
+        def model_info(self, *a, **k):
+            return _Info(filenames)
+
+    monkeypatch.setattr(huggingface_hub, "HfApi", _Api)
+
+
+# ----- Finding P: variant-aware weight-format selection -----
+
+
+def test_variant_keeps_bin_when_only_default_safetensors(monkeypatch):
+    """A default model.safetensors must not prove a variant .bin redundant; without a variant it does."""
+    _install_fake_model_info(monkeypatch, ["model.safetensors", "pytorch_model.fp16.bin"])
+    ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True)
+    assert "*.bin" not in ig
+    ig_default = U._prefetch_ignore_patterns("org/repo", weights_at_root = True)
+    assert "*.bin" in ig_default
+
+
+def test_variant_drops_bin_when_variant_safetensors_present(monkeypatch):
+    """A variant-matching safetensors makes the variant .bin redundant, so .bin is dropped."""
+    _install_fake_model_info(monkeypatch, ["model.fp16.safetensors", "pytorch_model.fp16.bin"])
+    ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True)
+    assert "*.bin" in ig
+
+
+def test_no_variant_keeps_bin_when_only_variant_safetensors(monkeypatch):
+    """For a no-variant load, only a canonical safetensors (not a lone variant) makes .bin redundant."""
+    _install_fake_model_info(monkeypatch, ["model.fp16.safetensors", "pytorch_model.bin"])
+    ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = True)
+    assert "*.bin" not in ig
+    _install_fake_model_info(monkeypatch, ["model.safetensors", "pytorch_model.bin"])
+    ig2 = U._prefetch_ignore_patterns("org/repo", weights_at_root = True)
+    assert "*.bin" in ig2
+
+
+def test_variant_keeps_bin_for_noncanonical_sidecar(monkeypatch):
+    """A non-canonical variant sidecar must not prove the variant .bin redundant; a canonical one does."""
+    _install_fake_model_info(
+        monkeypatch, ["consolidated.fp16.safetensors", "pytorch_model.fp16.bin"]
+    )
+    ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True)
+    assert "*.bin" not in ig
+    _install_fake_model_info(monkeypatch, ["model.fp16.safetensors", "pytorch_model.fp16.bin"])
+    ig2 = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True)
+    assert "*.bin" in ig2
+
+
+def test_is_canonical_model_weight_safetensors():
+    """The canonical detector matches only non-variant model-weight safetensors names."""
+    assert U._is_canonical_model_weight_safetensors("model.safetensors") is True
+    assert U._is_canonical_model_weight_safetensors("model-00001-of-00002.safetensors") is True
+    assert U._is_canonical_model_weight_safetensors("model.safetensors.index.json") is True
+    assert U._is_canonical_model_weight_safetensors("model.fp16.safetensors") is False
+    assert (
+        U._is_canonical_model_weight_safetensors("model.fp16-00001-of-00002.safetensors") is False
+    )
+    assert U._is_canonical_model_weight_safetensors("adapter_model.safetensors") is False
+
+
+def test_st_prefetch_resolves_env_cache_and_runs_after_validation():
+    """The ST prefetch must resolve SENTENCE_TRANSFORMERS_HOME and run after load-mode validation."""
+    import ast
+    import os
+
+    src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py")
+    with open(src_path, "r", encoding = "utf-8") as f:
+        src = f.read()
+    tree = ast.parse(src)
+    prefetch_calls = [
+        n
+        for n in ast.walk(tree)
+        if isinstance(n, ast.Call)
+        and isinstance(n.func, ast.Name)
+        and n.func.id == "maybe_prefetch_hf_snapshot"
+    ]
+    assert len(prefetch_calls) == 1, "expected exactly one ST prefetch call"
+    call = prefetch_calls[0]
+    # cache_dir kwarg resolves SENTENCE_TRANSFORMERS_HOME.
+    cache_dir_kw = next((kw for kw in call.keywords if kw.arg == "cache_dir"), None)
+    assert cache_dir_kw is not None, "ST prefetch must pass cache_dir"
+    assert "SENTENCE_TRANSFORMERS_HOME" in ast.dump(
+        cache_dir_kw.value
+    ), "ST prefetch cache_dir must resolve SENTENCE_TRANSFORMERS_HOME"
+    # Load-mode validation runs before the prefetch (fewer source lines = earlier).
+    val_lineno = src[: src.index("Can only load in 4bit or 8bit or 16bit")].count("\n")
+    assert val_lineno < call.lineno, "load-mode validation must precede the ST prefetch"
+
+
+def test_st_cache_resolutions_honor_explicit_hf_cache_dir():
+    """Every ST cache resolution falling back to SENTENCE_TRANSFORMERS_HOME must first honor an explicit HF cache_dir."""
+    import ast
+    import os
+
+    src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py")
+    with open(src_path, "r", encoding = "utf-8") as f:
+        tree = ast.parse(f.read())
+    resolutions = [
+        kw
+        for kw in ast.walk(tree)
+        if isinstance(kw, ast.keyword)
+        and kw.arg == "cache_dir"
+        and "SENTENCE_TRANSFORMERS_HOME" in ast.dump(kw.value)
+    ]
+    assert resolutions, "expected cache_dir resolutions referencing SENTENCE_TRANSFORMERS_HOME"
+    for kw in resolutions:
+        assert "'cache_dir'" in ast.dump(
+            kw.value
+        ), "an ST cache_dir resolution must read an explicit kwargs.get('cache_dir') first"
+
+
+def test_st_native_loads_map_hf_cache_dir_to_cache_folder():
+    """Native SentenceTransformer loads take cache_folder, so an explicit HF cache_dir must be mapped onto it."""
+    import ast
+    import os
+
+    src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py")
+    with open(src_path, "r", encoding = "utf-8") as f:
+        src = f.read()
+    tree = ast.parse(src)
+    # Every native SentenceTransformer(...) forwarding cache_folder must read cache_dir.
+    st_calls = [
+        n
+        for n in ast.walk(tree)
+        if isinstance(n, ast.Call)
+        and isinstance(n.func, ast.Name)
+        and n.func.id == "SentenceTransformer"
+    ]
+    cache_folder_kws = [kw for call in st_calls for kw in call.keywords if kw.arg == "cache_folder"]
+    assert cache_folder_kws, "expected a native SentenceTransformer call forwarding cache_folder"
+    for kw in cache_folder_kws:
+        assert "'cache_dir'" in ast.dump(
+            kw.value
+        ), "a native SentenceTransformer cache_folder must map the explicit HF cache_dir first"
+    # for_inference feeds cache_folder via st_kwargs; both native branches map cache_dir -> cache_folder.
+    normalized = "".join(src.split())
+    assert (
+        'st_kwargs["cache_folder"]=' in normalized
+    ), "for_inference must set st_kwargs cache_folder"
+    assert (
+        normalized.count('kwargs.get("cache_dir")orkwargs.get("cache_folder")') >= 2
+    ), "both native ST branches (for_inference, fast-encoder) must map cache_dir -> cache_folder"
+
+
+def test_vision_warms_vllm_tokenizer_after_remap():
+    """On the vLLM path the tokenizer warm is deferred until after the fast_inference_setup remap."""
+    import os
+
+    src_path = os.path.join(os.path.dirname(U.__file__), "vision.py")
+    with open(src_path, "r", encoding = "utf-8") as f:
+        src = f.read()
+    guard = "if _vllm_owns_weights and isinstance(tokenizer_name"
+    assert guard in src, "expected a vLLM-gated tokenizer warm"
+    assert src.index(guard) > src.index(
+        "fast_inference_setup("
+    ), "the vLLM tokenizer warm must run after the fast_inference_setup remap"
+
+
+def test_diffusion_forwards_variant_to_real_load():
+    """FastDiffusionModel must forward variant to the real model_cls.from_pretrained load, not just the prefetch."""
+    import os
+
+    src_path = os.path.join(os.path.dirname(U.__file__), "diffusion.py")
+    with open(src_path, "r", encoding = "utf-8") as f:
+        src = f.read()
+    assert (
+        'load_kwargs["variant"] = kwargs["variant"]' in src
+    ), "the diffusion load must forward variant to model_cls.from_pretrained"
+
+
+def test_vision_prefetch_runs_after_load_mode_validation():
+    """The FastBaseModel (vision) prefetch must run after the load-mode validation."""
+    import ast
+    import os
+
+    src_path = os.path.join(os.path.dirname(U.__file__), "vision.py")
+    with open(src_path, "r", encoding = "utf-8") as f:
+        src = f.read()
+    tree = ast.parse(src)
+    prefetch_calls = [
+        n
+        for n in ast.walk(tree)
+        if isinstance(n, ast.Call)
+        and isinstance(n.func, ast.Name)
+        and n.func.id == "maybe_prefetch_hf_snapshot"
+    ]
+    assert prefetch_calls, "expected a vision prefetch call"
+    first_prefetch = min(call.lineno for call in prefetch_calls)
+    val_lineno = src[: src.index("Can only load in 4bit or 8bit or 16bit")].count("\n")
+    assert val_lineno < first_prefetch, "load-mode validation must precede the vision prefetch"
+
+
+def test_llama_prefetch_skips_only_real_vllm_loads():
+    """The llama prefetch's fast_inference skip must be gated on num_labels is None (a classification load still downloads)."""
+    import ast
+    import os
+
+    src_path = os.path.join(os.path.dirname(U.__file__), "llama.py")
+    with open(src_path, "r", encoding = "utf-8") as f:
+        tree = ast.parse(f.read())
+    gated = False
+    for n in ast.walk(tree):
+        if not (
+            isinstance(n, ast.Call)
+            and isinstance(n.func, ast.Name)
+            and n.func.id == "maybe_prefetch_hf_snapshot"
+        ):
+            continue
+        fi_kw = next((kw for kw in n.keywords if kw.arg == "fast_inference"), None)
+        if fi_kw is None:
+            continue
+        dumped = ast.dump(fi_kw.value)
+        if "fast_inference" in dumped and "num_labels" in dumped:
+            gated = True
+    assert gated, "llama prefetch fast_inference must be gated on num_labels is None"
+
+
+def test_st_fallback_module_loads_resolve_env_cache():
+    """Fallback module loads deriving cache_dir from cache_folder must also fall back to SENTENCE_TRANSFORMERS_HOME."""
+    import ast
+    import os
+
+    src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py")
+    with open(src_path, "r", encoding = "utf-8") as f:
+        src = f.read()
+    tree = ast.parse(src)
+
+    # Fallback sites (cache_dir derived from cache_folder) must resolve SENTENCE_TRANSFORMERS_HOME.
+    checked = 0
+    for node in ast.walk(tree):
+        if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)):
+            continue
+        if node.func.attr not in ("_module_path", "_load_modules"):
+            continue
+        cache_dir_kw = next((kw for kw in node.keywords if kw.arg == "cache_dir"), None)
+        if cache_dir_kw is None:
+            continue
+        dumped = ast.dump(cache_dir_kw.value)
+        if "cache_folder" not in dumped:
+            continue  # internal pass-through, not a resolution site
+        checked += 1
+        assert (
+            "SENTENCE_TRANSFORMERS_HOME" in dumped
+        ), f"{node.func.attr} cache_dir resolves cache_folder but not SENTENCE_TRANSFORMERS_HOME"
+    assert (
+        checked >= 2
+    ), "expected the fallback _module_path and _load_modules calls to resolve the env cache"
+
+
+def test_st_fallback_module_loads_forward_revision():
+    """The fallback module loads must forward revision so module files match the revision-pinned weights.
+    Guards: (a) helpers accept revision, (b) every download primitive forwards it, (c) _load_modules
+    threads it into internal calls, (d) the from_pretrained fallback sites forward it."""
+    import ast
+    import os
+
+    src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py")
+    with open(src_path, "r", encoding = "utf-8") as f:
+        tree = ast.parse(f.read())
+
+    funcs = {
+        n.name: n
+        for n in ast.walk(tree)
+        if isinstance(n, ast.FunctionDef)
+        and n.name in ("_module_path", "_read_pooling_mode", "_load_modules")
+    }
+    assert set(funcs) == {"_module_path", "_read_pooling_mode", "_load_modules"}
+
+    # (a) each helper takes a revision parameter.
+    for name, fn in funcs.items():
+        arg_names = {a.arg for a in fn.args.args + fn.args.kwonlyargs}
+        assert "revision" in arg_names, f"{name} must accept a revision argument"
+
+    # (b) every download primitive inside the helpers forwards revision.
+    downloads = 0
+    for name, fn in funcs.items():
+        for node in ast.walk(fn):
+            if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Name)):
+                continue
+            if node.func.id not in ("hf_hub_download", "load_dir_path"):
+                continue
+            downloads += 1
+            assert any(
+                kw.arg == "revision" for kw in node.keywords
+            ), f"{node.func.id} in {name} must forward revision"
+    assert downloads >= 3, "expected the module-download primitives to be revision-guarded"
+
+    # (c) _load_modules threads revision into its internal _module_path / _read_pooling_mode calls.
+    internal = 0
+    for node in ast.walk(funcs["_load_modules"]):
+        if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)):
+            continue
+        if node.func.attr not in ("_module_path", "_read_pooling_mode"):
+            continue
+        internal += 1
+        assert any(
+            kw.arg == "revision" for kw in node.keywords
+        ), f"_load_modules must forward revision to {node.func.attr}"
+    assert internal >= 2, "expected _load_modules to call _module_path and _read_pooling_mode"
+
+    # (d) the from_pretrained fallback _module_path / _load_modules sites forward revision.
+    checked = 0
+    for node in ast.walk(tree):
+        if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)):
+            continue
+        if node.func.attr not in ("_module_path", "_load_modules"):
+            continue
+        cache_dir_kw = next((kw for kw in node.keywords if kw.arg == "cache_dir"), None)
+        if cache_dir_kw is None or "cache_folder" not in ast.dump(cache_dir_kw.value):
+            continue  # internal pass-through, not a fallback site
+        checked += 1
+        rev_kw = next((kw for kw in node.keywords if kw.arg == "revision"), None)
+        assert rev_kw is not None and "revision" in ast.dump(
+            rev_kw.value
+        ), f"{node.func.attr} fallback call must forward revision"
+    assert (
+        checked >= 2
+    ), "expected the fallback _module_path and _load_modules calls to forward revision"
+
+
+def test_st_fallback_model_load_resolves_env_cache():
+    """from_pretrained must resolve the warmed ST cache into kwargs['cache_dir'] before the FastModel weight load."""
+    import ast
+    import os
+
+    src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py")
+    with open(src_path, "r", encoding = "utf-8") as f:
+        tree = ast.parse(f.read())
+
+    def _resolves_st_cache(value_node):
+        # Resolution may be inline or in the assignment to an intermediate variable the value references.
+        dumped = ast.dump(value_node)
+        if "cache_folder" in dumped and "SENTENCE_TRANSFORMERS_HOME" in dumped:
+            return True
+        if isinstance(value_node, ast.Name):
+            for n in ast.walk(tree):
+                if isinstance(n, ast.Assign) and any(
+                    isinstance(t, ast.Name) and t.id == value_node.id for t in n.targets
+                ):
+                    d = ast.dump(n.value)
+                    if "cache_folder" in d and "SENTENCE_TRANSFORMERS_HOME" in d:
+                        return True
+        return False
+
+    resolved_lines = []
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.Assign):
+            continue
+        for tgt in node.targets:
+            if (
+                isinstance(tgt, ast.Subscript)
+                and isinstance(tgt.value, ast.Name)
+                and tgt.value.id == "kwargs"
+                and isinstance(tgt.slice, ast.Constant)
+                and tgt.slice.value == "cache_dir"
+                and _resolves_st_cache(node.value)
+            ):
+                resolved_lines.append(node.lineno)
+    assert resolved_lines, "from_pretrained must resolve the ST cache into kwargs['cache_dir']"
+
+    fastmodel_calls = [
+        n.lineno
+        for n in ast.walk(tree)
+        if isinstance(n, ast.Call)
+        and isinstance(n.func, ast.Attribute)
+        and n.func.attr == "from_pretrained"
+        and isinstance(n.func.value, ast.Name)
+        and n.func.value.id == "FastModel"
+    ]
+    assert fastmodel_calls, "expected a FastModel.from_pretrained call"
+    assert min(resolved_lines) < min(
+        fastmodel_calls
+    ), "kwargs['cache_dir'] must be resolved before the fallback FastModel weight load"
+
+
+def test_canonical_variant_model_weight_matches_transformers_names():
+    """The variant safetensors detector matches only canonical variant names, rejecting sidecars and wrong variants."""
+    f = U._is_canonical_variant_model_weight_safetensors
+    assert f("model.fp16.safetensors", "fp16") is True
+    assert f("model.fp16-00001-of-00002.safetensors", "fp16") is True
+    assert f("model-00001-of-00002.fp16.safetensors", "fp16") is True
+    assert f("model.safetensors.index.fp16.json", "fp16") is True
+    assert f("consolidated.fp16.safetensors", "fp16") is False
+    assert f("model.safetensors", "fp16") is False
+    assert f("model-00001-of-00002.safetensors", "fp16") is False
+    assert f("model.bf16.safetensors", "fp16") is False
+
+
+def test_variant_is_forwarded_to_downloader(capture):
+    """maybe_prefetch_hf_snapshot must forward variant to the downloader (absent a variant, nothing is forwarded)."""
+    _, st = capture(weights_at_root = True, use_safetensors = True, variant = "fp16")
+    assert st["variant"] == "fp16"
+    _, st = capture(weights_at_root = True, use_safetensors = True)
+    assert st["variant"] is None
+
+
+def test_variant_drops_bin_for_sharded_variant_safetensors(monkeypatch):
+    """A sharded variant safetensors is recognized, so its redundant variant .bin is dropped."""
+    _install_fake_model_info(
+        monkeypatch,
+        [
+            "model.fp16-00001-of-00002.safetensors",
+            "model.fp16-00002-of-00002.safetensors",
+            "pytorch_model.fp16-00001-of-00002.bin",
+        ],
+    )
+    ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True)
+    assert "*.bin" in ig
+
+
+def test_tokenizer_only_warms_extra_vocab_files(capture):
+    """tokenizer_only must warm SentencePiece / vocab / processor files, including a named jinja template."""
+    _, st = capture(tokenizer_only = True)
+    allow = st["allow_patterns"]
+    for name in (
+        "spm.model",
+        "normalizer.json",
+        "video_preprocessor_config.json",
+        "tokenizer.model.v3",
+    ):
+        assert name in allow, name
+    sample = [
+        "spm.model",
+        "normalizer.json",
+        "video_preprocessor_config.json",
+        "tokenizer.model.v3",
+        "additional_chat_templates/custom.jinja",
+    ]
+    kept = _filter(sample, allow, st["ignore_patterns"])
+    assert set(kept) == set(sample)
+
+
+def test_format_probe_runs_even_when_config_cached(capture, monkeypatch):
+    """A cached config.json must not skip the weight-format probe; model_info still drops the redundant .bin."""
+    import huggingface_hub
+
+    # Pretend config.json is cached (the AutoConfig side effect); this must not gate the probe.
+    monkeypatch.setattr(
+        huggingface_hub, "try_to_load_from_cache", lambda *a, **k: "/cache/config.json"
+    )
+    _install_fake_model_info(monkeypatch, ["model.safetensors", "pytorch_model.bin"])
+    _, st = capture(weights_at_root = True)
+    ig = st["ignore_patterns"] or []
+    assert "*.bin" in ig
+
+
+def test_optimizer_safetensors_does_not_drop_bin(monkeypatch):
+    """An optimizer.safetensors sidecar must not count as model safetensors, so the real .bin weights are kept."""
+    _install_fake_model_info(monkeypatch, ["pytorch_model.bin", "optimizer.safetensors"])
+    ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = True)
+    assert "*.bin" not in ig
+
+
+def test_model_safetensors_still_drops_bin(monkeypatch):
+    """Control for the optimizer case: a real model.safetensors next to pytorch_model.bin still drops the .bin."""
+    _install_fake_model_info(
+        monkeypatch, ["model.safetensors", "pytorch_model.bin", "optimizer.safetensors"]
+    )
+    ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = True)
+    assert "*.bin" in ig
+
+
+def test_whole_multi_component_snapshot_keeps_subdir_bin(monkeypatch):
+    """A whole multi-component snapshot must not drop *.bin (it would strip a subdir module's weight); a root load still does."""
+    _install_fake_model_info(monkeypatch, ["model.safetensors", "1_Dense/pytorch_model.bin"])
+    ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = False)
+    assert "*.bin" not in ig
+    ig_root = U._prefetch_ignore_patterns("org/repo", weights_at_root = True)
+    assert "*.bin" in ig_root
+
+
+def test_is_model_weight_safetensors_classification():
+    """Real model weights count; adapter / trainer-state sidecars do not."""
+    assert U._is_model_weight_safetensors("model.safetensors") is True
+    assert U._is_model_weight_safetensors("model-00001-of-00002.safetensors") is True
+    assert U._is_model_weight_safetensors("model.safetensors.index.json") is True
+    assert U._is_model_weight_safetensors("consolidated.safetensors") is True
+    assert U._is_model_weight_safetensors("adapter_model.safetensors") is False
+    assert U._is_model_weight_safetensors("optimizer.safetensors") is False
+    assert U._is_model_weight_safetensors("scheduler.safetensors") is False
+    assert U._is_model_weight_safetensors("rng_state_0.safetensors") is False
+
+
+def test_tokenizer_only_warms_slow_sentencepiece_vocab(capture):
+    """tokenizer_only must warm the slow-tokenizer SentencePiece / BPE vocab files AutoTokenizer fetches first."""
+    _, st = capture(tokenizer_only = True)
+    allow = st["allow_patterns"]
+    for name in (
+        "sentencepiece.bpe.model",
+        "source.spm",
+        "target.spm",
+        "bpe.codes",
+        "vocab.bpe",
+        "sentencepiece.model",
+        "vocab-src.json",
+        "vocab-tgt.json",
+    ):
+        assert name in allow, name
+
+
+def test_adapter_safetensors_check_scoped_to_root(monkeypatch):
+    """_adapter_repo_has_safetensors must only count a root adapter_model*.safetensors, not a subdir one."""
+    import huggingface_hub
+
+    class _Sib:
+        def __init__(self, name):
+            self.rfilename = name
+
+    class _Api:
+        def __init__(self, names):
+            self._names = names
+
+        def model_info(self, *a, **k):
+            return type("MI", (), {"siblings": [_Sib(n) for n in self._names]})()
+
+    # Subdir safetensors only -> not reported present.
+    monkeypatch.setattr(
+        huggingface_hub,
+        "HfApi",
+        lambda: _Api(
+            ["adapter_config.json", "adapter_model.bin", "checkpoint-5/adapter_model.safetensors"]
+        ),
+    )
+    assert U._adapter_repo_has_safetensors("org/repo") is False
+    # Root safetensors -> reported present.
+    monkeypatch.setattr(
+        huggingface_hub,
+        "HfApi",
+        lambda: _Api(["adapter_config.json", "adapter_model.safetensors"]),
+    )
+    assert U._adapter_repo_has_safetensors("org/repo") is True
+
+
+def test_gguf_file_warm_keeps_gguf(capture):
+    """A gguf_file load allow-lists that GGUF while not pulling other quants the repo publishes."""
+    _, st = capture(weights_at_root = True, gguf_file = "model-Q4_K_M.gguf")
+    allow = st["allow_patterns"]
+    ig = st["ignore_patterns"]
+    assert allow is not None and "model-Q4_K_M.gguf" in allow
+    sample = [
+        "model-Q4_K_M.gguf",
+        "model-Q8_0.gguf",
+        "config.json",
+        "tokenizer.json",
+    ]
+    kept = _filter(sample, allow, ig)
+    assert "model-Q4_K_M.gguf" in kept
+    assert "config.json" in kept
+    assert "model-Q8_0.gguf" not in kept
+
+
+# ----- Finding Q: adapter weight-format selection -----
+
+
+def test_adapter_only_prefers_safetensors_over_bin(capture, monkeypatch):
+    """A mixed-format adapter repo warms only the safetensors PeftModel reads, not both formats."""
+    _install_fake_model_info(
+        monkeypatch, ["adapter_config.json", "adapter_model.safetensors", "adapter_model.bin"]
+    )
+    _, st = capture(adapter_only = True)
+    ig = st["ignore_patterns"]
+    assert ig is not None and "adapter_model*.bin" in ig
+    kept = _filter(
+        ["adapter_config.json", "adapter_model.safetensors", "adapter_model.bin"],
+        st["allow_patterns"],
+        ig,
+    )
+    assert "adapter_model.safetensors" in kept
+    assert "adapter_model.bin" not in kept
+
+
+def test_adapter_only_bin_only_keeps_bin(capture, monkeypatch):
+    """A .bin-only adapter repo must keep adapter_model.bin (no safetensors found -> both formats eligible)."""
+    _install_fake_model_info(monkeypatch, ["adapter_config.json", "adapter_model.bin"])
+    _, st = capture(adapter_only = True)
+    kept = _filter(
+        ["adapter_config.json", "adapter_model.bin"], st["allow_patterns"], st["ignore_patterns"]
+    )
+    assert "adapter_model.bin" in kept
+
+
+def test_adapter_only_explicit_use_safetensors_false_keeps_bin(capture):
+    """An explicit use_safetensors=False forces the .bin form without a model_info call."""
+    _, st = capture(adapter_only = True, use_safetensors = False)
+    ig = st["ignore_patterns"]
+    assert ig is not None and "adapter_model*.safetensors" in ig
+    kept = _filter(
+        ["adapter_config.json", "adapter_model.safetensors", "adapter_model.bin"],
+        st["allow_patterns"],
+        ig,
+    )
+    assert "adapter_model.bin" in kept
+    assert "adapter_model.safetensors" not in kept
+
+
+def test_gguf_file_with_subfolder_warms_subfolder_path(capture):
+    """gguf_file + subfolder: the warm allow-lists <subfolder>/<gguf_file>, not the bare root name."""
+    _, st = capture(weights_at_root = True, gguf_file = "model-Q4_K_M.gguf", subfolder = "gguf")
+    allow = st["allow_patterns"]
+    assert "gguf/model-Q4_K_M.gguf" in allow
+    kept = _filter(["gguf/model-Q4_K_M.gguf", "config.json"], allow, st["ignore_patterns"])
+    assert "gguf/model-Q4_K_M.gguf" in kept and "config.json" in kept
+
+
+def test_from_tf_root_load_ignores_nested_h5(capture):
+    """A from_tf root load keeps the root .h5 but drops nested .h5 / .msgpack checkpoints."""
+    _, st = capture(weights_at_root = True, from_tf = True)
+    ig = st["ignore_patterns"]
+    assert "*/*.h5" in ig and "*/*.msgpack" in ig
+    kept = _filter(["model.h5", "checkpoint-1/model.h5", "config.json"], st["allow_patterns"], ig)
+    assert "model.h5" in kept
+    assert "checkpoint-1/model.h5" not in kept
+
+
+def test_sentence_transformer_from_pretrained_is_prefetch_wired():
+    """from_pretrained must call maybe_prefetch_hf_snapshot as an unconditional top-level statement before any return."""
+    import ast
+    import os
+
+    src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py")
+    with open(src_path, "r", encoding = "utf-8") as f:
+        tree = ast.parse(f.read())
+    cls = next(
+        n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "FastSentenceTransformer"
+    )
+    fp = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "from_pretrained")
+
+    def _prefetch_call(node):
+        # a bare call statement, or one whose return is captured (e.g. _st_prefetched = ...)
+        value = node.value if isinstance(node, (ast.Expr, ast.Assign)) else None
+        if (
+            isinstance(value, ast.Call)
+            and isinstance(value.func, ast.Name)
+            and value.func.id == "maybe_prefetch_hf_snapshot"
+        ):
+            return value
+        return None
+
+    prefetch_pos = next((i for i, n in enumerate(fp.body) if _prefetch_call(n)), None)
+    return_pos = next((i for i, n in enumerate(fp.body) if isinstance(n, ast.Return)), len(fp.body))
+    assert (
+        prefetch_pos is not None
+    ), "from_pretrained must call maybe_prefetch_hf_snapshot at top level"
+    assert prefetch_pos < return_pos, "prefetch must run before any top-level return"
+    # local_files_only must be forwarded so an offline load does not start a Hub download.
+    prefetch_call = _prefetch_call(fp.body[prefetch_pos])
+    assert "local_files_only" in {
+        kw.arg for kw in prefetch_call.keywords
+    }, "prefetch must forward local_files_only"
+
+
+def test_st_module_download_forwards_cache_folder():
+    """_load_modules must forward the custom cache_folder into load_dir_path so per-module subdirs read the warmed cache."""
+    import ast
+    import os
+
+    src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py")
+    with open(src_path, "r", encoding = "utf-8") as f:
+        tree = ast.parse(f.read())
+    calls = [
+        n
+        for n in ast.walk(tree)
+        if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == "load_dir_path"
+    ]
+    assert calls, "expected a load_dir_path call in sentence_transformer.py"
+    assert all(
+        "cache_folder" in {kw.arg for kw in c.keywords} for c in calls
+    ), "every load_dir_path call must forward cache_folder"
+
+
+def test_st_native_sentence_transformer_calls_forward_cache_folder():
+    """Every native SentenceTransformer(model_name, ...) load must forward cache_folder; a modules-based build needs none."""
+    import ast
+    import os
+
+    src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py")
+    with open(src_path, "r", encoding = "utf-8") as f:
+        tree = ast.parse(f.read())
+    weight_loading_calls = []
+    for n in ast.walk(tree):
+        if not (
+            isinstance(n, ast.Call)
+            and isinstance(n.func, ast.Name)
+            and n.func.id == "SentenceTransformer"
+        ):
+            continue
+        kw_names = {kw.arg for kw in n.keywords}
+        # A modules-based build downloads nothing; only a repo-name load reads the cache.
+        if "modules" in kw_names:
+            continue
+        weight_loading_calls.append(n)
+    assert (
+        weight_loading_calls
+    ), "expected a repo-name SentenceTransformer load in sentence_transformer.py"
+    # cache_folder is forwarded explicitly or via a **kwargs unpacking (kw.arg == None).
+    for c in weight_loading_calls:
+        kw_names = {kw.arg for kw in c.keywords}
+        forwards = "cache_folder" in kw_names or None in kw_names
+        assert forwards, (
+            "a repo-name SentenceTransformer load must forward cache_folder "
+            f"(explicitly or via **kwargs) at line {c.lineno}"
+        )

tests/test_synthetic_chunk_data.py

@@ -104,10 +104,36 @@ def test_chunk_data_rejects_overlap_not_smaller_than_chunk():
         os.unlink(path)
 
 
+def test_chunk_data_uninitialized_error_names_real_class():
+    # Without max_seq_length the guard tells the user which method to call first.
+    # The message must name the real class (SyntheticDataKit) so copying it works;
+    # a misspelling would raise NameError when the user follows it verbatim.
+    kit = SyntheticDataKit.__new__(SyntheticDataKit)
+    kit.tokenizer = _MockTokenizer()  # max_seq_length intentionally unset
+    with tempfile.NamedTemporaryFile("w", suffix = ".txt", delete = False) as f:
+        f.write("word " * 50)
+        path = f.name
+    try:
+        try:
+            kit.chunk_data(filename = path)
+            raise AssertionError("expected RuntimeError when max_seq_length is unset")
+        except RuntimeError as e:
+            msg = str(e)
+            assert (
+                "SyntheticDataKit.from_pretrained" in msg
+            ), f"error must name SyntheticDataKit.from_pretrained, got: {msg}"
+            assert (
+                "SynthetidDataKit" not in msg
+            ), f"error must not misspell the class name, got: {msg}"
+    finally:
+        os.unlink(path)
+
+
 if __name__ == "__main__":
     test_chunk_data_keeps_single_chunk_document()
     test_chunk_data_still_splits_long_document()
     test_chunk_data_empty_document_yields_no_chunks()
     test_chunk_data_short_document_is_not_split_into_fragments()
     test_chunk_data_rejects_overlap_not_smaller_than_chunk()
+    test_chunk_data_uninitialized_error_names_real_class()
     print("OK")

tests/utils/test_rope_scaling_drift.py

@@ -31,6 +31,7 @@ requires_cuda = pytest.mark.skipif(
 
 REPO_ROOT = Path(__file__).resolve().parents[2]
 LLAMA_PY = REPO_ROOT / "unsloth" / "models" / "llama.py"
+LOADER_PY = REPO_ROOT / "unsloth" / "models" / "loader.py"
 
 CLASS_NAME = "LlamaRotaryEmbedding"
 
@@ -78,42 +79,88 @@ def _config_branch(init_fn):
     return None
 
 
+def _iter_names_and_calls(node):
+    """(attribute/string names, bare-name calls, method-call attrs) under node."""
+    names, calls, call_attrs = set(), set(), set()
+    for sub in ast.walk(node):
+        if isinstance(sub, ast.Attribute):
+            names.add(sub.attr)
+        elif isinstance(sub, ast.Constant) and isinstance(sub.value, str):
+            names.add(sub.value)
+        elif isinstance(sub, ast.Call):
+            if isinstance(sub.func, ast.Name):
+                calls.add(sub.func.id)
+            elif isinstance(sub.func, ast.Attribute):
+                call_attrs.add(sub.func.attr)
+    return names, calls, call_attrs
+
+
+def _find_method(source_path, class_name, method_name):
+    for node in ast.walk(ast.parse(source_path.read_text())):
+        if isinstance(node, ast.ClassDef) and node.name == class_name:
+            for sub in node.body:
+                if isinstance(sub, ast.FunctionDef) and sub.name == method_name:
+                    return sub
+    return None
+
+
+def _find_function(source_path, function_name):
+    for node in ast.walk(ast.parse(source_path.read_text())):
+        if isinstance(node, ast.FunctionDef) and node.name == function_name:
+            return node
+    return None
+
+
 def test_config_path_inspects_rope_scaling():
     init_fn = _load_class_init()
-    branch = _config_branch(init_fn)
-    assert branch is not None, (
-        f"{CLASS_NAME}.__init__ no longer has an `if config is not None:` "
-        "branch; the config constructor path must read config.rope_scaling so "
-        "scaled models (llama3/linear/longrope) are not silently unscaled "
-        "(issue #2405)"
-    )
+    # inv_freq is derived through the shared _unsloth_recompute_inv_freq helper
+    # (or still inlined in the config branch on older layouts); whichever scope
+    # holds the scaling must read config.rope_scaling and call
+    # _compute_config_rope_inv_freq, else scaled models run unscaled (#2405).
+    _, _, init_call_attrs = _iter_names_and_calls(init_fn)
+    scope = _find_method(LLAMA_PY, CLASS_NAME, "_unsloth_recompute_inv_freq")
+    if scope is not None:
+        assert "_unsloth_recompute_inv_freq" in init_call_attrs, (
+            f"{CLASS_NAME}.__init__ no longer derives inv_freq via "
+            "_unsloth_recompute_inv_freq; keep the constructor wired to the "
+            "shared scaling helper or scaled configs silently lose RoPE scaling "
+            "(issue #2405)."
+        )
+    else:
+        scope = _config_branch(init_fn)
+        assert scope is not None, (
+            f"{CLASS_NAME}.__init__ has neither a _unsloth_recompute_inv_freq "
+            "helper nor an `if config is not None:` branch; the config path must "
+            "apply llama3/linear/longrope scaling (issue #2405)."
+        )
 
-    names = set()
-    for stmt in branch.body:
-        for sub in ast.walk(stmt):
-            if isinstance(sub, ast.Attribute):
-                names.add(sub.attr)
-            elif isinstance(sub, ast.Constant) and isinstance(sub.value, str):
-                names.add(sub.value)
+    names, called, _ = _iter_names_and_calls(scope)
     assert "rope_scaling" in names, (
-        f"{CLASS_NAME}.__init__ config path does not reference `rope_scaling`. "
-        "When a rotary class is built straight from a config (the path modern "
-        "transformers takes, since rotary moved to LlamaModel), the llama3 / "
-        "linear / longrope scaling must still be applied; otherwise long inputs "
-        "produce repeated-pattern gibberish (issue #2405)."
+        f"{CLASS_NAME} inv_freq computation does not reference `rope_scaling`; "
+        "scaled models (llama3/linear/longrope) would run unscaled and produce "
+        "repeated-pattern gibberish past the original context (issue #2405)."
+    )
+    assert "_compute_config_rope_inv_freq" in called, (
+        f"{CLASS_NAME} inv_freq computation no longer calls "
+        "_compute_config_rope_inv_freq; keep it wired or scaled configs silently "
+        "lose RoPE scaling again (issue #2405)."
     )
 
-    called = {
-        sub.func.id
-        for stmt in branch.body
-        for sub in ast.walk(stmt)
-        if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)
-    }
-    assert "_compute_config_rope_inv_freq" in called, (
-        f"{CLASS_NAME}.__init__ config path no longer calls "
-        "_compute_config_rope_inv_freq; the CPU behavioral tests below cover "
-        "that helper directly, so the constructor must stay wired to it or "
-        "scaled configs silently lose RoPE scaling again (issue #2405)."
+
+def test_v5_repair_reuses_recompute():
+    # transformers v5 blanks non-persistent buffers on load, so
+    # loader._fix_rope_inv_freq rebuilds inv_freq; it must reuse the scaled
+    # recompute, since an unscaled rebuild re-drops llama3 scaling (#2405).
+    fix_fn = _find_function(LOADER_PY, "_fix_rope_inv_freq")
+    assert fix_fn is not None, (
+        "loader._fix_rope_inv_freq not found; if it was renamed, update this "
+        "guard so the v5 rope repair keeps applying config scaling (issue #2405)."
+    )
+    _, _, call_attrs = _iter_names_and_calls(fix_fn)
+    assert "_unsloth_recompute_inv_freq" in call_attrs, (
+        "loader._fix_rope_inv_freq no longer rebuilds inv_freq via "
+        "_unsloth_recompute_inv_freq; transformers v5 blanks the buffer on load "
+        "and an unscaled rebuild re-drops llama3 scaling (issue #2405)."
     )
 
 
@@ -189,6 +236,27 @@ def test_default_rope_type_matches_vanilla_inv_freq():
     )
 
 
+def test_recompute_helper_scales_on_cpu():
+    # Exercise the exact method loader._fix_rope_inv_freq calls, without CUDA.
+    from unsloth.models.llama import LlamaRotaryEmbedding, _get_rope_theta
+
+    def recompute(config):
+        rot = object.__new__(LlamaRotaryEmbedding)
+        rot.attention_scaling = 1.0
+        rot.base = _get_rope_theta(config, 10000.0)
+        rot.dim = config.head_dim
+        rot._unsloth_rope_config = config
+        return rot._unsloth_recompute_inv_freq().float().cpu()
+
+    config = _make_config(LLAMA3_ROPE_SCALING)
+    assert torch.allclose(
+        recompute(config), _reference_inv_freq(config, "llama3"), rtol = 1e-4, atol = 1e-6
+    ), "_unsloth_recompute_inv_freq dropped llama3 scaling (issue #2405)."
+    assert torch.allclose(
+        recompute(_make_config(None)), _vanilla_inv_freq(), rtol = 1e-4, atol = 1e-6
+    ), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled."
+
+
 def _cos_at_position(rot, position):
     """cos row at one position, built like _set_cos_sin_cache but CPU-only."""
     inv_freq = rot.inv_freq.float().cpu()

unsloth/_compressed_quantize.py

@@ -203,6 +203,7 @@ def main():
     ap.add_argument("--max-seq-length", type = int, default = 2048)
     ap.add_argument("--is-vlm", action = "store_true")
     ap.add_argument("--trust-remote-code", action = "store_true")
+    ap.add_argument("--trust-remote-code-tokenizer", action = "store_true")
     ap.add_argument("--variant", default = "", help = "weight-filename variant for the output shards")
     args = ap.parse_args()
 
@@ -232,7 +233,11 @@ def main():
     model.eval()
     # A tokenizer may be absent if the caller saved it separately; only calibration needs one.
     try:
-        tokenizer = auto_proc.from_pretrained(args.model, trust_remote_code = args.trust_remote_code)
+        # The tokenizer/processor has its own trust flag: consent for one component must not
+        # let the other's custom code run.
+        tokenizer = auto_proc.from_pretrained(
+            args.model, trust_remote_code = args.trust_remote_code_tokenizer
+        )
     except Exception:
         if args.needs_calibration:
             raise RuntimeError(

unsloth/dataprep/synthetic.py

@@ -391,7 +391,7 @@ class SyntheticDataKit:
         assert os.path.exists(filename)
         assert hasattr(self, "tokenizer")
         if not hasattr(self, "max_seq_length"):
-            raise RuntimeError("Please use SynthetidDataKit.from_pretrained(...) first!")
+            raise RuntimeError("Please use SyntheticDataKit.from_pretrained(...) first!")
         if not hasattr(self, "overlap") or not hasattr(self, "max_generation_tokens"):
             raise RuntimeError("Please use prepare_qa_generation first!")
 

unsloth/kernels/fp8.py

@@ -68,7 +68,9 @@ def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr):
     n = tl.cdiv(N, BLOCK_SIZE)
     offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
     offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    offs = offs_m[:, None] * N + offs_n[None, :]
+    # tl.arange is int32, so offs_m * N overflows for tensors with more than
+    # 2**31 elements (e.g. flattened MoE expert stacks); index in int64.
+    offs = offs_m[:, None].to(tl.int64) * N + offs_n[None, :].to(tl.int64)
     mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
     x = tl.load(x_ptr + offs, mask = mask).to(tl.float32)
     s = tl.load(s_ptr + pid_m * n + pid_n)
@@ -327,11 +329,42 @@ fp8_block_matmul = (
 )
 
 
+def _blockwise_weight_dequant_any_shape(weight, weight_scale, block_size, out_dtype):
+    """Blockwise fp8 weight dequant for any shape: triton when the weight tiles
+    evenly into block_size, else a torch-native per-block scale expansion."""
+    m, n = weight.shape
+    if weight_scale.dtype not in (torch.float32, torch.float16, torch.bfloat16):
+        weight_scale = weight_scale.to(torch.float32)  # e.g. float8_e8m0fnu scales break triton
+    if weight_scale.numel() == 1:
+        # Per-tensor scale: the normal forward stashes the un-expanded scalar,
+        # which repeat_interleave cannot grow to (m, n). Scale directly.
+        return (weight.to(torch.float32) * weight_scale.float()).to(out_dtype)
+    if m % block_size[0] != 0 or n % block_size[1] != 0 or block_size[0] != block_size[1]:
+        # Uneven tiling, or rectangular blocks. The triton kernel uses a single
+        # BLOCK_SIZE for both axes and derives the column scale stride from it, so
+        # it mis-indexes the scale when block_size[0] != block_size[1]. Expand the
+        # per-block scales in torch, which handles both dimensions independently.
+        s_full = weight_scale.repeat_interleave(block_size[0], 0)[:m]
+        s_full = s_full.repeat_interleave(block_size[1], 1)[:, :n]
+        return (weight.to(torch.float32) * s_full).to(out_dtype)
+    # Even tiling with square blocks: block-quant dequant with the real block size
+    # (weight_dequant would silently default to 128 and dequantize wrongly).
+    return weight_dequant_block(weight, weight_scale, block_size = block_size[0], dtype = out_dtype)
+
+
 class FP8BlockQuantLinear(torch.autograd.Function):
     @staticmethod
     def forward(ctx, X, weight, weight_scale):
         m, n = weight.shape
 
+        if weight_scale.dtype not in (torch.float32, torch.float16, torch.bfloat16):
+            # Upcast (e.g. e8m0) returns a fresh tensor and drops any Python
+            # attribute, so carry block_size across the cast for the lookup below.
+            _scale_block_size = getattr(weight_scale, "block_size", None)
+            weight_scale = weight_scale.to(torch.float32)  # e8m0 scales break triton dtype mapping
+            if _scale_block_size is not None:
+                weight_scale.block_size = _scale_block_size
+
         # Original scale, saved for backward before any transformation
         original_weight_scale = weight_scale
 
@@ -360,6 +393,18 @@ class FP8BlockQuantLinear(torch.autograd.Function):
         if not weight.is_contiguous():
             weight = weight.contiguous()
 
+        if X.shape[-1] % block_size[1] != 0:
+            # Hidden dim not divisible by the activation block: dequant + plain matmul.
+            # Use the original (un-expanded) scale so a scalar per-tensor scale keeps
+            # the fast scalar path in both forward and backward.
+            W_deq = _blockwise_weight_dequant_any_shape(
+                weight, original_weight_scale, block_size, X.dtype
+            )
+            ctx.weight = weight
+            ctx.weight_scale = original_weight_scale
+            ctx.block_size = block_size
+            return torch_matmul(X, W_deq.T).to(X.dtype)
+
         qinput, scale = act_quant(X, block_size[1])
         output = fp8_block_matmul(
             qinput,
@@ -371,11 +416,14 @@ class FP8BlockQuantLinear(torch.autograd.Function):
         )
         ctx.weight = weight
         ctx.weight_scale = original_weight_scale  # Save original for backward
+        ctx.block_size = block_size
         return output.to(X.dtype)
 
     @staticmethod
     def backward(ctx, grad_output):
-        W_deq = weight_dequant(ctx.weight, ctx.weight_scale)
+        W_deq = _blockwise_weight_dequant_any_shape(
+            ctx.weight, ctx.weight_scale, ctx.block_size, grad_output.dtype
+        )
         grad_X = torch_matmul(grad_output, W_deq)
         del W_deq
         return grad_X, None, None

unsloth/models/_utils.py

@@ -83,8 +83,10 @@ __all__ = [
     "verify_fp8_support_if_applicable",
     "_get_inference_mode_context_manager",
     "hf_login",
+    "maybe_prefetch_hf_snapshot",
     "is_moe_model",
     "get_moe_target_parameters",
+    "_select_moe_detection_targets",
     "make_fast_generate_wrapper",
     "_mark_unsloth_disable_data_parallel",
     "_patch_transformers_trainer_data_parallel",
@@ -421,6 +423,18 @@ def apply_unsloth_gradient_checkpointing(use_gradient_checkpointing, max_seq_len
 _FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h", "modernbert")
 _FLEX_PREFERRED_MODELS = ("gemma3", "gemma3_text", "shieldgemma2")
 _SDPA_EXCLUDED_MODELS = ("gpt_oss",)
+# The loader (loader.py) forces supports_sdpa=False for these because their bundled
+# SDPA modules are wrong. Kept here, not in loader.py, so _is_sdpa_excluded can honor
+# them without a loader -> _utils import cycle (loader.py already imports from _utils
+# and re-exports this name for callers like sentence_transformer.py). Entries are matched
+# as substrings against a comma-joined model_types string ending in a comma, so "gemma3,"
+# matches a distinct "gemma3" entry but not "gemma3n", and "gemma3_text" matches the
+# EmbeddingGemma text model.
+DISABLE_SDPA_MODEL_NAMES = [
+    "gemma3,",  # Add comma bc gemma3 will match gemma3n
+    "gemma3_text",  # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore
+    "gpt_oss",
+]
 _FLASH_EXCLUDED_MODELS = ("gpt_oss",)
 _EAGER_ONLY_PREFIXES = ("gemma3n",)
 _FLASH_ATTENTION_MAX_HEAD_DIM = 256
@@ -431,8 +445,23 @@ def _is_flex_excluded(model_type):
     return model_type in _FLEX_EXCLUDED_MODELS
 
 
+def _is_sdpa_disabled_by_name(model_type):
+    # Mirror the loader's DISABLE_SDPA_MODEL_NAMES check: loader.py builds
+    # model_types_all = ",".join(model_types) + "," and tests `name in model_types_all`.
+    # Rebuild the same trailing-comma form for a single model_type so the match is
+    # identical (e.g. "gemma3," matches "gemma3" but not "gemma3n", and "gemma3_text"
+    # still matches "gemma3_text").
+    model_types_all = model_type.lower() + ","
+    return any(name.lower() in model_types_all for name in DISABLE_SDPA_MODEL_NAMES)
+
+
 def _is_sdpa_excluded(model_type):
-    return model_type in _SDPA_EXCLUDED_MODELS
+    # SDPA is known-broken for these models, so an explicit sdpa request must not
+    # re-enable it. Two sources: _SDPA_EXCLUDED_MODELS (resolver-level, e.g. gpt_oss)
+    # and DISABLE_SDPA_MODEL_NAMES (loader-level, e.g. gemma3 / gemma3_text, which the
+    # loader also forces to supports_sdpa=False).
+    lowered = model_type.lower()
+    return lowered in _SDPA_EXCLUDED_MODELS or _is_sdpa_disabled_by_name(lowered)
 
 
 def _is_flash_excluded(model_type):
@@ -608,6 +637,12 @@ def _disable_flash_attention_if_needed(
     if disable_reason is None:
         return attn_implementation
 
+    # Only an implementation passed by the caller counts as an explicit request.
+    # Values read from the config are synthesized by the loaders (the language path
+    # seeds the config with attn_implementation="sdpa") or come from Transformers
+    # defaults, so they must not be treated as a deliberate user choice.
+    explicit_request = attn_implementation
+
     requested_attn_implementation = attn_implementation
     if requested_attn_implementation is None:
         requested_attn_implementation = _config_get(config, "_attn_implementation", None)
@@ -617,6 +652,20 @@ def _disable_flash_attention_if_needed(
     if requested_attn_implementation == "eager":
         return _set_attn_impl(config, "eager")
 
+    model_type = _config_get(config, "model_type", "")
+
+    # The disable reason is flash-specific: honor an explicit non-flash request from
+    # the caller instead of downgrading it. SDPA is honored unless the model's SDPA is
+    # known-broken - _SDPA_EXCLUDED_MODELS (e.g. gpt_oss) or DISABLE_SDPA_MODEL_NAMES
+    # (e.g. gemma3 / gemma3_text); flex_attention
+    # is honored only when it is actually usable, since supports_flex_attention already
+    # rejects the excluded/broken/unavailable configs. This keeps an explicit request
+    # from selecting a backend the repo marks as wrong.
+    if explicit_request == "sdpa" and not _is_sdpa_excluded(model_type.lower()):
+        return _set_attn_impl(config, "sdpa")
+    if explicit_request == "flex_attention" and supports_flex_attention:
+        return _set_attn_impl(config, "flex_attention")
+
     if supports_sdpa:
         fallback_attn_implementation = "sdpa"
     elif supports_flex_attention:
@@ -629,7 +678,6 @@ def _disable_flash_attention_if_needed(
             if _is_flash_attention_requested(requested_attn_implementation)
             else "flash_attention_2"
         )
-        model_type = _config_get(config, "model_type", "")
         warning_key = (
             model_type,
             logged_attn_implementation,
@@ -843,7 +891,19 @@ def resolve_attention_implementation(
         final_attn_impl = requested_attn_implementation
         _set_attn_impl(config, final_attn_impl)
 
-    if not supports_sdpa and final_attn_impl == "sdpa":
+    # A caller who explicitly passes requested_attn_implementation="sdpa" keeps it even
+    # on a conservatively unsupported model, mirroring _disable_flash_attention_if_needed
+    # which honors an explicit sdpa request. The exception is a model whose SDPA is
+    # known-broken - _SDPA_EXCLUDED_MODELS (e.g. gpt_oss) or DISABLE_SDPA_MODEL_NAMES
+    # (e.g. gemma3 / gemma3_text, which the loader also forces to supports_sdpa=False):
+    # an explicit request must not re-enable it, so it still downgrades to eager, just
+    # like flex falls back for _FLEX_EXCLUDED_MODELS. A synthesized/default sdpa
+    # (requested is None, so the value came from the model resolution above or the
+    # config) also downgrades.
+    honor_explicit_sdpa = requested_attn_implementation == "sdpa" and not _is_sdpa_excluded(
+        model_type
+    )
+    if not supports_sdpa and final_attn_impl == "sdpa" and not honor_explicit_sdpa:
         print(
             f"Unsloth: {(model_type_name or 'model').title()} does not support SDPA - switching to fast eager."
         )
@@ -905,6 +965,411 @@ logging.getLogger("transformers.tokenization_utils_base").setLevel(logging.CRITI
 TORCHAO_MSG = "Error: torchao not found, please install with `pip install torchao`"
 
 
+# Artifacts a Transformers/PEFT load never reads (ONNX/TF/Flax/CoreML/GGUF/training state), skipped
+# when prewarming so a mixed-format repo is not pulled in full.
+_PREFETCH_IGNORE_PATTERNS = (
+    "*.onnx",
+    "onnx/*",
+    "*.h5",
+    "*.msgpack",
+    "*.tflite",
+    "coreml/*",
+    "*.mlpackage/*",
+    "*.mlmodel",
+    "*.gguf",
+    # Training / checkpoint formats from_pretrained never reads.
+    "*.pt",
+    "*.pth",
+    "*.ckpt",
+    "optimizer.*",
+    "scheduler.*",
+    "rng_state*",
+    "trainer_state.json",
+    "events.out.tfevents*",
+    "checkpoint-*/*",
+)
+
+
+# Repo-root tokenizer / config / processor files from_pretrained reads from root even when weights
+# load from a subfolder. Exact names (no wildcard) so they match only root-level files.
+_ROOT_AUX_PREFETCH_PATTERNS = (
+    "config.json",
+    "generation_config.json",
+    "tokenizer_config.json",
+    "tokenizer.json",
+    "tokenizer.model",
+    "special_tokens_map.json",
+    "added_tokens.json",
+    "vocab.json",
+    "vocab.txt",
+    "merges.txt",
+    "spiece.model",
+    # More VOCAB_FILES_NAMES the slow tokenizer may fetch (DeBERTa-v2, Whisper, Mistral, XLM-R/mBART, Marian, FSMT/XLM, GPT-2).
+    "spm.model",
+    "normalizer.json",
+    "tokenizer.model.v3",
+    "sentencepiece.bpe.model",
+    "source.spm",
+    "target.spm",
+    "bpe.codes",
+    "vocab.bpe",
+    # More VOCAB_FILES_NAMES (RemBERT, FSMT) a distinct-tokenizer-repo warm must cache too.
+    "sentencepiece.model",
+    "vocab-src.json",
+    "vocab-tgt.json",
+    "chat_template.jinja",
+    "chat_template.json",
+    # chat_template="<name>" fetches additional_chat_templates/<name>.jinja.
+    "additional_chat_templates/*.jinja",
+    "preprocessor_config.json",
+    "processor_config.json",
+    "video_preprocessor_config.json",  # Qwen2.5-VL-style video processors
+    # trust_remote_code auto_map can name any module, so warm every *.py (tiny; none in a non-remote repo).
+    "*.py",
+    "*.tiktoken",  # tiktoken vocab (e.g. Qwen's qwen.tiktoken)
+)
+
+
+# Files a PEFT adapter load reads: config + weights (glob covers sharded adapters). Any merged
+# full-model weights the repo also ships match none of these.
+_ADAPTER_PREFETCH_PATTERNS = (
+    "adapter_config.json",
+    "adapter_model*",
+)
+
+
+# Weight files in a SUBDIRECTORY. A bare root load reads only root weights, so ignoring these drops
+# alternate-precision/experimental dirs (fp16/, experimental/). "*/*" spans "/" (HF fnmatch), so nested
+# weights match while root "model.safetensors" is kept. Only applied when weights_at_root (diffusion
+# keeps weights in subfolders).
+_SUBDIR_WEIGHT_IGNORE_PATTERNS = (
+    "*/*.safetensors",
+    "*/*.bin",
+    "*/*.h5",
+    "*/*.msgpack",
+    "*/*.pt",
+    "*/*.pth",
+)
+
+
+def _in_requested_load_scope(filename, subfolder):
+    """True if *filename* is in the location being loaded (*subfolder*, else root). Scopes the ".bin is
+    redundant when safetensors exist" test so a .bin-only subfolder keeps its .bin."""
+    filename = filename.replace("\\", "/")
+    if isinstance(subfolder, str) and subfolder.strip("/"):
+        return filename.startswith(subfolder.strip("/") + "/")
+    return "/" not in filename  # root load: no directory component
+
+
+# .safetensors training-state files that are NOT model weights (e.g. optimizer.safetensors next to a
+# real pytorch_model.bin); counting them as "model safetensors present" would drop the needed .bin.
+_NON_MODEL_WEIGHT_STEMS = frozenset(
+    {
+        "optimizer",
+        "scheduler",
+        "scaler",
+        "rng_state",
+        "training_args",
+    }
+)
+
+
+def _is_model_weight_safetensors(filename):
+    """True if *filename* is a model-weights safetensors, not a PEFT adapter/sidecar
+    (adapter_model.safetensors) or trainer-state (optimizer.safetensors). Only a real one proves the
+    .bin redundant; counting a sidecar would wrongly drop the needed .bin (fetched then without Xet fallback)."""
+    name = filename.replace("\\", "/").rsplit("/", 1)[-1]
+    if not name.endswith((".safetensors", ".safetensors.index.json")):
+        return False
+    if name.startswith("adapter_"):
+        return False
+    # Stem before first dot: "optimizer.safetensors" -> "optimizer" (real shards kept); rng_state via prefix.
+    stem = name.split(".", 1)[0].lower()
+    if stem in _NON_MODEL_WEIGHT_STEMS or stem.startswith("rng_state"):
+        return False
+    return True
+
+
+def _is_canonical_variant_model_weight_safetensors(filename, variant):
+    """True for a canonical model-weights safetensors carrying the requested *variant*, in the forms
+    transformers reads (single, either numbered-shard layout, or the index). Strict (base must be
+    "model"): a sidecar like consolidated.<variant>.safetensors does not prove the variant .bin redundant."""
+    base = filename.replace("\\", "/").rsplit("/", 1)[-1]
+    v = re.escape(variant)
+    return bool(
+        re.match(
+            rf"^(?:model\.{v}\.safetensors"
+            rf"|model\.{v}-\d{{5}}-of-\d{{5}}\.safetensors"
+            rf"|model-\d{{5}}-of-\d{{5}}\.{v}\.safetensors"
+            rf"|model\.safetensors\.index\.{v}\.json)$",
+            base,
+        )
+    )
+
+
+_CANONICAL_MODEL_WEIGHT_SAFETENSORS_RE = re.compile(
+    r"^(?:model\.safetensors|model-\d{5}-of-\d{5}\.safetensors|model\.safetensors\.index\.json)$"
+)
+
+
+def _is_canonical_model_weight_safetensors(filename):
+    """True for a canonical (non-variant) model-weights safetensors a default load reads (model.safetensors,
+    a numbered shard, or the index). Strict: an unrecognized name keeps both formats, so a variant-only
+    safetensors + pytorch_model.bin repo never has its .bin dropped for a no-variant load."""
+    name = filename.replace("\\", "/").rsplit("/", 1)[-1]
+    return bool(_CANONICAL_MODEL_WEIGHT_SAFETENSORS_RE.match(name))
+
+
+def _adapter_repo_has_safetensors(
+    model_name,
+    *,
+    token = None,
+    revision = None,
+):
+    """Best-effort: does the adapter repo ship a root safetensors adapter weight (making the .bin
+    redundant)? Scoped to root adapter_model* files; any failure returns False."""
+    try:
+        from huggingface_hub import HfApi
+        siblings = HfApi().model_info(model_name, revision = revision, token = token).siblings or []
+        return any(
+            "/" not in sibling.rfilename.replace("\\", "/")  # root only
+            and sibling.rfilename.startswith("adapter_model")
+            and sibling.rfilename.endswith(".safetensors")
+            for sibling in siblings
+        )
+    except Exception:
+        return False
+
+
+def _prefetch_ignore_patterns(
+    model_name,
+    *,
+    token = None,
+    revision = None,
+    subfolder = None,
+    use_safetensors = None,
+    from_tf = False,
+    from_flax = False,
+    variant = None,
+    weights_at_root = False,
+):
+    """ignore_patterns for the prewarm snapshot: the static skip list, minus the checkpoint guard when
+    loading from a checkpoint-* subfolder, minus the weight format the load will not read. use_safetensors
+    is a format allowlist (True -> skip *.bin, False -> skip *.safetensors); auto (None) skips *.bin only
+    when in-scope safetensors are shipped. from_tf/from_flax keep *.h5/*.msgpack.
+
+    Suppressed for a whole multi-component snapshot (weights_at_root=False, no subfolder: ST/diffusers
+    repos with per-subfolder weights, each in its own format), since "*" spans "/" so dropping "*.bin"
+    would strip a module's only weight."""
+    # Keep checkpoint-*/* under a checkpoint-* subfolder; keep *.h5 / *.msgpack under from_tf/flax.
+    ignore_patterns = [
+        pattern
+        for pattern in _PREFETCH_IGNORE_PATTERNS
+        if not (
+            (
+                pattern == "checkpoint-*/*"
+                and isinstance(subfolder, str)
+                and subfolder.startswith("checkpoint-")
+            )
+            or (from_tf and pattern == "*.h5")
+            or (from_flax and pattern == "*.msgpack")
+        )
+    ]
+    # Drop the format the load will not read (the other doubles the download); skipped for a whole
+    # multi-component snapshot (see docstring).
+    whole_multi_component = not weights_at_root and not (
+        isinstance(subfolder, str) and subfolder.strip("/")
+    )
+    if whole_multi_component:
+        pass
+    elif from_tf or from_flax:
+        # TF / Flax loads never read the PyTorch formats; drop safetensors and .bin.
+        ignore_patterns.extend(
+            (
+                "*.safetensors",
+                "*.safetensors.index.json",
+                "*.bin",
+                "*.bin.index.json",
+            )
+        )
+    elif use_safetensors is True:
+        # Explicit safetensors: load never reads .bin (no model_info call needed).
+        ignore_patterns.extend(("*.bin", "*.bin.index.json"))
+    elif use_safetensors is False:
+        # Explicit .bin: load never reads safetensors.
+        ignore_patterns.extend(("*.safetensors", "*.safetensors.index.json"))
+    else:
+        # Auto: skip .bin only once in-scope safetensors are confirmed (best-effort; any failure keeps both).
+        try:
+            from huggingface_hub import HfApi
+
+            siblings = (
+                HfApi()
+                .model_info(
+                    model_name,
+                    revision = revision,
+                    token = token,
+                )
+                .siblings
+                or []
+            )
+            # Count only in-scope model-weights safetensors (not adapters/sidecars): variant-matching if
+            # a variant is requested, else canonical, proving the .bin redundant.
+            has_safetensors = any(
+                _is_model_weight_safetensors(sibling.rfilename)
+                and _in_requested_load_scope(sibling.rfilename, subfolder)
+                and (
+                    _is_canonical_variant_model_weight_safetensors(sibling.rfilename, variant)
+                    if variant
+                    else _is_canonical_model_weight_safetensors(sibling.rfilename)
+                )
+                for sibling in siblings
+            )
+            if has_safetensors:
+                ignore_patterns.extend(("*.bin", "*.bin.index.json"))
+        except Exception:
+            pass
+    return ignore_patterns
+
+
+def maybe_prefetch_hf_snapshot(
+    model_name,
+    token = None,
+    *,
+    revision = None,
+    cache_dir = None,
+    local_files_only = False,
+    fast_inference = False,
+    subfolder = None,
+    force_download = False,
+    use_safetensors = None,
+    from_tf = False,
+    from_flax = False,
+    tokenizer_only = False,
+    adapter_only = False,
+    weights_at_root = False,
+    variant = None,
+    gguf_file = None,
+):
+    """Warm the HF cache for a remote repo before the in-process load.
+
+    Xet can hang on a blob with no progress or exception, and a blocked native Xet thread cannot be
+    killed in-process. So pull the snapshot first in a killable subprocess that falls back Xet -> HTTP
+    on a stall (unsloth_zoo.hf_xet_fallback), making from_pretrained a cache hit.
+
+    Returns True iff warmed (caller can clear force_download), else False (skipped: local/offline/
+    local_files_only/fast_inference/old unsloth_zoo, or failed). Only a both-transports-stalled
+    DownloadStallError is raised; other failures are left for from_pretrained to surface.
+    """
+    try:
+        from unsloth_zoo.hf_xet_fallback import (
+            snapshot_download_with_xet_fallback,
+            DownloadStallError,
+        )
+    except Exception:
+        return False  # older unsloth_zoo without the helper: load normally
+
+    if not isinstance(model_name, str) or not model_name:
+        return False
+    # Local path: nothing to download. Expand ~ first (os.path.exists does not).
+    model_path = os.path.expanduser(model_name)
+    if os.path.isdir(model_path) or os.path.exists(model_path):
+        return False
+    # Looks local but not yet on disk (e.g. an uncreated output dir): not a Hub repo id, so leave it
+    # for from_pretrained rather than download it.
+    if (
+        os.path.isabs(model_path)
+        or model_name.startswith(("~", "./", "../", ".\\", "..\\"))
+        or "\\" in model_name
+    ):
+        return False
+    if local_files_only:  # cache-only: never reach out
+        return False
+    if any(
+        os.environ.get(flag, "0").lower() in ("1", "true", "yes", "on")
+        for flag in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")
+    ):
+        return False
+    if fast_inference:  # vLLM has its own download path
+        return False
+
+    # tokenizer-only / adapter-only warms allow-list exact files below, so the weight-format ignore
+    # list (and its auto-branch model_info call) is skipped.
+    ignore_patterns = (
+        None
+        if tokenizer_only or adapter_only or gguf_file
+        else _prefetch_ignore_patterns(
+            model_name,
+            token = token,
+            revision = revision,
+            subfolder = subfolder,
+            use_safetensors = use_safetensors,
+            from_tf = from_tf,
+            from_flax = from_flax,
+            variant = variant,
+            weights_at_root = weights_at_root,
+        )
+    )
+    # Narrow the warm to what the load reads (skip extra checkpoints/precisions); every branch still warms
+    # root tokenizer/config/custom-code so those never fall in-process.
+    allow_patterns = None
+    if gguf_file:
+        # gguf_file=NAME reads exactly that GGUF, but the static ignore list drops *.gguf; so warm just
+        # that file (plus root aux), under <subfolder>/ if set.
+        _gguf_path = (
+            f"{subfolder.strip('/')}/{gguf_file}"
+            if isinstance(subfolder, str) and subfolder.strip("/")
+            else gguf_file
+        )
+        allow_patterns = [_gguf_path, *_ROOT_AUX_PREFETCH_PATTERNS]
+    elif tokenizer_only:
+        # A distinct tokenizer repo: warm only tokenizer / config / vocab files, never its weights.
+        allow_patterns = list(_ROOT_AUX_PREFETCH_PATTERNS)
+    elif adapter_only:
+        # A PEFT adapter load reads only adapter_config.json + adapter_model.* (plus root aux), not any
+        # merged weights the repo may also publish.
+        allow_patterns = [*_ADAPTER_PREFETCH_PATTERNS, *_ROOT_AUX_PREFETCH_PATTERNS]
+        # PeftModel reads one format (safetensors when present): explicit use_safetensors wins, else
+        # prefer safetensors when shipped (best-effort; any failure keeps both).
+        if use_safetensors is False:
+            ignore_patterns = [
+                "adapter_model*.safetensors",
+                "adapter_model*.safetensors.index.json",
+            ]
+        elif use_safetensors is True or _adapter_repo_has_safetensors(
+            model_name, token = token, revision = revision
+        ):
+            ignore_patterns = ["adapter_model*.bin", "adapter_model*.bin.index.json"]
+    elif isinstance(subfolder, str) and subfolder.strip("/"):
+        # subfolder=X: load resolves every weight under X/, so warm that subfolder (plus root aux).
+        allow_patterns = [f"{subfolder.strip('/')}/*", *_ROOT_AUX_PREFETCH_PATTERNS]
+    elif weights_at_root:
+        # A bare load reads only root weights: drop subdir weights (fp16/, checkpoint dirs) while keeping
+        # subdir configs. Diffusion leaves weights_at_root False.
+        ignore_patterns = [*(ignore_patterns or []), *_SUBDIR_WEIGHT_IGNORE_PATTERNS]
+    try:
+        snapshot_download_with_xet_fallback(
+            model_name,
+            token = token,
+            revision = revision,
+            cache_dir = cache_dir,
+            allow_patterns = allow_patterns,
+            ignore_patterns = ignore_patterns,
+            force_download = force_download,
+            variant = variant,
+        )
+        return True
+    except DownloadStallError:
+        # Both transports stalled: surface a clear network error, not a silent in-process hang.
+        raise
+    except Exception as exception:
+        logger.warning_once(
+            f"Unsloth: Could not pre-download {model_name} "
+            f"({type(exception).__name__}: {exception}); continuing with the normal load."
+        )
+        return False
+
+
 # Ignore logging messages
 class HideLoggingMessage(logging.Filter):
     __slots__ = ("text",)
@@ -3507,8 +3972,25 @@ def _moe_target_set_from_string(target_modules: str) -> set[str]:
         return {target_modules}
 
     is_regex = re.search(r"[*+?()[\]{}|\\^$]", target_modules) is not None
-    targets_mlp = "mlp" in target_modules or "ffn" in target_modules
-    if is_regex and "proj" in target_modules and targets_mlp:
+    # Key detection on the mlp/ffn/experts path segment (absent from an
+    # attention-only regex), never on q/k/v/o leaves alone.
+    targets_mlp_path = any(
+        tag in target_modules for tag in ("mlp", "ffn", "feed_forward", "experts")
+    )
+    if not is_regex or not targets_mlp_path:
+        return set()
+    # Explicit expert leaves scope the target set to exactly those leaves.
+    named = {name for name in _MOE_BROAD_MLP_TARGETS if name in target_modules}
+    if named:
+        return named
+    # A generic projection under an mlp path (e.g. ".*mlp.*proj"): any proj
+    # occurrence that is not an attention leaf name.
+    if re.search(r"(?<![qkvo]_)(?<!out_)(?<!in_)proj", target_modules):
+        return set(_MOE_BROAD_MLP_TARGETS)
+    # The auto regex on fused-expert models lists only attention Linears as
+    # leaves; its mlp tag block is the remaining MLP-intent signal. A regex
+    # like "(mlp|self_attn).(q_proj|o_proj)" has neither and stays attention-only.
+    if "mlp|feed_forward|ffn|dense" in target_modules:
         return set(_MOE_BROAD_MLP_TARGETS)
 
     return set()
@@ -3594,6 +4076,31 @@ def get_moe_target_parameters(model, target_modules = None) -> Optional[List[str
     return None
 
 
+def _select_moe_detection_targets(
+    original_target_modules,
+    scoped_target_modules,
+    finetune_mlp_modules = True,
+    finetune_language_layers = True,
+):
+    """Pick what get_moe_target_parameters keys expert detection on.
+
+    Prefer the caller's ORIGINAL explicit leaf list over the scoped regex so an
+    attention-only request is not pushed into the experts by get_peft_regex's
+    ``mlp|feed_forward|ffn|dense`` component block (which the string fallback
+    cannot tell apart from a fused-expert auto regex).
+
+    But only when the MLP and language families are BOTH still in scope. If the
+    caller scoped MLP or language OFF (``finetune_mlp_modules=False`` or
+    ``finetune_language_layers=False``) the scoped regex already drops the MoE
+    experts, and reusing the original list -- which may still name gate/up/down
+    leaves -- would wrongly re-introduce them. In that case honor the scoped
+    result so the frozen-MLP / vision-only request is respected.
+    """
+    if original_target_modules is not None and finetune_mlp_modules and finetune_language_layers:
+        return original_target_modules
+    return scoped_target_modules
+
+
 def make_fast_generate_wrapper(original_generate):
     """
     Creates a wrapper around model.generate that checks for incorrect

unsloth/models/cohere.py

@@ -22,6 +22,7 @@ from ..utils.attention_dispatch import (
     AttentionContext,
     run_attention,
     select_attention_backend,
+    resolve_prefix_seg_info,
 )
 
 try:
@@ -151,6 +152,9 @@ def CohereAttention_fast_forward(
             "softmax_scale": getattr(self, "softmax_scale", None),
         },
     )
+    # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse
+    # (KV cache / padding mask) raises. None => byte-identical default.
+    _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask)
     context = AttentionContext(
         bsz = bsz,
         q_len = q_len,
@@ -161,6 +165,7 @@ def CohereAttention_fast_forward(
         seq_info = seq_info,
         attention_mask = attention_mask,
         causal_mask = causal_mask,
+        prefix_seg_info = _pg_seg,
     )
 
     A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V)

unsloth/models/diffusion.py

@@ -24,7 +24,7 @@ import os
 import torch
 from transformers import AutoConfig, AutoProcessor, AutoTokenizer
 
-from ._utils import is_bfloat16_supported
+from ._utils import is_bfloat16_supported, maybe_prefetch_hf_snapshot
 from .llama import logger
 
 __all__ = ["FastDiffusionModel", "DIFFUSION_MODEL_TYPES", "is_diffusion_model_type"]
@@ -79,7 +79,14 @@ def _resolve_diffusion_model_class(config):
     )
 
 
-def _load_diffusion_config(model_name, token, trust_remote_code, revision, local_files_only):
+def _load_diffusion_config(
+    model_name,
+    token,
+    trust_remote_code,
+    revision,
+    local_files_only,
+    cache_dir = None,
+):
     """Load the config, aliasing the legacy ``diffusion_gemma`` model_type to the ``diffusion_gemma4``
     classes current transformers ships. AutoConfig raises on the legacy type; catch that, rewrite the
     type/arch names in-memory, and rebuild."""
@@ -90,6 +97,7 @@ def _load_diffusion_config(model_name, token, trust_remote_code, revision, local
             trust_remote_code = trust_remote_code,
             revision = revision,
             local_files_only = local_files_only,
+            cache_dir = cache_dir,
         )
     except ValueError as e:
         if "diffusion_gemma" not in str(e):
@@ -103,6 +111,7 @@ def _load_diffusion_config(model_name, token, trust_remote_code, revision, local
             token = token,
             revision = revision,
             local_files_only = local_files_only,
+            cache_dir = cache_dir,
         )
         with open(cfg_path, encoding = "utf-8") as f:
             cd = json.load(f)
@@ -152,12 +161,16 @@ class FastDiffusionModel:
                 os.environ.get("HF_HUB_OFFLINE", "0") == "1"
                 or os.environ.get("TRANSFORMERS_OFFLINE", "0") == "1"
             )
+
+        cache_dir = kwargs.get("cache_dir")
+
         config = _load_diffusion_config(
             model_name,
             token,
             trust_remote_code,
             revision,
             local_files_only,
+            cache_dir = cache_dir,
         )
         model_type = getattr(config, "model_type", None)
         if not is_diffusion_model_type(model_type):
@@ -168,6 +181,21 @@ class FastDiffusionModel:
 
         model_cls = _resolve_diffusion_model_class(config)
 
+        # Prefetch the whole repo root so the weight load is a cache hit. No subfolder: the pipeline
+        # loads every component subfolder, so narrowing would leave unet/vae/text_encoder to Xet.
+        maybe_prefetch_hf_snapshot(
+            model_name,
+            token = token,
+            revision = revision,
+            cache_dir = cache_dir,
+            local_files_only = local_files_only,
+            fast_inference = False,
+            force_download = kwargs.get("force_download", False),
+            use_safetensors = kwargs.get("use_safetensors"),
+            # Forward variant (e.g. "fp16") so the warm keeps variant weights.
+            variant = kwargs.get("variant"),
+        )
+
         load_kwargs = dict(
             dtype = dtype,
             device_map = device_map,
@@ -176,7 +204,14 @@ class FastDiffusionModel:
             attn_implementation = attn_implementation,
             revision = revision,
             local_files_only = local_files_only,
+            cache_dir = cache_dir,
         )
+        # Match the load's weight format to the warm (None/auto already matches).
+        if kwargs.get("use_safetensors") is not None:
+            load_kwargs["use_safetensors"] = kwargs["use_safetensors"]
+        # Forward variant to the real load so it reads the warmed variant weights.
+        if kwargs.get("variant") is not None:
+            load_kwargs["variant"] = kwargs["variant"]
 
         # Optional bitsandbytes quant. The MoE experts (3D Parameters) are not nn.Linear so bnb skips
         # them; only attention + dense MLP Linears quantize, lm_head/embeddings stay full precision.
@@ -222,6 +257,7 @@ class FastDiffusionModel:
                 trust_remote_code = trust_remote_code,
                 revision = revision,
                 local_files_only = local_files_only,
+                cache_dir = cache_dir,
             )
         except Exception:
             tokenizer = AutoTokenizer.from_pretrained(
@@ -230,6 +266,7 @@ class FastDiffusionModel:
                 trust_remote_code = trust_remote_code,
                 revision = revision,
                 local_files_only = local_files_only,
+                cache_dir = cache_dir,
             )
 
         return model, tokenizer

unsloth/models/gemma2.py

@@ -22,6 +22,7 @@ from ..utils.attention_dispatch import (
     AttentionContext,
     run_attention,
     select_attention_backend,
+    resolve_prefix_seg_info,
     SDPA,
 )
 from .gemma import (
@@ -168,6 +169,11 @@ def Gemma2Attention_fast_forward(
             },
         )
 
+        # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse
+        # (KV cache / padding mask) raises. None => byte-identical default. gemma2 is
+        # sliding-window and softcapped: the engage gate caps spans at the window and
+        # excludes softcap models entirely, so PG never engages here.
+        _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask)
         context = AttentionContext(
             bsz = bsz,
             q_len = q_len,
@@ -179,6 +185,7 @@ def Gemma2Attention_fast_forward(
             attention_mask = attention_mask,
             causal_mask = causal_mask,
             sliding_window = sliding_window,
+            prefix_seg_info = _pg_seg,
         )
 
         A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V)

unsloth/models/granite.py

@@ -23,6 +23,7 @@ from ..utils.attention_dispatch import (
     AttentionContext,
     run_attention,
     select_attention_backend,
+    resolve_prefix_seg_info,
     SDPA,
 )
 from .llama import (
@@ -159,6 +160,9 @@ def GraniteAttention_fast_forward(
         },
     )
 
+    # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse
+    # (KV cache / padding mask) raises. None => byte-identical default.
+    _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask)
     context = AttentionContext(
         bsz = bsz,
         q_len = q_len,
@@ -169,6 +173,7 @@ def GraniteAttention_fast_forward(
         seq_info = seq_info,
         attention_mask = attention_mask,
         causal_mask = causal_mask,
+        prefix_seg_info = _pg_seg,
     )
 
     A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V)

unsloth/models/llama.py

@@ -39,6 +39,7 @@ from ..utils.attention_dispatch import (
     run_attention,
     SDPA,
     select_attention_backend,
+    resolve_prefix_seg_info,
 )
 from torch.nn.functional import scaled_dot_product_attention
 from transformers import __version__ as transformers_version
@@ -738,6 +739,10 @@ def LlamaAttention_fast_forward(
         flash_dense_kwargs = {"causal": True},
         flash_varlen_kwargs = {"dropout_p": 0.0, "causal": True},
     )
+    # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward (same route
+    # as packed_seq_lengths); misuse (KV cache / padding mask) raises. None => byte-identical
+    # default. Reuse of this forward also carries the branch to qwen2 & gemma.
+    _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask)
     context = AttentionContext(
         bsz = bsz,
         q_len = q_len,
@@ -748,6 +753,7 @@ def LlamaAttention_fast_forward(
         seq_info = seq_info,
         attention_mask = attention_mask,
         causal_mask = causal_mask,
+        prefix_seg_info = _pg_seg,
     )
 
     A = run_attention(config = config, context = context, Q = Q, K = K, V = V)
@@ -895,8 +901,10 @@ def LlamaModel_fast_forward(
     seq_length_with_past = seq_length
 
     # Fix out of bounds tokenization unless we were given packed metadata
-    allow_overlength = getattr(self, "_unsloth_allow_packed_overlength", False) or (
-        "packed_seq_lengths" in kwargs
+    allow_overlength = (
+        getattr(self, "_unsloth_allow_packed_overlength", False)
+        or ("packed_seq_lengths" in kwargs)
+        or ("prefix_seg_info" in kwargs and kwargs["prefix_seg_info"] is not None)
     )
     if hasattr(self, "max_seq_length") and not allow_overlength:
         if seq_length > self.max_seq_length:
@@ -1748,7 +1756,6 @@ class LlamaRotaryEmbedding(torch.nn.Module):
         # Base-class-from-config path (modern transformers): derive inv_freq like
         # transformers so config.rope_scaling is not dropped (#2405). Scaled
         # subclasses are excluded to avoid double-scaling.
-        config_inv_freq = None
         if config is not None:
             # [TODO] Hack to pass in config - need to remove later
             base = _get_rope_theta(config, default = base)
@@ -1761,32 +1768,17 @@ class LlamaRotaryEmbedding(torch.nn.Module):
             device = DEVICE_TYPE_TORCH
             max_position_embeddings = config.max_position_embeddings
 
-            rope_scaling = getattr(config, "rope_scaling", None)
-            if rope_scaling is not None and type(self) is LlamaRotaryEmbedding:
-                config_inv_freq, self.attention_scaling = _compute_config_rope_inv_freq(
-                    config,
-                    rope_scaling,
-                )
-
         self.dim = dim
         self.max_position_embeddings = max_position_embeddings
         self.base = base
+        # Kept so the v5 rope repair can rebuild the scaled inv_freq (#2405).
+        self._unsloth_rope_config = config
         # Dynamic RoPE we first set it to a max of 4 * 8192 tokens then we iteratively grow this
         self.current_rope_size = min(4 * 8192, self.max_position_embeddings)
         self.multi_gpu_cos_cached = [None] * DEVICE_COUNT
         self.multi_gpu_sin_cached = [None] * DEVICE_COUNT
 
-        if config_inv_freq is not None:
-            inv_freq = config_inv_freq  # already scaled; skip subclass scaling
-        else:
-            # Normal Llama-3 RoPE
-            inv_freq = 1.0 / (
-                self.base
-                ** (
-                    torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim
-                )
-            )
-            inv_freq = self._apply_inv_freq_scaling(inv_freq)
+        inv_freq = self._unsloth_recompute_inv_freq()
         self.register_buffer("inv_freq", inv_freq, persistent = False)
 
         # Build here to make `torch.jit.trace` work.
@@ -1809,6 +1801,25 @@ class LlamaRotaryEmbedding(torch.nn.Module):
         """Override to apply custom inv_freq scaling (e.g., extended RoPE)."""
         return inv_freq
 
+    def _unsloth_recompute_inv_freq(self):
+        # Config scaling (llama3/yarn) first, else vanilla + subclass scaling.
+        # Shared by __init__ and the v5 rope repair so they cannot diverge.
+        config = getattr(self, "_unsloth_rope_config", None)
+        config_inv_freq = None
+        rope_scaling = getattr(config, "rope_scaling", None) if config is not None else None
+        if rope_scaling is not None and type(self) is LlamaRotaryEmbedding:
+            config_inv_freq, self.attention_scaling = _compute_config_rope_inv_freq(
+                config,
+                rope_scaling,
+            )
+        if config_inv_freq is not None:
+            return config_inv_freq
+        inv_freq = 1.0 / (
+            self.base
+            ** (torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim)
+        )
+        return self._apply_inv_freq_scaling(inv_freq)
+
     def _apply_time_scaling(self, t):
         """Override to apply custom time scaling (e.g., linear scaling)."""
         return t
@@ -2420,6 +2431,73 @@ class FastLlamaModel:
 
         preferred_attn_impl = resolve_attention_implementation(model_function, model_config)
 
+        # Prefetch the repo (killable child) so the weight load is a cache hit. Runs after the
+        # AutoConfig/model-class check so an unsupported repo fails on its small config fetch. No
+        # revision: the load resolves model_name (maybe a remapped prequant repo) on its default branch.
+        _prefetched = maybe_prefetch_hf_snapshot(
+            model_name,
+            token = token,
+            cache_dir = kwargs.get("cache_dir"),
+            local_files_only = kwargs.get("local_files_only", False),
+            # Skip the warm only for a real vLLM load; a num_labels classification load still goes
+            # in-process below, so it must be warmed even under fast_inference.
+            fast_inference = fast_inference and num_labels is None,
+            subfolder = kwargs.get("subfolder"),
+            force_download = kwargs.get("force_download", False),
+            use_safetensors = kwargs.get("use_safetensors"),
+            from_tf = kwargs.get("from_tf", False),
+            from_flax = kwargs.get("from_flax", False),
+            # Bare load reads only ROOT weights; skip subdir weights. Ignored when a subfolder is set.
+            weights_at_root = True,
+            variant = kwargs.get("variant"),  # forward so the warm keeps the variant .bin
+            gguf_file = kwargs.get(
+                "gguf_file"
+            ),  # forward so the warm fetches the GGUF (else ignored)
+        )
+        # Child did the forced download; clear the flag so the load reuses the warm cache.
+        if _prefetched and kwargs.get("force_download", False):
+            kwargs["force_download"] = False
+
+        # Tokenizer always loads in-process. Resolve the cache_dir the tokenizer load will actually
+        # use, mirroring load_correct_tokenizer: without an explicit cache_dir, Colab/Kaggle route to
+        # a special tokenizer cache (huggingface_tokenizers_cache / Kaggle tmp), NOT the HF-default
+        # cache the base snapshot warmed. So the base warm does not cover the tokenizer there.
+        from ..tokenizer_utils import (
+            IS_COLAB_ENVIRONMENT,
+            IS_KAGGLE_ENVIRONMENT,
+            KAGGLE_TMP,
+        )
+
+        _tokenizer_repo = (
+            tokenizer_name if (isinstance(tokenizer_name, str) and tokenizer_name) else model_name
+        )
+        _tokenizer_cache_dir = kwargs.get("cache_dir")
+        if _tokenizer_cache_dir is None:
+            if IS_COLAB_ENVIRONMENT:
+                _tokenizer_cache_dir = "huggingface_tokenizers_cache"
+            elif IS_KAGGLE_ENVIRONMENT:
+                _tokenizer_cache_dir = os.path.join(KAGGLE_TMP, "huggingface_tokenizers_cache")
+        # Warm the tokenizer repo into the cache the load will use whenever the base warm did not
+        # cover it: a distinct tokenizer repo, fast_inference (base warm skipped), or a tokenizer
+        # cache_dir that differs from the base-warm cache_dir (Colab/Kaggle special cache).
+        _warm_tokenizer_repo = (
+            isinstance(_tokenizer_repo, str)
+            and bool(_tokenizer_repo)
+            and (
+                _tokenizer_repo != model_name
+                or fast_inference
+                or _tokenizer_cache_dir != kwargs.get("cache_dir")
+            )
+        )
+        if _warm_tokenizer_repo:
+            maybe_prefetch_hf_snapshot(
+                _tokenizer_repo,
+                token = token,
+                cache_dir = _tokenizer_cache_dir,
+                local_files_only = kwargs.get("local_files_only", False),
+                tokenizer_only = True,
+            )
+
         has_rope_scaling = False
         try:
             with open(inspect.getfile(model_function), "r", encoding = "utf-8") as file:
@@ -2672,6 +2750,10 @@ class FastLlamaModel:
 
         # Counteract saved tokenizers
         tokenizer_name = model_name if tokenizer_name is None else tokenizer_name
+        # Route the tokenizer load to the custom cache_dir the prefetch warmed.
+        _tokenizer_cache_kwargs = {}
+        if kwargs.get("cache_dir") is not None:
+            _tokenizer_cache_kwargs["cache_dir"] = kwargs["cache_dir"]
         tokenizer = load_correct_tokenizer(
             tokenizer_name = tokenizer_name,
             model_max_length = max_position_embeddings,
@@ -2679,6 +2761,7 @@ class FastLlamaModel:
             token = token,
             trust_remote_code = trust_remote_code,
             fix_tokenizer = fix_tokenizer,
+            **_tokenizer_cache_kwargs,
         )
 
         model, tokenizer = patch_tokenizer(model, tokenizer)
@@ -2805,6 +2888,7 @@ class FastLlamaModel:
                 model_max_length = max_position_embeddings,
                 padding_side = "right",
                 token = token,
+                cache_dir = kwargs.get("cache_dir"),
             )
         patch_saving_functions(tokenizer)
 
@@ -3743,4 +3827,17 @@ class FastLlamaModel:
 
 from .rl import PatchFastRL
 
+# Auto-enable grouped-GEMM MoE (tf<5 ModuleList experts) on built / PEFT'd models. Wrap the
+# loader leaves before PatchFastRL so downstream patchers see the wrapped versions. Guarded.
+try:
+    from unsloth_zoo.temporary_patches.moe_grouped_modulelist import wrap_loader_for_grouped_moe
+    FastLlamaModel.from_pretrained = staticmethod(
+        wrap_loader_for_grouped_moe(FastLlamaModel.from_pretrained)
+    )
+    FastLlamaModel.get_peft_model = staticmethod(
+        wrap_loader_for_grouped_moe(FastLlamaModel.get_peft_model)
+    )
+except Exception:
+    pass
+
 PatchFastRL(FastLanguageModel = FastLlamaModel)

unsloth/models/loader.py

@@ -21,6 +21,10 @@ from ._utils import (
     USE_MODELSCOPE,
     get_transformers_model_type,
     hf_login,
+    # Single source of truth is _utils.py; re-exported here so callers doing
+    # `from unsloth.models.loader import DISABLE_SDPA_MODEL_NAMES` keep working and so
+    # _is_sdpa_excluded (in _utils) can honor it without a loader -> _utils cycle.
+    DISABLE_SDPA_MODEL_NAMES,
 )
 from .granite import FastGraniteModel
 from .llama import FastLlamaModel, logger
@@ -106,24 +110,31 @@ from ._utils import (
     _is_family_text_decoder,
     _apply_text_only_key_mapping,
     set_task_config_attr,
+    maybe_prefetch_hf_snapshot,
 )
 
-# Single source of truth is unsloth_zoo.model_lists. Re-exported so callers
-# doing `from unsloth.models.loader import FORCE_FLOAT32` keep working.
-# Fallback list mirrors zoo for users who upgrade unsloth without upgrading
-# unsloth_zoo (so this module never fails at import).
+# Source of truth is unsloth_zoo.model_lists. Re-exported so callers doing
+# `from unsloth.models.loader import FORCE_FLOAT32` keep working. The fallback
+# list is also unioned in so a newer unsloth still forces float32 for these
+# archs when paired with an older unsloth_zoo that predates them (upgrade skew).
+_FORCE_FLOAT32_FALLBACK = [
+    "gemma3,",  # Add comma bc gemma3 will match gemma3n
+    "gemma3text",  # Gemma3TextModel (EmbeddingGemma, standalone text-only Gemma3)
+    "gemma3n",
+    "gemma4",  # Gemma4 (gemma4 / gemma4_text): float16 NaNs grad norms in the backward
+    "glm4_moe",  # GLM-4.x MoE (glm4_moe / glm4_moe_lite): float16 NaNs grad norms
+    "gpt_oss",
+    "qwen3_5",  # Qwen3.5 GDN layers produce NaN grad norms in float16 training
+    "qwen3_moe",  # Qwen3-MoE (Qwen3-30B-A3B): float16 NaNs grad norms in the backward
+]
 try:
-    from unsloth_zoo import FORCE_FLOAT32  # noqa: F401
+    from unsloth_zoo import FORCE_FLOAT32 as _ZOO_FORCE_FLOAT32
+    FORCE_FLOAT32 = list(_ZOO_FORCE_FLOAT32)
 except ImportError:
-    global FORCE_FLOAT32
-    # Forces float32 precision since float16 goes to infinity
-    FORCE_FLOAT32 = [
-        "gemma3,",  # Add comma bc gemma3 will match gemma3n
-        "gemma3text",  # Gemma3TextModel (EmbeddingGemma, standalone text-only Gemma3)
-        "gemma3n",
-        "gpt_oss",
-        "qwen3_5",  # Qwen3.5 GDN layers produce NaN grad norms in float16 training
-    ]
+    FORCE_FLOAT32 = []
+for _mt in _FORCE_FLOAT32_FALLBACK:
+    if not any(_mt in _entry for _entry in FORCE_FLOAT32):
+        FORCE_FLOAT32.append(_mt)
 
 global DISABLE_COMPILE_MODEL_NAMES
 # Must be alphabetically sorted for each entry
@@ -195,13 +206,44 @@ DISABLE_COMPILE_MODEL_NAMES = [
     "granite,llava_next",  # Granite-vision 3
 ]
 
-global DISABLE_SDPA_MODEL_NAMES
-# Disables some SDPA modules since it's wrong
-DISABLE_SDPA_MODEL_NAMES = [
-    "gemma3,",  # Add comma bc gemma3 will match gemma3n
-    "gemma3_text",  # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore
-    "gpt_oss",
-]
+# Architectures with gated-deltanet (linear attention) layers. Unsloth bundles the
+# flash-linear-attention Triton kernels (unsloth_zoo/_vendored/fla), so no install is
+# needed; transformers uses the much slower pure PyTorch path only when they can't be enabled.
+FLA_MODEL_TYPE_PREFIXES = ("qwen3_next", "qwen3_5", "kimi_linear", "olmo_hybrid")
+_fla_advised = False
+
+
+def _maybe_advise_fla_install(model_types):
+    """One-time note when a gated-deltanet model loads without the fast kernels.
+
+    The kernels ship with Unsloth (no install needed); this fires only when they
+    could not be enabled on this platform (e.g. no CUDA, torch < 2.7 or
+    triton < 3.3), i.e. exactly when transformers uses the slow pure PyTorch path.
+    """
+    global _fla_advised
+    if _fla_advised:
+        return
+    if model_types is None:
+        return
+    if isinstance(model_types, str):
+        model_types = [model_types]  # a lone string would otherwise iterate chars
+    try:
+        if not any(
+            isinstance(t, str) and t.startswith(FLA_MODEL_TYPE_PREFIXES) for t in model_types
+        ):
+            return
+        from transformers.utils.import_utils import is_flash_linear_attention_available
+        if is_flash_linear_attention_available():
+            return  # bundled (or user-installed) fast kernels are active
+    except Exception:
+        return
+    _fla_advised = True
+    print(
+        "Unsloth: This model uses gated-deltanet linear attention layers. Unsloth\n"
+        "bundles the flash-linear-attention kernels, but they could not be enabled\n"
+        "on this setup (they need CUDA with torch >= 2.7 and triton >= 3.3), so\n"
+        "transformers will use a slower pure PyTorch path."
+    )
 
 
 def _fix_rope_inv_freq(model):
@@ -227,14 +269,18 @@ def _fix_rope_inv_freq(model):
             and hasattr(module, "_apply_inv_freq_scaling")
             and hasattr(module, "multi_gpu_cos_cached")
         ):
-            inv_freq = 1.0 / (
-                module.base
-                ** (
-                    torch.arange(0, module.dim, 2, dtype = torch.int64, device = "cpu").float()
-                    / module.dim
+            if hasattr(module, "_unsloth_recompute_inv_freq"):
+                # Restore config scaling (llama3/yarn); unscaled here broke v5.
+                inv_freq = module._unsloth_recompute_inv_freq()
+            else:
+                inv_freq = 1.0 / (
+                    module.base
+                    ** (
+                        torch.arange(0, module.dim, 2, dtype = torch.int64, device = "cpu").float()
+                        / module.dim
+                    )
                 )
-            )
-            inv_freq = module._apply_inv_freq_scaling(inv_freq)
+                inv_freq = module._apply_inv_freq_scaling(inv_freq)
             module.inv_freq = inv_freq
             for device_idx in range(len(module.multi_gpu_cos_cached)):
                 if module.multi_gpu_cos_cached[device_idx] is not None:
@@ -469,8 +515,10 @@ class FastLanguageModel(FastLlamaModel):
             ("-unsloth-bnb-4bit", "-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"):
+        # '-bf16' hub repos load bf16; a local dir keeps the requested quant unless 16bit is set
+        if model_name.lower().endswith("-bf16") and (
+            load_in_16bit or not os.path.isdir(os.path.expanduser(model_name))
+        ):
             load_in_4bit = False
             load_in_8bit = False
             load_in_fp8 = False
@@ -628,8 +676,10 @@ class FastLanguageModel(FastLlamaModel):
                 ("-unsloth-bnb-4bit", "-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"):
+            # '-bf16' hub repos load bf16; a local dir keeps the requested quant unless 16bit is set
+            if model_name.lower().endswith("-bf16") and (
+                load_in_16bit or not os.path.isdir(os.path.expanduser(model_name))
+            ):
                 load_in_4bit = False
                 load_in_8bit = False
                 load_in_fp8 = False
@@ -865,6 +915,28 @@ class FastLanguageModel(FastLlamaModel):
         if is_peft:
             # From https://github.com/huggingface/peft/issues/184
             # Now add PEFT adapters
+            # Warm the adapter repo: PeftModel downloads it in-process and can hang on Xet.
+            _prefetched = maybe_prefetch_hf_snapshot(
+                old_model_name,
+                token = token,
+                revision = revision,
+                cache_dir = kwargs.get("cache_dir"),
+                local_files_only = local_files_only,
+                # Adapter always loads in-process via PeftModel, so warm it even under fast_inference.
+                fast_inference = False,
+                force_download = kwargs.get("force_download", False),
+                # Leave use_safetensors auto (inheriting base format could skip a safetensors-only
+                # adapter). adapter_only restricts the warm to the adapter files + root aux.
+                adapter_only = True,
+            )
+            # Child did the forced download; clear the flag so the load reuses the warm cache.
+            if _prefetched and kwargs.get("force_download", False):
+                kwargs["force_download"] = False
+            # Forward cache_dir so the load reads the warmed adapter. No subfolder (that targets the
+            # base checkpoint; adapters live at the root).
+            peft_load_kwargs = {}
+            if kwargs.get("cache_dir") is not None:
+                peft_load_kwargs["cache_dir"] = kwargs["cache_dir"]
             model = PeftModel.from_pretrained(
                 model,
                 old_model_name,
@@ -873,9 +945,19 @@ class FastLanguageModel(FastLlamaModel):
                 local_files_only = local_files_only,
                 is_trainable = True,
                 trust_remote_code = trust_remote_code,
+                **peft_load_kwargs,
             )
             # Patch it as well!
             model = dispatch_model.patch_peft_model(model, use_gradient_checkpointing)
+            # Re-evaluate grouped MoE now the adapter is attached: an expert-LoRA block falls back
+            # to the original loop, an attention-only adapter keeps the grouped path. Guarded.
+            try:
+                from unsloth_zoo.temporary_patches.moe_grouped_modulelist import (
+                    auto_enable_grouped_moe,
+                )
+                auto_enable_grouped_moe(model)
+            except Exception:
+                pass  # optional speedup; never block model loading
 
         # Patch Tiled MLP
         # to turn on set UNSLOTH_TILED_MLP to "arctic", "target", or "target:{GB}""
@@ -1116,8 +1198,10 @@ class FastModel(FastBaseModel):
             ("-unsloth-bnb-4bit", "-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"):
+        # '-bf16' hub repos load bf16; a local dir keeps the requested quant unless 16bit is set
+        if model_name.lower().endswith("-bf16") and (
+            load_in_16bit or not os.path.isdir(os.path.expanduser(model_name))
+        ):
             load_in_4bit = False
             load_in_8bit = False
             load_in_fp8 = False
@@ -1263,6 +1347,7 @@ class FastModel(FastBaseModel):
             trust_remote_code = trust_remote_code,
         )
         model_types_all = ",".join(model_types) + ","
+        _maybe_advise_fla_install(model_types)
 
         # ---- Text-diffusion models (e.g. DiffusionGemma) take a transformers-only slow path. ----
         # These use a custom block-diffusion `generate` and a novel backbone, so we skip Unsloth's
@@ -1474,8 +1559,10 @@ class FastModel(FastBaseModel):
                 ("-unsloth-bnb-4bit", "-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"):
+            # '-bf16' hub repos load bf16; a local dir keeps the requested quant unless 16bit is set
+            if model_name.lower().endswith("-bf16") and (
+                load_in_16bit or not os.path.isdir(os.path.expanduser(model_name))
+            ):
                 load_in_4bit = False
                 load_in_8bit = False
                 load_in_fp8 = False
@@ -1790,6 +1877,28 @@ class FastModel(FastBaseModel):
 
                 _LoraModel._create_and_replace = _patched_car
 
+            # Warm the adapter repo: PeftModel downloads it in-process and can hang on Xet.
+            _prefetched = maybe_prefetch_hf_snapshot(
+                old_model_name,
+                token = token,
+                revision = revision,
+                cache_dir = kwargs.get("cache_dir"),
+                local_files_only = local_files_only,
+                # Adapter always loads in-process via PeftModel, so warm it even under fast_inference.
+                fast_inference = False,
+                force_download = kwargs.get("force_download", False),
+                # Leave use_safetensors auto (inheriting base format could skip a safetensors-only
+                # adapter). adapter_only restricts the warm to the adapter files + root aux.
+                adapter_only = True,
+            )
+            # Child did the forced download; clear the flag so the load reuses the warm cache.
+            if _prefetched and kwargs.get("force_download", False):
+                kwargs["force_download"] = False
+            # Forward cache_dir so the load reads the warmed adapter. No subfolder (that targets the
+            # base checkpoint; adapters live at the root).
+            peft_load_kwargs = {}
+            if kwargs.get("cache_dir") is not None:
+                peft_load_kwargs["cache_dir"] = kwargs["cache_dir"]
             try:
                 model = PeftModel.from_pretrained(
                     model,
@@ -1799,6 +1908,7 @@ class FastModel(FastBaseModel):
                     local_files_only = local_files_only,
                     is_trainable = True,
                     trust_remote_code = trust_remote_code,
+                    **peft_load_kwargs,
                 )
             finally:
                 # Always restore original PEFT method, even if loading fails
@@ -1809,6 +1919,15 @@ class FastModel(FastBaseModel):
             model = FastBaseModel.post_patch_model(
                 model, use_gradient_checkpointing, trust_remote_code = trust_remote_code
             )
+            # Re-evaluate grouped MoE now the adapter is attached: an expert-LoRA block falls back
+            # to the original loop, an attention-only adapter keeps the grouped path. Guarded.
+            try:
+                from unsloth_zoo.temporary_patches.moe_grouped_modulelist import (
+                    auto_enable_grouped_moe,
+                )
+                auto_enable_grouped_moe(model)
+            except Exception:
+                pass  # optional speedup; never block model loading
 
         # Apply QAT if specified
         if qat_scheme is not None:

unsloth/models/mistral.py

@@ -27,6 +27,7 @@ from ..utils.attention_dispatch import (
     run_attention,
     SDPA,
     select_attention_backend,
+    resolve_prefix_seg_info,
 )
 from .llama import (
     LlamaRotaryEmbedding,
@@ -124,6 +125,9 @@ def MistralAttention_fast_forward(
             "softmax_scale": getattr(self, "softmax_scale", None),
         },
     )
+    # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse
+    # (KV cache / padding mask) raises. None => byte-identical default.
+    _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask)
     context = AttentionContext(
         bsz = bsz,
         q_len = q_len,
@@ -134,6 +138,7 @@ def MistralAttention_fast_forward(
         seq_info = seq_info,
         attention_mask = attention_mask,
         causal_mask = causal_mask,
+        prefix_seg_info = _pg_seg,
     )
 
     A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V)
@@ -161,7 +166,13 @@ def MistralForCausalLM_fast_forward(
     *args,
     **kwargs,
 ) -> Union[Tuple, CausalLMOutputWithPast]:
-    if causal_mask is None and past_key_values is None:
+    # PrefixGrouper brings its own mask: a synthesized causal attention_mask would trip
+    # resolve_prefix_seg_info on the no-xFormers path and force a fallback.
+    if (
+        causal_mask is None
+        and past_key_values is None
+        and kwargs.get("prefix_seg_info", None) is None
+    ):
         bsz, q_len = input_ids.shape
         sliding_window = getattr(self.config, "sliding_window", None)
 

unsloth/models/qwen3.py

@@ -23,6 +23,7 @@ from ..utils.attention_dispatch import (
     run_attention,
     SDPA,
     select_attention_backend,
+    resolve_prefix_seg_info,
 )
 from .llama import (
     LlamaRotaryEmbedding,
@@ -146,6 +147,9 @@ def Qwen3Attention_fast_forward(
             "softmax_scale": getattr(self, "softmax_scale", None),
         },
     )
+    # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse
+    # (KV cache / padding mask) raises. None => byte-identical default.
+    _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask)
     context = AttentionContext(
         bsz = bsz,
         q_len = q_len,
@@ -156,6 +160,7 @@ def Qwen3Attention_fast_forward(
         seq_info = seq_info,
         attention_mask = attention_mask,
         causal_mask = causal_mask,
+        prefix_seg_info = _pg_seg,
     )
 
     A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V)

unsloth/models/rl_replacements.py

@@ -29,7 +29,9 @@ from collections import defaultdict
 from unsloth_zoo.rl_replacements import (
     RL_REPLACEMENTS,
     left_pack_padding,
+    create_completion_attention_mask,
     chunked_selective_log_softmax,
+    chunked_hidden_states_selective_log_softmax,
     _unsloth_get_mm_token_id,
     _unsloth_fix_mm_token_type_ids,
 )
@@ -48,7 +50,41 @@ from ..device_type import (
     ALLOW_PREQUANTIZED_MODELS,
 )
 import textwrap
-from ._utils import _get_inference_mode_context_manager
+from ._utils import _get_inference_mode_context_manager, UNSLOTH_ENABLE_LOGGING
+
+# One-time GRPO sequence-packing gates; mirrored into the generated trainer cache via RL_PRE_ITEMS.
+UNSLOTH_GRPO_SEQ_PACKING_ON = os.environ.get("UNSLOTH_GRPO_SEQ_PACKING", "1").lower() not in (
+    "0",
+    "false",
+    "no",
+    "off",
+)
+# Packing needs zoo#840's masked-column guard in grpo_compute_loss (installed zoo is fixed per-process).
+try:
+    UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = "torch.where(_keep, new" in inspect.getsource(
+        RL_REPLACEMENTS["grpo_compute_loss"]
+    )
+except Exception:
+    UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = False
+# One-time PrefixGrouper gate; any import failure degrades to "PrefixGrouper off".
+_pg_build_layout = _pg_enabled_fn = _pg_verify_on = _pg_tol_ok = _PG_TOL_KILL = None
+UNSLOTH_GRPO_PREFIX_GROUPER_ON = os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER", "1").lower() not in (
+    "0",
+    "false",
+    "no",
+    "off",
+)
+if UNSLOTH_GRPO_PREFIX_GROUPER_ON:
+    try:
+        from ..utils.prefix_grouper import (
+            build_group_layout as _pg_build_layout,
+            prefix_grouper_enabled as _pg_enabled_fn,
+            verify_on as _pg_verify_on,
+            tol_ok as _pg_tol_ok,
+            TOL_KILL as _PG_TOL_KILL,
+        )
+    except Exception:
+        UNSLOTH_GRPO_PREFIX_GROUPER_ON = False
 
 RL_EXTRA_ARGS = defaultdict(list)
 RL_FUNCTIONS = defaultdict(list)
@@ -1359,6 +1395,439 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
             )
             os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
 
+            # ---- Sequence packing (default-on; disable with UNSLOTH_GRPO_SEQ_PACKING=0) ----
+            # One varlen [1, sum L] forward replaces the padded [B, Lmax] loop (also fixes the
+            # left-pad RoPE error). Self-verified against the per-row forward, re-checked as T
+            # grows; falls back if a backend ignores packed_seq_lengths.
+            logprobs = None
+
+            # ---- PrefixGrouper (GRPO shared-prompt dedup; default ON, exact + self-verified) ----
+            # G completions per prompt share the prefix; the packed path forwards it G times,
+            # PrefixGrouper stores it once (FlexAttention shared-prefix mask), cutting the trunk
+            # forward from G*(P+R) to P+G*R tokens. Gated by UNSLOTH_GRPO_PREFIX_GROUPER (needs
+            # seq-packing), tok_r auto-gate, and first-use self-verify vs the packed path
+            # (mismatch => fall back + mark unsafe), so a mask/isolation regression cannot ship
+            # silently. When off / ungrouped / unverified, the packed path below runs as before.
+            _pg_result = None
+            _pg_use = False
+            _pg_skip_pk = False  # once a shape is PG-verified, skip the full-row forward
+            _pg_forward_fn = None  # deferred PG forward (runs at the verify site below)
+            _pg_num_gen = getattr(self, "num_generations", None)
+            # Env gate hoisted to module level (mirrored via RL_PRE_ITEMS). Skip PG under vLLM
+            # (fast_inference=True): the rollout dominates the step, so PG saves little and its
+            # first-use self-verify is net overhead.
+            _pg_engage = (
+                UNSLOTH_GRPO_PREFIX_GROUPER_ON
+                and not getattr(self, "use_vllm", False)
+                and not getattr(unwrapped_model, "_unsloth_prefix_grouper_nograd_disabled", False)
+            )
+            if _pg_engage:
+                try:
+                    # Skip softcap models (the flex kernel never applies attn_logit_softcapping)
+                    # and hybrid SSM / MoE models: only the threaded attention forwards get the
+                    # shared-prefix isolation, so a Mamba or MoE decoder that does not forward
+                    # prefix_seg_info would leak suffixes across completions. PG also rides on
+                    # sequence packing, so it needs the same zoo masked-column guard.
+                    _pg_cfg = getattr(unwrapped_model, "config", None)
+                    _pg_engage = (
+                        _pg_enabled_fn()
+                        and UNSLOTH_ZOO_HAS_MASKED_COL_GUARD
+                        and pixel_values is None
+                        and token_type_ids is None
+                        and mm_token_type_ids is None
+                        and _pg_num_gen is not None
+                        and _pg_num_gen >= 2
+                        and not getattr(_pg_cfg, "attn_logit_softcapping", None)
+                        # normal backends apply config.attention_dropout in training; the flex
+                        # path is deterministic, so skip PG when it is set.
+                        and not getattr(_pg_cfg, "attention_dropout", 0)
+                        and not any(
+                            getattr(_pg_cfg, _pg_a, None) is not None
+                            for _pg_a in (
+                                "mamba_d_ssm",
+                                "mamba_d_state",
+                                "mamba_expand",
+                                "num_experts",
+                                "num_local_experts",
+                                "n_routed_experts",
+                                "moe_intermediate_size",
+                            )
+                        )
+                    )
+                except Exception:
+                    _pg_engage = False
+            if _pg_engage:
+                try:
+                    _pg_pad = self.processing_class.pad_token_id
+                    # cap the PG span (P+max(R)) at the sliding window, like the packed _pk_sw guard.
+                    _pg_sw = getattr(
+                        getattr(unwrapped_model, "config", None), "sliding_window", None
+                    )
+                    if not (isinstance(_pg_sw, int) and _pg_sw > 0):
+                        _pg_sw = None
+                    _pg_layout = _pg_build_layout(
+                        input_ids,
+                        logits_to_keep,
+                        _pg_pad,
+                        _pg_num_gen,
+                        left_pad_tokens_per_prompt,
+                        max_segment_cap = _pg_sw,
+                    )
+                    _pg_unsafe = getattr(
+                        unwrapped_model, "_unsloth_prefix_grouper_nograd_unsafe", None
+                    )
+                    if _pg_unsafe is None:
+                        _pg_unsafe = set()
+                    if _pg_layout is not None and _pg_layout.signature not in _pg_unsafe:
+                        _pg_sig = _pg_layout.signature
+                        _pg_verified = getattr(
+                            unwrapped_model, "_unsloth_prefix_grouper_nograd_verified", None
+                        )
+                        if _pg_verified is None:
+                            _pg_verified = set()
+                        _pg_chunks = max(1, total_rows * multiplier)
+
+                        def _pg_run_forward(_pg_layout = _pg_layout, _pg_chunks = _pg_chunks):
+                            with _get_inference_mode_context_manager(model):
+                                with torch.amp.autocast(
+                                    device_type = "cuda", dtype = self._autocast_dtype
+                                ):
+                                    _pg_hidden = unwrapped_model(
+                                        input_ids = _pg_layout.flat_ids,
+                                        position_ids = _pg_layout.position_ids,
+                                        prefix_seg_info = _pg_layout.prefix_seg_info,
+                                        use_cache = False,
+                                    ).logits
+                                    _pg_r = _pg_layout.extract_logps(
+                                        _pg_hidden,
+                                        lm_head,
+                                        chunked_hidden_states_selective_log_softmax,
+                                        _pg_chunks,
+                                        logit_scale_multiply,
+                                        logit_scale_divide,
+                                        logit_softcapping,
+                                        temperature,
+                                    )
+                                    _pg_hidden = None  # release before any verify forward
+                            device_synchronize()
+                            # clip to the loss window [B, logits_to_keep+max_left_pad]
+                            _pg_w = logits_to_keep + max_left_pad
+                            if _pg_r.shape[1] > _pg_w:
+                                _pg_r = _pg_r[:, -_pg_w:]
+                            return _pg_r
+
+                        # trust only within the verified envelope: re-verify when T or the
+                        # longest segment grows, like the packed path
+                        _pg_T = int(_pg_layout.flat_ids.shape[1])
+                        _pg_maxseg = int(_pg_layout.position_ids.max()) + 1
+                        _pg_env = (
+                            _pg_verified.get(_pg_sig) if isinstance(_pg_verified, dict) else None
+                        )
+                        if (not _pg_verify_on()) or (
+                            _pg_env is not None and _pg_T <= _pg_env[0] and _pg_maxseg <= _pg_env[1]
+                        ):
+                            # trusted shape: run PG now and skip the full-row forward below
+                            _pg_result = _pg_run_forward()
+                            _pg_use = True
+                            _pg_skip_pk = True
+                        else:
+                            # unverified shape: defer the forward until the packed reference
+                            # exists (verify site below), so a declined packed path never wastes
+                            # a whole-batch PG forward
+                            _pg_forward_fn = _pg_run_forward
+                except Exception as _pg_err:
+                    _pg_result = None
+                    _pg_use = False
+                    _pg_skip_pk = False
+                    _pg_forward_fn = None
+                    # A FlexAttention/Triton compile failure or OOM here is GPU-wide, not
+                    # layout-specific, so retrying the same PG forward every step just re-pays
+                    # the failure. Persistently disable PG (mirrors the seq-packing handler
+                    # setting _unsloth_seq_packing_nograd_ok = False); the packed/padded path
+                    # below still produces the exact result.
+                    unwrapped_model._unsloth_prefix_grouper_nograd_disabled = True
+                    if isinstance(_pg_err, torch.cuda.OutOfMemoryError):
+                        torch.cuda.empty_cache()
+                    os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
+                    if UNSLOTH_ENABLE_LOGGING:
+                        print(
+                            f"[Unsloth] GRPO PrefixGrouper (no-grad) disabled (fell back to packed): {_pg_err!r}",
+                            flush = True,
+                        )
+
+            # ---- Sequence packing (default-on; disable with UNSLOTH_GRPO_SEQ_PACKING=0) ----
+            # One varlen [1, sum L] block-diagonal forward replaces the padded [B, Lmax] loop
+            # (exact per-row result; also fixes the padded path's left-pad RoPE error).
+            # Self-verified vs the per-row forward, re-checked as T grows; falls back if a
+            # backend ignores packed_seq_lengths. lm_head runs on completion positions only.
+            _pk_result = None
+            _pk_use = False
+            _pk_enabled = UNSLOTH_GRPO_SEQ_PACKING_ON
+            # Without zoo#840's masked-column guard, zeroed prompt/pad columns turn NaN in exp().
+            _pk_enabled = _pk_enabled and UNSLOTH_ZOO_HAS_MASKED_COL_GUARD
+            _pk_ok = getattr(unwrapped_model, "_unsloth_seq_packing_nograd_ok", None)
+            if (
+                _pk_enabled
+                and not _pg_skip_pk
+                and pixel_values is None
+                and token_type_ids is None
+                and mm_token_type_ids is None
+                and _pk_ok is not False
+            ):
+                try:
+                    _pk_pad = self.processing_class.pad_token_id
+                    _pk_keep = input_ids != _pk_pad
+                    _pk_len = _pk_keep.sum(dim = 1)
+                    _pk_len_cpu = _pk_len.tolist()  # single GPU->CPU sync, reused below
+                    _pk_nz_cpu = [_n for _n in _pk_len_cpu if _n > 0]
+                    _pk_flat = input_ids[_pk_keep].unsqueeze(0)
+                    _pk_T = _pk_flat.shape[1]
+                    _pk_L = input_ids.shape[1]
+                    _pk_W = logits_to_keep + max_left_pad
+                    _pk_maxseg = max(_pk_nz_cpu) if _pk_nz_cpu else 0
+                    # sliding-window models lose the per-sequence local window in a packed stream
+                    _pk_sw = getattr(
+                        getattr(unwrapped_model, "config", None), "sliding_window", None
+                    )
+                    _pk_sw_ok = not (isinstance(_pk_sw, int) and _pk_sw > 0 and _pk_maxseg > _pk_sw)
+                    # per-row completion mask (same as the loss); prompt-only rows count as inactive
+                    _pk_cmask = create_completion_attention_mask(
+                        input_ids[:, -_pk_W:], left_pad_tokens_per_prompt, max_left_pad, _pk_pad
+                    )
+                    _pk_active = int(_pk_cmask.any(dim = 1).sum())
+                    # skip the packed forward entirely at known-unsafe lengths (avoids a wasted pass / OOM)
+                    _pk_unsafe = getattr(
+                        unwrapped_model, "_unsloth_seq_packing_nograd_unsafe_T", None
+                    )
+                    # cap the flattened forward at one padded [batch_size, seq_len] mini-batch's
+                    # token budget; anything larger uses the chunked padded loop
+                    _pk_cap = batch_size * seq_len
+                    if (
+                        _pk_T >= 2
+                        and _pk_T <= _pk_cap
+                        and len(_pk_nz_cpu) > 0
+                        and _pk_sw_ok
+                        and not (_pk_unsafe is not None and _pk_T >= _pk_unsafe)
+                        and (_pk_ok is True or _pk_active >= 2)
+                    ):
+                        # reset 0-based position_ids per segment
+                        _pk_pos = (_pk_keep.cumsum(dim = 1) - 1)[_pk_keep].unsqueeze(0)
+                        _pk_chunks = max(1, total_rows * multiplier)
+                        _pk_nz_idx = _pk_keep.nonzero(
+                            as_tuple = False
+                        )  # [T, 2] = (row, col), row-major
+                        _pk_within = _pk_nz_idx[1:, 0] == _pk_nz_idx[:-1, 0]  # [T-1]
+                        # per-row completion start after left-packing (matches create_completion_attention_mask)
+                        _pk_cstart = (_pk_L - logits_to_keep) - left_pad_tokens_per_prompt  # [rows]
+                        _pk_ctgt = (_pk_nz_idx[1:, 1] >= _pk_cstart[_pk_nz_idx[1:, 0]]) & _pk_within
+                        with _get_inference_mode_context_manager(model):
+                            with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype):
+                                # use_cache=False: a KV cache silently disables varlen packing
+                                _pk_hidden = unwrapped_model(
+                                    input_ids = _pk_flat,
+                                    position_ids = _pk_pos,
+                                    packed_seq_lengths = torch.tensor(
+                                        _pk_nz_cpu, dtype = torch.int32, device = input_ids.device
+                                    ),
+                                    use_cache = False,
+                                ).logits
+                                _pk_sel = chunked_hidden_states_selective_log_softmax(
+                                    _pk_hidden[0, :-1, :][_pk_ctgt].unsqueeze(0),
+                                    lm_head,
+                                    _pk_flat[0, 1:][_pk_ctgt].unsqueeze(0),
+                                    _pk_chunks,
+                                    logit_scale_multiply,
+                                    logit_scale_divide,
+                                    logit_softcapping,
+                                    temperature,
+                                )[0]
+                        # GPT-OSS offload race guard (matches the padded loop)
+                        device_synchronize()
+                        # scatter each logprob back to its (row, col) so [:, -_pk_W:] matches padded
+                        _pk_tgt = (_pk_nz_idx[1:, 0] * _pk_L + _pk_nz_idx[1:, 1])[_pk_ctgt]
+                        _pk_result = (
+                            torch.zeros(
+                                total_rows * _pk_L,
+                                dtype = torch.float32,
+                                device = input_ids.device,
+                            )
+                            .index_put((_pk_tgt,), _pk_sel.to(torch.float32))
+                            .view(total_rows, _pk_L)[:, -_pk_W:]
+                        )
+                        # re-verify when T or the longest segment grows past what was verified
+                        # (a LongRoPE cache switch can change the result)
+                        _pk_vT = int(
+                            getattr(unwrapped_model, "_unsloth_seq_packing_nograd_verified_T", 0)
+                        )
+                        _pk_vS = int(
+                            getattr(unwrapped_model, "_unsloth_seq_packing_nograd_verified_seg", 0)
+                        )
+                        # debug: hand-edit this condition to force re-verify every step
+                        if _pk_ok is True and _pk_T <= _pk_vT and _pk_maxseg <= _pk_vS:
+                            _pk_use = True  # already verified for this shape
+                        else:
+                            # verify against the per-row forward (ground truth)
+                            _pk_ref = torch.zeros_like(_pk_result)
+                            with _get_inference_mode_context_manager(model):
+                                with torch.amp.autocast(
+                                    device_type = "cuda", dtype = self._autocast_dtype
+                                ):
+                                    for _pk_i in range(total_rows):
+                                        _pk_ni = _pk_len_cpu[_pk_i]
+                                        if _pk_ni < 2:
+                                            continue
+                                        _pk_rmask = _pk_keep[_pk_i]
+                                        _pk_real = input_ids[_pk_i][_pk_rmask].unsqueeze(0)
+                                        _pk_rpos = torch.arange(
+                                            _pk_ni, device = input_ids.device
+                                        ).unsqueeze(0)
+                                        _pk_rh = unwrapped_model(
+                                            input_ids = _pk_real,
+                                            position_ids = _pk_rpos,
+                                            use_cache = False,
+                                        ).logits
+                                        _pk_rsel = chunked_hidden_states_selective_log_softmax(
+                                            _pk_rh[:, :-1, :],
+                                            lm_head,
+                                            _pk_real[:, 1:],
+                                            1,
+                                            logit_scale_multiply,
+                                            logit_scale_divide,
+                                            logit_softcapping,
+                                            temperature,
+                                        )[0]
+                                        _pk_rcols = _pk_rmask.nonzero(as_tuple = False).squeeze(1)[
+                                            1:
+                                        ] - (_pk_L - _pk_W)
+                                        _pk_rkeep = _pk_rcols >= 0
+                                        _pk_ref[_pk_i, _pk_rcols[_pk_rkeep]] = _pk_rsel[
+                                            _pk_rkeep
+                                        ].to(torch.float32)
+                            device_synchronize()
+                            # compare over the loss-mask region only
+                            _pk_cm = _pk_cmask.float()
+                            _pk_diff = float(((_pk_result - _pk_ref).abs() * _pk_cm).max())
+                            if UNSLOTH_ENABLE_LOGGING:
+                                print(
+                                    f"[Unsloth] GRPO seq-packing (no-grad) verify: T={_pk_T} maxseg={_pk_maxseg} packed-vs-perrow max|d|={_pk_diff:.4f}",
+                                    flush = True,
+                                )
+                            # kernel-noise floor ~0.25; cross-sample contamination is >= 2.4
+                            if _pk_diff < 7e-1:
+                                unwrapped_model._unsloth_seq_packing_nograd_ok = True
+                                # widen the trusted shape only when >= 2 completion rows exercised
+                                # cross-sample packing; single-row passes prove nothing
+                                if _pk_active >= 2:
+                                    unwrapped_model._unsloth_seq_packing_nograd_verified_T = max(
+                                        _pk_vT, _pk_T
+                                    )
+                                    unwrapped_model._unsloth_seq_packing_nograd_verified_seg = max(
+                                        _pk_vS, _pk_maxseg
+                                    )
+                                _pk_ok = True
+                                _pk_use = True
+                            else:
+                                _pk_use = False
+                                if _pk_diff >= 1.5:
+                                    # contamination (attention ignores the packed mask): disable packing
+                                    unwrapped_model._unsloth_seq_packing_nograd_ok = False
+                                else:
+                                    # likely a length boundary (LongRoPE): mark unsafe, keep smaller shapes
+                                    unwrapped_model._unsloth_seq_packing_nograd_unsafe_T = (
+                                        _pk_T if _pk_unsafe is None else min(_pk_unsafe, _pk_T)
+                                    )
+                                if UNSLOTH_ENABLE_LOGGING:
+                                    print(
+                                        f"[Unsloth] GRPO seq-packing (no-grad) fell back at T={_pk_T} (diff={_pk_diff:.3f})",
+                                        flush = True,
+                                    )
+                except Exception as _pk_err:
+                    # any failure: drop intermediates, use the padded loop, do not retry
+                    _pk_hidden = None
+                    _pk_sel = None
+                    _pk_result = None
+                    _pk_use = False
+                    if isinstance(_pk_err, torch.cuda.OutOfMemoryError):
+                        torch.cuda.empty_cache()
+                    unwrapped_model._unsloth_seq_packing_nograd_ok = False
+                    if UNSLOTH_ENABLE_LOGGING:
+                        print(
+                            f"[Unsloth] GRPO sequence-packing (no-grad) disabled (fell back to padded): {_pk_err!r}",
+                            flush = True,
+                        )
+            # ---- PrefixGrouper first-use self-verify (no-grad) ----
+            # Compare the untrusted PG result to the full-row packed result (itself verified vs
+            # per-row) over the completion mask: < tol_ok -> trust the structure; >= TOL_KILL ->
+            # unsafe forever; borderline -> fall back this shape.
+            if _pg_forward_fn is not None and not _pg_use:
+                if _pk_use and _pk_result is not None:
+                    try:
+                        # deferred PG forward, run only now that the packed reference exists
+                        _pg_result = _pg_forward_fn()
+                        _pg_W2 = logits_to_keep + max_left_pad
+                        _pg_cm = create_completion_attention_mask(
+                            input_ids[:, -_pg_W2:],
+                            left_pad_tokens_per_prompt,
+                            max_left_pad,
+                            self.processing_class.pad_token_id,
+                        ).float()
+                        _pg_a = _pg_result[:, -_pg_W2:].float()
+                        _pg_b = _pk_result[:, -_pg_W2:].float()
+                        _pg_diff = float(((_pg_a - _pg_b).abs() * _pg_cm).max())
+                        if UNSLOTH_ENABLE_LOGGING:
+                            print(
+                                f"[Unsloth] GRPO PrefixGrouper (no-grad) verify: sig={_pg_layout.signature} "
+                                f"shared-prefix vs full-row-packed max|d|={_pg_diff:.4f}",
+                                flush = True,
+                            )
+                        if _pg_diff < _pg_tol_ok():
+                            _pg_v = getattr(
+                                unwrapped_model, "_unsloth_prefix_grouper_nograd_verified", None
+                            )
+                            if not isinstance(_pg_v, dict):
+                                _pg_v = {}
+                            _pg_vT = int(_pg_layout.flat_ids.shape[1])
+                            _pg_vS = int(_pg_layout.position_ids.max()) + 1
+                            _pg_old = _pg_v.get(_pg_layout.signature, (0, 0))
+                            _pg_v[_pg_layout.signature] = (
+                                max(_pg_vT, _pg_old[0]),
+                                max(_pg_vS, _pg_old[1]),
+                            )
+                            unwrapped_model._unsloth_prefix_grouper_nograd_verified = _pg_v
+                            _pg_use = True
+                        else:
+                            _pg_u = getattr(
+                                unwrapped_model, "_unsloth_prefix_grouper_nograd_unsafe", None
+                            )
+                            if _pg_u is None:
+                                _pg_u = set()
+                            if _pg_diff >= _PG_TOL_KILL:
+                                _pg_u.add(_pg_layout.signature)
+                                unwrapped_model._unsloth_prefix_grouper_nograd_unsafe = _pg_u
+                            _pg_use = False
+                    except Exception as _pg_err3:
+                        _pg_result = None
+                        _pg_use = False
+                        if isinstance(_pg_err3, torch.cuda.OutOfMemoryError):
+                            torch.cuda.empty_cache()
+                        os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
+                        if UNSLOTH_ENABLE_LOGGING:
+                            print(
+                                f"[Unsloth] GRPO PrefixGrouper (no-grad) verify failed (fell back to packed): {_pg_err3!r}",
+                                flush = True,
+                            )
+                # else: no packed reference (packing off/failed) -> cannot verify; fall back.
+
+            if _pg_use and _pg_result is not None:
+                logprobs = _pg_result  # PrefixGrouper verified/trusted -> skip the loop
+                zipped_inputs = []
+            elif _pk_use and _pk_result is not None:
+                logprobs = _pk_result  # verified -> skip the loop
+                zipped_inputs = []
+            else:
+                # free packed intermediates before running the padded loop
+                _pk_hidden = _pk_sel = _pk_result = _pk_ref = None
+
             with _get_inference_mode_context_manager(model):
                 for (
                     input_ids_chunk,
@@ -1443,7 +1912,8 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
                     # However, it seems that this line does not slow down or disrupt models.
                     device_synchronize()
                     all_logprobs_list.append(logprobs_chunk)
-                logprobs = torch.cat(all_logprobs_list, dim = 0)
+                if logprobs is None:  # padded fallback when packing was not used
+                    logprobs = torch.cat(all_logprobs_list, dim = 0)
                 entropies = None
 
             os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0"
@@ -1523,6 +1993,34 @@ RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_accumulated_loss))
 RL_PRE_ITEMS["grpo_trainer"].append(grpo_compute_loss_slow)
 RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_update_SamplingParams))
 RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_get_inference_mode_context_manager))
+# inspect.getsource inlines function bodies but not module imports, so constants the inlined
+# grpo functions reference (e.g. UNSLOTH_ENABLE_LOGGING) must be redefined in the generated cache.
+RL_PRE_ITEMS["grpo_trainer"].append(
+    "import os as _unsloth_os\n"
+    "UNSLOTH_ENABLE_LOGGING = _unsloth_os.environ.get('UNSLOTH_ENABLE_LOGGING', '0') in ('1', 'True', 'true')\n"
+)
+# Sequence-packing gates, same values as the module-top constants.
+RL_PRE_ITEMS["grpo_trainer"].append(
+    "UNSLOTH_GRPO_SEQ_PACKING_ON = _unsloth_os.environ.get('UNSLOTH_GRPO_SEQ_PACKING', '1').lower() not in ('0', 'false', 'no', 'off')\n"
+)
+RL_PRE_ITEMS["grpo_trainer"].append(
+    "try:\n"
+    "    import inspect as _unsloth_inspect\n"
+    "    from unsloth_zoo.rl_replacements import RL_REPLACEMENTS as _unsloth_zoo_RL\n"
+    "    UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = 'torch.where(_keep, new' in _unsloth_inspect.getsource(_unsloth_zoo_RL['grpo_compute_loss'])\n"
+    "except Exception:\n"
+    "    UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = False\n"
+)
+# PrefixGrouper gate, same shape as the module-top constants.
+RL_PRE_ITEMS["grpo_trainer"].append(
+    "_pg_build_layout = _pg_enabled_fn = _pg_verify_on = _pg_tol_ok = _PG_TOL_KILL = None\n"
+    "UNSLOTH_GRPO_PREFIX_GROUPER_ON = _unsloth_os.environ.get('UNSLOTH_GRPO_PREFIX_GROUPER', '1').lower() not in ('0', 'false', 'no', 'off')\n"
+    "if UNSLOTH_GRPO_PREFIX_GROUPER_ON:\n"
+    "    try:\n"
+    "        from unsloth.utils.prefix_grouper import build_group_layout as _pg_build_layout, prefix_grouper_enabled as _pg_enabled_fn, verify_on as _pg_verify_on, tol_ok as _pg_tol_ok, TOL_KILL as _PG_TOL_KILL\n"
+    "    except Exception:\n"
+    "        UNSLOTH_GRPO_PREFIX_GROUPER_ON = False\n"
+)
 
 
 # Edit _get_per_token_logps to handle mixed precision

unsloth/models/sentence_transformer.py

@@ -19,6 +19,7 @@ from ._utils import (
     SUPPORTS_BFLOAT16,
     resolve_model_class,
     resolve_encoder_attention_implementation,
+    maybe_prefetch_hf_snapshot,
 )
 import inspect
 import json
@@ -541,7 +542,12 @@ class FastSentenceTransformer(FastModel):
         return transformer_module
 
     @staticmethod
-    def _read_pooling_mode(model_name, token):
+    def _read_pooling_mode(
+        model_name,
+        token,
+        cache_dir = None,
+        revision = None,
+    ):
         """Read the pooling mode from modules.json, else return "mean"."""
         try:
             if os.path.exists(model_name) and os.path.exists(
@@ -549,7 +555,13 @@ class FastSentenceTransformer(FastModel):
             ):
                 modules_json_path = os.path.join(model_name, "modules.json")
             else:
-                modules_json_path = hf_hub_download(model_name, "modules.json", token = token)
+                modules_json_path = hf_hub_download(
+                    model_name,
+                    "modules.json",
+                    token = token,
+                    cache_dir = cache_dir,
+                    revision = revision,
+                )
 
             with open(modules_json_path, "r", encoding = "utf-8") as f:
                 modules_config = json.load(f)
@@ -571,6 +583,8 @@ class FastSentenceTransformer(FastModel):
                                 model_name,
                                 os.path.join(pooling_path, "config.json"),
                                 token = token,
+                                cache_dir = cache_dir,
+                                revision = revision,
                             )
                         break
 
@@ -950,7 +964,12 @@ class FastSentenceTransformer(FastModel):
             f.write(content)
 
     @staticmethod
-    def _module_path(model_name, token = None):
+    def _module_path(
+        model_name,
+        token = None,
+        cache_dir = None,
+        revision = None,
+    ):
         """Return the path to the modules.json file, or None."""
         try:
             if os.path.exists(model_name) and os.path.isdir(model_name):
@@ -958,7 +977,13 @@ class FastSentenceTransformer(FastModel):
                 return path if os.path.exists(path) else None
             else:
                 try:
-                    return hf_hub_download(model_name, "modules.json", token = token)
+                    return hf_hub_download(
+                        model_name,
+                        "modules.json",
+                        token = token,
+                        cache_dir = cache_dir,
+                        revision = revision,
+                    )
                 except:
                     return None
         except:
@@ -1135,6 +1160,8 @@ class FastSentenceTransformer(FastModel):
         max_seq_length,
         pooling_mode,
         trust_remote_code = False,
+        cache_dir = None,
+        revision = None,
     ) -> tuple[OrderedDict, bool]:
         """Load modules from modules.json, else fall back to hard-coded modules.
 
@@ -1145,7 +1172,9 @@ class FastSentenceTransformer(FastModel):
         from sentence_transformers.models import Pooling, Normalize
 
         modules = OrderedDict()
-        modules_json_path = FastSentenceTransformer._module_path(model_name, token)
+        modules_json_path = FastSentenceTransformer._module_path(
+            model_name, token, cache_dir = cache_dir, revision = revision
+        )
 
         if modules_json_path:
             with open(modules_json_path, encoding = "utf8") as f:
@@ -1171,7 +1200,13 @@ class FastSentenceTransformer(FastModel):
                         load_path = os.path.join(model_name, module_path)
                     else:
                         try:
-                            load_path = load_dir_path(model_name, module_path, token = token)
+                            load_path = load_dir_path(
+                                model_name,
+                                module_path,
+                                token = token,
+                                cache_folder = cache_dir,
+                                revision = revision,
+                            )
                         except Exception as e:
                             print(f"Unsloth Warning: Could not download module {module_path}: {e}")
                             continue
@@ -1198,7 +1233,9 @@ class FastSentenceTransformer(FastModel):
         hidden_size = getattr(model.config, "hidden_size", 768)
 
         if pooling_mode == "mean":
-            pooling_mode = FastSentenceTransformer._read_pooling_mode(model_name, token)
+            pooling_mode = FastSentenceTransformer._read_pooling_mode(
+                model_name, token, cache_dir = cache_dir, revision = revision
+            )
 
         modules["1"] = Pooling(word_embedding_dimension = hidden_size, pooling_mode = pooling_mode)
         modules["2"] = Normalize()
@@ -1386,6 +1423,45 @@ class FastSentenceTransformer(FastModel):
                 "Run `pip install sentence-transformers` to install it."
             )
 
+        # Validate the load modes BEFORE the prefetch so a bad config fails without downloading weights.
+        # Guard on not for_inference: that branch below never used these flags.
+        if not for_inference:
+            # sanity check, thanks Etherl:
+            if full_finetuning and (load_in_4bit or load_in_8bit):
+                print(
+                    "Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA."
+                )
+                load_in_4bit = False
+                load_in_8bit = False
+                load_in_fp8 = False
+                load_in_16bit = False
+
+            if int(load_in_4bit) + int(load_in_8bit) + int(load_in_16bit) >= 2:
+                raise RuntimeError(
+                    "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!\n"
+                    "Also, we by default set `load_in_16bit = True`.\n"
+                    "If you want 4bit LoRA finetuning, set `load_in_16bit = False` and `load_in_4bit = True`\n"
+                    "If you want 8bit finetuning, set both `load_in_16bit = False` and `load_in_8bit = True`"
+                )
+
+        # Prefetch so the ST load below is a cache hit. weights_at_root stays False (ST component
+        # weights live in per-module subfolders). Resolve the same cache the load uses: HF cache_dir,
+        # else cache_folder, else SENTENCE_TRANSFORMERS_HOME, else default -- a wrong cache misses the warm.
+        _st_prefetched = maybe_prefetch_hf_snapshot(
+            model_name,
+            token = token,
+            revision = revision,
+            cache_dir = kwargs.get("cache_dir")
+            or kwargs.get("cache_folder")
+            or os.environ.get("SENTENCE_TRANSFORMERS_HOME"),
+            local_files_only = kwargs.get("local_files_only", False),
+            # Forward force_download so the refresh happens in the killable child, then clear it so the
+            # in-process ST load reuses the warm cache instead of re-downloading over unguarded Xet.
+            force_download = kwargs.get("force_download", False),
+        )
+        if _st_prefetched and kwargs.get("force_download", False):
+            kwargs["force_download"] = False
+
         # if for_inference == True, skip Unsloth optimizations to avoid torch compile issues
         if for_inference:
             st_device = device_map
@@ -1416,27 +1492,16 @@ class FastSentenceTransformer(FastModel):
                 if k in kwargs:
                     st_kwargs[k] = kwargs[k]
 
+            # ST takes cache_folder, not cache_dir: map cache_dir onto it so this load hits the warm
+            # (None lets ST honor SENTENCE_TRANSFORMERS_HOME, matching the prefetch).
+            _st_cache = kwargs.get("cache_dir") or kwargs.get("cache_folder")
+            if _st_cache is not None:
+                st_kwargs["cache_folder"] = _st_cache
+
             st_model = SentenceTransformer(model_name, **st_kwargs)
             return st_model
 
-        # sanity check, thanks Etherl:
-        if full_finetuning and (load_in_4bit or load_in_8bit):
-            print(
-                "Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA."
-            )
-            load_in_4bit = False
-            load_in_8bit = False
-            load_in_fp8 = False
-            load_in_16bit = False
-
-        if int(load_in_4bit) + int(load_in_8bit) + int(load_in_16bit) >= 2:
-            raise RuntimeError(
-                "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!\n"
-                "Also, we by default set `load_in_16bit = True`.\n"
-                "If you want 4bit LoRA finetuning, set `load_in_16bit = False` and `load_in_4bit = True`\n"
-                "If you want 8bit finetuning, set both `load_in_16bit = False` and `load_in_8bit = True`"
-            )
-
+        # Load-mode validation already ran before the prefetch above.
         if "auto_model" not in kwargs:
             kwargs["auto_model"] = AutoModel
 
@@ -1533,7 +1598,8 @@ class FastSentenceTransformer(FastModel):
                 elif is_mpnet:
                     FastSentenceTransformer._patch_mpnet_v5()
 
-            # Load via native SentenceTransformer (bypasses Unsloth patching)
+            # ST takes cache_folder, not cache_dir: map cache_dir onto it so this load hits the warm
+            # (None lets ST honor SENTENCE_TRANSFORMERS_HOME, matching the prefetch).
             st_model = SentenceTransformer(
                 model_name,
                 device = st_device,
@@ -1541,6 +1607,7 @@ class FastSentenceTransformer(FastModel):
                 token = token,
                 revision = revision,
                 model_kwargs = model_kwargs,
+                cache_folder = kwargs.get("cache_dir") or kwargs.get("cache_folder"),
             )
 
             # Store metadata for get_peft_model
@@ -1646,7 +1713,18 @@ class FastSentenceTransformer(FastModel):
 
         # No modules.json -> force 16-bit: saving is custom for these models and
         # 4-bit would need dequant in save_pretrained_merged, not worth it.
-        has_modules_json = FastSentenceTransformer._module_path(model_name, token) is not None
+        # Resolve the warmed cache: hf_hub_download ignores SENTENCE_TRANSFORMERS_HOME, so pass it as cache_dir.
+        has_modules_json = (
+            FastSentenceTransformer._module_path(
+                model_name,
+                token,
+                cache_dir = kwargs.get("cache_dir")
+                or kwargs.get("cache_folder")
+                or os.environ.get("SENTENCE_TRANSFORMERS_HOME"),
+                revision = revision,
+            )
+            is not None
+        )
 
         if not has_modules_json and load_in_4bit:
             print(
@@ -1656,6 +1734,12 @@ class FastSentenceTransformer(FastModel):
             load_in_4bit = False
             load_in_16bit = True
 
+        # The fallback FastModel load reads HF cache_dir, not ST's cache_folder/SENTENCE_TRANSFORMERS_HOME.
+        # Point it at the warmed cache, but only when no explicit cache_dir was passed (which wins).
+        _st_cache_dir = kwargs.get("cache_folder") or os.environ.get("SENTENCE_TRANSFORMERS_HOME")
+        if _st_cache_dir is not None and "cache_dir" not in kwargs:
+            kwargs["cache_dir"] = _st_cache_dir
+
         try:
             model, tokenizer = FastModel.from_pretrained(
                 model_name = model_name,
@@ -1697,6 +1781,12 @@ class FastSentenceTransformer(FastModel):
             max_seq_length,
             pooling_mode,
             trust_remote_code = trust_remote_code,
+            # Same resolved cache as above so the fallback module loads hit the warm, not Xet.
+            cache_dir = kwargs.get("cache_dir")
+            or kwargs.get("cache_folder")
+            or os.environ.get("SENTENCE_TRANSFORMERS_HOME"),
+            # Same revision as the weight load so modules hit the warm (None = default branch).
+            revision = revision,
         )
 
         st_device = device_map

unsloth/models/vision.py

@@ -37,6 +37,7 @@ from ._utils import (
     _get_text_only_config,
     _is_family_text_decoder,
     _apply_text_only_key_mapping,
+    _select_moe_detection_targets,
     set_task_config_attr,
 )
 from ._utils import *
@@ -551,6 +552,7 @@ def _construct_vlm_processor_fallback(
     model_type,
     token,
     trust_remote_code,
+    cache_dir = None,
     local_files_only = False,
 ):
     """Build a VLM processor manually when AutoProcessor.from_pretrained fails (some VLMs
@@ -568,6 +570,7 @@ def _construct_vlm_processor_fallback(
             tokenizer_name,
             token = token,
             trust_remote_code = trust_remote_code,
+            cache_dir = cache_dir,
             local_files_only = local_files_only,
         )
         # Load tokenizer via PreTrainedTokenizerFast (bypasses tokenizer_class check)
@@ -576,6 +579,7 @@ def _construct_vlm_processor_fallback(
             padding_side = "left",
             token = token,
             trust_remote_code = trust_remote_code,
+            cache_dir = cache_dir,
             local_files_only = local_files_only,
         )
         # Read tokenizer_config.json for special tokens: prefer the local file (offline
@@ -601,6 +605,7 @@ def _construct_vlm_processor_fallback(
                     tokenizer_name,
                     "tokenizer_config.json",
                     token = token,
+                    cache_dir = cache_dir,
                     local_files_only = local_files_only,
                 )
                 with open(config_path, "r", encoding = "utf-8") as f:
@@ -632,6 +637,7 @@ def _construct_vlm_processor_fallback(
                     tokenizer_name,
                     token = token,
                     trust_remote_code = trust_remote_code,
+                    cache_dir = cache_dir,
                     local_files_only = local_files_only,
                 )
                 proc_class_name = PROCESSOR_MAPPING_NAMES.get(config.model_type)
@@ -872,6 +878,9 @@ class FastBaseModel:
         # For debugging - we use a download counter to see if environments are not breaking or if HF is down
         get_statistics(kwargs.get("local_files_only", False))
 
+        # The base + tokenizer prefetch runs AFTER the load-mode validation below, so an invalid
+        # load_in_* combination fails without first downloading a snapshot.
+
         if dtype is None:
             dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16
         elif os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
@@ -968,6 +977,53 @@ class FastBaseModel:
             raise RuntimeError(
                 "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!"
             )
+
+        # Prefetch the repo (killable child) so the in-process load below is a cache hit. vLLM owns the
+        # weight download only when actually available; if fast_inference was requested but vLLM is
+        # missing, the load falls through in-process, so weights must still be warmed here.
+        _vllm_owns_weights = fast_inference and is_vLLM_available()
+        _prefetched = maybe_prefetch_hf_snapshot(
+            model_name,
+            token = token,
+            revision = kwargs.get("revision"),
+            cache_dir = kwargs.get("cache_dir"),
+            local_files_only = kwargs.get("local_files_only", False),
+            fast_inference = _vllm_owns_weights,
+            subfolder = kwargs.get("subfolder"),
+            force_download = kwargs.get("force_download", False),
+            use_safetensors = kwargs.get("use_safetensors"),
+            from_tf = kwargs.get("from_tf", False),
+            from_flax = kwargs.get("from_flax", False),
+            # Bare load reads only ROOT weights; skip subdir weights. Ignored when a subfolder is set.
+            weights_at_root = True,
+            variant = kwargs.get("variant"),  # forward so the warm keeps the variant .bin
+            gguf_file = kwargs.get(
+                "gguf_file"
+            ),  # forward so the warm fetches the GGUF (else ignored)
+        )
+        # Child did the forced download; clear the flag so the load reuses the warm cache.
+        if _prefetched and kwargs.get("force_download", False):
+            kwargs["force_download"] = False
+
+        # Warm a SEPARATE tokenizer repo only (model_name is covered above). Not model_name here: this
+        # runs before fast_inference_setup may remap the repo, so it would warm the wrong one.
+        _tokenizer_repo = (
+            tokenizer_name if (isinstance(tokenizer_name, str) and tokenizer_name) else model_name
+        )
+        _warm_tokenizer_repo = (
+            isinstance(_tokenizer_repo, str)
+            and bool(_tokenizer_repo)
+            and _tokenizer_repo != model_name
+        )
+        if _warm_tokenizer_repo:
+            maybe_prefetch_hf_snapshot(
+                _tokenizer_repo,
+                token = token,
+                cache_dir = kwargs.get("cache_dir"),
+                local_files_only = kwargs.get("local_files_only", False),
+                tokenizer_only = True,
+            )
+
         _skip_modules = SKIP_QUANTIZATION_MODULES.copy()
         # Nemotron-H uses 'mixer' (not 'mamba') for Mamba layers.
         # Mamba fused kernels pass out_proj.weight directly to F.linear,
@@ -1278,6 +1334,18 @@ class FastBaseModel:
         # Counteract saved tokenizers
         tokenizer_name = model_name if tokenizer_name is None else tokenizer_name
 
+        # On the vLLM path the tokenizer warm was deferred (fast_inference_setup may remap model_name).
+        # Warm the now-final tokenizer repo so the load below hits the cache (a cached/local repo is a no-op).
+        if _vllm_owns_weights and isinstance(tokenizer_name, str) and tokenizer_name:
+            maybe_prefetch_hf_snapshot(
+                tokenizer_name,
+                token = token,
+                revision = kwargs.get("revision"),
+                cache_dir = kwargs.get("cache_dir"),
+                local_files_only = kwargs.get("local_files_only", False),
+                tokenizer_only = True,
+            )
+
         # Fix _Unsloth_Patched_ prefix in local config files from old saves (issue #4085)
         if os.path.isdir(tokenizer_name):
             import json as _json
@@ -1315,6 +1383,7 @@ class FastBaseModel:
                         language = whisper_language,
                         task = whisper_task,
                         trust_remote_code = trust_remote_code,
+                        cache_dir = kwargs.get("cache_dir"),
                         local_files_only = lfo,
                     )
                 except Exception as _e:
@@ -1327,6 +1396,7 @@ class FastBaseModel:
                         padding_side = "left",
                         token = token,
                         trust_remote_code = trust_remote_code,
+                        cache_dir = kwargs.get("cache_dir"),
                         local_files_only = lfo,
                     )
                 except Exception as _e:
@@ -1337,6 +1407,7 @@ class FastBaseModel:
                             padding_side = "left",
                             token = token,
                             trust_remote_code = trust_remote_code,
+                            cache_dir = kwargs.get("cache_dir"),
                             local_files_only = lfo,
                         )
                     except Exception:
@@ -1355,6 +1426,7 @@ class FastBaseModel:
                         model_type_arch,
                         token,
                         trust_remote_code,
+                        cache_dir = kwargs.get("cache_dir"),
                         local_files_only = lfo,
                     )
                 except Exception as _fe:
@@ -1440,6 +1512,7 @@ class FastBaseModel:
                     padding_side = "left",
                     token = token,
                     trust_remote_code = trust_remote_code,
+                    cache_dir = kwargs.get("cache_dir"),
                     local_files_only = local_files_only,
                 )
                 model, _fallback_tok = patch_tokenizer(model, _fallback_tok)
@@ -1469,6 +1542,7 @@ class FastBaseModel:
                         padding_side = "left",
                         token = token,
                         trust_remote_code = trust_remote_code,
+                        cache_dir = kwargs.get("cache_dir"),
                         local_files_only = lfo,
                     )
                 except Exception:
@@ -1478,6 +1552,7 @@ class FastBaseModel:
                         padding_side = "left",
                         token = token,
                         trust_remote_code = trust_remote_code,
+                        cache_dir = kwargs.get("cache_dir"),
                         local_files_only = lfo,
                     )
 
@@ -1629,6 +1704,16 @@ class FastBaseModel:
             )
         else:
             _audio_kwargs = {}
+        # Remember the caller's ORIGINAL explicit leaf list for MoE expert
+        # detection. When an explicit list is routed through get_peft_regex for
+        # family scoping below, the generated regex carries get_peft_regex's full
+        # "mlp|feed_forward|ffn|dense" component block even when the caller named
+        # only attention leaves (q/k/v/o_proj). Keying expert detection on that
+        # regex would train the experts for an attention-only request. The
+        # original list carries the true leaf intent, so use it for MoE detection;
+        # only the auto (None / "all-linear") path relies on the regex, whose mlp
+        # block is the sole remaining MLP-intent signal on fused-expert models.
+        _moe_detect_target = target_modules if type(target_modules) in (list, tuple) else None
         if target_modules is None or target_modules == "all-linear":
             target_modules = get_peft_regex(
                 model,
@@ -1706,9 +1791,21 @@ class FastBaseModel:
             loftq_config, lora_dropout, bias, init_lora_weights, model
         )
 
-        # Auto-detect MoE models and populate target_parameters for expert layers
+        # Auto-detect MoE models and populate target_parameters for expert layers.
+        # Prefer the caller's ORIGINAL explicit leaf list over the scoped regex so an
+        # attention-only request does not train experts via get_peft_regex's mlp block,
+        # but only when MLP and language families are both still in scope. If the caller
+        # scoped MLP or language OFF (finetune_mlp_modules / finetune_language_layers
+        # False), the scoped regex already dropped the experts, so honor it instead of
+        # re-introducing the original list's gate/up/down leaves.
         if target_parameters is None:
-            target_parameters = get_moe_target_parameters(model, target_modules)
+            _moe_targets = _select_moe_detection_targets(
+                _moe_detect_target,
+                target_modules,
+                finetune_mlp_modules = finetune_mlp_modules,
+                finetune_language_layers = finetune_language_layers,
+            )
+            target_parameters = get_moe_target_parameters(model, _moe_targets)
 
         if finetune_last_n_layers is not None and layers_to_transform is None:
             _total_layers = _get_total_transformer_layers(model)
@@ -2207,3 +2304,16 @@ def check_dataset_for_missing_videos(
         warnings.warn(error_msg, stacklevel = 2)
 
     return missing
+
+
+# Auto-enable grouped-GEMM MoE (transformers<5 ModuleList experts); see llama.py.
+try:
+    from unsloth_zoo.temporary_patches.moe_grouped_modulelist import wrap_loader_for_grouped_moe
+    FastBaseModel.from_pretrained = staticmethod(
+        wrap_loader_for_grouped_moe(FastBaseModel.from_pretrained)
+    )
+    FastBaseModel.get_peft_model = staticmethod(
+        wrap_loader_for_grouped_moe(FastBaseModel.get_peft_model)
+    )
+except Exception:
+    pass

unsloth/save.py

@@ -223,6 +223,49 @@ def _normalize_torchao_method(save_method):
     return TORCHAO_EXPORT_SCHEMES.get(key)
 
 
+def _loaded_via_remote_code(obj):
+    """True if `obj`'s class comes from downloaded custom code (an auto_map module).
+
+    Transformers loads auto_map code into the ``transformers_modules`` package, so a
+    ``transformers_modules`` class proves the original load actually ran that remote code
+    (which the caller's / Studio's consent gate scans at load time). Export paths derive their
+    reload trust_remote_code from this - the already approved load decision - instead of from a
+    checkpoint's static ``auto_map``: a model that loads with built-in classes must not have its
+    unvetted remote code run when it is re-read during quantization export. Walks PEFT / wrapper
+    layers so a LoRA over a custom-code base is still detected, and processor components so a
+    custom tokenizer held inside a built-in processor keeps its approved trust.
+    """
+    seen = set()
+    queue = [obj]
+    while queue and len(seen) < 16:
+        node = queue.pop(0)
+        if node is None or id(node) in seen:
+            continue
+        seen.add(id(node))
+        # __module__ can be None/absent on some dynamically created or C-extension classes;
+        # treat anything non-string as "not remote code" rather than crashing the export.
+        module = getattr(type(node), "__module__", None)
+        if isinstance(module, str) and module.startswith("transformers_modules"):
+            return True
+        if hasattr(node, "get_base_model"):
+            try:
+                queue.append(node.get_base_model())
+            except Exception:
+                pass
+        # PEFT / trainer wrappers hold the real model in base_model / model; a built-in
+        # ProcessorMixin holds its (possibly custom-code) components as attributes.
+        for attr in (
+            "base_model",
+            "model",
+            "tokenizer",
+            "image_processor",
+            "feature_extractor",
+            "video_processor",
+        ):
+            queue.append(getattr(node, attr, None))
+    return False
+
+
 def _normalize_compressed_method(save_method):
     """Return (scheme, needs_calibration, suffix) if `save_method` is an FP8/FP4 compressed
     export, else None (so normal lora / merged_16bit / merged_4bit handling proceeds).
@@ -2883,15 +2926,16 @@ def unsloth_save_pretrained_gguf(
                 "Unsloth: quantization_method can only be a string or a list of strings"
             )
         for i, quant_method in enumerate(quantization_method):
-            quant_method = quant_method.lower()
+            if quant_method is None:
+                quant_method = "q8_0"
+            else:
+                quant_method = quant_method.lower()
             if quant_method == "not_quantized":
                 quant_method = "f16"
             elif quant_method == "fast_quantized":
                 quant_method = "q8_0"
             elif quant_method == "quantized":
                 quant_method = "q4_k_m"
-            elif quant_method is None:
-                quant_method = "q8_0"
             quantization_methods.append(quant_method.lower())
 
     try:
@@ -3532,7 +3576,9 @@ def _unsloth_save_lora_gguf(
             cmd += ["--base", base_model_id]
         else:
             cmd += ["--base-model-id", base_model_id]
-        if bool(getattr(model.config, "auto_map", None)):
+        # Only pass --trust-remote-code when the loaded model actually came from custom code (the
+        # approved load decision), not merely because its config carries an auto_map entry.
+        if _loaded_via_remote_code(model):
             cmd.append("--trust-remote-code")
 
         # Expose the token to the converter so it can fetch a gated/private base config from the Hub.
@@ -3682,15 +3728,16 @@ def save_to_gguf_generic(
                 "Unsloth: quantization_method can only be a string or a list of strings"
             )
         for i, quant_method in enumerate(quantization_method):
-            quant_method = quant_method.lower()
+            if quant_method is None:
+                quant_method = "q8_0"
+            else:
+                quant_method = quant_method.lower()
             if quant_method == "not_quantized":
                 quant_method = "f16"
             elif quant_method == "fast_quantized":
                 quant_method = "q8_0"
             elif quant_method == "quantized":
                 quant_method = "q4_k_m"
-            elif quant_method is None:
-                quant_method = "q8_0"
             new_quantization_methods.append(quant_method.lower())
     else:
         new_quantization_methods.append(quantization_type.lower())
@@ -4387,8 +4434,8 @@ def _unsloth_save_compressed_tensors(
         )
         unsloth_generic_save(**merge_args)
 
-        # 4) Detect VLM + trust_remote_code from the in-memory model config. A vision/multimodal
-        #    model exposes a vision_config or an explicitly vision-named architecture; a bare
+        # 4) Detect VLM from the in-memory model config. A vision/multimodal model exposes a
+        #    vision_config or an explicitly vision-named architecture; a bare
         #    *ForConditionalGeneration also matches text seq2seq models (T5/BART/Whisper), so it
         #    is not treated as a VLM on its own.
         is_vlm = False
@@ -4402,9 +4449,13 @@ def _unsloth_save_compressed_tensors(
                 "Unsloth: FP8/FP4 compressed export for vision / multimodal models is "
                 "experimental; vision-tower layers may be affected."
             )
-        trust_remote_code = (
-            bool(getattr(model.config, "auto_map", None)) if hasattr(model, "config") else False
-        )
+        # trust_remote_code must reflect the approved load decision (whether the model / tokenizer
+        # was actually loaded from custom code), not the config's static auto_map, so a
+        # built-in-loadable model carrying auto_map cannot run unvetted code in the subprocess.
+        # Model and tokenizer trust stay separate, like the torchao path: an approved custom
+        # tokenizer must not enable an unapproved model's code in the subprocess (or vice versa).
+        model_trust = _loaded_via_remote_code(model)
+        tok_trust = _loaded_via_remote_code(tokenizer)
 
         # 5) Marshal the calibration dataset for the subprocess: None -> ultrachat default; a
         #    str/PathLike is a local save_to_disk dir if it exists else a Hub id; Dataset -> temp.
@@ -4479,8 +4530,10 @@ def _unsloth_save_compressed_tensors(
             cmd += ["--calibration-dataset", calib_value]
         if is_vlm:
             cmd.append("--is-vlm")
-        if trust_remote_code:
+        if model_trust:
             cmd.append("--trust-remote-code")
+        if tok_trust:
+            cmd.append("--trust-remote-code-tokenizer")
         if variant:
             cmd += ["--variant", variant]
 
@@ -4679,35 +4732,19 @@ def _unsloth_save_torchao(
         )
         unsloth_generic_save(**merge_args)
 
-        # 2) Detect VLM + trust_remote_code so the right auto class reloads the staged checkpoint.
-        #    A bare *ForConditionalGeneration also matches text seq2seq (T5/BART/Whisper), so key off
-        #    vision_config / a vision-named architecture only, like the compressed path.
+        # 2) Detect VLM + reload class. A bare *ForConditionalGeneration also matches text seq2seq
+        #    (T5/BART/Whisper), so key off vision_config / a vision-named architecture only.
         is_vlm = False
-        trust_remote_code = False
         if hasattr(model, "config"):
             archs = getattr(model.config, "architectures", None) or []
             is_vlm = hasattr(model.config, "vision_config") or any(
                 x.endswith("ForVisionText2Text") for x in archs
             )
-            trust_remote_code = bool(getattr(model.config, "auto_map", None))
-        # Custom code can be declared only in the tokenizer/processor config, so also honor an
-        # auto_map in any staged config (the original load already had the user's consent).
-        if not trust_remote_code:
-            for _cfg in (
-                "config.json",
-                "tokenizer_config.json",
-                "processor_config.json",
-                "preprocessor_config.json",
-            ):
-                try:
-                    _p = os.path.join(staging, _cfg)
-                    if os.path.exists(_p):
-                        with open(_p, "r", encoding = "utf-8") as _f:
-                            if "auto_map" in json.load(_f):
-                                trust_remote_code = True
-                                break
-                except Exception:
-                    pass
+        # trust_remote_code must reflect the approved load decision - whether the in-memory model /
+        # tokenizer was itself loaded from custom code - not the staged config's auto_map, which an
+        # attacker can set on a built-in-loadable model to run unvetted code past the consent gate.
+        model_trust = _loaded_via_remote_code(model)
+        tok_trust = _loaded_via_remote_code(tokenizer)
         # Reload with the class that matches the checkpoint: an image-text VLM class (with a
         # fallback for older Transformers that lack AutoModelForImageTextToText); the model's own
         # architecture class for encoder-decoder seq2seq (T5/BART/Whisper are not causal LMs, and
@@ -4766,12 +4803,10 @@ def _unsloth_save_torchao(
             staging,
             device_map = "auto",
             quantization_config = TorchAoConfig(quant_type = quant_type),
-            trust_remote_code = trust_remote_code,
+            trust_remote_code = model_trust,
             **dtype_kw,
         )
-        staged_tokenizer = auto_processor.from_pretrained(
-            staging, trust_remote_code = trust_remote_code
-        )
+        staged_tokenizer = auto_processor.from_pretrained(staging, trust_remote_code = tok_trust)
 
         quantized_model.save_pretrained(out_dir, safe_serialization = safe_serialization)
         staged_tokenizer.save_pretrained(out_dir)

unsloth/tokenizer_utils.py

@@ -563,8 +563,11 @@ def _load_correct_tokenizer(
         # /tmp of Kaggle seems has a 80GB limit!
         # Let's utilize them
         cache_dir = os.path.join(KAGGLE_TMP, cache_dir)
-    else:
+    elif cache_dir == "huggingface_tokenizers_cache":
+        # This default name is Colab/Kaggle-only; elsewhere use the HF default cache.
         cache_dir = None
+    # else: keep a caller-supplied cache_dir so the tokenizer loads from the prefetch-warmed dir instead
+    # of risking an in-process Hub/Xet transfer.
 
     # Try loading the slow tokenizer. If it fails, then try Fast only
     # Mainly to solve Deepseek models with no tokenizer.model file
@@ -1323,6 +1326,7 @@ def check_tokenizer(
     padding_side = "right",
     token = None,
     _reload = True,
+    cache_dir = None,
 ):
     # Checks tokenizer for out of bounds ids.
     # Mainly a fix for https://huggingface.co/berkeley-nest/Starling-LM-7B-alpha
@@ -1413,10 +1417,11 @@ def check_tokenizer(
                     f"Fix your tokenizer since it'll perform out of bounds memory accesses."
                 )
 
-            if IS_COLAB_ENVIRONMENT or IS_KAGGLE_ENVIRONMENT:
-                cache_dir = "huggingface_tokenizers_cache"
-            else:
-                cache_dir = None
+            # Reuse a caller-supplied cache_dir (warmed cache) for the repair reload; else the
+            # Colab/Kaggle sentinel (HF default elsewhere), as load_correct_tokenizer does.
+            reload_cache_dir = cache_dir
+            if reload_cache_dir is None and (IS_COLAB_ENVIRONMENT or IS_KAGGLE_ENVIRONMENT):
+                reload_cache_dir = "huggingface_tokenizers_cache"
 
             # Sometimes slow tokenizer does not work like Deepseek
             try:
@@ -1430,7 +1435,7 @@ def check_tokenizer(
                     use_fast = False,
                     legacy = False,
                     from_slow = True,
-                    cache_dir = cache_dir,
+                    cache_dir = reload_cache_dir,
                 )
                 return check_tokenizer(
                     model = model,
@@ -1440,6 +1445,7 @@ def check_tokenizer(
                     padding_side = padding_side,
                     token = token,
                     _reload = False,
+                    cache_dir = cache_dir,
                 )
                 break
             except:

unsloth/utils/attention_dispatch.py

@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass
 from typing import Any, Optional, Tuple
 
@@ -42,6 +43,17 @@ if HAS_XFORMERS and torch.cuda.is_available():
         HAS_XFORMERS = False
 SDPA_HAS_GQA = "enable_gqa" in (scaled_dot_product_attention.__doc__ or "")
 
+# PrefixGrouper kernel, resolved once when the env gate is on so PG-off users never load
+# torch flex_attention.
+_flex_shared_prefix_attention = None
+if os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER", "1").lower() not in ("0", "false", "no", "off"):
+    try:
+        from .prefix_grouper_kernel import (
+            flex_shared_prefix_attention as _flex_shared_prefix_attention,
+        )
+    except Exception:
+        _flex_shared_prefix_attention = None
+
 FLASH_VARLEN = "flash_varlen"
 FLASH_DENSE = "flash_dense"
 XFORMERS = "xformers"
@@ -84,6 +96,9 @@ class AttentionContext:
     attention_mask: Optional[Tensor]
     causal_mask: Optional[Any]
     sliding_window: Optional[int] = None
+    # PrefixGrouper: non-None routes Q/K/V through the FlexAttention shared-prefix kernel;
+    # None leaves every existing construction/behavior unchanged.
+    prefix_seg_info: Optional[Any] = None
 
 
 def select_attention_backend(use_varlen: bool = False) -> str:
@@ -99,6 +114,33 @@ def select_attention_backend(use_varlen: bool = False) -> str:
     return SDPA
 
 
+def resolve_prefix_seg_info(kwargs, past_key_value, attention_mask):
+    """PrefixGrouper shared-prefix segment table resolver for the arch attention forwards.
+
+    The GRPO PrefixGrouper packed path rides a ``PrefixSegInfo`` in through ``**kwargs``
+    (same route as ``packed_seq_lengths``). When present, the forward must route Q/K/V
+    through the FlexAttention shared-prefix kernel via ``AttentionContext.prefix_seg_info``.
+
+    Returns the seg table (or ``None`` when PrefixGrouper did not group this batch -- the
+    unchanged path). Hardened: the shared-prefix stream is NOT a plain causal sequence, so running
+    it under a KV cache or an explicit padding mask would silently produce wrong logprobs.
+    That combination can only arise from misuse (PrefixGrouper only rides in via the GRPO
+    logprob forward, which is mask-free prefill), so we RAISE loudly instead of degrading
+    to a wrong result.
+
+    Factored here so every arch (llama/mistral/qwen3/gemma2/cohere/granite/falcon_h1)
+    shares one implementation and cannot drift.
+    """
+    seg = kwargs.get("prefix_seg_info", None)
+    if seg is not None and (past_key_value is not None or attention_mask is not None):
+        raise RuntimeError(
+            "PrefixGrouper: prefix_seg_info requires prefill with no KV cache and no "
+            f"attention_mask (got past_key_value={past_key_value is not None}, "
+            f"attention_mask={attention_mask is not None})."
+        )
+    return seg
+
+
 def run_attention(
     *, config: AttentionConfig, context: AttentionContext, Q: Tensor, K: Tensor, V: Tensor
 ) -> Tensor:
@@ -111,6 +153,28 @@ def run_attention(
     and SDPA handle packing via a block-diagonal mask.
     """
 
+    # PrefixGrouper shared-prefix attention (GRPO dedup). Q/K/V here are [bsz, H, T, D];
+    # the kernel takes/returns [1, T, H, D], matching the other backends. The field is
+    # only set when the env gate is on and grouping succeeded; None keeps every backend
+    # byte-identical.
+    if context.prefix_seg_info is not None:
+        flex_shared_prefix_attention = _flex_shared_prefix_attention
+        if flex_shared_prefix_attention is None:
+            # gate flipped on after import (or one-time load failed): resolve lazily.
+            from ..utils.prefix_grouper_kernel import flex_shared_prefix_attention
+
+        scale = None
+        if config.flash_varlen_kwargs:
+            scale = config.flash_varlen_kwargs.get("softmax_scale")
+        A = flex_shared_prefix_attention(
+            Q.transpose(1, 2),
+            K.transpose(1, 2),
+            V.transpose(1, 2),
+            context.prefix_seg_info,
+            scale = scale,
+        )
+        return A  # [1, T, n_heads, head_dim]
+
     backend = config.backend
     if backend == FLASH_VARLEN and context.seq_info is None:
         backend = FLASH_DENSE if HAS_FLASH_ATTENTION else SDPA
@@ -337,5 +401,6 @@ __all__ = [
     "AttentionConfig",
     "AttentionContext",
     "select_attention_backend",
+    "resolve_prefix_seg_info",
     "run_attention",
 ]

unsloth/utils/prefix_grouper.py

@@ -0,0 +1,351 @@
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""PrefixGrouper layout builder + completion-logprob extraction for the Unsloth GRPO
+packed path (all archs that route through the varlen attention dispatch).
+
+Given the de-padded, LEFT-PACKED input_ids the packed GRPO path already works with, this
+module:
+
+  1. Detects consecutive ``num_generations`` rows that share a prompt prefix (byte-
+     identical prompt precondition; falls back / returns None otherwise).
+  2. Builds ONE flat shared-prefix stream across all groups
+     ``[ prefix_g0, suf_g0_0 .. suf_g0_{G-1}, prefix_g1, ... ]`` with position_ids that
+     continue each prefix positionally, plus a ``PrefixSegInfo`` segment table for the
+     FlexAttention shared-prefix kernel.
+  3. Extracts completion logprobs via the index map (completion pos ``j==0`` predicted
+     from the shared prefix's last token; ``j>=1`` from the preceding suffix token) and
+     scatters them back into ``[total_rows, W]`` EXACTLY where the full-row packed path
+     puts them (dest = ``orig_row*L + orig_col``), so grpo_compute_loss / completion_mask
+     / TIS / metrics are byte-untouched.
+
+The flat stream is built by GATHERING original (row, col) coordinates out of input_ids,
+so the grad path's autograd flows to the same embedding rows as today (the shared prefix
+now contributes grad once = the sum of the G repeats, which is mathematically identical).
+
+``chunked_hidden_states_selective_log_softmax`` (from unsloth_zoo, passed in) is reused
+verbatim over the gathered predicting-position hidden states, so fp32 accumulation,
+logit_scale/softcapping/temperature are all preserved.
+
+Env:
+  UNSLOTH_GRPO_PREFIX_GROUPER=1            engage (default ON; set 0 to disable). Auto-off under vLLM.
+  UNSLOTH_GRPO_PREFIX_GROUPER_TOKR=1.3     tok_r auto-gate threshold (env-overridable)
+  UNSLOTH_GRPO_PREFIX_GROUPER_VERIFY=1     first-step self-verify (default ON)
+  UNSLOTH_GRPO_PREFIX_GROUPER_TOL=0.7      self-verify PASS band (nats)
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import torch
+
+from .prefix_grouper_kernel import build_seg_info_multigroup, PrefixSegInfo
+
+
+# ---------------------------------------------------------------------------
+# Env helpers
+# ---------------------------------------------------------------------------
+def env_on(name: str, default: str = "0") -> bool:
+    return os.environ.get(name, default).lower() not in ("0", "false", "no", "off")
+
+
+# One-time env reads; the helpers stay callable since unsloth_zoo imports and calls them.
+_ENABLED = env_on("UNSLOTH_GRPO_SEQ_PACKING", "1") and env_on("UNSLOTH_GRPO_PREFIX_GROUPER", "1")
+_VERIFY_ON = env_on("UNSLOTH_GRPO_PREFIX_GROUPER_VERIFY", "1")
+_TOKR_THRESHOLD = float(os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER_TOKR", "1.3"))
+_TOL_OK = float(os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER_TOL", "0.7"))
+
+
+def prefix_grouper_enabled() -> bool:
+    """PrefixGrouper requires seq-packing on (it reuses its de-pad + scatter machinery)."""
+    return _ENABLED
+
+
+def verify_on() -> bool:
+    return _VERIFY_ON
+
+
+def tokr_threshold() -> float:
+    return _TOKR_THRESHOLD
+
+
+def tol_ok() -> float:
+    return _TOL_OK
+
+
+# diff >= TOL_KILL = broken mask/isolation -> structure permanently unsafe; between
+# tol_ok and TOL_KILL -> fall back for this shape but keep trying others.
+TOL_KILL = 1.5
+
+
+@dataclass
+class GroupLayout:
+    """Everything the GRPO forward needs to run + extract the shared-prefix path."""
+
+    flat_ids: torch.Tensor  # [1, T]  (T == seg.T)
+    position_ids: torch.Tensor  # [1, T]
+    prefix_seg_info: PrefixSegInfo
+    # per completion target token, aligned 1:1:
+    tgt_rows: torch.Tensor  # [N] original row index
+    tgt_cols: torch.Tensor  # [N] original padded column in that row
+    tgt_pred: torch.Tensor  # [N] flat predicting index (into the T stream)
+    tgt_flat: torch.Tensor  # [N] flat index of the target token itself (into T)
+    total_rows: int
+    L: int  # original padded seq length (input_ids.shape[1])
+    W: int  # logits_to_keep + max_left_pad (scatter width)
+    tok_r: float
+    signature: Tuple
+
+    def extract_logps(
+        self,
+        hidden,
+        lm_head,
+        chunked_fn,
+        chunks,
+        logit_scale_multiply,
+        logit_scale_divide,
+        logit_softcapping,
+        temperature,
+    ) -> torch.Tensor:
+        """hidden: [1, T, Hdim] (pre-lm_head hidden states, UNSLOTH_RETURN_HIDDEN_STATES=1).
+        Returns [total_rows, W] float32, byte-compatible with the packed path result."""
+        # In a sharded model hidden may live on the lm-head device; move the small index
+        # maps to hidden.device before indexing.
+        device = hidden.device
+        pred_h = hidden[0, self.tgt_pred.to(device), :].unsqueeze(0)  # [1, N, Hdim]
+        tgt_ids = self.flat_ids[0, self.tgt_flat].to(device).unsqueeze(0)  # [1, N]
+        sel = chunked_fn(
+            pred_h,
+            lm_head,
+            tgt_ids,
+            chunks,
+            logit_scale_multiply,
+            logit_scale_divide,
+            logit_softcapping,
+            temperature,
+        )[0]  # [N] logprobs
+        dest = self.tgt_rows.to(device) * self.L + self.tgt_cols.to(device)
+        result = (
+            torch.zeros(self.total_rows * self.L, dtype = torch.float32, device = device)
+            .index_put((dest,), sel.to(torch.float32))
+            .view(self.total_rows, self.L)[:, -self.W :]
+        )
+        return result
+
+
+def _build_groups(ids_cpu, real_cols_cpu, cstart_cpu, num_generations, total_rows):
+    """CPU-side grouping. Returns group dicts or None. Mirrors the packed _pk_* partition.
+
+    A row's REAL tokens are the columns where input != pad. Its completion region (what
+    the packed path scatters, then completion_mask masks) is the real columns with
+    original col >= cstart_r, where cstart_r = (L - logits_to_keep) - left_pad_r. The
+    prompt is the real columns < cstart_r. Within a GRPO group all G rows share the same
+    prompt => same left_pad => same cstart => the prompt real columns are BYTE-IDENTICAL
+    across the group (the shared prefix). We require that byte-identity (falls back
+    otherwise). No prompt-tail special-casing: every suffix token is scattered exactly
+    like the packed path; completion_mask masks the leading prompt-tail positions.
+    """
+    G = num_generations
+    if G is None or G < 2 or total_rows % G != 0:
+        return None
+    groups = []
+    for g0 in range(0, total_rows, G):
+        rows = list(range(g0, g0 + G))
+        prompt_cols_per_row = []  # real cols < cstart
+        prompt_toks_per_row = []
+        comp_cols_per_row = []  # real cols >= cstart  (the completion region packed scatters)
+        for r in rows:
+            cs = cstart_cpu[r]
+            rc = real_cols_cpu[r]
+            p_cols = [c for c in rc if c < cs]
+            c_cols = [c for c in rc if c >= cs]
+            prompt_cols_per_row.append(p_cols)
+            prompt_toks_per_row.append([ids_cpu[r][c] for c in p_cols])
+            comp_cols_per_row.append(c_cols)
+        if any(len(p) == 0 for p in prompt_toks_per_row):
+            return None
+        # require BYTE-IDENTICAL prompts across the group (shared-prefix precondition).
+        P = len(prompt_toks_per_row[0])
+        if any(len(prompt_toks_per_row[k]) != P for k in range(1, G)):
+            return None
+        p0 = prompt_toks_per_row[0]
+        if any(prompt_toks_per_row[k] != p0 for k in range(1, G)):
+            return None
+        if P == 0:
+            return None
+        R_list = [len(c) for c in comp_cols_per_row]
+        if sum(R_list) == 0:
+            return None
+        groups.append(
+            dict(
+                rows = rows,
+                P = P,
+                prefix_cols = prompt_cols_per_row[0],  # shared prompt real columns (row0)
+                prefix_row = rows[0],
+                R_list = R_list,
+                suf_cols = comp_cols_per_row,  # per-row completion-region real columns
+            )
+        )
+    return groups
+
+
+def _tok_r(groups) -> float:
+    tok_full = 0
+    tok_sp = 0
+    for gm in groups:
+        P = gm["P"]
+        Rs = gm["R_list"]
+        tok_full += sum(P + r for r in Rs)  # G*P + sumR
+        tok_sp += P + sum(Rs)  # P + sumR
+    return (tok_full / tok_sp) if tok_sp else 1.0
+
+
+def build_group_layout(
+    input_ids,
+    logits_to_keep,
+    pad_id,
+    num_generations,
+    left_pad_tokens_per_prompt,
+    *,
+    apply_tokr_gate = True,
+    max_segment_cap = None,
+):
+    """Build the shared-prefix GroupLayout, or return None to fall back to the packed path.
+
+    input_ids : [B, L]. GRPO's layout is left-padded in the prompt and right-padded in
+        the completion. Real tokens of a row are a contiguous run not necessarily
+        starting at column 0.
+    logits_to_keep : int
+    left_pad_tokens_per_prompt : [B] long tensor (per-row left-pad count in the prompt).
+    """
+    device = input_ids.device
+    total_rows, L = input_ids.shape
+    keep = input_ids != pad_id
+    # completion start column per row (matches create_completion_attention_mask / _pk_cstart).
+    cstart = ((L - logits_to_keep) - left_pad_tokens_per_prompt).to(torch.long)
+    cstart_cpu = cstart.tolist()
+    ids_cpu = input_ids.tolist()
+    # per-row real (non-pad) columns. GRPO rows are one contiguous real run, so derive
+    # [first, first+n) on GPU; the O(B*L) scan is only a non-contiguous fallback.
+    n_real = keep.sum(dim = 1)
+    first = torch.argmax(keep.to(torch.int8), dim = 1)
+    ar = torch.arange(L, device = device)
+    contiguous = bool(
+        (keep == ((ar >= first.unsqueeze(1)) & (ar < (first + n_real).unsqueeze(1)))).all()
+    )
+    if contiguous:
+        real_cols_cpu = [list(range(f, f + n)) for f, n in zip(first.tolist(), n_real.tolist())]
+    else:
+        keep_cpu = keep.tolist()
+        real_cols_cpu = [[c for c in range(L) if keep_cpu[r][c]] for r in range(total_rows)]
+
+    groups = _build_groups(ids_cpu, real_cols_cpu, cstart_cpu, num_generations, total_rows)
+    if groups is None:
+        return None
+
+    # sliding-window guard: a group's PG span is P + max(R); fall back if it exceeds the window.
+    if max_segment_cap is not None:
+        for gm in groups:
+            if gm["P"] + max(gm["R_list"]) > max_segment_cap:
+                return None
+
+    tok_r = _tok_r(groups)
+    if apply_tokr_gate and tok_r < tokr_threshold():
+        return None  # low reuse -> not worth it; use the full-row packed path
+
+    # Build flat stream by gathering original (row, col) coordinates.
+    group_specs = [(gm["P"], gm["R_list"]) for gm in groups]
+    seg, group_meta = build_seg_info_multigroup(group_specs, device)
+
+    flat_src_rows: List[int] = []
+    flat_src_cols: List[int] = []
+    pos_list: List[int] = []
+    tgt_rows: List[int] = []
+    tgt_cols: List[int] = []
+    tgt_pred: List[int] = []
+    tgt_flat: List[int] = []
+
+    for gm, meta in zip(groups, group_meta):
+        rows = gm["rows"]
+        P = gm["P"]
+        r0 = gm["prefix_row"]
+        prefix_cols = gm["prefix_cols"]  # ORIGINAL real prompt columns (len P) of row0
+        plast = meta["prefix_last_index"]  # base + P - 1
+        # gather the shared prefix once, from row0.
+        flat_src_rows.extend([r0] * P)
+        flat_src_cols.extend(prefix_cols)
+        pos_list.extend(range(P))
+        # suffixes: every suffix token is a completion-region target (scattered like the
+        # packed path; completion_mask hides prompt-tail positions).
+        for i, r in enumerate(rows):
+            cols = gm["suf_cols"][i]
+            r_i = len(cols)
+            s, e = meta["suffix_slices"][i]  # flat offsets [s, e)
+            flat_src_rows.extend([r] * r_i)
+            flat_src_cols.extend(cols)
+            pos_list.extend(range(P, P + r_i))
+            for j in range(r_i):
+                # pos 0 is predicted from the prefix's last token; j>=1 from the previous suffix token.
+                pred = plast if j == 0 else (s + j - 1)
+                tgt_rows.append(r)
+                tgt_cols.append(cols[j])  # ORIGINAL padded column in row r
+                tgt_pred.append(pred)
+                tgt_flat.append(s + j)  # flat index of the target token itself
+
+    T = len(flat_src_rows)
+    assert T == seg.T, f"flat stream len {T} != seg.T {seg.T}"
+    fr = torch.tensor(flat_src_rows, device = device, dtype = torch.long)
+    fc = torch.tensor(flat_src_cols, device = device, dtype = torch.long)
+    flat_ids = input_ids[fr, fc].unsqueeze(0)  # [1, T] (grad-safe gather)
+    position_ids = torch.tensor(pos_list, device = device, dtype = torch.long).unsqueeze(0)
+
+    max_left_pad = int(left_pad_tokens_per_prompt.max().item()) if total_rows else 0
+    W = logits_to_keep + max_left_pad
+
+    # self-verify cache key: the mask/index-map/scatter logic is structural, so key on
+    # (num_groups, group_sizes), not exact lengths -- GRPO lengths change every step and
+    # keying on T would re-verify forever ("verify once, then trust", like the packed path).
+    grp_sizes = tuple(sorted(len(gm["R_list"]) for gm in groups))
+    sig = (len(groups), grp_sizes)
+
+    return GroupLayout(
+        flat_ids = flat_ids,
+        position_ids = position_ids,
+        prefix_seg_info = seg,
+        tgt_rows = torch.tensor(tgt_rows, device = device, dtype = torch.long),
+        tgt_cols = torch.tensor(tgt_cols, device = device, dtype = torch.long),
+        tgt_pred = torch.tensor(tgt_pred, device = device, dtype = torch.long),
+        tgt_flat = torch.tensor(tgt_flat, device = device, dtype = torch.long),
+        total_rows = total_rows,
+        L = L,
+        W = W,
+        tok_r = tok_r,
+        signature = sig,
+    )
+
+
+__all__ = [
+    "GroupLayout",
+    "build_group_layout",
+    "prefix_grouper_enabled",
+    "verify_on",
+    "tokr_threshold",
+    "tol_ok",
+    "TOL_KILL",
+    "env_on",
+]

unsloth/utils/prefix_grouper_kernel.py

@@ -0,0 +1,436 @@
+# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""FlexAttention shared-prefix kernel for PrefixGrouper (GRPO shared-prompt dedup).
+
+In GRPO every prompt spawns ``G = num_generations`` completions that share the same
+prompt prefix. The full-row packed path forwards the identical prefix ``G`` times.
+PrefixGrouper stores the prefix ONCE and concatenates only the ``G`` suffixes, with an
+attention layout where each suffix token attends to ``[the single shared prefix] +
+[causal within its own suffix]``. This kernel expresses that one-prefix -> many-suffix
+fan-out via a ``torch.nn.attention.flex_attention`` block mask, so the masked-out
+cross-suffix / cross-group blocks are never computed and the ``P + G*R`` FLOP saving is
+realised (not merely a masked dense ``O(T^2)``).
+
+Mask semantics (identical to the certified SDPA oracle):
+
+    keep(q_idx, kv_idx) = same_group(q, kv) AND
+        ( is_prefix[kv_idx]                                  # full prefix visibility
+          OR ( suffix_of_kv[kv_idx] == suffix_of_kv[q_idx]   # same suffix ...
+               AND kv_idx <= q_idx ) )                       # ... causal within it
+
+This module is self-contained (no dependency on any temp/ scratch dir) so PrefixGrouper
+works from the installed source after a fresh compile. It is only imported lazily from
+``attention_dispatch.run_attention`` when ``prefix_seg_info`` is present, which itself is
+only ever set when ``UNSLOTH_GRPO_PREFIX_GROUPER`` is on and grouping succeeded, so the
+default (off) path never touches this file.
+
+Provided entry points:
+  * ``PrefixSegInfo``            : per-flat-token segment metadata + cache signature.
+  * ``build_seg_info_multigroup``: build PrefixSegInfo for many groups packed flat.
+  * ``build_seg_info_from_layout``: build PrefixSegInfo for ONE group (test helper).
+  * ``get_block_mask``          : cached create_block_mask keyed on the signature.
+  * ``flex_shared_prefix_attention(Q, K, V, prefix_seg_info)``
+        Q/K/V of shape [1, T, n_heads, head_dim]; returns [1, T, n_heads, head_dim],
+        IDENTICAL semantics to the SDPA oracle.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import torch
+from torch.nn.attention.flex_attention import (
+    BlockMask,
+    create_block_mask,
+    flex_attention,
+)
+
+# GRPO feeds many distinct segment lengths; at dynamo's default recompile_limit (8) the
+# compiled kernel silently reuses a mismatched specialisation (wrong results). Raise it.
+torch._dynamo.config.recompile_limit = max(getattr(torch._dynamo.config, "recompile_limit", 8), 256)
+torch._dynamo.config.accumulated_recompile_limit = max(
+    getattr(torch._dynamo.config, "accumulated_recompile_limit", 256), 2048
+)
+
+
+# Compiled kernels: torch.compile fuses the sparse mask into one kernel. dynamic=True is
+# required: T changes almost every GRPO batch and dynamic=False recompiles per T (~14s
+# each). T is still padded to a multiple of 128 (_pad_len) for the backward kernel.
+_flex_attention_compiled = torch.compile(flex_attention, dynamic = True)
+_create_block_mask_compiled = torch.compile(create_block_mask, dynamic = True)
+
+# Flash block sizes by Q dtype (env-overridable). The two disjoint key runs (prefix +
+# own-suffix) stress online-softmax accumulation: fp32 needs 32/32 for a ~1e-6 floor;
+# bf16 passes parity at 128/64 and is ~5x faster (128/128 OOMs Triton on B200).
+_FP32_BLOCK_M = int(os.environ.get("PG_FLEX_BLOCK_M", "32"))
+_FP32_BLOCK_N = int(os.environ.get("PG_FLEX_BLOCK_N", "32"))
+_BF16_BLOCK_M = int(os.environ.get("PG_FLEX_BF16_BLOCK_M", "128"))
+_BF16_BLOCK_N = int(os.environ.get("PG_FLEX_BF16_BLOCK_N", "64"))
+
+
+def _kernel_options_for_dtype(dtype):
+    """Pick the numerically-safe flash block sizes for the Q dtype."""
+    if dtype == torch.bfloat16 or dtype == torch.float16:
+        return {"BLOCK_M": _BF16_BLOCK_M, "BLOCK_N": _BF16_BLOCK_N}
+    return {"BLOCK_M": _FP32_BLOCK_M, "BLOCK_N": _FP32_BLOCK_N}
+
+
+# Backward-compat constant (fp32 default).
+_FLEX_KERNEL_OPTIONS = {"BLOCK_M": _FP32_BLOCK_M, "BLOCK_N": _FP32_BLOCK_N}
+
+# The compiled backward trips an Inductor assertion when T is not a multiple of 128, so
+# pad the flat sequence. Pad tokens form a group that attends to / is attended by nothing
+# (all-masked rows return 0, not NaN) and are sliced off the output.
+_PAD_MULTIPLE = 128
+_PAD_GROUP = -99  # sentinel group id / suffix id for pad tokens
+
+
+def _pad_len(T: int) -> int:
+    return ((T + _PAD_MULTIPLE - 1) // _PAD_MULTIPLE) * _PAD_MULTIPLE
+
+
+# ---------------------------------------------------------------------------
+# Segment metadata
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class PrefixSegInfo:
+    """Per-flat-token segment metadata driving the shared-prefix block mask.
+
+    The label tensors are 1-D of length ``T_pad`` (>= real ``T``, padded up to a multiple
+    of 128 so the backward kernel compiles). Positions ``[T:T_pad)`` are pad tokens
+    (group/suffix == _PAD_GROUP) that attend to nothing.
+
+    Attributes
+    ----------
+    group_of_kv : LongTensor [T_pad]
+        Group id per flat token (0..num_groups-1); _PAD_GROUP for pad tokens.
+    is_prefix : BoolTensor [T_pad]
+        True iff the token is a prefix token of its group (False for pad).
+    suffix_of_kv : LongTensor [T_pad]
+        Suffix id per flat token; -1 for prefix, _PAD_GROUP for pad. Suffix ids are
+        globally unique across groups.
+    signature : hashable
+        Cache key for the block mask (depends only on the labels + T_pad).
+    T : int
+        Real flat sequence length (Q/K/V of this length are padded internally).
+    T_pad : int
+        Padded length (multiple of 128) at which the block mask is built.
+    """
+
+    group_of_kv: torch.Tensor
+    is_prefix: torch.Tensor
+    suffix_of_kv: torch.Tensor
+    signature: Tuple
+    T: int
+    T_pad: int
+
+
+def _pad_labels(group_of_kv, is_prefix, suffix_of_kv, device):
+    """Pad the label tensors up to a multiple of 128 with pad-token sentinels."""
+    T = int(group_of_kv.numel())
+    T_pad = _pad_len(T)
+    if T_pad == T:
+        return group_of_kv, is_prefix, suffix_of_kv, T, T_pad
+    pad = T_pad - T
+    group_of_kv = torch.cat(
+        [group_of_kv, torch.full((pad,), _PAD_GROUP, dtype = torch.long, device = device)]
+    )
+    is_prefix = torch.cat([is_prefix, torch.zeros(pad, dtype = torch.bool, device = device)])
+    suffix_of_kv = torch.cat(
+        [suffix_of_kv, torch.full((pad,), _PAD_GROUP, dtype = torch.long, device = device)]
+    )
+    return group_of_kv, is_prefix, suffix_of_kv, T, T_pad
+
+
+def build_seg_info_from_layout(layout, device: Optional[torch.device] = None) -> PrefixSegInfo:
+    """Build PrefixSegInfo for ONE group from an object with ``.flat_ids``, ``.P`` and
+    ``.suffix_slices`` (used by the parity test / oracle helpers)."""
+    if device is None:
+        device = layout.flat_ids.device
+    T = int(layout.flat_ids.shape[1])
+    P = int(layout.P)
+
+    group_of_kv = torch.zeros(T, dtype = torch.long, device = device)  # single group -> 0
+    is_prefix = torch.zeros(T, dtype = torch.bool, device = device)
+    is_prefix[:P] = True
+    suffix_of_kv = torch.full((T,), -1, dtype = torch.long, device = device)
+    for i, (s, e) in enumerate(layout.suffix_slices):
+        suffix_of_kv[s:e] = i
+
+    group_of_kv, is_prefix, suffix_of_kv, T, T_pad = _pad_labels(
+        group_of_kv, is_prefix, suffix_of_kv, device
+    )
+    sig = ("single", T_pad, P, tuple((s, e) for (s, e) in layout.suffix_slices))
+    return PrefixSegInfo(
+        group_of_kv = group_of_kv,
+        is_prefix = is_prefix,
+        suffix_of_kv = suffix_of_kv,
+        signature = sig,
+        T = T,
+        T_pad = T_pad,
+    )
+
+
+def build_seg_info_multigroup(
+    group_specs: List[Tuple[int, List[int]]], device: torch.device
+) -> Tuple[PrefixSegInfo, List[dict]]:
+    """Build PrefixSegInfo for several shared-prefix groups packed block-diagonally.
+
+    Parameters
+    ----------
+    group_specs : list of (P_g, [R_{g,0}, R_{g,1}, ...])
+        For each group: prefix length and the list of suffix lengths.
+
+    Returns
+    -------
+    seg : PrefixSegInfo
+    group_meta : list of dicts with 'base', 'P', 'prefix_last_index', 'suffix_slices'
+        (flat offsets), enough to build the completion index map.
+    """
+    group_of_list = []
+    is_prefix_list = []
+    suffix_of_list = []
+    group_meta = []
+
+    base = 0
+    suffix_counter = 0
+    sig_parts = []
+    for gid, (P, R_list) in enumerate(group_specs):
+        # prefix
+        group_of_list.append(torch.full((P,), gid, dtype = torch.long, device = device))
+        is_prefix_list.append(torch.ones(P, dtype = torch.bool, device = device))
+        suffix_of_list.append(torch.full((P,), -1, dtype = torch.long, device = device))
+        prefix_last_index = base + P - 1
+        suffix_slices = []
+        cursor = base + P
+        for r in R_list:
+            group_of_list.append(torch.full((r,), gid, dtype = torch.long, device = device))
+            is_prefix_list.append(torch.zeros(r, dtype = torch.bool, device = device))
+            suffix_of_list.append(torch.full((r,), suffix_counter, dtype = torch.long, device = device))
+            suffix_slices.append((cursor, cursor + r))
+            cursor += r
+            suffix_counter += 1
+        group_meta.append(
+            {
+                "base": base,
+                "P": P,
+                "prefix_last_index": prefix_last_index,
+                "suffix_slices": suffix_slices,
+            }
+        )
+        sig_parts.append((P, tuple(R_list)))
+        base = cursor
+
+    group_of_kv = torch.cat(group_of_list)
+    is_prefix = torch.cat(is_prefix_list)
+    suffix_of_kv = torch.cat(suffix_of_list)
+    group_of_kv, is_prefix, suffix_of_kv, T, T_pad = _pad_labels(
+        group_of_kv, is_prefix, suffix_of_kv, device
+    )
+    sig = ("multi", T_pad, tuple(sig_parts))
+    seg = PrefixSegInfo(
+        group_of_kv = group_of_kv,
+        is_prefix = is_prefix,
+        suffix_of_kv = suffix_of_kv,
+        signature = sig,
+        T = T,
+        T_pad = T_pad,
+    )
+    return seg, group_meta
+
+
+# ---------------------------------------------------------------------------
+# Block-mask builder + cache, keyed on (signature, device): the mask depends only on the
+# per-token labels and T, so it is reused across layers and steps.
+
+_BLOCK_MASK_CACHE: Dict[Tuple, BlockMask] = {}
+
+
+def _make_mask_mod(group_of_kv, is_prefix, suffix_of_kv):
+    """Return a mask_mod closure over the (device) label tensors.
+
+    keep(q, kv) = same_group AND
+        ( is_prefix[kv]  AND kv <= q                     # causal within/ into prefix
+          OR ( suffix_of_kv[kv] == suffix_of_kv[q]       # same suffix ...
+               AND (not is_prefix[q])                    # q is a suffix token ...
+               AND kv <= q ) )                           # ... causal within it
+
+    The single ``kv <= q`` guard on the is_prefix branch gives BOTH prefix-causal
+    behaviour (a prefix q sees only earlier prefix tokens) AND full-prefix-visibility for
+    suffixes (every prefix index < every suffix index in a group, so kv <= q always holds
+    for a suffix q vs a prefix kv of its group), matching the SDPA oracle exactly.
+    """
+
+    def mask_mod(b, h, q_idx, kv_idx):
+        same_group = group_of_kv[q_idx] == group_of_kv[kv_idx]
+        kv_is_prefix = is_prefix[kv_idx]
+        causal = kv_idx <= q_idx
+        same_suffix = (suffix_of_kv[kv_idx] == suffix_of_kv[q_idx]) & (~is_prefix[q_idx])
+        keep = same_group & ((kv_is_prefix & causal) | (same_suffix & causal))
+        return keep
+
+    return mask_mod
+
+
+def get_block_mask(
+    seg: PrefixSegInfo,
+    device: torch.device,
+    compile_mask: bool = True,
+) -> BlockMask:
+    """Return a cached BlockMask for the segment signature (built once, reused).
+
+    CRITICAL: the block mask is cached and shared across BOTH the no-grad old/ref logprob
+    forward (which runs under torch.inference_mode) and the grad training forward. If the
+    mask were first built under inference_mode, its tensors would be INFERENCE tensors that
+    "cannot be saved for backward" when reused in the grad forward. We therefore build the
+    mask with inference mode explicitly DISABLED, so the same cached BlockMask is a normal
+    tensor usable by autograd. (The mask depends only on integer labels; it needs no grad.)
+    """
+    key = (seg.signature, str(device))
+    bm = _BLOCK_MASK_CACHE.get(key)
+    if bm is not None:
+        return bm
+
+    # Move labels to the consumer (Q) device: with a sharded model the seg tensors live on
+    # input_ids.device and would index cross-device. Copies once per (signature, device).
+    # These copies must also run with inference mode DISABLED (same reason as the mask build):
+    # when this entry is first built under the no-grad old/ref forward's inference_mode and
+    # device != seg.device, a .to(device) copy would be an inference tensor that mask_mod
+    # captures, which then cannot be saved for backward when the grad training forward reuses
+    # the cached mask.
+    builder = _create_block_mask_compiled if compile_mask else create_block_mask
+    with torch.inference_mode(False):
+        mask_mod = _make_mask_mod(
+            seg.group_of_kv.to(device), seg.is_prefix.to(device), seg.suffix_of_kv.to(device)
+        )
+        bm = builder(
+            mask_mod,
+            B = 1,
+            H = None,
+            Q_LEN = seg.T_pad,
+            KV_LEN = seg.T_pad,
+            device = device,
+        )
+    # FIFO bound: GRPO lengths change nearly every step, so evict the oldest to cap GPU pins.
+    if len(_BLOCK_MASK_CACHE) >= 8:
+        _BLOCK_MASK_CACHE.pop(next(iter(_BLOCK_MASK_CACHE)))
+    _BLOCK_MASK_CACHE[key] = bm
+    return bm
+
+
+def clear_block_mask_cache():
+    _BLOCK_MASK_CACHE.clear()
+
+
+def _pad_qkv_seq(x: torch.Tensor, T_pad: int) -> torch.Tensor:
+    """Zero-pad a [B, H, T, D] tensor along the sequence dim up to T_pad."""
+    T = x.shape[2]
+    if T_pad == T:
+        return x
+    pad = torch.zeros(x.shape[0], x.shape[1], T_pad - T, x.shape[3], device = x.device, dtype = x.dtype)
+    return torch.cat([x, pad], dim = 2)
+
+
+def _run_flex(q, k, v, block_mask, enable_gqa, scale, compiled, T, T_pad):
+    """Pad q/k/v to T_pad, run flex, slice the output back to T. q/k/v: [B,H,T,D]."""
+    qp = _pad_qkv_seq(q, T_pad)
+    kp = _pad_qkv_seq(k, T_pad)
+    vp = _pad_qkv_seq(v, T_pad)
+    if compiled:
+        out = _flex_attention_compiled(
+            qp,
+            kp,
+            vp,
+            block_mask = block_mask,
+            enable_gqa = enable_gqa,
+            scale = scale,
+            kernel_options = _kernel_options_for_dtype(qp.dtype),
+        )
+    else:
+        # eager path (fp64 parity): dense scores, no kernel_options.
+        out = flex_attention(
+            qp,
+            kp,
+            vp,
+            block_mask = block_mask,
+            enable_gqa = enable_gqa,
+            scale = scale,
+        )
+    return out[:, :, :T, :]
+
+
+# ---------------------------------------------------------------------------
+# The kernel entry point
+# ---------------------------------------------------------------------------
+
+
+def flex_shared_prefix_attention(
+    Q: torch.Tensor,
+    K: torch.Tensor,
+    V: torch.Tensor,
+    prefix_seg_info: PrefixSegInfo,
+    scale: Optional[float] = None,
+    block_mask: Optional[BlockMask] = None,
+    compiled: bool = True,
+) -> torch.Tensor:
+    """Shared-prefix attention via FlexAttention.
+
+    Parameters
+    ----------
+    Q, K, V : Tensor [1, T, n_heads, head_dim]
+        (Q has n_heads, K/V have n_kv_heads for GQA).
+    prefix_seg_info : PrefixSegInfo
+    scale : optional float, softmax scale (defaults to 1/sqrt(head_dim)).
+    block_mask : optional precomputed BlockMask (else built/cached from seg info).
+
+    Returns
+    -------
+    Tensor [1, T, n_heads, head_dim], identical semantics to the SDPA oracle branch.
+    """
+    assert Q.dim() == 4 and Q.shape[0] == 1, f"expected [1,T,H,D], got {tuple(Q.shape)}"
+    device = Q.device
+    # FlexAttention wants [B, H, T, D].
+    q = Q.transpose(1, 2)  # [1, n_heads, T, D]
+    k = K.transpose(1, 2)  # [1, n_kv_heads, T, D]
+    v = V.transpose(1, 2)
+
+    n_heads = q.shape[1]
+    n_kv = k.shape[1]
+    enable_gqa = n_heads != n_kv
+    T = q.shape[2]
+    T_pad = prefix_seg_info.T_pad
+    assert T == prefix_seg_info.T, f"Q length {T} != seg.T {prefix_seg_info.T}"
+
+    if block_mask is None:
+        block_mask = get_block_mask(prefix_seg_info, device, compile_mask = compiled)
+
+    out = _run_flex(q, k, v, block_mask, enable_gqa, scale, compiled, T, T_pad)
+    # back to [1, T, n_heads, D]
+    return out.transpose(1, 2).contiguous()
+
+
+__all__ = [
+    "PrefixSegInfo",
+    "build_seg_info_multigroup",
+    "build_seg_info_from_layout",
+    "get_block_mask",
+    "clear_block_mask_cache",
+    "flex_shared_prefix_attention",
+]