ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: reset quantized KV cache to f16 when the flash-attn-off crash-recovery fallback fires (#7390)

Browse source 
* Studio: reset quantized KV cache to f16 when flash-attn-off fallback fires

Studio force-enables --flash-attn on for GGUF launches. On a hard startup
or first-decode crash it retries via _with_flash_attn_off, which flipped FA
off but left --cache-type-k/-v untouched. A quantized KV cache (q8_0, q4_0,
q4_1, q5_0, q5_1, iq4_nl) requires flash attention in llama.cpp, so the retry
itself aborted at init with 'V cache quantization requires flash_attn' instead
of recovering.

Reset any quantized --cache-type-k/-v to f16 in the FA-off fallback path so
the retry can actually launch. Non-quantized types (f16, bf16, f32) run fine
without flash attention and are left unchanged. Handles long and short flag
forms and both space and equals syntax, rewriting in place to preserve list
length. Adds pytest coverage.

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

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

* Studio: FA-off fallback resets only the quantized V cache and drops env-only V cache

Only the V cache requires flash attention in llama.cpp; a quantized K cache
runs fine without it. Restrict the FA-off crash-recovery reset to the V axis
(main and draft) so a memory-constrained config keeps its quantized K cache
instead of risking an OOM on the recovery. Also drop an inherited quantized V
cache set purely through the environment (LLAMA_ARG_CACHE_TYPE_V /
LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V) at the FA-off retry sites, which the argv
rewrite cannot reach, so the child falls back to the f16 default rather than
aborting.

* Studio: normalize underscore V-cache aliases in the FA-off fallback

llama.cpp rewrites '_' to '-' for any '--' long option before matching,
so a pass-through --cache_type_v q8_0 enables a quantized V cache just
like --cache-type-v. The FA-off crash-recovery reset only matched the
hyphenated spelling, so the underscore alias slipped through and the
retry still aborted with "V cache quantization requires flash_attn".
Canonicalize the flag name the same way before matching (short flags and
the type value are untouched).

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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-25 04:10:44 -07:00 • committed by GitHub
commit 2d026a1184
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 301 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

106
studio/backend/core/inference/llama_cpp.py
View file

@ -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 <path>' 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")

195
studio/backend/tests/test_llama_cpp_mmproj_fallback.py
View file

@ -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