Studio: retry as one sequence when llama.cpp refuses a unified KV cache (#10371)

* Studio: retry as one sequence when llama.cpp refuses a unified KV cache

Studio appends --kv-unified whenever it asks for more than one slot, so
llama.cpp does not split -c into per-slot windows. Some architectures need one
sequence per stream and refuse to build a context that way:

  llama_init_from_model: failed to initialize the context: glm5next: the pooled
  indexer needs one sequence per stream, so a unified KV cache is only supported
  with a single sequence

The flag is Studio's, not the user's, so nothing they can change in the UI
reaches it. A Strix Halo report shows GLM-5.3-Flash tried four times over two
days after a 146 GB download and never once ran; the second attempt cut the
context 15x, which could not have helped because it was never the context.

Retry once with the geometry that would never have added the flag: one slot,
no unified cache, the requested context untouched. Matched on llama.cpp's own
wording rather than on an architecture name, so a second model with the same
constraint needs no list entry. The rung sits ahead of the flash-attention
retry, which only fires on a signal crash, and this is a clean refusal.

Every --parallel occurrence is rewritten, not just the emitted one: extras are
appended after Unsloth's flags and llama.cpp is last-wins. LLAMA_ARG_KV_UNIFIED
is dropped from the child environment for the same reason the flash-attn retry
drops LLAMA_ARG_FLASH_ATTN - llama.cpp applies the environment before argv, and
there is no negated flag to emit against it.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: the single-sequence retry commits the one-slot geometry it launched

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Mark the single-sequence retry's slot clamp for PR #10371

* Keep the single-sequence retry ahead of the fit rungs and cover every slot alias

* Tighten the comments on the single-sequence retry

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: danielhanchen

studio/backend/core/inference/llama_cpp.py

@@ -77,6 +77,7 @@
     _GPU_LAYER_FLAGS,
     _LAYER_OFFLOAD_FLAGS,
     _MOE_OFFLOAD_FLAGS,
+    _PARALLEL_FLAGS,
     _SPLIT_MODE_FLAGS,
     _TENSOR_SPLIT_FLAGS,
     _effective_tensor_parallel,
@@ -17693,6 +17694,79 @@ def _drop_env_flash_attn(env: MutableMapping[str, str]) -> bool:
         """
         return env.pop("LLAMA_ARG_FLASH_ATTN", None) is not None
 
+    # On the wording, not an architecture name, so the next model needs no list.
+    # glm5next spells it as the GGML_ASSERT of ggml-org/llama.cpp#27754.
+    _KV_UNIFIED_REFUSED_MARKERS = (
+        "a unified kv cache is only supported with a single sequence",
+        "needs one sequence per stream",
+    )
+
+    @staticmethod
+    def _is_kv_unified_refused(output: str) -> bool:
+        low = (output or "").lower()
+        return any(marker in low for marker in LlamaCppBackend._KV_UNIFIED_REFUSED_MARKERS)
+
+    @staticmethod
+    def _with_single_sequence(cmd: list[str]) -> Optional[list[str]]:
+        """Return cmd re-run as one sequence, or None when it already is one.
+
+        Studio adds --kv-unified itself above one slot, so this reverses Studio's
+        choice, not the user's. Every alias goes, from _PARALLEL_FLAGS so it cannot
+        drift from the denylist; llama.cpp is last-wins.
+        """
+        out: list[str] = []
+        saw_kv_unified = False
+        multi_slot = False
+        skip_value = False
+        for i, tok in enumerate(cmd):
+            if skip_value:
+                skip_value = False
+                continue
+            name = _flag_name(tok)
+            if name in ("--kv-unified", "-kvu"):
+                saw_kv_unified = True
+                # The flag is bare, but an inline "=1" or a separate "1" reach us.
+                if (
+                    "=" not in tok
+                    and cmd[i + 1 : i + 2]
+                    # Case-sensitive, like llama.cpp's own bool parse.
+                    and cmd[i + 1] in _LLAMA_ARG_TRUE_FALSE_AUTO_VALUES
+                ):
+                    skip_value = True
+                continue
+            if name in _PARALLEL_FLAGS:
+                if "=" in tok:
+                    value = tok.partition("=")[2]
+                elif tok != name:
+                    # The attached short, -np8, which _flag_name peels to -np.
+                    value = tok[len(name) :]
+                else:
+                    # Valueless --parallel: eating the next token would delete a flag.
+                    value = cmd[i + 1] if i + 1 < len(cmd) else ""
+                    skip_value = _flag_name(value) is None
+                try:
+                    multi_slot = multi_slot or int(value.strip()) > 1
+                except (AttributeError, TypeError, ValueError):
+                    pass
+                out.extend([name, "1"])
+                continue
+            out.append(tok)
+        if not saw_kv_unified and not multi_slot:
+            return None
+        if not any(_flag_name(tok) in _PARALLEL_FLAGS for tok in out):
+            out.extend(["--parallel", "1"])
+        return out
+
+    @staticmethod
+    def _drop_env_single_sequence(env: MutableMapping[str, str]) -> bool:
+        """Drop inherited unified-cache and slot-count env before that retry.
+
+        llama.cpp reads its environment before argv. Dropped, not negated with
+        --no-kv-unified, which needs every build reaching here to know it.
+        """
+        dropped = env.pop("LLAMA_ARG_KV_UNIFIED", None) is not None
+        return env.pop("LLAMA_ARG_N_PARALLEL", None) is not None or dropped
+
     @staticmethod
     def _strip_mmproj_args(cmd: list[str]) -> list[str]:
         """Return cmd without the '--mmproj <path>' pair (text-only retry).
@@ -23692,6 +23766,11 @@ def _spawn_and_wait(run_cmd, *, label = ""):
                             _startup_output
                         ) or self._is_tensor_quant_kv_unsupported(_startup_output)
                         _hip_rocr_mismatch = self._is_bundled_hip_rocr_mismatch(_startup_output)
