diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 147174451e..d46e5a0fe1 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -33,6 +33,7 @@ from typing import ( List, Literal, Mapping, + MutableMapping, Optional, Union, ) @@ -3696,6 +3697,14 @@ class LlamaCppBackend: # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + # V cache types that llama.cpp can run WITHOUT flash attention. Only the V + # axis has the dependency: a quantized V cache (q8_0/q4_0/q4_1/q5_0/q5_1/ + # iq4_nl) aborts init with "V cache quantization requires flash_attn", while + # a quantized K cache runs fine without FA. So the flash-attn-off crash- + # recovery fallback must reset a quantized V cache to f16 before it can + # launch (and leaves K alone). These three are the only non-quantized types. + _NON_QUANTIZED_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + # Main-model placement settings that Manual mode owns. They must not leak # from Studio's parent environment into llama-server and silently override # the command assembled from the current request. Draft-model placement is @@ -6149,6 +6158,21 @@ class LlamaCppBackend: cls._is_signal_crash(returncode) or cls._is_abort_exit(returncode) ) + @staticmethod + def _canonical_long_flag(name: str) -> str: + """Return ``name`` with llama.cpp's long-option underscore normalization. + + llama.cpp runs ``std::replace(arg.begin(), arg.end(), '_', '-')`` on any + argv token that starts with ``--`` before looking it up, so a legal + pass-through spelling like ``--cache_type_v`` parses as + ``--cache-type-v``. Mirror that here so managed-flag matching sees the + same canonical name. Short flags (``-ctv``) never carry underscores and + keep their exact spelling; pass only the flag name (no attached value). + """ + if name.startswith("--"): + return name.replace("_", "-") + return name + @staticmethod def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]: """Return cmd with flash attention forced off, or None when its effective @@ -6181,8 +6205,76 @@ class LlamaCppBackend: out[i + 1] = "off" elif explicit(i) is None: # bare flag (reads as on) -> explicit off out[i] = f"{tok}=off" + + # A quantized V cache requires flash attention in llama.cpp: the init + # aborts with "V cache quantization requires flash_attn". A quantized K + # cache has no such requirement and runs fine without FA, so it is left + # untouched -- resetting it would needlessly enlarge the K cache and can + # OOM a memory-constrained config. Studio launches with FA on, so a + # quantized --cache-type-v is legal at launch but would make THIS FA-off + # retry crash on init instead of recovering. Reset a quantized V cache -- + # main and draft (the draft context shares the global --flash-attn flag, + # so its V cache aborts too) -- to f16 (the llama.cpp default); + # non-quantized types -- f16/bf16/f32 -- run fine without FA and are left + # untouched. The value is rewritten in place so the list length is + # preserved for downstream slices, matching the flash-attn flip above. + _v_cache_flags = ( + "--cache-type-v", + "-ctv", + "--cache-type-v-draft", + "--spec-draft-type-v", + "-ctvd", + ) + _cache_reset = False + for i, tok in enumerate(out): + # llama.cpp rewrites '_' to '-' for any argv token starting with + # '--' before matching, so a legal pass-through spelling such as + # --cache_type_v parses as --cache-type-v and still enables a + # quantized V cache. Canonicalize the flag name the same way so the + # reset recognizes the underscore aliases too; short flags (-ctv) + # and the type value are left untouched. + name = LlamaCppBackend._canonical_long_flag(tok.partition("=")[0]) + if name not in _v_cache_flags: + continue + if "=" in tok: + flag, _, value = tok.partition("=") + if value.strip().lower() not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES: + out[i] = f"{flag}=f16" + _cache_reset = True + elif i + 1 < len(out): + if out[i + 1].strip().lower() not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES: + out[i + 1] = "f16" + _cache_reset = True + if _cache_reset: + logger.info( + "V cache dtype reset to f16 because flash attention was disabled " + "by the crash-recovery fallback (quantized V cache requires flash " + "attention in llama.cpp; the K cache is left untouched)." + ) return out + @staticmethod + def _drop_env_quantized_v_cache(env: MutableMapping[str, str]) -> bool: + """Drop an inherited quantized V-cache env var (main or draft) in place + before a flash-attn-off retry, returning True if anything was removed. + + The argv rewrite in ``_with_flash_attn_off`` only reaches flags on the + command line. Studio deliberately lets an env-only cache type reach the + child untouched (an asymmetric K/V env must survive), so a quantized V + cache set purely through ``LLAMA_ARG_CACHE_TYPE_V`` (or the draft + ``LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V``) would still abort the FA-off retry + with "V cache quantization requires flash_attn". Dropping it lets + llama.cpp fall back to the f16 default. Only V is dropped: a quantized K + cache runs fine without flash attention, so its env var is preserved. + """ + dropped = False + for var in ("LLAMA_ARG_CACHE_TYPE_V", "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V"): + value = (env.get(var) or "").strip().lower() + if value and value not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES: + env.pop(var, None) + dropped = True + return dropped + @staticmethod def _strip_mmproj_args(cmd: list[str]) -> list[str]: """Return cmd without the '--mmproj ' pair (text-only retry). @@ -8339,6 +8431,13 @@ class LlamaCppBackend: _fa_rc, ) self._kill_process() + # The argv rewrite can't reach an env-only quantized V + # cache; drop it so the FA-off child doesn't abort on it. + if self._drop_env_quantized_v_cache(env): + logger.info( + "Dropped inherited quantized V-cache env for the " + "--flash-attn off retry (requires flash attention)." + ) cmd = _fa_cmd healthy = _spawn_and_wait(_fa_cmd, label = "-noflash") @@ -8384,6 +8483,13 @@ class LlamaCppBackend: _probe_rc, ) self._kill_process() + # The argv rewrite can't reach an env-only quantized V + # cache; drop it so the FA-off child doesn't abort on it. + if self._drop_env_quantized_v_cache(env): + logger.info( + "Dropped inherited quantized V-cache env for the " + "--flash-attn off retry (requires flash attention)." + ) cmd = _fa_cmd healthy = ( _spawn_and_wait(_fa_cmd, label = "-noflash-mtp") diff --git a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py index 4332a440a5..45c8bcb032 100644 --- a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py +++ b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py @@ -250,6 +250,201 @@ class TestFlashAttnOff: assert _flash_off(["llama-server", "-fa"]) == ["llama-server", "-fa=off"] +_drop_env_v = LlamaCppBackend._drop_env_quantized_v_cache + + +class TestFlashAttnOffQuantizedKvCache: + """Only the V cache requires flash attention in llama.cpp (init aborts with + "V cache quantization requires flash_attn"); a quantized K cache runs fine + without FA. Studio launches FA on, so a quantized --cache-type-v is legal at + launch but would make the FA-off crash-recovery retry crash on init. The + fallback must reset a quantized V cache (main and draft) to f16 while leaving + the K cache and non-quantized (f16/bf16/f32) types unchanged -- resetting K + would needlessly enlarge it and can OOM a memory-constrained config.""" + + _QUANTIZED = ["q8_0", "q4_0", "q4_1", "q5_0", "q5_1", "iq4_nl"] + _NON_QUANTIZED = ["f16", "bf16", "f32"] + + @pytest.mark.parametrize("qtype", _QUANTIZED) + def test_quantized_v_reset_k_preserved(self, qtype): + cmd = [ + "llama-server", + "--flash-attn", + "on", + "--cache-type-k", + qtype, + "--cache-type-v", + qtype, + ] + out = _flash_off(cmd) + assert out is not None + # FA flipped off AND the V axis reset to f16; the K axis is preserved so + # the FA-off retry keeps its memory budget (quantized K is FA-independent). + assert out[out.index("--flash-attn") + 1] == "off" + assert out[out.index("--cache-type-k") + 1] == qtype + assert out[out.index("--cache-type-v") + 1] == "f16" + assert len(out) == len(cmd) + + @pytest.mark.parametrize("qtype", _QUANTIZED) + def test_quantized_draft_v_reset(self, qtype): + # The draft context shares the global --flash-attn flag, so its quantized + # V cache aborts too and must be reset; the draft K cache is preserved. + for v_flag, k_flag in ( + ("--cache-type-v-draft", "--cache-type-k-draft"), + ("--spec-draft-type-v", "--spec-draft-type-k"), + ("-ctvd", "-ctkd"), + ): + cmd = ["llama-server", "-fa", "on", k_flag, qtype, v_flag, qtype] + out = _flash_off(cmd) + assert out is not None + assert out[out.index(v_flag) + 1] == "f16" + assert out[out.index(k_flag) + 1] == qtype + + @pytest.mark.parametrize("ntype", _NON_QUANTIZED) + def test_nonquantized_cache_left_unchanged(self, ntype): + cmd = [ + "llama-server", + "--flash-attn", + "on", + "--cache-type-k", + ntype, + "--cache-type-v", + ntype, + ] + out = _flash_off(cmd) + assert out is not None + # Only FA flips; the non-quantized cache type is preserved verbatim. + assert out[out.index("--flash-attn") + 1] == "off" + assert out[out.index("--cache-type-k") + 1] == ntype + assert out[out.index("--cache-type-v") + 1] == ntype + + def test_equals_form_quantized_v_reset(self): + out = _flash_off(["llama-server", "--flash-attn=on", "--cache-type-v=q8_0"]) + assert out == ["llama-server", "--flash-attn=off", "--cache-type-v=f16"] + + def test_equals_form_quantized_k_preserved(self): + out = _flash_off(["llama-server", "--flash-attn=on", "--cache-type-k=q8_0"]) + assert out == ["llama-server", "--flash-attn=off", "--cache-type-k=q8_0"] + + def test_short_alias_v_reset_k_preserved(self): + out = _flash_off(["llama-server", "-fa", "on", "-ctk", "q4_0", "-ctv", "q4_0"]) + assert out == ["llama-server", "-fa", "off", "-ctk", "q4_0", "-ctv", "f16"] + + def test_asymmetric_cache_only_v_reset(self): + # Quantized V, non-quantized K: reset V, keep K untouched. + out = _flash_off( + [ + "llama-server", + "--flash-attn", + "on", + "--cache-type-k", + "f16", + "--cache-type-v", + "q8_0", + ] + ) + assert out[out.index("--cache-type-k") + 1] == "f16" + assert out[out.index("--cache-type-v") + 1] == "f16" + + def test_no_cache_flags_still_flips_fa(self): + out = _flash_off(["llama-server", "--flash-attn", "on", "-c", "4096"]) + assert out == ["llama-server", "--flash-attn", "off", "-c", "4096"] + + def test_quantized_k_only_still_flips_fa_but_keeps_k(self): + # A quantized K cache with no V flag is a valid FA-off launch; the retry + # must not touch the K cache (it would waste memory for nothing). + out = _flash_off(["llama-server", "--flash-attn", "on", "--cache-type-k", "q8_0"]) + assert out == ["llama-server", "--flash-attn", "off", "--cache-type-k", "q8_0"] + + def test_input_not_mutated(self): + cmd = ["llama-server", "--flash-attn", "on", "--cache-type-v", "q8_0"] + _flash_off(cmd) + assert cmd[-1] == "q8_0" + + @pytest.mark.parametrize( + "flag", + ["--cache_type_v", "--cache-type_v", "--cache_type-v"], + ) + def test_underscore_alias_v_reset(self, flag): + # llama.cpp normalizes '_' to '-' in any '--' long option before + # matching, so a pass-through --cache_type_v enables a quantized V cache + # and must be reset by the FA-off retry too (else init aborts). + out = _flash_off(["llama-server", "--flash-attn", "on", flag, "q8_0"]) + assert out is not None + assert out[out.index("--flash-attn") + 1] == "off" + # The user's flag spelling is preserved; llama.cpp normalizes it anyway. + assert out[out.index(flag) + 1] == "f16" + + def test_underscore_alias_draft_v_reset(self): + out = _flash_off(["llama-server", "-fa", "on", "--spec_draft_type_v", "q4_0"]) + assert out is not None + assert out[out.index("--spec_draft_type_v") + 1] == "f16" + + def test_underscore_alias_equals_form_v_reset(self): + out = _flash_off(["llama-server", "--flash-attn=on", "--cache_type_v=q8_0"]) + assert out == ["llama-server", "--flash-attn=off", "--cache_type_v=f16"] + + def test_underscore_value_not_normalized_for_nonquantized(self): + # Only the flag name is canonicalized; a non-quantized type value is + # matched verbatim and left untouched (no spurious reset). + out = _flash_off(["llama-server", "--flash-attn", "on", "--cache_type_v", "f16"]) + assert out[out.index("--cache_type_v") + 1] == "f16" + assert out[out.index("--flash-attn") + 1] == "off" + + def test_short_alias_underscore_not_applied(self): + # Short flags are never underscore-normalized by llama.cpp; -ctv still + # matches and resets, and an unrelated short token is left alone. + out = _flash_off(["llama-server", "-fa", "on", "-ctv", "q8_0"]) + assert out == ["llama-server", "-fa", "off", "-ctv", "f16"] + + +class TestDropEnvQuantizedVCache: + """The argv rewrite can't reach a cache type set purely through the + environment (Studio deliberately lets an env-only type reach the child), so + the FA-off retry separately drops a quantized V-cache env var. Only V is + dropped: a quantized K cache is FA-independent and must survive.""" + + _QUANTIZED = ["q8_0", "q4_0", "q4_1", "q5_0", "q5_1", "iq4_nl"] + + @pytest.mark.parametrize("qtype", _QUANTIZED) + def test_drops_quantized_main_v_env(self, qtype): + env = {"LLAMA_ARG_CACHE_TYPE_V": qtype, "PATH": "/usr/bin"} + assert _drop_env_v(env) is True + assert "LLAMA_ARG_CACHE_TYPE_V" not in env + assert env["PATH"] == "/usr/bin" + + @pytest.mark.parametrize("qtype", _QUANTIZED) + def test_drops_quantized_draft_v_env(self, qtype): + env = {"LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V": qtype} + assert _drop_env_v(env) is True + assert "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V" not in env + + def test_preserves_quantized_k_env(self): + # A quantized K cache runs without FA, so its env must not be dropped. + env = {"LLAMA_ARG_CACHE_TYPE_K": "q8_0", "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K": "q4_0"} + assert _drop_env_v(env) is False + assert env["LLAMA_ARG_CACHE_TYPE_K"] == "q8_0" + assert env["LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K"] == "q4_0" + + @pytest.mark.parametrize("ntype", ["f16", "bf16", "f32", "F16", " q8_0 "]) + def test_preserves_nonquantized_v_env(self, ntype): + # Non-quantized V env values (and whitespace/case variants of them) run + # fine without FA; only a genuinely quantized value is dropped. + if ntype.strip().lower() in ("q8_0",): + env = {"LLAMA_ARG_CACHE_TYPE_V": ntype} + assert _drop_env_v(env) is True + assert "LLAMA_ARG_CACHE_TYPE_V" not in env + else: + env = {"LLAMA_ARG_CACHE_TYPE_V": ntype} + assert _drop_env_v(env) is False + assert env["LLAMA_ARG_CACHE_TYPE_V"] == ntype + + def test_noop_on_empty_env(self): + env = {} + assert _drop_env_v(env) is False + assert env == {} + + class TestNonProjectorDiagnostic: """_output_has_nonprojector_diagnostic gates the signal-only text-only retry: a hard crash that already names OOM / a bad arch / a TP limit must surface