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

Studio: match llama.cpp SWA cache sizing (#7530)

Browse source 
* Studio: match llama.cpp SWA cache sizing

* Studio: account for batch-capped SWA ubatch

* Studio: match llama.cpp KV stream padding

* Match llama.cpp batch and FA-off cache sizing

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

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

* Skip unusable compact SWA slot saves

* Align KV planning with launched server

* Match cache type casing and narrow the compact SWA slot-save skip

The launcher tested the requested cache type case-sensitively while the budget
lowercases it via _planned_main_cache_types, so a Q8_0 request emitted no
--cache-type flag and llama.cpp ran f16 while the estimate priced q8_0 (1.01 GiB
under-reserved on a 27B SWA model at ctx 32768 with 4 slots).

The compact SWA slot-save skip keyed on the sliding window alone, but the
estimator's SWA path also requires key/value length. phi3 GGUFs report a window
without those dimensions and llama.cpp runs them non-SWA, so their slots restore
fine and were being skipped.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
oobabooga 2026-07-28 09:18:15 -03:00 • committed by GitHub
commit 7b048168c8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 1331 additions and 294 deletions
Download patch file Download diff file Expand all files Collapse all files

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

File diff suppressed because it is too large Load diff

9
studio/backend/core/inference/llama_server_args.py
View file

@ -80,9 +80,10 @@ _DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)
def _flag_name(token: str) -> Optional[str]:
"""Flag name for ``token``, or None if it isn't a flag.
Peels `--key=value` to `--key`, treats `-1`/`-0.5` as values (shorts
always start with a letter), and normalises attached `-np8` / `-np-1` /
`-np8x` to `-np`. Mirrors the CLI's `_expand_attached_np_short`.
Peels `--key=value` to `--key`, normalises long-option underscores like
llama.cpp, treats `-1`/`-0.5` as values (shorts always start with a letter),
and normalises attached `-np8` / `-np-1` / `-np8x` to `-np`. Mirrors the
CLI's `_expand_attached_np_short`.
"""
token = token.strip()
if not token.startswith("-") or token in {"-", "--"}:
@ -90,6 +91,8 @@ def _flag_name(token: str) -> Optional[str]:
if len(token) >= 2 and (token[1].isdigit() or token[1] == "."):
return None
name = token.split("=", 1)[0]
if name.startswith("--"):
name = name.replace("_", "-")
if len(name) > 3 and name.startswith("-np"):
suffix = name[3:]
if suffix[0].isdigit() or (

2
studio/backend/models/inference.py
View file

@ -254,6 +254,8 @@ class ValidateModelRequest(BaseModel):
# /load; defaults preserve old behavior for callers that omit them.
max_seq_length: int = Field(0, ge = 0, le = 1048576)
load_in_4bit: bool = Field(True)
cache_type_kv: Optional[str] = Field(None)
tensor_parallel: bool = Field(False)
gpu_ids: Optional[List[int]] = Field(None)
gpu_memory_mode: Literal["auto", "manual"] = Field(
"auto",

78
studio/backend/routes/inference.py
View file

@ -1004,8 +1004,13 @@ try:
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_STREAM_STALL_TIMEOUT_S,
_canonicalize_spec_mode,
_extra_args_n_ubatch,
_extra_args_set_spec_type,
_hf_offline_if_dns_dead,
_kv_bytes_per_elem,
_kv_unified_from_args,
_planned_main_cache_types,
_swa_full_from_args_or_env,
detect_reasoning_flags,
)
from core.inference.llama_server_args import (
@ -1043,8 +1048,13 @@ except ImportError:
_DEFAULT_MAX_TOKENS_FLOOR,
_DEFAULT_STREAM_STALL_TIMEOUT_S,
_canonicalize_spec_mode,
_extra_args_n_ubatch,
_extra_args_set_spec_type,
_hf_offline_if_dns_dead,
_kv_bytes_per_elem,
_kv_unified_from_args,
_planned_main_cache_types,
_swa_full_from_args_or_env,
detect_reasoning_flags,
)
from core.inference.llama_server_args import (
@ -3320,6 +3330,10 @@ def _request_matches_loaded_settings(
strip_offload = request.gpu_memory_mode == "manual",
)
)
if not llama_backend.is_diffusion and llama_backend.swa_full != _swa_full_from_args_or_env(
effective_extra
):
return False
if not _tensor_parallel_matches_loaded(
effective_extra, request.tensor_parallel, llama_backend.tensor_parallel
):
@ -4435,10 +4449,12 @@ def _estimate_gguf_kv_gb(
max_seq_length: int,
llama_extra_args: Optional[list[str]] = None,
n_parallel: int = 1,
cache_type_kv: Optional[str] = None,
tensor_parallel: bool = False,
) -> float:
"""KV-cache VRAM (GB) at the larger of max_seq_length and any `--ctx-size`/`-c`
override, over n_parallel slots, with the default f16 cache so the estimate is
never below what the server allocates. 0 if metadata is unreadable."""
override, over n_parallel slots, using the effective cache settings and managed
launcher defaults. 0 if metadata is unreadable."""
try:
from core.inference.llama_server_args import parse_ctx_override
@ -4453,7 +4469,43 @@ def _estimate_gguf_kv_gb(
ctx = max(max_seq_length or 0, ctx_override) or (probe._context_length or 0)
if ctx <= 0:
return 0.0
kv = probe._estimate_kv_cache_bytes(ctx, n_parallel = max(1, n_parallel or 1))
slots = max(1, n_parallel or 1)
managed_kv_unified = bool(
slots > 1
and LlamaCppBackend.probe_server_capabilities().get("supports_kv_unified", False)
)
planned_cache_types = _planned_main_cache_types(
cache_type_kv,
llama_extra_args,
)
if tensor_parallel and any(
cache_type not in LlamaCppBackend._TENSOR_PARALLEL_KV_TYPES
for cache_type in planned_cache_types
):
# Tensor mode strips quantized axes, but a layer fallback restores
# the original settings. Size for the larger successful outcome.
tensor_cache_types = _planned_main_cache_types(None, None)
cache_type_for_budget = max(
(*planned_cache_types, *tensor_cache_types, "f16"),
key = _kv_bytes_per_elem,
)
else:
cache_type_for_budget = max(
planned_cache_types,
key = _kv_bytes_per_elem,
)
kv = probe._estimate_kv_cache_bytes(
ctx,
cache_type_for_budget,
n_parallel = slots,
swa_full = _swa_full_from_args_or_env(llama_extra_args),
kv_unified = _kv_unified_from_args(
llama_extra_args,
default = managed_kv_unified,
),
n_ubatch = _extra_args_n_ubatch(llama_extra_args, n_ctx = ctx),
flash_attn = False,
)
return kv / (1024**3)
except Exception as e:
logger.warning(f"Could not size GGUF KV cache for training guard: {e}")
@ -4466,6 +4518,8 @@ def _estimate_gguf_required_gb(
max_seq_length: int = 0,
llama_extra_args: Optional[list[str]] = None,
n_parallel: int = 1,
cache_type_kv: Optional[str] = None,
tensor_parallel: bool = False,
) -> Optional[float]:
"""Approximate GGUF VRAM (GB): quantized weights + companions, plus the KV
cache for local files (unreadable pre-download for remote). None when nothing
@ -4481,7 +4535,12 @@ def _estimate_gguf_required_gb(
total_bytes += Path(f).stat().st_size
if total_bytes > 0:
return total_bytes / (1024**3) + _estimate_gguf_kv_gb(
main, max_seq_length, llama_extra_args, n_parallel
main,
max_seq_length,
llama_extra_args,
n_parallel,
cache_type_kv,
tensor_parallel,
)
repo = getattr(config, "gguf_hf_repo", None)
@ -4622,6 +4681,8 @@ def _guard_chat_load_against_training(
requested_gpu_ids: Optional[List[int]],
llama_extra_args: Optional[list[str]] = None,
n_parallel: int = 1,
cache_type_kv: Optional[str] = None,
tensor_parallel: bool = False,
gpu_memory_mode: Literal["auto", "manual"] = "auto",
) -> None:
"""Protect active training from automatically placed chat-model loads.
@ -4676,6 +4737,11 @@ def _guard_chat_load_against_training(
max_seq_length = max_seq_length,
llama_extra_args = llama_extra_args,
n_parallel = n_parallel,
cache_type_kv = cache_type_kv,
tensor_parallel = (
_effective_tensor_parallel(llama_extra_args, tensor_parallel)
and (is_vulkan or LlamaCppBackend._effective_gpu_count(requested_gpu_ids) >= 2)
),
)
if is_gguf
else None
@ -5416,6 +5482,8 @@ async def _load_model_impl(
requested_gpu_ids = effective_gpu_ids,
llama_extra_args = extra_llama_args,
n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1),
cache_type_kv = request.cache_type_kv,
tensor_parallel = bool(request.tensor_parallel),
gpu_memory_mode = request.gpu_memory_mode,
)
@ -6092,6 +6160,8 @@ async def validate_model(
if fastapi_request is not None
else 1
),
cache_type_kv = request.cache_type_kv,
tensor_parallel = request.tensor_parallel,
gpu_memory_mode = request.gpu_memory_mode,
)

103
studio/backend/tests/test_chat_load_during_training.py
View file

@ -451,6 +451,9 @@ class TestChatLoadGuardRoute(unittest.TestCase):
decision,
gpu_memory_mode = "auto",
requested_gpu_ids = None,
llama_extra_args = None,
cache_type_kv = None,
tensor_parallel = False,
):
config = config or SimpleNamespace(is_gguf = False, is_lora = False, path = None)
with _stub_guard_deps(
@ -463,6 +466,9 @@ class TestChatLoadGuardRoute(unittest.TestCase):
load_in_4bit = True,
max_seq_length = 0,
requested_gpu_ids = requested_gpu_ids,
llama_extra_args = llama_extra_args,
cache_type_kv = cache_type_kv,
tensor_parallel = tensor_parallel,
gpu_memory_mode = gpu_memory_mode,
)
@ -597,6 +603,32 @@ class TestChatLoadGuardRoute(unittest.TestCase):
self.assertEqual(captured[0]["is_gguf"], True)
self.assertEqual(captured[0]["required_override_gb"], 12.5)
def test_vulkan_gguf_estimate_keeps_tensor_cache_coercion(self):
config = SimpleNamespace(is_gguf = True)
estimate_kwargs = {}
with (
patch.object(
self.route,
"_estimate_gguf_required_gb",
side_effect = lambda *args, **kwargs: estimate_kwargs.update(kwargs) or 12.5,
),
patch.object(
self.route.LlamaCppBackend,
"_effective_gpu_count",
return_value = 0,
),
patch.object(self.route.LlamaCppBackend, "_is_vulkan_backend", return_value = True),
):
self._guard(
config = config,
training_active = True,
decision = (True, {}),
llama_extra_args = ["--split-mode", "tensor"],
cache_type_kv = "q4_0",
)
self.assertEqual(estimate_kwargs["cache_type_kv"], "q4_0")
self.assertTrue(estimate_kwargs["tensor_parallel"])
class TestEffectiveLoadIn4bit(unittest.TestCase):
@classmethod
@ -745,7 +777,12 @@ class TestValidateRefusesDuringTraining(unittest.TestCase):
# /load then 409s after the frontend has already unloaded.
from models.inference import ValidateModelRequest
request = ValidateModelRequest(model_path = "unsloth/Qwen3-1.7B", max_seq_length = 4096)
request = ValidateModelRequest(
model_path = "unsloth/Qwen3-1.7B",
max_seq_length = 4096,
cache_type_kv = "f32",
tensor_parallel = True,
)
cfg = SimpleNamespace(
identifier = "unsloth/Qwen3-1.7B",
display_name = "Qwen3-1.7B",
@ -774,6 +811,8 @@ class TestValidateRefusesDuringTraining(unittest.TestCase):
asyncio.run(self.route.validate_model(request, current_subject = "u"))
self.assertEqual(captured.get("llama_extra_args"), ["-c", "32768"])
self.assertIn("n_parallel", captured)
self.assertEqual(captured.get("cache_type_kv"), "f32")
self.assertTrue(captured.get("tensor_parallel"))
def test_metadata_probe_skips_training_guard(self):
# A header-only probe (include_context_length) allocates no VRAM, so the
@ -985,6 +1024,8 @@ class TestEstimateGgufRequiredGb(unittest.TestCase):
class _FakeBackend:
_context_length = 2048
_TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"})
supports_kv_unified = True
def _read_gguf_metadata(self, path):
pass
@ -992,13 +1033,27 @@ class TestEstimateGgufRequiredGb(unittest.TestCase):
def _can_estimate_kv(self):
return True
@classmethod
def probe_server_capabilities(cls):
return {"supports_kv_unified": cls.supports_kv_unified}
def _estimate_kv_cache_bytes(
self,
ctx,
cache_type = None,
n_parallel = 1,
swa_full = False,
kv_unified = False,
n_ubatch = None,
flash_attn = True,
):
seen["ctx"] = ctx
seen["cache_type"] = cache_type
seen["n_parallel"] = n_parallel
seen["swa_full"] = swa_full
seen["kv_unified"] = kv_unified
seen["n_ubatch"] = n_ubatch
seen["flash_attn"] = flash_attn
return ctx * n_parallel * (1024**2) # 1 MiB per ctx unit per slot
with patch.object(self.route, "LlamaCppBackend", _FakeBackend):
@ -1009,6 +1064,8 @@ class TestEstimateGgufRequiredGb(unittest.TestCase):
)
self.assertEqual(seen["ctx"], 131072)
self.assertEqual(seen["n_parallel"], 1) # default single slot
self.assertFalse(seen["swa_full"])
self.assertFalse(seen["flash_attn"])
# override below max_seq_length -> larger (max_seq_length) wins
self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "1024"]), 4.0)
self.assertEqual(seen["ctx"], 4096)
@ -1020,6 +1077,50 @@ class TestEstimateGgufRequiredGb(unittest.TestCase):
# --parallel slots scale the cache the same way the launcher does
self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, None, 4), 16.0)
self.assertEqual(seen["n_parallel"], 4)
self.assertTrue(seen["kv_unified"])
# User extras are appended after Studio's managed default.
r._estimate_gguf_kv_gb("m", 4096, ["--no-kv-unified"], 4)
self.assertFalse(seen["kv_unified"])
# An older binary without the flag keeps separate KV streams.
_FakeBackend.supports_kv_unified = False
r._estimate_gguf_kv_gb("m", 4096, None, 4)
self.assertFalse(seen["kv_unified"])
r._estimate_gguf_kv_gb("m", 4096, None, 1, "f32")
self.assertEqual(seen["cache_type"], "f32")
r._estimate_gguf_kv_gb("m", 4096, ["--cache-type-v", "f32"])
self.assertEqual(seen["cache_type"], "f32")
with patch.dict(self.route.os.environ, {"LLAMA_ARG_CACHE_TYPE_K": "f32"}):
r._estimate_gguf_kv_gb("m", 4096)
self.assertEqual(seen["cache_type"], "f32")
with patch.dict(
self.route.os.environ,
{
"LLAMA_ARG_CACHE_TYPE_K": "q4_0",
"LLAMA_ARG_CACHE_TYPE_V": "q4_0",
},
):
r._estimate_gguf_kv_gb("m", 4096)
self.assertEqual(seen["cache_type"], "q4_0")
r._estimate_gguf_kv_gb(
"m",
4096,
["--cache-type-k", "q4_0", "--cache-type-v", "q4_0"],
tensor_parallel = True,
)
self.assertEqual(seen["cache_type"], "f16")
r._estimate_gguf_kv_gb(
"m",
4096,
["--cache-type-k", "f32", "--cache-type-v", "q4_0"],
tensor_parallel = True,
)
self.assertEqual(seen["cache_type"], "f32")
# Full SWA mode follows the same pass-through args as the launcher.
r._estimate_gguf_kv_gb("m", 4096, ["--swa_full"])
self.assertTrue(seen["swa_full"])
r._estimate_gguf_kv_gb("m", 4096, ["--kv_unified", "--ubatch_size", "256"])
self.assertTrue(seen["kv_unified"])
self.assertEqual(seen["n_ubatch"], 256)
# ── load_model integration: authoritative 409, and no unload before refusal ──

3
studio/backend/tests/test_gpu_memory_mode.py
View file

@ -183,11 +183,12 @@ def test_already_in_target_state_reloads_on_mode_change(loaded, requested):
assert _target_state(_loaded_backend(loaded), requested) is False
def test_already_in_target_state_ignores_mode_for_diffusion():
def test_already_in_target_state_ignores_mode_for_diffusion(monkeypatch):
# The diffusion runner is mode-agnostic (always "auto"), so a standing manual
# preference must not force a needless reload.
backend = _loaded_backend("auto")
backend._is_diffusion = True
monkeypatch.setenv("LLAMA_ARG_SWA_FULL", "1")
assert _target_state(backend, "manual") is True

355
studio/backend/tests/test_kv_cache_estimation.py
View file

@ -76,6 +76,39 @@ from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend
# Helpers
def _runtime_kv_cells(
n_ctx: int,
*,
slots: int = 1,
unified: bool = True,
) -> int:
"""Total KV cells allocated by llama.cpp across all streams."""
slots = max(1, slots)
padded_ctx = ((n_ctx + 255) // 256) * 256
streams = 1 if unified else slots
cells_per_stream = padded_ctx if unified else ((max(1, padded_ctx // slots) + 255) // 256) * 256
return cells_per_stream * streams
def _runtime_swa_cells(
n_ctx: int,
sliding_window: int,
*,
slots: int = 1,
unified: bool = True,
n_ubatch: int = 512,
) -> tuple[int, int]:
"""Return total non-SWA and compact-SWA cells allocated by llama.cpp."""
slots = max(1, slots)
streams = 1 if unified else slots
base_cells = _runtime_kv_cells(n_ctx, slots = slots, unified = unified)
cells_per_stream = base_cells // streams
swa_limit = sliding_window * (slots if unified else 1) + n_ubatch
swa_cells_per_stream = min(cells_per_stream, swa_limit)
swa_cells_per_stream = ((swa_cells_per_stream + 255) // 256) * 256
return base_cells, swa_cells_per_stream * streams
def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes:
"""Build a minimal GGUF v3 blob with the given KV metadata.
@ -789,7 +822,7 @@ class TestMLAEstimation:
b = self._mla_backend()
result = b._estimate_kv_cache_bytes(1000, "f16")
# n_layers * ctx * 1 * key_len(576) * 2
expected = 61 * 1000 * 1 * 576 * 2
expected = 61 * _runtime_kv_cells(1000) * 1 * 576 * 2
assert result == expected
def test_mla_fallback_when_no_key_length(self):
@ -797,14 +830,14 @@ class TestMLAEstimation:
b = self._mla_backend(_kv_key_length = None)
# default _key_length_mla=192, so rope_dim=192
result = b._estimate_kv_cache_bytes(1000, "f16")
expected = 61 * 1000 * 1 * (512 + 192) * 2 # 704
expected = 61 * _runtime_kv_cells(1000) * 1 * (512 + 192) * 2 # 704
assert result == expected
def test_mla_fallback_no_key_length_mla(self):
"""No key_length and no key_length_mla: fall back to +64."""
b = self._mla_backend(_kv_key_length = None, _key_length_mla = None)
result = b._estimate_kv_cache_bytes(1000, "f16")
expected = 61 * 1000 * 1 * (512 + 64) * 2 # 576
expected = 61 * _runtime_kv_cells(1000) * 1 * (512 + 64) * 2 # 576
assert result == expected
def test_mla_defaults_n_kv_to_1_when_heads_absent(self):
@ -812,7 +845,7 @@ class TestMLAEstimation:
b = self._mla_backend(_n_kv_heads = None) # n_heads=128 still set
result = b._estimate_kv_cache_bytes(1000, "f16")
# Uses n_kv_mla=1, NOT n_heads=128
expected = 61 * 1000 * 1 * 576 * 2
expected = 61 * _runtime_kv_cells(1000) * 1 * 576 * 2
assert result == expected
def test_mla_q4_quantization(self):
@ -821,7 +854,7 @@ class TestMLAEstimation:
result_q4 = b._estimate_kv_cache_bytes(1000, "q4_0")
assert result_q4 < result_f16
# q4_0 bpe = 0.5625, f16 bpe = 2.0
assert result_q4 == int(61 * 1000 * 1 * 576 * 0.5625)
assert result_q4 == int(61 * _runtime_kv_cells(1000) * 1 * 576 * 0.5625)
# D. Path 2: Hybrid Mamba Estimation
@ -910,9 +943,8 @@ class TestSlidingWindowEstimation:
n_global = max(1, 62 // 4) # 15
n_swa = 62 - n_global # 47
kv_per = 16 * (128 + 128) * 2
# SWA cache is double-buffered: 2 * sliding_window cells, capped at n_ctx.
swa_cells = min(131072, 2 * 1024)
expected = int(n_global * 131072 * kv_per + n_swa * swa_cells * kv_per)
base_cells, swa_cells = _runtime_swa_cells(131072, 1024)
expected = int(n_global * base_cells * kv_per + n_swa * swa_cells * kv_per)
assert b._estimate_kv_cache_bytes(131072, "f16") == expected
def test_gpt_oss(self):
@ -929,8 +961,8 @@ class TestSlidingWindowEstimation:
n_global = max(1, 24 // 4) # 6
n_swa = 24 - n_global # 18
kv_per = 8 * (64 + 64) * 2
swa_cells = min(131072, 2 * 128)
expected = int(n_global * 131072 * kv_per + n_swa * swa_cells * kv_per)
base_cells, swa_cells = _runtime_swa_cells(131072, 128)
expected = int(n_global * base_cells * kv_per + n_swa * swa_cells * kv_per)
assert b._estimate_kv_cache_bytes(131072, "f16") == expected
def test_gemma4_per_layer_swa_metadata(self):
@ -952,21 +984,67 @@ class TestSlidingWindowEstimation:
sliding_layers = 25
def expected(ctx):
full = full_layers * ctx * 2 * (512 + 512) * 2
sliding = sliding_layers * min(ctx, 2 * 1024) * 8 * (256 + 256) * 2
base_cells, swa_cells = _runtime_swa_cells(ctx, 1024)
full = full_layers * base_cells * 2 * (512 + 512) * 2
sliding = sliding_layers * swa_cells * 8 * (256 + 256) * 2
return int(full + sliding)
for ctx in (4096, 46500, 262144):
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected(ctx)
def test_gemma4_flash_attn_off_pads_v_to_model_max(self):
b = self._swa_backend(
_n_layers = 35,
_n_kv_heads = 1,
_n_heads = 8,
_embedding_length = 1536,
_kv_key_length = 512,
_kv_value_length = 512,
_sliding_window = 512,
_sliding_window_pattern = [True, True, True, True, False] * 7,
_kv_key_length_swa = 256,
_kv_value_length_swa = 256,
_shared_kv_layers = 20,
)
ctx = 5000
slots = 3
base_cells, swa_cells = _runtime_swa_cells(ctx, 512, slots = slots, unified = True)
max_v_width = 512
expected = (
3 * base_cells * (512 + max_v_width) * 2 + 12 * swa_cells * (256 + max_v_width) * 2
)
actual = b._estimate_kv_cache_bytes(
ctx,
"f16",
n_parallel = slots,
flash_attn = False,
)
assert actual == expected
assert actual == 66 * 1024**2
assert actual > b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots)
def test_flash_attn_off_prices_quantized_v_retry_as_f16(self):
b = self._swa_backend(
_n_layers = 2,
_n_kv_heads = None,
_n_kv_heads_by_layer = [8, 2],
_sliding_window_pattern = [True, False],
_kv_key_length_swa = 64,
_kv_value_length_swa = 64,
)
off = b._estimate_kv_cache_bytes(4096, "q4_0", flash_attn = False)
on = b._estimate_kv_cache_bytes(4096, "q4_0")
assert off > on
def test_ctx_smaller_than_window(self):
"""When ctx < 2 * sliding_window, SWA cache caps at ctx."""
"""When context is smaller than the compact allowance, SWA caps at context."""
b = self._swa_backend(_sliding_window = 8192)
n_global = max(1, 62 // 4) # 15
n_swa = 62 - n_global # 47
kv_per = 16 * (128 + 128) * 2
ctx = 4096
expected = int(n_global * ctx * kv_per + n_swa * min(ctx, 2 * 8192) * kv_per)
base_cells, swa_cells = _runtime_swa_cells(ctx, 8192)
expected = int(n_global * base_cells * kv_per + n_swa * swa_cells * kv_per)
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
def test_odd_layer_count(self):
@ -974,7 +1052,8 @@ class TestSlidingWindowEstimation:
n_global = max(1, 63 // 4) # 15
n_swa = 63 - n_global # 48
kv_per = 16 * (128 + 128) * 2
expected = int(n_global * 1000 * kv_per + n_swa * min(1000, 2 * 1024) * kv_per)
base_cells, swa_cells = _runtime_swa_cells(1000, 1024)
expected = int(n_global * base_cells * kv_per + n_swa * swa_cells * kv_per)
assert b._estimate_kv_cache_bytes(1000, "f16") == expected
@ -1086,8 +1165,7 @@ class TestPathPriority:
b._full_attention_interval = 4
b._sliding_window = 1024 # Would trigger SWA
# MLA: 61 * 1000 * 1 * 576 * 2
expected_mla = int(61 * 1000 * 1 * 576 * 2)
expected_mla = int(61 * _runtime_kv_cells(1000) * 1 * 576 * 2)
assert b._estimate_kv_cache_bytes(1000, "f16") == expected_mla
def test_hybrid_over_swa(self):
@ -1104,7 +1182,7 @@ class TestPathPriority:
b._sliding_window = 1024 # Would trigger SWA
n_attn = 64 // 4
expected_hybrid = int(n_attn * 1000 * 4 * (256 + 256) * 2)
expected_hybrid = int(n_attn * _runtime_kv_cells(1000) * 4 * (256 + 256) * 2)
assert b._estimate_kv_cache_bytes(1000, "f16") == expected_hybrid
def test_all_paths_produce_different_values(self):
@ -1192,7 +1270,7 @@ class TestQuantization:
b._kv_key_length = 64
b._kv_value_length = 64
result = b._estimate_kv_cache_bytes(1000, cache_type)
expected = int(10 * 1000 * 1 * (64 + 64) * expected_bpe)
expected = int(10 * _runtime_kv_cells(1000) * 1 * (64 + 64) * expected_bpe)
assert result == expected
@ -1221,7 +1299,7 @@ class TestEdgeCases:
b._kv_key_length = 64
b._kv_value_length = 64
result = b._estimate_kv_cache_bytes(1, "f16")
assert result == int(10 * 1 * 1 * (64 + 64) * 2)
assert result == int(10 * _runtime_kv_cells(1) * 1 * (64 + 64) * 2)
def test_very_large_context(self):
"""1M context should not overflow or crash."""
@ -1242,7 +1320,7 @@ class TestEdgeCases:
b._kv_key_length = 64
b._kv_value_length = 64
result = b._estimate_kv_cache_bytes(100, "f16")
expected = int(10 * 100 * 8 * (64 + 64) * 2)
expected = int(10 * _runtime_kv_cells(100) * 8 * (64 + 64) * 2)
assert result == expected
def test_both_heads_none_falls_to_one(self):
@ -1253,7 +1331,7 @@ class TestEdgeCases:
b._kv_key_length = 64
b._kv_value_length = 64
result = b._estimate_kv_cache_bytes(100, "f16")
expected = int(10 * 100 * 1 * (64 + 64) * 2)
expected = int(10 * _runtime_kv_cells(100) * 1 * (64 + 64) * 2)
assert result == expected
@ -1335,12 +1413,21 @@ class TestServerFlags:
assert with_cp_full == no_cp_full
assert with_cp > b._estimate_kv_cache_bytes(8192, "f16")
def test_compact_swa_includes_ubatch_headroom_and_padding(self):
b = self._swa_backend(_sliding_window = 128)
ctx = 8192
result = b._estimate_kv_cache_bytes(ctx, "f16", n_ubatch = 512)
per_token = 4 * (256 + 256) * 2
n_swa = sum(b._sliding_window_pattern)
n_global = b._n_layers - n_swa
expected = n_global * ctx * per_token + n_swa * 768 * per_token
assert result == expected
# ── --parallel + --kv-unified ──────────────────────────────────
# Verified against llama-server: non-SWA caches partition n_ctx across
# slots (total memory constant); only SWA layers scale with --parallel.
# --kv-unified is a no-op for memory math (kept for API forward-compat).
# non-unified streams. Compact SWA sizing depends on the stream layout.
def test_gqa_kv_constant_across_parallel(self):
def test_gqa_kv_constant_for_aligned_stream_divisions(self):
b = self._gqa_backend()
baseline = b._estimate_kv_cache_bytes(4096, "f16")
for slots in (1, 2, 4, 8):
@ -1359,7 +1446,7 @@ class TestServerFlags:
== baseline
)
def test_swa_path_scales_only_swa_portion(self):
def test_swa_path_matches_aligned_stream_layout(self):
b = self._swa_backend()
ctx = 8192
baseline = b._estimate_kv_cache_bytes(ctx, "f16")
@ -1367,27 +1454,27 @@ class TestServerFlags:
swa = b._sliding_window
per_token_global = 4 * (256 + 256) * 2 # n_kv * (k+v) * f16
per_token_swa = 4 * (256 + 256) * 2 # k_swa/val_swa fall back
per_slot_swa_cells = min(ctx, 2 * swa) # not clamped at parallel=1
base_cells, swa_cells = _runtime_swa_cells(ctx, swa)
global_bytes = sum(
ctx * per_token_global for f in b._sliding_window_pattern[: b._n_layers] if not f
base_cells * per_token_global for f in b._sliding_window_pattern[: b._n_layers] if not f
)
swa_bytes_per_slot = sum(
per_slot_swa_cells * per_token_swa
for f in b._sliding_window_pattern[: b._n_layers]
if f
swa_bytes = sum(
swa_cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f
)
# Sanity: parallel=1 reproduces baseline exactly
assert global_bytes + swa_bytes_per_slot == baseline
# Only the SWA portion scales by parallel
assert global_bytes + swa_bytes == baseline
for slots in (1, 2, 3, 4):
scaled = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False)
# SWA cells clamp to per_slot_ctx when ctx/slots < 2*swa
per_slot_ctx = max(1, ctx // slots)
cells = min(ctx, 2 * swa, per_slot_ctx)
swa_bps = sum(
cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f
base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = False)
expected_global = sum(
base_cells * per_token_global
for f in b._sliding_window_pattern[: b._n_layers]
if not f
)
assert scaled == global_bytes + slots * swa_bps
expected_swa = sum(
swa_cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f
)
assert scaled == expected_global + expected_swa
def test_mla_kv_constant_across_parallel(self):
b = LlamaCppBackend()
@ -1444,19 +1531,17 @@ class TestServerFlags:
ctx = 8192
swa = b._sliding_window
per_token = 4 * (256 + 256) * 2
global_bytes = sum(
ctx * per_token for f in b._sliding_window_pattern[: b._n_layers] if not f
)
n_swa_layers = sum(1 for f in b._sliding_window_pattern[: b._n_layers] if f)
slots = 3
per_slot_ctx = max(1, ctx // slots)
swa_cells = min(ctx, 2 * swa, per_slot_ctx)
swa_bytes_per_slot = n_swa_layers * swa_cells * per_token
base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = False)
n_global_layers = b._n_layers - n_swa_layers
global_bytes = n_global_layers * base_cells * per_token
swa_bytes = n_swa_layers * swa_cells * per_token
cp_extra_per_slot = n_swa_layers * 4 * swa * per_token # 4 checkpoints
flagged = b._estimate_kv_cache_bytes(
ctx, "f16", ctx_checkpoints = 4, n_parallel = slots, kv_unified = False
)
assert flagged == global_bytes + slots * (swa_bytes_per_slot + cp_extra_per_slot)
assert flagged == global_bytes + swa_bytes + slots * cp_extra_per_slot
# ── --kv-offload (kv_on_gpu) ───────────────────────────────────
@ -1535,22 +1620,40 @@ class TestServerFlags:
assert fitted_default == ctx
assert fitted_full < ctx
def test_tensor_planner_threads_swa_full_through_estimator(self):
b = self._swa_backend()
estimate = b._estimate_kv_cache_bytes
calls = []
def record(*args, **kwargs):
calls.append(kwargs)
return estimate(*args, **kwargs)
b._estimate_kv_cache_bytes = record
b._plan_tensor_parallel(
[(0, 32768), (1, 32768)],
1024**3,
8192,
cache_type_kv = "f16",
swa_full = True,
flash_attn = False,
)
assert calls
assert all(call["swa_full"] is True for call in calls)
assert all(call["flash_attn"] is False for call in calls)
# J2.5. --parallel N memory accounting (per-layer-type scaling rule)
class TestParallelSWAScaling:
"""Per-layer-type scaling rule vs the closed form measured from
llama-server. Empirical formula on Gemma-3 270m at ctx=8192:
total_kv = 24 + parallel * 15 (MiB).
"""Per-layer-type scaling rule measured from llama-server.
Rule (verified vs ``llama-server`` log on real GGUFs):
* non-SWA layers: total cells = n_ctx, partitioned across slots,
memory CONSTANT in n_parallel.
* SWA layers: per-slot cells = 2 * sliding_window (clamped at
n_ctx and at per_slot_ctx); memory LINEAR in n_parallel.
* --kv-unified is a no-op for memory math; both modes give the
same total in measured cases.
* non-SWA layers use the padded per-stream context.
* compact SWA adds ubatch headroom and pads to 256 cells.
* unified mode uses one stream with all slot windows.
* non-unified mode allocates one stream per slot.
"""
def _gqa_backend(self, **overrides):
@ -1586,7 +1689,7 @@ class TestParallelSWAScaling:
setattr(b, k, v)
return b
# ── non-SWA paths: constant ────────────────────────────────────
# ── non-SWA paths: constant when stream divisions are aligned ──
def test_pure_gqa_constant_across_parallel(self):
b = self._gqa_backend()
@ -1633,25 +1736,53 @@ class TestParallelSWAScaling:
for slots in (1, 2, 4, 8):
assert b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots) == baseline
# ── SWA paths: scale only the SWA portion ──────────────────────
def test_non_swa_paths_follow_unaligned_stream_padding(self):
mla = LlamaCppBackend()
mla._n_layers = 60
mla._n_kv_heads = 1
mla._kv_lora_rank = 512
mla._key_length_mla = 64
mla._kv_key_length = 576
def test_swa_pattern_scales_only_swa_portion(self):
hybrid = LlamaCppBackend()
hybrid._n_layers = 64
hybrid._n_kv_heads = 16
hybrid._n_heads = 32
hybrid._embedding_length = 4096
hybrid._kv_key_length = 128
hybrid._kv_value_length = 128
hybrid._ssm_inner_size = 4096
hybrid._full_attention_interval = 4
legacy = LlamaCppBackend()
legacy._n_layers = 32
legacy._n_kv_heads = 8
legacy._n_heads = 8
legacy._embedding_length = 4096
for backend in (self._gqa_backend(), mla, hybrid, legacy):
bytes_per_cell = backend._estimate_kv_cache_bytes(256, "f16") // 256
unified = backend._estimate_kv_cache_bytes(5000, "f16", n_parallel = 3, kv_unified = True)
separate = backend._estimate_kv_cache_bytes(5000, "f16", n_parallel = 3, kv_unified = False)
assert unified == 5120 * bytes_per_cell
assert separate == 5376 * bytes_per_cell
# ── SWA paths: aligned stream scaling ──────────────────────────
def test_swa_pattern_matches_aligned_stream_layout(self):
b = self._swa_backend()
ctx = 8192
swa = b._sliding_window
per_token = 1 * (256 + 256) * 2 # n_kv * (k+v) * f16
n_global = sum(1 for f in b._sliding_window_pattern if not f)
n_swa = sum(1 for f in b._sliding_window_pattern if f)
global_bytes = n_global * ctx * per_token
for slots in (1, 2, 4, 8):
per_slot_ctx = max(1, ctx // slots)
cells = min(ctx, 2 * swa, per_slot_ctx)
swa_bps = n_swa * cells * per_token
for unified in (True, False):
base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = unified)
got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = unified)
assert got == global_bytes + slots * swa_bps
assert got == (n_global * base_cells * per_token + n_swa * swa_cells * per_token)
def test_swa_fallback_scales_only_swa_portion(self):
def test_swa_fallback_matches_aligned_stream_layout(self):
# No per-layer pattern -> 1/4-global heuristic.
b = self._swa_backend(_sliding_window_pattern = None)
ctx = 8192
@ -1660,34 +1791,28 @@ class TestParallelSWAScaling:
n_global = max(1, n_layers // 4)
n_swa = n_layers - n_global
per_token = 1 * (256 + 256) * 2
global_bytes = n_global * ctx * per_token
for slots in (1, 2, 4, 8):
per_slot_ctx = max(1, ctx // slots)
cells = min(ctx, 2 * swa, per_slot_ctx)
swa_bps = n_swa * cells * per_token
got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots)
assert got == global_bytes + slots * swa_bps
for unified in (True, False):
base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = unified)
got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = unified)
assert got == (n_global * base_cells * per_token + n_swa * swa_cells * per_token)
def test_swa_per_slot_clamped_when_ctx_lt_slots_x_2window(self):
# ctx=4096 / slots=8 -> per_slot_ctx=512, but 2*sliding=1024.
# SWA cells clamp at per_slot_ctx (512), not 2*sliding.
# ctx=4096 / slots=8 gives a 512-cell stream, which caps compact SWA.
b = self._swa_backend()
ctx = 4096
per_slot_ctx_at_8 = ctx // 8
assert per_slot_ctx_at_8 < 2 * b._sliding_window
# Build expected with the clamped formula
n_swa = sum(1 for f in b._sliding_window_pattern if f)
n_global = sum(1 for f in b._sliding_window_pattern if not f)
per_token = 1 * (256 + 256) * 2
global_bytes = n_global * ctx * per_token
cells = min(ctx, 2 * b._sliding_window, per_slot_ctx_at_8)
assert cells == per_slot_ctx_at_8
expected = global_bytes + 8 * (n_swa * cells * per_token)
assert b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = 8) == expected
base_cells, swa_cells = _runtime_swa_cells(ctx, b._sliding_window, slots = 8, unified = False)
assert swa_cells == 8 * per_slot_ctx_at_8
expected = n_global * base_cells * per_token + n_swa * swa_cells * per_token
assert b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = 8, kv_unified = False) == expected
def test_swa_full_does_not_scale_under_parallel(self):
# swa_full forces every layer to n_ctx -> all-global GQA-style
# total, constant in parallel.
def test_swa_full_constant_for_aligned_stream_divisions(self):
# swa_full forces every layer to n_ctx. This aligned context remains
# constant across the tested stream divisions.
b = self._swa_backend()
ctx = 8192
baseline = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
@ -1696,25 +1821,32 @@ class TestParallelSWAScaling:
b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True, n_parallel = slots) == baseline
)
# ── kv_unified: no-op for memory math ──────────────────────────
# ── kv_unified stream layout ────────────────────────────────────
def test_kv_unified_is_no_op_for_memory_math(self):
# unified=True and unified=False must give the same total bytes
# for every backend type and parallel value.
backends = [
("gqa", self._gqa_backend()),
("swa", self._swa_backend()),
]
for label, b in backends:
for slots in (1, 2, 4, 8):
u = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = True)
nu = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = False)
assert u == nu, f"{label} parallel={slots} unified-mismatch"
def test_kv_unified_changes_only_compact_swa_for_aligned_context(self):
gqa = self._gqa_backend()
swa = self._swa_backend()
for slots in (1, 2, 4, 8):
gqa_unified = gqa._estimate_kv_cache_bytes(
8192, "f16", n_parallel = slots, kv_unified = True
)
gqa_separate = gqa._estimate_kv_cache_bytes(
8192, "f16", n_parallel = slots, kv_unified = False
)
assert gqa_unified == gqa_separate
swa_unified = swa._estimate_kv_cache_bytes(
8192, "f16", n_parallel = slots, kv_unified = True
)
swa_separate = swa._estimate_kv_cache_bytes(
8192, "f16", n_parallel = slots, kv_unified = False
)
assert (swa_unified == swa_separate) is (slots == 1)
# ── Empirical Gemma-3 270m formula ─────────────────────────────
def test_matches_empirical_gemma3_270m_formula(self):
"""Exact match against the formula measured from llama-server:
"""Exact match against the non-unified formula measured from llama-server:
total_kv = 24 + parallel * 15 (MiB) at ctx=8192.
Geometry: 18 layers (3 global + 15 SWA), n_kv=1, head_dim=256,
@ -1736,12 +1868,16 @@ class TestParallelSWAScaling:
# Confirm pattern shape
assert sum(b._sliding_window_pattern) == n_swa
for slots, expected_mib in [(1, 39), (2, 54), (4, 84)]:
got_bytes = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots)
got_bytes = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = False)
got_mib = got_bytes / (1024 * 1024)
assert (
got_mib == expected_mib
), f"slots={slots}: got {got_mib} MiB, expected {expected_mib} MiB"
for slots, expected_mib in [(1, 39), (2, 46.5), (4, 61.5)]:
got_bytes = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = True)
assert got_bytes / (1024 * 1024) == expected_mib
# J3. shared_kv_layers (Gemma 3n / Gemma 4)
@ -1844,8 +1980,8 @@ class TestSharedKVLayers:
assert sliding_in_unshared == 16
assert full_in_unshared == 4
kv_per = 4 * (256 + 256) * 2
swa_cells = min(ctx, 2 * 1024)
expected = full_in_unshared * ctx * kv_per + sliding_in_unshared * swa_cells * kv_per
base_cells, swa_cells = _runtime_swa_cells(ctx, 1024)
expected = full_in_unshared * base_cells * kv_per + sliding_in_unshared * swa_cells * kv_per
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
def test_shared_layers_reduces_estimate(self):
@ -1875,8 +2011,8 @@ class TestSharedKVLayers:
n_global = max(1, n_layers_kv // 4) # 5
n_swa = n_layers_kv - n_global # 15
kv_per = 4 * (256 + 256) * 2
swa_cells = min(ctx, 2 * 1024)
expected = n_global * ctx * kv_per + n_swa * swa_cells * kv_per
base_cells, swa_cells = _runtime_swa_cells(ctx, 1024)
expected = n_global * base_cells * kv_per + n_swa * swa_cells * kv_per
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
def test_shared_floors_at_one_layer(self):
@ -1896,13 +2032,12 @@ class TestSharedKVLayers:
unshared_pattern = b._sliding_window_pattern[:20] # 35 - 15 shared
sliding_in_unshared = sum(unshared_pattern)
global_in_unshared = len(unshared_pattern) - sliding_in_unshared
global_bytes = global_in_unshared * ctx * per_token
slots = 3
per_slot_ctx = max(1, ctx // slots)
swa_cells = min(ctx, 2 * swa, per_slot_ctx)
swa_bytes_per_slot = sliding_in_unshared * swa_cells * per_token
base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = False)
global_bytes = global_in_unshared * base_cells * per_token
swa_bytes = sliding_in_unshared * swa_cells * per_token
flagged = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False)
assert flagged == global_bytes + slots * swa_bytes_per_slot
assert flagged == global_bytes + swa_bytes
def test_composes_with_ctx_checkpoints(self):
b = self._gemma3n_backend()
@ -2036,14 +2171,14 @@ class TestLifecycle:
)
assert b._can_estimate_kv()
result = b._estimate_kv_cache_bytes(131072, "f16")
# gemma3 -> period 6 from bootstrap; SWA cache double-buffered to
# 2 * sliding_window cells.
# gemma3 uses period 6 from the bootstrap resolver.
period = 6
kv_per = 16 * 256 * 2
base_cells, swa_cells = _runtime_swa_cells(131072, 1024)
expected = 0
for i in range(62):
is_swa = (i + 1) % period != 0
layer_ctx = min(131072, 2 * 1024) if is_swa else 131072
layer_ctx = swa_cells if is_swa else base_cells
expected += layer_ctx * kv_per
assert result == expected

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

@ -221,6 +221,18 @@ class TestFlashAttnOff:
assert _flash_off(["llama-server", "-fa", "auto"]) == ["llama-server", "-fa", "off"]
assert _flash_off(["llama-server", "-fa=on"]) == ["llama-server", "-fa=off"]
@pytest.mark.parametrize("value", ["on", "enabled", "true", "1", "auto", "-1"])
def test_flips_every_enabled_value(self, value):
assert _flash_off(["llama-server", "--flash-attn", value]) == [
"llama-server",
"--flash-attn",
"off",
]
@pytest.mark.parametrize("value", ["off", "disabled", "false", "0"])
def test_none_for_every_disabled_value(self, value):
assert _flash_off(["llama-server", "--flash-attn", value]) is None
def test_flips_every_occurrence_last_wins(self):
# extra_args can re-enable FA after Unsloth's flag; llama.cpp is last-wins,
# so one leftover 'on' would re-crash the retry. Every enable must flip.
@ -384,6 +396,10 @@ class TestFlashAttnOffQuantizedKvCache:
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_alias_flash_attn_is_disabled(self):
out = _flash_off(["llama-server", "--flash_attn=on"])
assert out == ["llama-server", "--flash_attn=off"]
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).

37
studio/backend/tests/test_llama_cpp_mtp_detection.py
View file

@ -63,7 +63,9 @@ from core.inference.llama_cpp import (
_extra_args_set_any_flag,
_extra_args_set_spec_type,
_is_mtp_model_name,
_kv_unified_from_args,
_mla_mtp_auto_enabled,
_swa_full_from_args_or_env,
)
@ -147,6 +149,41 @@ def test_is_mtp_model_name_handles_none():
assert _is_mtp_model_name("", "") is False
@pytest.mark.parametrize("flag", ["--swa-full", "--swa_full"])
def test_swa_full_detects_llama_cpp_long_flag_spellings(flag):
assert _swa_full_from_args_or_env([flag], {}) is True
@pytest.mark.parametrize("value", ["on", "enabled", "true", "1"])
def test_swa_full_detects_llama_cpp_env_truth_values(value):
assert _swa_full_from_args_or_env([], {"LLAMA_ARG_SWA_FULL": value}) is True
@pytest.mark.parametrize("value", ["", "off", "yes", "TRUE", " true ", "0"])
def test_swa_full_rejects_values_llama_cpp_treats_as_false(value):
assert _swa_full_from_args_or_env([], {"LLAMA_ARG_SWA_FULL": value}) is False
def test_swa_full_cli_wins_when_env_is_false():
assert _swa_full_from_args_or_env(["--swa-full"], {"LLAMA_ARG_SWA_FULL": "0"}) is True
@pytest.mark.parametrize("flag", ["--kv-unified", "--kv_unified", "-kvu"])
def test_kv_unified_detects_enable_aliases(flag):
assert _kv_unified_from_args([flag]) is True
@pytest.mark.parametrize("flag", ["--no-kv-unified", "--no_kv_unified", "-no-kvu"])
def test_kv_unified_detects_disable_aliases(flag):
assert _kv_unified_from_args(["--kv-unified", flag]) is False
def test_kv_unified_uses_environment_before_cli():
assert _kv_unified_from_args([], env = {"LLAMA_ARG_KV_UNIFIED": "true"}) is True
assert _kv_unified_from_args([], default = True, env = {"LLAMA_ARG_KV_UNIFIED": "false"}) is True
assert _kv_unified_from_args(["--kv-unified"], env = {"LLAMA_ARG_KV_UNIFIED": "false"}) is True
def test_is_mtp_model_name_detects_marker_in_filename(tmp_path):
gguf = tmp_path / "Qwen3.6-27B-MTP-Q4_K_M.gguf"
gguf.write_bytes(b"")

28
studio/backend/tests/test_llama_cpp_props_readback.py
View file

@ -104,6 +104,9 @@ def _make_backend(effective_ctx = 98304, port = 51234):
inst._port = port
inst._effective_context_length = effective_ctx
inst._context_length = 262144
inst._effective_parallel_slots = 1
inst._kv_cache_unified = False
inst._kv_cache_context_total = None
return inst
@ -173,6 +176,31 @@ def test_fit_shrunk_ctx_overwrites_advertised_value(monkeypatch):
assert inst.context_length == 67584
def test_props_keeps_total_cache_context_for_slot_preflight(monkeypatch):
inst = _make_backend(effective_ctx = 32768)
inst._effective_parallel_slots = 4
_stub_props(
monkeypatch,
body = {"default_generation_settings": {"n_ctx": 8192}},
)
inst._reconcile_effective_ctx_with_server()
assert inst._effective_context_length == 8192
assert inst._kv_cache_context_total == 32768
def test_props_does_not_multiply_unified_cache_context(monkeypatch):
inst = _make_backend(effective_ctx = 32768)
inst._effective_parallel_slots = 4
inst._kv_cache_unified = True
_stub_props(
monkeypatch,
body = {"default_generation_settings": {"n_ctx": 32768}},
)
inst._reconcile_effective_ctx_with_server()
assert inst._effective_context_length == 32768
assert inst._kv_cache_context_total == 32768
def test_matching_ctx_is_left_alone(monkeypatch):
inst = _make_backend(effective_ctx = 98304)
_stub_props(

103
studio/backend/tests/test_llama_cpp_slot_resume.py
View file

@ -221,6 +221,34 @@ def test_fingerprint_tracks_effective_context_length(tmp_path):
assert backend._slot_launch_fingerprint() != before
def test_fingerprint_tracks_swa_full_mode(tmp_path):
backend = _resume_backend(tmp_path)
before = backend._slot_launch_fingerprint()
backend._swa_full = True
assert backend._slot_launch_fingerprint() != before
def test_fingerprint_tracks_unified_cache_mode(tmp_path):
backend = _resume_backend(tmp_path)
before = backend._slot_launch_fingerprint()
backend._kv_cache_unified = True
assert backend._slot_launch_fingerprint() != before
def test_fingerprint_tracks_flash_attention_mode(tmp_path):
backend = _resume_backend(tmp_path)
before = backend._slot_launch_fingerprint()
backend._flash_attn_enabled = False
assert backend._slot_launch_fingerprint() != before
def test_fingerprint_tracks_effective_cache_types(tmp_path):
backend = _resume_backend(tmp_path)
before = backend._slot_launch_fingerprint()
backend._effective_cache_types = ("f32", "f16")
assert backend._slot_launch_fingerprint() != before
def test_gguf_file_identity_covers_split_shards(tmp_path):
backend = _resume_backend(tmp_path)
first = tmp_path / "m-00001-of-00002.gguf"
@ -444,6 +472,81 @@ def test_save_skipped_when_estimate_exceeds_cap(monkeypatch, tmp_path):
assert backend.save_slots_for_resume() is None
def test_save_estimate_uses_total_context_and_active_cache_settings(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path, n_slots = 4)
backend._effective_context_length = 8192
backend._kv_cache_context_total = 32768
backend._sliding_window = 4096
backend._swa_full = True
backend._flash_attn_enabled = False
backend._effective_cache_types = ("f32", "f16")
calls = []
def estimate(ctx, cache_type, **kwargs):
calls.append((ctx, cache_type, kwargs))
return 0
backend._estimate_kv_cache_bytes = estimate
_fake_disk(monkeypatch)
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: _Resp(200, {"n_saved": 1, "n_written": 1}),
raising = False,
)
assert backend.save_slots_for_resume() is not None
assert calls == [
(
32768,
"f32",
{
"n_parallel": 4,
"swa_full": True,
"kv_unified": False,
"n_ubatch": 512,
"flash_attn": False,
},
)
]
def test_compact_swa_slot_save_is_skipped(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
backend._sliding_window = 4096
backend._kv_key_length = 256
backend._kv_value_length = 256
backend._swa_full = False
backend._estimate_kv_cache_bytes = lambda *a, **k: (_ for _ in ()).throw(AssertionError)
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
def test_window_without_kv_dims_still_saves(monkeypatch, tmp_path):
# phi3 reports a window but no key/value length, and llama.cpp runs it
# non-SWA, so the compact-SWA skip must not catch it.
backend = _resume_backend(tmp_path)
backend._sliding_window = 262144
backend._kv_key_length = None
backend._kv_value_length = None
backend._swa_full = False
posted = []
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: posted.append(a)
or SimpleNamespace(status_code = 200, json = lambda: {"filename": "slot.bin"}),
raising = False,
)
backend.save_slots_for_resume()
assert posted
def test_save_skipped_when_model_file_changed_since_load(monkeypatch, tmp_path):
# The GGUF/sidecars were swapped on disk after the server loaded them, so the
# live KV belongs to the old weights: refuse to persist it (no POST at all).

5
studio/backend/tests/test_llama_server_args.py
View file

@ -112,6 +112,11 @@ def test_value_with_equals_form_passes_through():
assert validate_extra_args(["--top-k=20"]) == ["--top-k=20"]
def test_managed_long_flag_underscore_alias_is_rejected():
with pytest.raises(ValueError, match = "slot-save-path"):
validate_extra_args(["--slot_save_path", "/tmp/slots"])
def test_non_flag_token_passes_through():
# Bare positionals are passed through; llama-server can reject them.
assert validate_extra_args(["foo"]) == ["foo"]

128
studio/backend/tests/test_mtp_vram_budget.py
View file

@ -76,7 +76,9 @@ from core.inference.llama_cpp import ( # noqa: E402
_extra_args_spec_draft_n_max,
_effective_tensor_parallel,
_env_main_cache_type_for_budget,
_effective_main_cache_types,
_extra_args_main_cache_type_for_budget,
_flash_attn_enabled_from_args,
_kv_bytes_per_elem,
_tensor_parallel_matches_loaded,
)
@ -132,6 +134,7 @@ class _StubDrafter:
def __init__(self, kv_per_token):
self._kv_per_token = kv_per_token
self._architecture = "gemma3"
def _can_estimate_kv(self):
return True
@ -177,6 +180,14 @@ class TestEmbeddedDraftKv:
two = _make_backend(nextn = 2)._mtp_draft_kv_bytes(65536)
assert two == pytest.approx(2 * one)
def test_unaligned_context_follows_runtime_stream_padding(self):
b = _make_backend()
bytes_per_cell = b._mtp_draft_kv_bytes(256) // 256
unified = b._mtp_draft_kv_bytes(5000, n_parallel = 3, kv_unified = True)
separate = b._mtp_draft_kv_bytes(5000, n_parallel = 3, kv_unified = False)
assert unified == 5120 * bytes_per_cell
assert separate == 5376 * bytes_per_cell
def test_embedded_draft_kv_floored_at_f16(self):
# The embedded MTP head is one layer, so llama.cpp's quantized-KV
# overhead is not amortized: a quantized draft KV fits LESS context than
@ -201,6 +212,15 @@ class TestEmbeddedDraftKv:
both_f16 = b._mtp_draft_kv_bytes(131072, draft_cache_type_k = "f16", draft_cache_type_v = "f16")
assert both_q4 == k_only == both_f16 # floored at f16, never under-reserved
def test_flash_attn_off_uses_model_wide_v_width(self):
b = _make_backend(n_layers = 2)
b._n_kv_heads_by_layer = [4, 1]
b._sliding_window_pattern = [False, True]
b._kv_value_length_swa = 2048
ctx = 4096
expected_per_cell = 4 * 256 * 2 + 1 * 2048 * 2
assert b._mtp_draft_kv_bytes(ctx, flash_attn = False) == ctx * expected_per_cell
def test_none_when_dims_missing(self):
assert _make_backend(nextn = 0)._mtp_draft_kv_bytes(65536) is None
assert _make_backend(kv_key_length = None)._mtp_draft_kv_bytes(65536) is None
@ -232,6 +252,30 @@ class TestSeparateDrafter:
c = b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf")
assert c == pytest.approx(4 * a)
def test_gemma4_assistant_shares_target_kv(self, monkeypatch):
b = _make_backend(nextn = None)
stub = _StubDrafter(kv_per_token = 2000)
stub._architecture = "gemma4-assistant"
monkeypatch.setattr(b, "_draft_backend_for", lambda path: stub)
assert (
b._mtp_draft_kv_bytes(
65536,
drafter_path = "/m/mtp-gemma4.gguf",
swa_full = True,
)
== 0
)
assert (
b._estimate_mtp_overhead_bytes(
65536,
drafter_path = "/m/mtp-gemma4.gguf",
draft_weights_bytes = GIB,
swa_full = True,
)
== GIB
)
def test_drafter_kv_scales_with_parallel_slots(self, monkeypatch):
# The drafter is served under the same --parallel slots as the main model,
# so a sliding-window drafter's KV grows per slot; the reserve must thread
@ -398,6 +442,7 @@ class TestExtraArgsMtpDetection:
(["--spec-type", "mtp"], True),
(["--spec-type", "ngram-mod,draft-mtp"], True),
(["--spec-type=draft-mtp"], True),
(["--spec_type=draft-mtp"], True),
(["--spec-type", "ngram-mod"], False),
(["--spec-default"], False),
(["-c", "131072"], False),
@ -579,6 +624,7 @@ class TestExtraArgsMtpDetection:
(["--spec-draft-ngl", "0"], True),
(["-ngld", "0"], True),
(["--spec-draft-ngl=0"], True),
(["--spec_draft_ngl=0"], True),
(["--n-gpu-layers-draft", "0"], True),
(["--spec-draft-ngl", "20"], False),
(["--spec-draft-device", "none"], True),
@ -623,6 +669,7 @@ class TestExtraArgsMtpDetection:
[
(["--spec-draft-n-max", "4"], 4),
(["--spec-draft-n-max=6"], 6),
(["--spec_draft_n_max=6"], 6),
(["--spec-type", "draft-mtp", "--spec-draft-n-max", "3"], 3),
(["--spec-draft-n-max", "2", "--spec-draft-n-max", "5"], 5), # last wins
(["--spec-draft-n-max", "notanint"], None),
@ -644,6 +691,7 @@ class TestExtraArgsMtpDetection:
(["--spec-draft-model", "/m/draft.gguf"], "/m/draft.gguf"),
(["-md", "/m/draft.gguf"], "/m/draft.gguf"),
(["--model-draft=/m/draft.gguf"], "/m/draft.gguf"),
(["--model_draft=/m/draft.gguf"], "/m/draft.gguf"),
(["--model-draft", "--spec-type"], None),
(["-c", "4096"], None),
(None, None),
@ -689,6 +737,7 @@ class TestExtraArgsMtpDetection:
(["--cache-type-v-draft", "q4_0"], (None, "q4_0")), # K stays f16, V only
(["--cache-type-k-draft", "q4_0", "--cache-type-v-draft", "q8_0"], ("q4_0", "q8_0")),
(["--cache-type-k-draft=q8_0"], ("q8_0", None)),
(["--cache_type_k_draft=q8_0"], ("q8_0", None)),
(["--cache-type-k", "q8_0"], (None, None)), # main type, not draft
(["-c", "4096"], (None, None)),
(None, (None, None)),
@ -717,8 +766,17 @@ class TestExtraArgsMtpDetection:
"args,expected",
[
(["--ubatch-size", "1024"], 1024),
(["-ub", "4096"], 4096),
(["-ub", "4096"], 2048),
(["--ubatch-size", "0"], 2048),
(["--batch-size", "256", "--ubatch-size", "0"], 256),
(["--batch-size", "-1"], 512),
(["--ubatch-size", "-1"], 2048),
(["--ubatch-size=512"], 512),
(["--ubatch_size=512"], 512),
(["--batch-size", "256"], 256),
(["--batch_size=256"], 256),
(["-b", "256", "-ub", "1024"], 256),
(["-b", "4096"], 512),
(["--ubatch", "2048"], None), # not a real llama-server flag; ignore it
(["-c", "4096"], None),
(None, None),
@ -727,12 +785,76 @@ class TestExtraArgsMtpDetection:
def test_n_ubatch(self, args, expected):
assert _extra_args_n_ubatch(args, env = {}) == expected
def test_n_ubatch_signed_values_cap_at_context(self):
assert (
_extra_args_n_ubatch(
["--batch-size", "-1", "--ubatch-size", "-1"],
env = {},
n_ctx = 4096,
)
== 4096
)
@pytest.mark.parametrize(
"args,expected",
[
(None, True),
(["--flash-attn", "off"], False),
(["--flash-attn", "disabled"], False),
(["--flash-attn", "false"], False),
(["--flash-attn", "0"], False),
(["--flash-attn=off"], False),
(["--flash-attn=disabled"], False),
(["--flash-attn=false"], False),
(["--flash-attn=0"], False),
(["--flash_attn", "off"], False),
(["-fa", "off", "--flash-attn", "auto"], True),
(["-fa", "off", "--flash-attn", "-1"], True),
(["-fa", "off", "--flash-attn", "enabled"], True),
(["-fa", "off", "--flash-attn=true"], True),
(["-fa", "off", "--flash-attn=1"], True),
(["--flash-attn", "off", "-fa"], True),
],
)
def test_flash_attn_last_value_wins(self, args, expected):
assert _flash_attn_enabled_from_args(args) is expected
def test_effective_main_cache_types_follow_env_then_cli(self):
env = {
"LLAMA_ARG_CACHE_TYPE_K": "f32",
"LLAMA_ARG_CACHE_TYPE_V": "q4_0",
}
assert _effective_main_cache_types([], env) == ("f32", "q4_0")
assert _effective_main_cache_types(["--cache-type-v", "f16"], env) == ("f32", "f16")
def test_n_ubatch_env_fallback(self):
# The child honors LLAMA_ARG_UBATCH; it must reach the compute-buffer reserve.
assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "4096"}) == 4096
# Environment values apply first, then each command-line option overrides
# its own axis before llama.cpp caps ubatch at batch size.
assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "4096"}) == 2048
assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_BATCH": "256"}) == 256
assert (
_extra_args_n_ubatch(
[],
env = {
"LLAMA_ARG_BATCH": "1024",
"LLAMA_ARG_UBATCH": "4096",
},
)
== 1024
)
assert (
_extra_args_n_ubatch(["-ub", "1024"], env = {"LLAMA_ARG_UBATCH": "4096"}) == 1024
) # CLI wins
assert (
_extra_args_n_ubatch(
["-b", "1024"],
env = {
"LLAMA_ARG_BATCH": "256",
"LLAMA_ARG_UBATCH": "4096",
},
)
== 1024
)
assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "notint"}) is None
def test_env_main_cache_type_for_budget(self):

30
studio/backend/tests/test_slot_offload_fit.py
View file

@ -36,6 +36,7 @@ def _backend(
vocab = 248320,
embd = 5120,
kv_fixed_mib = 0,
kv_calls = None,
):
"""Backend with the dims the compute buffer reads; KV mocked to a fixed size so the
only slot-dependent term is the compute buffer (485 MiB/slot f32 output x 1.15)."""
@ -43,7 +44,17 @@ def _backend(
b._vocab_size = vocab
b._embedding_length = embd
b._key_length_mla = None
b._estimate_kv_cache_bytes = lambda ctx, t = None, **k: kv_fixed_mib * MIB
def estimate(
ctx,
t = None,
**kwargs,
):
if kv_calls is not None:
kv_calls.append(kwargs)
return kv_fixed_mib * MIB
b._estimate_kv_cache_bytes = estimate
b._can_estimate_kv = lambda: True
return b
@ -55,6 +66,7 @@ def _run(
gpus,
total_by_idx,
overhead_mib = 0,
swa_full = False,
):
return b._slots_that_fit_on_gpu(
n_parallel,
@ -66,7 +78,8 @@ def _run(
FRAC,
int(overhead_mib * MIB),
1,
512,
n_ubatch = 512,
swa_full = swa_full,
)
@ -113,3 +126,16 @@ class TestSlotsThatFitOnGpu:
# base 19500 (= 22500 total at par-independent terms) the same par3 fit holds.
gi, use_fit, slots = _run(_backend(kv_fixed_mib = 3000), 4, 19500, [(0, 24576)], {0: 24576})
assert use_fit is False and slots == 3
def test_swa_full_is_used_for_every_candidate(self):
calls = []
_run(
_backend(kv_calls = calls),
4,
22500,
[(0, 24576)],
{0: 24576},
swa_full = True,
)
assert calls
assert all(call["swa_full"] is True for call in calls)

7
studio/backend/tests/test_tensor_parallel.py
View file

@ -209,6 +209,13 @@ def test_already_in_target_state_reloads_on_tensor_parallel_change(loaded, reque
assert _target_state(_loaded_backend(loaded), requested) is False
def test_already_in_target_state_reloads_when_swa_full_env_changes(monkeypatch):
backend = _loaded_backend(False)
backend._swa_full = False
monkeypatch.setenv("LLAMA_ARG_SWA_FULL", "1")
assert _target_state(backend, False) is False
def test_already_in_target_state_reconciles_split_mode_extras():
# Tensor engaged via --split-mode in extras (boolean omitted/default False)
# must match a server already running tensor mode -- no spurious reload.

23
studio/backend/tests/test_tp_vision_regression.py
View file

@ -663,6 +663,29 @@ def test_tensor_off_echo_preserves_multi_gpu_fallback():
)
def test_route_dedupe_reloads_when_swa_full_env_changes(monkeypatch):
from models.inference import LoadRequest
inference_routes = _load_inference_routes_module()
backend = _fallback_loaded_backend(layer_preserves_tensor_intent = False)
monkeypatch.setenv("LLAMA_ARG_SWA_FULL", "1")
request = LoadRequest(model_path = "owner/repo")
assert inference_routes._request_matches_loaded_settings(request, backend) is False
def test_route_dedupe_ignores_swa_full_for_diffusion(monkeypatch):
from models.inference import LoadRequest
inference_routes = _load_inference_routes_module()
backend = _fallback_loaded_backend(layer_preserves_tensor_intent = False)
backend._is_diffusion = True
monkeypatch.setenv("LLAMA_ARG_SWA_FULL", "1")
request = LoadRequest(model_path = "owner/repo")
assert inference_routes._request_matches_loaded_settings(request, backend) is True
def test_explicit_split_mode_layer_extras_reloads_after_multi_gpu_fallback():
"""Tensor intent can be dropped via extras too: an explicit --split-mode layer
matches the stored fallback extras but must still reload (reviewer.py P1, #6659)."""

4
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -1507,6 +1507,8 @@ async function autoLoadSmallestModel(): Promise<{
// The safetensors fallback omits both fields and uses HF auto-placement.
gpu_ids?: number[];
gpu_memory_mode?: "auto" | "manual";
cache_type_kv?: string | null;
tensor_parallel?: boolean | null;
}): Promise<boolean> {
const validation = await validateModel({
...payload,
@ -1595,6 +1597,8 @@ async function autoLoadSmallestModel(): Promise<{
max_seq_length: fitMaxSeqLength,
is_lora: false,
gguf_variant: candidate.ggufVariant,
cache_type_kv: config.kvCacheDtype,
tensor_parallel: config.tensorParallel,
// The same remembered-derived GPU pick the load below sends.
...(candidate.kind === "gguf"
? {

2
studio/frontend/src/features/chat/api/chat-api.ts
View file

@ -185,6 +185,8 @@ export async function validateModel(
// /load. Default placement is sized against the selected GPUs.
max_seq_length: payload.max_seq_length,
load_in_4bit: payload.load_in_4bit,
cache_type_kv: payload.cache_type_kv ?? null,
tensor_parallel: payload.tensor_parallel ?? false,
gpu_ids: payload.gpu_ids,
// Manual placement is an explicit override: Auto layers use llama.cpp
// --fit, while a pinned layer count is owned by the user. Tell validate

2
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -817,6 +817,8 @@ export function useChatModelRuntime() {
load_in_4bit: true,
is_lora: isLora,
gguf_variant: ggufVariant ?? null,
cache_type_kv: loadKvCacheDtype,
tensor_parallel: loadTensorParallel,
gpu_ids: validateGpuIds ?? undefined,
...(isGguf ? { gpu_memory_mode: loadGpuMemoryMode } : {}),
});

2
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -1122,6 +1122,8 @@ export function SharedComposer({
gguf_variant: sel.ggufVariant ?? null,
trust_remote_code: loadTrustRemoteCode,
chat_template_override: effectiveChatTemplateOverride,
cache_type_kv: ownConfig.kvCacheDtype ?? null,
tensor_parallel: effectiveTensorParallel,
// Scope the validate to the picked GPUs. GGUF-only, like the load
// below: a non-GGUF target must not inherit a hidden GGUF GPU pick.
...(targetIsGguf