+                        # No fit retry reaches it, and the rung below needs this
+                        # launch's argv. Whole buffer: it arrives with a backtrace.
+                        _capability_crash = _tensor_capability_crash or self._is_kv_unified_refused(
+                            "\n".join(self._stdout_lines)
+                        )
                         if (
                             not _did_rocm_retry
                             and _startup_crashed
@@ -23722,7 +23801,7 @@ def _spawn_and_wait(run_cmd, *, label = ""):
                         if (
                             not _did_fit_retry
                             and _startup_crashed
-                            and not _tensor_capability_crash
+                            and not _capability_crash
                             and not _hip_rocr_mismatch
                         ):
                             # A spill-planned launch that crashed on startup. The
@@ -23754,7 +23833,7 @@ def _spawn_and_wait(run_cmd, *, label = ""):
                             not _did_fit_retry
                             and fully_gpu_offloaded
                             and _startup_crashed
-                            and not _tensor_capability_crash
+                            and not _capability_crash
                             and not _hip_rocr_mismatch
                         ):
                             # We forced --fit off because Unsloth's (conservative) VRAM
@@ -23822,7 +23901,7 @@ def _spawn_and_wait(run_cmd, *, label = ""):
                             not _did_fit_retry
                             and _fit_retry_allowed
                             and _startup_crashed
