oobabooga · Oct 16, 2026 · Oct 15, 2026 · Oct 15, 2026 · Oct 15, 2026 · Oct 15, 2026
@@ -436,6 +436,11 @@ def _require_supported(self) -> None:
         hardware = self.hardware()
         if not hardware.get("present"):
             raise NpuError("No AMD Ryzen AI NPU was found on this machine.")
+        if hardware.get("family") == "XDNA2" and hardware.get("enabled") is False:
+            raise NpuError(
+                "The AMD NPU is off by default. Start Unsloth Studio with "
+                "UNSLOTH_STUDIO_ENABLE_NPU=1 to use it."
+            )
         if not hardware.get("supported"):
             raise NpuError(
                 "FastFlowLM needs an XDNA 2 NPU (Ryzen AI 300/400 or Ryzen AI Max); "
@@ -444,6 +449,8 @@ def _require_supported(self) -> None:
 
     def _ensure_running(self) -> LemonadeServer:
         """Start the installed runtime if needed and return its server."""
+        # A runtime enabled before the flag was turned off must not start again.
+        self._require_supported()
         with self._lock:
             if self._server is not None and self._server.is_alive():
                 return self._server

@@ -683,14 +683,15 @@ def test_every_local_exit_without_a_tool_loop_refuses_the_image():
         "_refuse_unused_mcp_image(_mcp_image, _catalog_names(_sf_tools_to_use) if _sf_use_tools else None)",
     ):
         assert guard in statements
-    # NPU, both speech models, audio input and the GGUF passthrough return before either loop.
+    # The NPU switch check, NPU, both speech models, audio input and the GGUF passthrough
+    # return before either loop.
     calls = [
         node
         for node in ast.walk(ast.parse(src))
         if isinstance(node, ast.Call)
         and getattr(node.func, "id", None) == "_refuse_unused_mcp_image"
     ]
-    assert len(calls) == 7
+    assert len(calls) == 8
     # The forced image approval parks the GGUF stream, so the slot must be tracked for reclaim.
     slot = [
         ast.unparse(node.value)

@@ -146,6 +146,23 @@ def test_unsupported_npu_is_refused_before_any_download(tmp_path, monkeypatch):
     assert installer.installs == 0
 
 
+@pytest.mark.parametrize("call", ["enable", "catalog"])
+def test_an_npu_left_off_is_refused_with_the_flag_named(tmp_path, monkeypatch, call):
+    # catalog: a runtime enabled before the flag went off must not start again.
+    installer = _Installer(_binary(tmp_path))
+    monkeypatch.setattr(nb, "_installer_module", lambda: installer)
+    monkeypatch.setattr(
+        nb,
+        "detect_amd_npu",
+        lambda: {"present": True, "enabled": False, "supported": False, "family": "XDNA2"},
+    )
+    backend = nb.LemonadeNpuBackend(root = tmp_path / "lemonade")
+    with pytest.raises(nb.NpuError, match = "UNSLOTH_STUDIO_ENABLE_NPU=1"):
+        getattr(backend, call)()
+    assert installer.installs == 0
+    assert backend.status()["supported"] is False
+
+
 def test_catalog_keeps_flm_chat_models_only(npu):
     npu.enable()
     ids = [model.id for model in npu.catalog()]

@@ -69,7 +69,38 @@ def _chunk(
 }
 
 
+# The end of a reply whose reasoning filled the loaded context, captured on Strix Halo: FastFlowLM
+# says "stop" with the KV cache full.
+_FULL_USAGE = {
+    "prompt_tokens": 45,
+    "completion_tokens": 8147,
+    "total_tokens": 8192,
+    "active_kv_tokens": 8192,
+    "max_kv_token_capacity": 8192,
+    "kv_token_occupancy_rate_percentage": 100.0,
+}
+
+
 def _script(prompt: str) -> str:
+    if "FULL" in prompt:
+        return (
+            _chunk({"reasoning_content": " Wait, no!"})
+            + _chunk({"content": None}, "stop", _FULL_USAGE)
+            + "data: [DONE]\n\n"
+        )
+    if "WSTOOL" in prompt:
+        call = {
+            "index": 0,
+            "id": "call_1",
+            "type": "function",
+            "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'},
+        }
+        return (
+            _chunk({"content": "\n\n\n"})
+            + _chunk({"tool_calls": [call]})
+            + _chunk({}, "tool_calls", _USAGE)
+            + "data: [DONE]\n\n"
+        )
     if "FAULT" in prompt:
         return (
             _chunk({"content": "par"})
@@ -826,6 +857,161 @@ async def _unload_gpu(_backend):
     assert teardowns == []
 
 
+def test_an_npu_load_drains_active_inference_before_unloading(monkeypatch):
+    # An auto-switch passes no on_reload_confirmed, so without the drain it cut a reply short.
+    import routes.inference as routes
+    from core.inference import npu_backend as nb
+    from models.inference import LoadRequest
+
+    order: list[str] = []
+    model = SimpleNamespace(id = "qwen3-0.6b-FLM", model_path = "lemonade:qwen3-0.6b-FLM")
+
+    class _Npu:
+        def resident(self):
+            return None
+
+        def loadable_model(self, model_id):
+            return model
+
+        def load(self, model_id, ctx):
+            order.append("load")
+
+    async def _drain(**kwargs):
+        order.append("drain")
+
+    async def _unload_gpu(_backend):
+        order.append("unload")
+
+    monkeypatch.setattr(nb, "get_npu_backend", lambda: _Npu())
+    monkeypatch.setattr(routes, "get_llama_cpp_backend", lambda: SimpleNamespace(is_active = True))
+    monkeypatch.setattr(routes, "_peek_inference_backend", lambda: None)
+    monkeypatch.setattr(routes, "_wait_for_model_switch_idle", _drain)
+    monkeypatch.setattr(routes, "_unload_llama_before_standard_load", _unload_gpu)
+    monkeypatch.setattr(routes, "release_chat_gpu_claim", lambda: True)
+    with pytest.raises(HTTPException):
+        # The stub never becomes resident, so the load ends in a 409 after the steps above.
+        asyncio.run(
+            routes._load_npu_model(
+                LoadRequest(model_path = model.model_path),
+                current_request_counted = True,
+                on_reload_confirmed = None,
+                load_cancel_event = None,
+            )
+        )
+    assert order == ["drain", "unload", "load"]
+
+
+@pytest.mark.parametrize("blocker", ["sidecar_swap", "shutdown"])
+def test_an_npu_load_refuses_before_unloading_during_a_swap_or_shutdown(monkeypatch, blocker):
+    # The NPU branch skipped the rechecks the other load paths make, so it unloaded the resident
+    # model under a transformers install or after shutdown had torn everything down.
+    import routes.inference as routes
+    from core.inference import npu_backend as nb
+    from models.inference import LoadRequest
+    from utils import transformers_version
+
+    order: list[str] = []
+    model = SimpleNamespace(id = "qwen3-0.6b-FLM", model_path = "lemonade:qwen3-0.6b-FLM")
+
+    class _Npu:
+        def resident(self):
+            return None
+
+        def loadable_model(self, model_id):
+            return model
+
+        def load(self, model_id, ctx):
+            order.append("load")
+
+    async def _drain(**kwargs):
+        order.append("drain")
+
+    async def _unload_gpu(_backend):
+        order.append("unload")
+
+    monkeypatch.setattr(nb, "get_npu_backend", lambda: _Npu())
+    monkeypatch.setattr(routes, "get_llama_cpp_backend", lambda: SimpleNamespace(is_active = True))
+    monkeypatch.setattr(routes, "_peek_inference_backend", lambda: None)
+    monkeypatch.setattr(routes, "_wait_for_model_switch_idle", _drain)
+    monkeypatch.setattr(routes, "_unload_llama_before_standard_load", _unload_gpu)
+    monkeypatch.setattr(routes, "release_chat_gpu_claim", lambda: True)
+    if blocker == "sidecar_swap":
+        monkeypatch.setattr(transformers_version, "sidecar_swap_in_progress", lambda: True)
+    else:
+        monkeypatch.setattr(routes, "_loads_shutting_down", True)
+    with pytest.raises(HTTPException) as excinfo:
+        asyncio.run(
+            routes._load_npu_model(
+                LoadRequest(model_path = model.model_path),
+                current_request_counted = True,
+                on_reload_confirmed = None,
+                load_cancel_event = None,
+            )
+        )
+    assert excinfo.value.status_code == 409
+    assert "unload" not in order and "load" not in order
+
+
+def test_a_stop_during_an_npu_load_cancels_it(monkeypatch):
+    # The cancel event was read only before and after the blocking load, so a Stop waited it out.
+    import routes.inference as routes
+    from core.inference import npu_backend as nb
+    from models.inference import LoadRequest
+
+    model = SimpleNamespace(id = "qwen3-0.6b-FLM", model_path = "lemonade:qwen3-0.6b-FLM")
+    loading = threading.Event()
+    stopped = threading.Event()
+    stop = threading.Event()
+
+    class _Npu:
+        def resident(self):
+            return None
+
+        def loadable_model(self, model_id):
+            return model
+
+        def load(self, model_id, ctx):
+            loading.set()
+            if not stopped.wait(10):
+                raise AssertionError("the load was never cancelled")
+            raise nb.NpuLoadCancelled("cancelled")
+
+        def cancel_load(self, model_id):
+            # The first ask lands before the load is registered, as when Stop beats npu.load.
+            asks.append(model_id)
+            if len(asks) == 1:
+                return False
+            stopped.set()
+            return True
+
+    asks: list[str] = []
+
+    def _stop_once_loading():
+        loading.wait(10)
+        stop.set()
+
+    async def _nothing(*_args, **_kwargs):
+        return None
+
+    monkeypatch.setattr(nb, "get_npu_backend", lambda: _Npu())
+    monkeypatch.setattr(routes, "get_llama_cpp_backend", lambda: SimpleNamespace(is_active = False))
+    monkeypatch.setattr(routes, "_peek_inference_backend", lambda: None)
+    monkeypatch.setattr(routes, "_unload_llama_before_standard_load", _nothing)
+    monkeypatch.setattr(routes, "release_chat_gpu_claim", lambda: True)
+    threading.Thread(target = _stop_once_loading, daemon = True).start()
+    with pytest.raises(HTTPException) as excinfo:
+        asyncio.run(
+            routes._load_npu_model(
+                LoadRequest(model_path = model.model_path),
+                current_request_counted = True,
+                on_reload_confirmed = None,
+                load_cancel_event = stop,
+            )
+        )
+    assert excinfo.value.status_code == 409
+    assert stopped.is_set() and len(asks) == 2
+
+
 def test_a_reply_cut_short_is_not_reported_complete(flm):
     flm()
     status, body = _call(stream = False, messages = [{"role": "user", "content": "CUT"}])
@@ -1278,3 +1464,57 @@ def test_a_tool_loop_turn_dropped_after_a_tool_call_is_cut_short():
     plain = _TurnFinish()
     plain.see(_chunk({}, "stop").strip())
     assert plain.finished is True
+
+
+def test_a_reply_that_fills_the_context_says_so(flm):
+    flm(reasoning = True)
+    full = [{"role": "user", "content": "FULL"}]
+    status, lines = _call(stream = True, messages = full)
+    assert status == 200
+    chunks = _data(lines)
+    assert [c["_toolEvent"] for c in chunks if "_toolEvent" in c] == [
+        {"type": "context_window_exceeded"}
+    ]
+    finishes = [
+        c["choices"][0]["finish_reason"] for c in chunks if c["choices"][0]["finish_reason"]
+    ]
+    assert finishes == ["length"]
+    _, body = _call(stream = False, messages = full)
+    assert body["choices"][0]["finish_reason"] == "length"
+
+
+@pytest.mark.parametrize(
+    ("usage", "full"),
+    [
+        (_FULL_USAGE, True),
+        # Thinking off: FastFlowLM decodes an empty think block, four tokens counted as one.
+        ({**_FULL_USAGE, "completion_tokens": 8144, "total_tokens": 8189}, True),
+        ({"total_tokens": 8192, "max_kv_token_capacity": 8192}, True),
+        (_FLM_USAGE, False),
+        ({"total_tokens": 8192}, False),
+        (None, False),
+    ],
+)
+def test_a_full_kv_cache_reads_the_kv_position(usage, full):
+    import routes.inference as ri
+    assert ri._npu_kv_cache_full(usage) is full
+
+
+def test_a_reply_with_room_left_keeps_its_stop(flm):
+    flm()
+    status, lines = _call(stream = True, messages = [{"role": "user", "content": "TIMED"}])
+    assert status == 200
+    chunks = _data(lines)
+    assert not any("_toolEvent" in c for c in chunks)
+    assert chunks[-1]["choices"][0]["finish_reason"] == "stop"
+
+
+def test_whitespace_beside_a_tool_call_is_no_content(flm):
+    flm(tools = True)
+    tools = [
+        {"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}
+    ]
+    _, body = _call(stream = False, tools = tools, messages = [{"role": "user", "content": "WSTOOL"}])
+    message = body["choices"][0]["message"]
+    assert message["content"] is None
+    assert message["tool_calls"][0]["function"]["name"] == "get_weather"
@@ -78,15 +78,39 @@ def test_only_xdna2_is_supported(tmp_path, monkeypatch, device, supported):
     monkeypatch.setattr(
         npu, "_linux_probe", lambda: probe(tmp_path / "devices", tmp_path / "module")
     )
+    monkeypatch.setenv(npu.ENABLE_ENV, "1")
     info = npu.detect_amd_npu()
     assert info["present"] is True
     assert info["supported"] is supported
 
 
+@pytest.mark.parametrize("value", [None, "", "0", "false"])
+def test_an_xdna2_npu_is_off_without_the_flag(tmp_path, monkeypatch, value):
+    _device(tmp_path, "0000:c4:00.1", "1022", "17f0", driver = "amdxdna")
+    probe = npu._linux_probe
+    monkeypatch.setattr(npu.sys, "platform", "linux")
+    monkeypatch.setattr(
+        npu, "_linux_probe", lambda: probe(tmp_path / "devices", tmp_path / "module")
+    )
+    if value is None:
+        monkeypatch.delenv(npu.ENABLE_ENV, raising = False)
+    else:
+        monkeypatch.setenv(npu.ENABLE_ENV, value)
+    info = npu.detect_amd_npu()
+    assert info["present"] is True and info["family"] == "XDNA2"
+    assert info["enabled"] is False and info["supported"] is False
+
+
 def test_detection_never_raises(monkeypatch):
     def broken():
         raise PermissionError("sysfs")
 
     monkeypatch.setattr(npu.sys, "platform", "linux")
     monkeypatch.setattr(npu, "_linux_probe", broken)
-    assert npu.detect_amd_npu() == {"present": False, "error": "sysfs", "supported": False}
+    monkeypatch.delenv(npu.ENABLE_ENV, raising = False)
+    assert npu.detect_amd_npu() == {
+        "present": False,
+        "error": "sysfs",
+        "enabled": False,
+        "supported": False,
+    }