-                            and not _tensor_capability_crash
+                            and not _capability_crash
                             and not _hip_rocr_mismatch
                         ):
                             logger.warning(
@@ -24517,6 +24596,34 @@ def _retry_is_oversized() -> bool:
                             self._memory_state = resolve_effective_memory_state(cmd, env)
                         healthy = _spawn_and_wait(cmd, label = "-archfallback")
 
+                # Studio adds --kv-unified itself above one slot, so nothing the user
+                # changes reaches it: retry at one slot, context intact. It aborts, so
+                # this MUST stay ahead of the flash-attn rung, which takes any signal
+                # crash; on the message, not the exit, so Windows lands here too.
+                if not healthy and not _load_cancelled():
+                    _kvu_cmd = (
+                        self._with_single_sequence(_last_spawn_cmd)
+                        if self._is_kv_unified_refused("\n".join(self._stdout_lines))
+                        else None
+                    )
+                    if _kvu_cmd is not None:
+                        logger.warning(
+                            "llama-server refused a unified KV cache with more than "
+                            "one sequence; retrying with one slot and --kv-unified "
+                            "dropped. Concurrent requests will queue."
+                        )
+                        self._kill_process()
+                        if self._drop_env_single_sequence(env):
+                            logger.info(
+                                "Dropped inherited LLAMA_ARG_KV_UNIFIED / "
+                                "LLAMA_ARG_N_PARALLEL for the single-sequence retry."
+                            )
+                        cmd = _kvu_cmd
+                        # Read by admission control; left as-is Studio over-admits.
+                        n_parallel = 1  # allow-slot-clamp: llama-server refused more
+                        kv_cache_unified = False
+                        healthy = _spawn_and_wait(_kvu_cmd, label = "-single-seq")
+
                 # Flash-attention kernels hard-crash at startup on some ROCm/GPU
                 # builds (frequently inside the vision tower). Disabling FA keeps
                 # both vision and MTP, so retry that way before dropping either.

studio/backend/core/inference/llama_server_args.py

@@ -35,12 +35,16 @@
 CTX_CHECKPOINTS_MAX = 256
 CACHE_RAM_MAX_MIB = 1024 * 1024
 
+# Slot-count aliases in one place: the denial below, its #9510 hint and the single-sequence retry must cover the same
+# set, or a spelling one of them misses reaches llama-server unnoticed.
+_PARALLEL_FLAGS: frozenset[str] = frozenset({"-np", "--parallel", "--n-parallel"})
+
 # Each group = every alias (short + long) of one hard-denied flag. Extend the matching group when llama.cpp adds a new
 # alias.
 _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
     # Parallel slots: owned by typer --parallel and LoadRequest.n_parallel; a pass-through would desync the slot
     # bookkeeping from llama-server.
-    frozenset({"-np", "--parallel", "--n-parallel"}),
+    _PARALLEL_FLAGS,
     # Model identity: a second -m would load a different model than Unsloth thinks it loaded
     # Model identity: Unsloth resolves it from LoadRequest; a second -m would load a different model than Unsloth thinks
     # it loaded.
@@ -275,7 +279,7 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
             # #9510: users reaching for `--parallel 1` hit this refusal with no pointer to the supported knob
             # Why (#9510): users reaching for `--parallel 1` to cap concurrent predictions on a local model hit this
             # refusal with no pointer to the supported knob; name it.
-            if flag in {"-np", "--parallel", "--n-parallel"}:
+            if flag in _PARALLEL_FLAGS:
                 message += "; set n_parallel on the load request (parallel decode slots) instead"
             raise ValueError(message)
         if flag is None:

studio/backend/tests/test_llama_cpp_single_sequence_retry.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""The single-sequence retry for architectures that refuse a unified KV cache.
+
+Studio appends ``--kv-unified`` itself above one slot, so nothing the user changes
+in the UI reaches the flag that stops the model loading.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from core.inference.llama_cpp import LlamaCppBackend
+
+# Verbatim from a Strix Halo report: llama-server refusing GLM-5.3-Flash.
+_REFUSAL = (
+    "0.56.012.623 E llama_init_from_model: failed to initialize the context: "
+    "glm5next: the pooled indexer needs one sequence per stream, so a unified "
+    "KV cache is only supported with a single sequence"
+)
+
+
+# What the binary prints: the full buffer is scanned, so the backtrace cannot bury it.
+_ASSERT = (
+    "/build/llama.cpp/src/models/glm5next.cpp:1018: GGML_ASSERT(n_ps == 1 && "
+    '"the per-cell pool view needs one sequence per stream") failed\n'
+    "[New LWP 4242]\n#0  0x00007f0000000000 in abort ()"
+)
+
+
+def test_the_reported_refusal_is_recognised():
+    assert LlamaCppBackend._is_kv_unified_refused(_REFUSAL)
+
+
+def test_the_assert_the_binary_actually_prints_is_recognised():
+    assert LlamaCppBackend._is_kv_unified_refused(_ASSERT)
+
+
+def test_an_unrelated_failure_is_not():
+    assert not LlamaCppBackend._is_kv_unified_refused(
+        "error loading model: unknown model architecture: 'qwen4exp'"
+    )
+    assert not LlamaCppBackend._is_kv_unified_refused("")
+    # minimax-m3 still loads: retrying it would cost three slots for nothing.
+    assert not LlamaCppBackend._is_kv_unified_refused(
+        "minimax_m3: unified KV cache with n_seq_max > 1; MSA needs per-sequence "
+        "streams -> running DENSE attention. Drop --kv-unified to enable MSA."
+    )
+
+
+def test_the_reported_launch_becomes_one_slot_with_no_unified_cache():
+    cmd = [
+        "llama-server",
+        "-m",
+        "GLM-5.3-Flash-UD-IQ4_XS-00001-of-00004.gguf",
+        "--parallel",
+        "4",
+        "--flash-attn",
+        "on",
+        "--no-context-shift",
+        "-c",
+        "128000",
+        "--gpu-layers",
+        "47",
+        "--fit",
+        "off",
+        "--kv-unified",
+        "--jinja",
+    ]
+
+    out = LlamaCppBackend._with_single_sequence(cmd)
+
+    assert out is not None
+    assert "--kv-unified" not in out
+    assert out[out.index("--parallel") + 1] == "1"
+    assert out[out.index("-c") + 1] == "128000"
+    assert out[out.index("--gpu-layers") + 1] == "47"
+    assert "--jinja" in out and "--no-context-shift" in out
+
+
+def test_a_parallel_surviving_in_the_extras_tail_is_rewritten_too():
+    """llama.cpp is last-wins and extras are appended after Unsloth's flags."""
+    cmd = ["llama-server", "--parallel", "4", "--kv-unified", "-np", "8"]
+
+    out = LlamaCppBackend._with_single_sequence(cmd)
+
+    assert out == ["llama-server", "--parallel", "1", "-np", "1"]
+
+
+def test_every_spelling_of_the_two_flags_is_handled():
+    cmd = ["llama-server", "--parallel=4", "-kvu", "--alias", "m"]
+    assert LlamaCppBackend._with_single_sequence(cmd) == [
+        "llama-server",
+        "--parallel",
+        "1",
+        "--alias",
+        "m",
+    ]
+
+    cmd = ["llama-server", "-np8", "--kv-unified", "1"]
+    assert LlamaCppBackend._with_single_sequence(cmd) == ["llama-server", "-np", "1"]
+
+
+def test_every_alias_of_the_slot_count_is_rewritten():
+    """--n-parallel is in the denylist group too."""
+    out = LlamaCppBackend._with_single_sequence(
+        ["llama-server", "--parallel", "4", "--kv-unified", "--n-parallel", "8"]
+    )
+
+    assert out == ["llama-server", "--parallel", "1", "--n-parallel", "1"]
+
+
+def test_a_valueless_slot_flag_does_not_swallow_the_next_flag():
+    """Malformed, but eating the token behind it would delete --kv-unified."""
+    out = LlamaCppBackend._with_single_sequence(["llama-server", "--parallel", "--kv-unified"])
+
+    assert out == ["llama-server", "--parallel", "1"]
+
+
+def test_a_command_already_running_one_sequence_has_nothing_to_retry():
+    assert LlamaCppBackend._with_single_sequence(["llama-server", "-c", "8192"]) is None
+    assert LlamaCppBackend._with_single_sequence(["llama-server", "--parallel", "1"]) is None
+
+
+def test_a_slot_count_is_added_when_the_command_carried_none():
+    out = LlamaCppBackend._with_single_sequence(["llama-server", "--kv-unified"])
+
+    assert out == ["llama-server", "--parallel", "1"]
+
+
+def test_the_inherited_environment_is_dropped_so_it_cannot_undo_the_retry():
+    """llama.cpp applies its environment before parsing argv."""
+    env = {"LLAMA_ARG_KV_UNIFIED": "1", "LLAMA_ARG_N_PARALLEL": "4", "PATH": "/usr/bin"}
+
+    assert LlamaCppBackend._drop_env_single_sequence(env) is True
+    assert "LLAMA_ARG_KV_UNIFIED" not in env
+    assert "LLAMA_ARG_N_PARALLEL" not in env
+    assert env["PATH"] == "/usr/bin"
+
+
+def test_a_clean_environment_reports_nothing_dropped():
+    env = {"PATH": "/usr/bin"}
+
+    assert LlamaCppBackend._drop_env_single_sequence(env) is False
+    assert env == {"PATH": "/usr/bin"}
+
+
+def test_the_fit_recovery_rungs_stand_down_for_this_refusal():
+    """A startup crash, so the --fit rungs would take it first and waste a load."""
+    import inspect
+
+    src = inspect.getsource(LlamaCppBackend.load_model)
+    assert "_capability_crash = _tensor_capability_crash or self._is_kv_unified_refused" in src
+    # Every --fit rung, and only those: the HIP rung keeps its narrower gate.
+    assert src.count("and not _capability_crash") == 3
+    assert src.count("and not _tensor_capability_crash") == 1
+
+
+def test_the_retry_commits_the_one_slot_geometry_it_launched():
+    # Committed after the ladder, so the retry must overwrite what it reverses.
+    import inspect
+
+    src = inspect.getsource(LlamaCppBackend.load_model)
+    start = src.index('label = "-single-seq"')
+    block = src[src.rindex("cmd = _kvu_cmd", 0, start) : start]
+    assert "n_parallel = 1" in block
+    assert "kv_cache_unified = False" in block

studio/backend/tests/test_tp_vision_regression.py

@@ -617,23 +617,23 @@ def test_fit_off_retry_skipped_on_a_tensor_capability_crash():
     """The fit-independent --fit off retry is skipped on the split-axis marker, else
     the model crashes a second time before the latch records it (reviewer.py, #6659).
 
-    The same skip now also covers the pre-b9455 refusal of a quantized KV cache in
-    tensor mode (ggml-org/llama.cpp#23792): both are capabilities the binary lacks,
-    so a second spawn to let it offload cannot help and costs a full model load.
-    Hence the guard's name is _tensor_capability_crash rather than the split-axis
-    one it started as.
+    It now also covers the pre-b9455 quantized-KV refusal in tensor mode
+    (ggml-org/llama.cpp#23792) and the unified-cache refusal: no second spawn helps
+    any of the three, and each costs a full model load. Hence _capability_crash,
+    with _tensor_capability_crash left as the half the ROCm rung gates on.
     """
     src = inspect.getsource(LlamaCppBackend.load_model)
     retry = src.find('run_cmd = [*run_cmd, "--fit", "off"]')
     assert retry != -1
     guard = src[max(0, retry - 1000) : retry]
     assert "_fit_retry_allowed" in guard and "_startup_crashed" in guard
     assert (
-        "not _tensor_capability_crash" in guard
-    ), "the fit-off retry must be skipped when the crash is a tensor capability limit"
-    # Both markers feed it, so neither can be dropped without this failing.
+        "not _capability_crash" in guard
+    ), "the fit-off retry must be skipped when the crash is a capability limit"
+    # All three markers feed it, so none can be dropped without this failing.
     assert "_is_tensor_split_assert" in src
     assert "_is_tensor_quant_kv_unsupported" in src
+    assert "_is_kv_unified_refused" in src
 
 
 def test_is_abort_exit_recognizes_windows_crt_abort():