# SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. """Tests for 5-path architecture-aware KV cache VRAM estimation. Covers the GGUF metadata parser, _can_estimate_kv gate, all 5 estimation paths (MLA, Hybrid Mamba, Sliding Window, Standard GQA, Legacy), KV cache quantization, edge cases, and lifecycle (init/unload/reparse). No GPU, network, or libraries beyond pytest. Cross-platform. """ import io import json import struct import sys import types as _types from pathlib import Path import pytest # Stub heavy / unavailable deps before importing the module under test. # Same pattern as test_native_context_length.py. _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) # loggers _loggers_stub = _types.ModuleType("loggers") _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) # structlog _structlog_stub = _types.ModuleType("structlog") sys.modules.setdefault("structlog", _structlog_stub) # httpx -- only stub when the real library is missing. Unconditional stubbing # shadows HTTPError/Response that huggingface_hub.errors imports at load time, # silently breaking the transformers introspection tier. try: import httpx as _httpx_real # noqa: F401 except ImportError: _httpx_stub = _types.ModuleType("httpx") for _exc_name in ( "ConnectError", "TimeoutException", "ReadTimeout", "ReadError", "RemoteProtocolError", "CloseError", "HTTPError", "RequestError", ): setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) class _FakeTimeout: def __init__(self, *a, **kw): pass _httpx_stub.Timeout = _FakeTimeout _httpx_stub.Response = type("Response", (), {}) _httpx_stub.Client = type( "Client", (), { "__init__": lambda self, **kw: None, "__enter__": lambda self: self, "__exit__": lambda self, *a: None, }, ) sys.modules["httpx"] = _httpx_stub from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend # Helpers def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes: """Build a minimal GGUF v3 blob with the given KV metadata. Supports the scalar and simple array metadata the parser uses. """ buf = io.BytesIO() # Header: magic, version, tensor_count, kv_count buf.write(struct.pack(" LlamaCppBackend: """Create a LlamaCppBackend with parsed GGUF metadata from given fields. `general` injects extra `general.*` metadata, to verify the dynamic SWA resolver picks up source-repo hints from GGUFs that ship them. """ kv = {"general.architecture": arch} for k, v in (general or {}).items(): kv[k] = v for k, v in fields.items(): kv[f"{arch}.{k}"] = v import tempfile, os data = _make_gguf_bytes(arch, kv) fd, path = tempfile.mkstemp(suffix = ".gguf") try: os.write(fd, data) os.close(fd) b = LlamaCppBackend() b._read_gguf_metadata(path) return b finally: os.unlink(path) # A. GGUF Parser Tests class TestGGUFParserNewFields: """Architecture-aware fields are parsed correctly.""" @pytest.mark.parametrize( "field,gguf_key,value", [ ("_kv_key_length", "attention.key_length", 128), ("_kv_value_length", "attention.value_length", 128), ("_sliding_window", "attention.sliding_window", 1024), ("_full_attention_interval", "full_attention_interval", 4), ("_kv_lora_rank", "attention.kv_lora_rank", 512), ("_key_length_mla", "attention.key_length_mla", 256), ("_ssm_inner_size", "ssm.inner_size", 6144), ("_ssm_state_size", "ssm.state_size", 128), ], ) def test_field_parsed(self, field, gguf_key, value): b = _backend_from_gguf("testarch", {gguf_key: value}) assert getattr(b, field) == value def test_missing_fields_are_none(self): b = _backend_from_gguf("testarch", {"block_count": 10}) for attr in [ "_kv_key_length", "_kv_value_length", "_sliding_window", "_sliding_window_pattern", "_full_attention_interval", "_kv_lora_rank", "_key_length_mla", "_kv_key_length_swa", "_kv_value_length_swa", "_ssm_inner_size", "_ssm_state_size", ]: assert getattr(b, attr) is None def test_array_fields_parsed(self): b = _backend_from_gguf( "gemma4", { "block_count": 6, "attention.head_count_kv": [8, 8, 8, 8, 8, 2], "attention.sliding_window_pattern": [ True, True, True, True, True, False, ], }, ) # Per-layer KV head count is preserved exactly... assert b._n_kv_heads_by_layer == [8, 8, 8, 8, 8, 2] # ...and mirrored into the scalar field as a conservative max, so # non-SWA paths and callers using `n_kv = self._n_kv_heads or ...` # get a safe upper bound. assert b._n_kv_heads == 8 assert b._sliding_window_pattern == [True, True, True, True, True, False] class TestArchSwaPatternDefaults: """Bootstrap arch table fires when GGUF reports `sliding_window` but no per-layer pattern (true for every Gemma 2/3/3n/gpt-oss GGUF today).""" @pytest.mark.parametrize( "arch,n_layers,expected_period", [ ("gemma2", 26, 2), ("gemma3", 18, 6), ("gemma3n", 35, 5), ("gpt_oss", 24, 2), ("cohere2", 32, 4), ], ) def test_arch_default_pattern_applied(self, arch, n_layers, expected_period): b = _backend_from_gguf( arch, { "block_count": n_layers, "attention.head_count": 4, "attention.head_count_kv": 1, "attention.key_length": 256, "attention.value_length": 256, "attention.sliding_window": 512, }, ) expected_pattern = [(i + 1) % expected_period != 0 for i in range(n_layers)] assert ( b._sliding_window_pattern == expected_pattern ), f"{arch} should expand to period={expected_period}" def test_unknown_arch_no_default(self): b = _backend_from_gguf( "totallymadeupv7", { "block_count": 24, "attention.head_count": 4, "attention.head_count_kv": 1, "attention.key_length": 128, "attention.value_length": 128, "attention.sliding_window": 1024, }, ) assert b._sliding_window_pattern is None def test_explicit_pattern_overrides_arch_default(self): # gemma3 default is period=6; the explicit array must win. b = _backend_from_gguf( "gemma3", { "block_count": 6, "attention.head_count": 4, "attention.head_count_kv": 1, "attention.key_length": 256, "attention.value_length": 256, "attention.sliding_window": 512, "attention.sliding_window_pattern": [ True, False, True, False, True, False, ], }, ) assert b._sliding_window_pattern == [True, False, True, False, True, False] def test_no_sliding_window_no_pattern(self): b = _backend_from_gguf( "gemma3", { "block_count": 18, "attention.head_count": 4, "attention.head_count_kv": 1, "attention.key_length": 256, "attention.value_length": 256, # no sliding_window key }, ) assert b._sliding_window_pattern is None @pytest.mark.parametrize( "arch", ["llama", "qwen2", "qwen3", "mistral", "mistral3", "glm4", "llama4"] ) def test_non_swa_arch_uses_full_attention_path(self, arch): # Pure-GQA arches: no sliding_window, no synthetic pattern, # estimator hits Path 4. b = _backend_from_gguf( arch, { "block_count": 32, "attention.head_count": 32, "attention.head_count_kv": 8, "attention.key_length": 128, "attention.value_length": 128, "embedding_length": 4096, }, ) assert b._sliding_window_pattern is None assert b._sliding_window is None kv = b._estimate_kv_cache_bytes(8192, "f16") gqa_expected = 32 * 8192 * 8 * (128 + 128) * 2 assert kv == gqa_expected def test_arch_default_reduces_kv_estimate_vs_legacy(self): common = { "block_count": 62, "attention.head_count": 32, "attention.head_count_kv": 16, "attention.key_length": 128, "attention.value_length": 128, "attention.sliding_window": 1024, "embedding_length": 5376, } with_default = _backend_from_gguf("gemma3", common) # Arch not in table -> legacy 1/4 path. without_default = _backend_from_gguf("totallymadeupv7", common) kv_default = with_default._estimate_kv_cache_bytes(131072, "f16") kv_legacy = without_default._estimate_kv_cache_bytes(131072, "f16") assert kv_default > 0 assert kv_legacy > 0 assert kv_default < kv_legacy, ( f"arch fallback should under-shoot legacy estimate: " f"{kv_default} >= {kv_legacy}" ) def test_scalar_sliding_window_pattern_expanded(self): block_count = 8 b = _backend_from_gguf( "gemma3", { "attention.sliding_window_pattern": 4, "block_count": block_count, "attention.head_count_kv": 4, "attention.key_length": 256, "attention.value_length": 256, "attention.sliding_window": 1024, }, ) expected = [(i + 1) % 4 != 0 for i in range(block_count)] assert isinstance(b._sliding_window_pattern, list) assert b._sliding_window_pattern == expected assert b._estimate_kv_cache_bytes(4096, "f16") > 0 def test_all_fields_parsed_together(self): fields = { "context_length": 131072, "block_count": 62, "attention.head_count_kv": 16, "attention.head_count": 32, "embedding_length": 5376, "attention.key_length": 128, "attention.value_length": 128, "attention.sliding_window": 1024, "attention.sliding_window_pattern": [True, False], "full_attention_interval": 6, "attention.kv_lora_rank": 512, "attention.key_length_mla": 256, "attention.key_length_swa": 64, "attention.value_length_swa": 64, "ssm.inner_size": 4096, "ssm.state_size": 128, } b = _backend_from_gguf("testarch", fields) assert b._context_length == 131072 assert b._n_layers == 62 assert b._n_kv_heads == 16 assert b._n_heads == 32 assert b._embedding_length == 5376 assert b._kv_key_length == 128 assert b._kv_value_length == 128 assert b._sliding_window == 1024 assert b._sliding_window_pattern == [True, False] assert b._full_attention_interval == 6 assert b._kv_lora_rank == 512 assert b._key_length_mla == 256 assert b._kv_key_length_swa == 64 assert b._kv_value_length_swa == 64 assert b._ssm_inner_size == 4096 assert b._ssm_state_size == 128 _SWA_FIELDS = { "block_count": 12, "attention.head_count": 4, "attention.head_count_kv": 1, "attention.key_length": 256, "attention.value_length": 256, "attention.sliding_window": 512, } class TestDynamicSwaResolver: """4-tier resolver: GGUF metadata, on-disk cache, bootstrap, HF fetch.""" def _isolate_cache(self, monkeypatch, tmp_path): from core.inference import llama_cpp as lc monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) monkeypatch.setattr(lc, "_SWA_CACHE", None) return tmp_path def test_period_from_layer_types_finds_smallest_period(self): from core.inference.llama_cpp import _period_from_layer_types # gemma3 (1 global/6), gpt-oss (alternating), gemma3n (1/5). assert _period_from_layer_types((["sliding_attention"] * 5 + ["full_attention"]) * 4) == 6 assert _period_from_layer_types(["sliding_attention", "full_attention"] * 12) == 2 assert _period_from_layer_types((["sliding_attention"] * 4 + ["full_attention"]) * 7) == 5 def test_period_from_layer_types_returns_none_for_aperiodic(self): from core.inference.llama_cpp import _period_from_layer_types lt = [ "sliding_attention", "full_attention", "sliding_attention", "sliding_attention", "full_attention", "sliding_attention", "sliding_attention", "sliding_attention", ] assert _period_from_layer_types(lt) is None def test_hf_repo_from_url(self): from core.inference.llama_cpp import _hf_repo_from_url assert ( _hf_repo_from_url("https://huggingface.co/google/gemma-3-1b-it") == "google/gemma-3-1b-it" ) assert ( _hf_repo_from_url("https://huggingface.co/google/gemma-3-1b-it/blob/main/config.json") == "google/gemma-3-1b-it" ) for bad in [ "https://huggingface.co/google", "https://example.com/foo/bar", None, "", ]: assert _hf_repo_from_url(bad) is None def test_bootstrap_tier_used_when_no_cache(self, monkeypatch, tmp_path): self._isolate_cache(monkeypatch, tmp_path) from core.inference import llama_cpp as lc def boom(*a, **kw): raise AssertionError("HF fetch must not run when bootstrap covers the arch") monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", boom) b = _backend_from_gguf("gemma3", dict(_SWA_FIELDS, block_count = 18)) assert b._sliding_window_pattern == [(i + 1) % 6 != 0 for i in range(18)] def test_disk_cache_takes_precedence_over_bootstrap(self, monkeypatch, tmp_path): self._isolate_cache(monkeypatch, tmp_path) # Cached period=3 overrides bootstrap=6. with open(tmp_path / "swa_cache.json", "w") as f: json.dump({"gemma3": 3}, f) b = _backend_from_gguf("gemma3", dict(_SWA_FIELDS, block_count = 18)) assert b._sliding_window_pattern == [(i + 1) % 3 != 0 for i in range(18)] def test_disk_cache_supports_array_entries(self, monkeypatch, tmp_path): # Aperiodic mask is tiled across n_layers. self._isolate_cache(monkeypatch, tmp_path) mask = [True, False, True, True, False, True, False, False] with open(tmp_path / "swa_cache.json", "w") as f: json.dump({"customarch": mask}, f) b = _backend_from_gguf("customarch", dict(_SWA_FIELDS, block_count = 16)) assert b._sliding_window_pattern == [bool(mask[i % 8]) for i in range(16)] def test_hf_fetch_populates_cache(self, monkeypatch, tmp_path): self._isolate_cache(monkeypatch, tmp_path) from core.inference import llama_cpp as lc calls = [] def fake_fetch(repo_id): calls.append(repo_id) return 4 if repo_id == "vendor/newmodel-1b-instruct" else None monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", fake_fetch) b = _backend_from_gguf( "newmodel", _SWA_FIELDS, general = {"general.source.huggingface.repository": "vendor/newmodel-1b-instruct"}, ) assert b._sliding_window_pattern == [(i + 1) % 4 != 0 for i in range(12)] assert calls == ["vendor/newmodel-1b-instruct"] with open(tmp_path / "swa_cache.json") as f: assert json.load(f) == {"newmodel": 4} def test_hf_fetch_falls_back_to_other_candidates(self, monkeypatch, tmp_path): self._isolate_cache(monkeypatch, tmp_path) from core.inference import llama_cpp as lc monkeypatch.setattr( lc, "_fetch_swa_entry_from_hf", lambda r: 6 if r == "vendor/newmodel-base" else None, ) b = _backend_from_gguf( "newmodel", _SWA_FIELDS, general = { "general.base_model.0.repo_url": "https://huggingface.co/vendor/newmodel-base" }, ) assert b._sliding_window_pattern == [(i + 1) % 6 != 0 for i in range(12)] def test_offline_env_skips_network(self, monkeypatch, tmp_path): self._isolate_cache(monkeypatch, tmp_path) monkeypatch.setenv("UNSLOTH_STUDIO_OFFLINE", "1") from core.inference import llama_cpp as lc def boom(*a, **kw): raise AssertionError("HF fetch must not run when offline=1") monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", boom) b = _backend_from_gguf( "newmodel", _SWA_FIELDS, general = {"general.source.huggingface.repository": "vendor/newmodel"}, ) assert b._sliding_window_pattern is None def test_hf_fetch_failure_falls_through_silently(self, monkeypatch, tmp_path): self._isolate_cache(monkeypatch, tmp_path) from core.inference import llama_cpp as lc monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", lambda repo_id: None) # Force failure into Tier 3; bypass Tier 2.5. monkeypatch.setattr(lc, "_resolve_swa_entry_from_transformers", lambda arch: None) b = _backend_from_gguf( "newmodel", _SWA_FIELDS, general = {"general.source.huggingface.repository": "vendor/does-not-exist"}, ) assert b._sliding_window_pattern is None assert not (tmp_path / "swa_cache.json").exists() class TestTransformersIntrospection: """Tier 2.5: default-init the matching Config; on failure, parse via inspect.""" def _isolate_cache(self, monkeypatch, tmp_path): from core.inference import llama_cpp as lc monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) monkeypatch.setattr(lc, "_SWA_CACHE", None) return tmp_path def test_arch_aliases_normalises_hyphen_underscore(self): from core.inference.llama_cpp import _arch_aliases aliases = _arch_aliases("falcon-h1") assert aliases[0] == "falcon-h1" and "falcon_h1" in aliases assert _arch_aliases("gemma3") == ("gemma3",) assert _arch_aliases("") == () def test_resolves_real_transformers_arches(self): from core.inference.llama_cpp import _resolve_swa_entry_from_transformers assert _resolve_swa_entry_from_transformers("gemma3") == 6 assert _resolve_swa_entry_from_transformers("gemma2") == 2 assert _resolve_swa_entry_from_transformers("cohere2") == 4 def test_falls_back_to_inspect_when_default_init_raises(self, monkeypatch): from core.inference import llama_cpp as lc class _FakeBrokenConfig: """Class with sliding_window_pattern: int = 7 in its docstring.""" def __init__(self, required_arg): raise TypeError("requires an argument") class _FakeLazyMapping(dict): def __getitem__(self, k): return _FakeBrokenConfig if k == "brokenarch" else super().__getitem__(k) import sys, types as _types fake_auto = _types.ModuleType("transformers.models.auto.configuration_auto") fake_auto.CONFIG_MAPPING_NAMES = {"brokenarch": "FakeBroken"} fake_auto.CONFIG_MAPPING = _FakeLazyMapping({"brokenarch": "FakeBroken"}) monkeypatch.setitem(sys.modules, "transformers.models.auto.configuration_auto", fake_auto) assert lc._resolve_swa_entry_from_transformers("brokenarch") == 7 def test_returns_none_when_transformers_unavailable(self, monkeypatch): from core.inference import llama_cpp as lc import sys orig_import = ( __builtins__["__import__"] if isinstance(__builtins__, dict) else __builtins__.__import__ ) def fake_import(name, *a, **kw): if name.startswith("transformers"): raise ImportError("transformers not installed") return orig_import(name, *a, **kw) monkeypatch.setattr("builtins.__import__", fake_import) for k in list(sys.modules): if k.startswith("transformers"): monkeypatch.delitem(sys.modules, k, raising = False) assert lc._resolve_swa_entry_from_transformers("gemma3") is None def test_returns_none_for_arch_unknown_to_transformers(self): from core.inference.llama_cpp import _resolve_swa_entry_from_transformers assert _resolve_swa_entry_from_transformers("totally-fake-arch-xyz") is None def test_full_resolver_uses_transformers_before_hf_fetch(self, monkeypatch, tmp_path): # Bootstrap empty: Tier 2.5 must answer before Tier 3 fires. self._isolate_cache(monkeypatch, tmp_path) from core.inference import llama_cpp as lc monkeypatch.setattr(lc, "_BOOTSTRAP_SWA_DEFAULTS", {}) def boom(repo_id): raise AssertionError("Tier 3 must not run when Tier 2.5 has the answer") monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", boom) b = _backend_from_gguf( "gemma3", dict(_SWA_FIELDS, block_count = 18), general = {"general.source.huggingface.repository": "google/gemma-3-1b-it"}, ) assert b._sliding_window_pattern == [(i + 1) % 6 != 0 for i in range(18)] with open(tmp_path / "swa_cache.json") as f: assert json.load(f) == {"gemma3": 6} class TestGGUFParserReset: """Fields are reset between parses.""" def test_reset_between_parses(self): # First parse: all fields set b = _backend_from_gguf( "arch1", { "block_count": 32, "attention.key_length": 128, "attention.kv_lora_rank": 512, "attention.head_count_kv": [8, 2], "attention.sliding_window_pattern": [True, False], "attention.key_length_swa": 64, "attention.value_length_swa": 64, "ssm.inner_size": 4096, }, ) assert b._kv_key_length == 128 assert b._kv_lora_rank == 512 assert b._n_kv_heads_by_layer == [8, 2] assert b._sliding_window_pattern == [True, False] assert b._kv_key_length_swa == 64 assert b._kv_value_length_swa == 64 assert b._ssm_inner_size == 4096 # Second parse without those fields -- they must be None kv = {"general.architecture": "arch2", "arch2.block_count": 64} import tempfile, os data = _make_gguf_bytes("arch2", kv) fd, path = tempfile.mkstemp(suffix = ".gguf") os.write(fd, data) os.close(fd) try: b._read_gguf_metadata(path) finally: os.unlink(path) assert b._kv_key_length is None assert b._kv_lora_rank is None assert b._n_kv_heads_by_layer is None assert b._sliding_window_pattern is None assert b._kv_key_length_swa is None assert b._kv_value_length_swa is None assert b._ssm_inner_size is None assert b._n_layers == 64 # B. _can_estimate_kv Gate Tests class TestCanEstimateKV: """Gate logic for all field combinations.""" def test_no_layers_returns_false(self): b = LlamaCppBackend() b._n_layers = None b._kv_key_length = 128 assert not b._can_estimate_kv() def test_explicit_both_dims_sufficient(self): b = LlamaCppBackend() b._n_layers = 32 b._kv_key_length = 128 b._kv_value_length = 128 assert b._can_estimate_kv() def test_key_length_alone_insufficient(self): """key_length without value_length is NOT enough.""" b = LlamaCppBackend() b._n_layers = 32 b._kv_key_length = 128 assert not b._can_estimate_kv() def test_kv_lora_rank_sufficient(self): b = LlamaCppBackend() b._n_layers = 61 b._kv_lora_rank = 512 assert b._can_estimate_kv() def test_legacy_embed_plus_heads(self): b = LlamaCppBackend() b._n_layers = 28 b._embedding_length = 1024 b._n_heads = 16 assert b._can_estimate_kv() def test_legacy_embed_plus_kv_heads(self): b = LlamaCppBackend() b._n_layers = 28 b._embedding_length = 1024 b._n_kv_heads = 8 assert b._can_estimate_kv() def test_legacy_no_embed_returns_false(self): b = LlamaCppBackend() b._n_layers = 28 b._n_heads = 16 # No embedding_length, no new-style fields assert not b._can_estimate_kv() def test_fresh_backend_returns_false(self): b = LlamaCppBackend() assert not b._can_estimate_kv() # C. Path 1: MLA Estimation class TestMLAEstimation: """MLA: K-only cache using compressed KV latent + RoPE.""" def _mla_backend(self, **overrides): defaults = { "_n_layers": 61, "_n_kv_heads": 1, "_n_heads": 128, "_embedding_length": 7168, "_kv_key_length": 576, "_kv_value_length": 512, "_kv_lora_rank": 512, "_key_length_mla": 192, } defaults.update(overrides) b = LlamaCppBackend() for k, v in defaults.items(): setattr(b, k, v) return b def test_deepseek_v3_f16(self): b = self._mla_backend() # 61 layers * 163840 ctx * 1 head * 576 key_len * 2 bpe expected = 61 * 163840 * 1 * 576 * 2 assert b._estimate_kv_cache_bytes(163840, "f16") == expected def test_mla_ignores_value_length(self): """MLA must NOT add value_length -- V is reconstructed from the latent.""" 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 assert result == expected def test_mla_fallback_when_no_key_length(self): """No key_length: fall back to kv_lora_rank + key_length_mla.""" 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 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 assert result == expected def test_mla_defaults_n_kv_to_1_when_heads_absent(self): """MLA uses n_kv=1 even if n_kv_heads is None (not n_heads).""" 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 assert result == expected def test_mla_q4_quantization(self): b = self._mla_backend() result_f16 = b._estimate_kv_cache_bytes(1000, "f16") 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) # D. Path 2: Hybrid Mamba Estimation class TestHybridMambaEstimation: """Hybrid Mamba: only attention layers (1 in N) need KV cache.""" def _hybrid_backend(self, **overrides): defaults = { "_n_layers": 64, "_n_kv_heads": 4, "_n_heads": 24, "_embedding_length": 5120, "_kv_key_length": 256, "_kv_value_length": 256, "_full_attention_interval": 4, "_ssm_inner_size": 6144, "_ssm_state_size": 128, } defaults.update(overrides) b = LlamaCppBackend() for k, v in defaults.items(): setattr(b, k, v) return b def test_qwen35_27b(self): b = self._hybrid_backend() # n_attn = 64 // 4 = 16 expected = 16 * 262144 * 4 * (256 + 256) * 2 assert b._estimate_kv_cache_bytes(262144, "f16") == expected def test_qwen35_35b_a3b(self): b = self._hybrid_backend( _n_layers = 40, _n_kv_heads = 2, _n_heads = 16, _embedding_length = 2048, _ssm_inner_size = 4096, ) # n_attn = 40 // 4 = 10 expected = 10 * 262144 * 2 * (256 + 256) * 2 assert b._estimate_kv_cache_bytes(262144, "f16") == expected def test_hybrid_without_explicit_dims(self): """Fall back to head_dim when key_length/value_length are missing.""" b = self._hybrid_backend(_kv_key_length = None, _kv_value_length = None) head_dim = 5120 // 24 # 213 expected = 16 * 4096 * 4 * 2 * head_dim * 2 assert b._estimate_kv_cache_bytes(4096, "f16") == expected def test_fai_zero_safety(self): """full_attention_interval=0 must not ZeroDivisionError.""" b = self._hybrid_backend(_full_attention_interval = 0) result = b._estimate_kv_cache_bytes(4096, "f16") # fai=0 -> n_attn = n_layers (all layers) expected = 64 * 4096 * 4 * (256 + 256) * 2 assert result == expected # E. Path 3: Sliding Window Estimation class TestSlidingWindowEstimation: """SWA: half global (full ctx) + half sliding window.""" def _swa_backend(self, **overrides): defaults = { "_n_layers": 62, "_n_kv_heads": 16, "_n_heads": 32, "_embedding_length": 5376, "_kv_key_length": 128, "_kv_value_length": 128, "_sliding_window": 1024, } defaults.update(overrides) b = LlamaCppBackend() for k, v in defaults.items(): setattr(b, k, v) return b def test_gemma3(self): b = self._swa_backend() # 1/4 heuristic: 62 // 4 = 15 global, 47 SWA 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) assert b._estimate_kv_cache_bytes(131072, "f16") == expected def test_gpt_oss(self): b = self._swa_backend( _n_layers = 24, _n_kv_heads = 8, _n_heads = 64, _embedding_length = 2880, _kv_key_length = 64, _kv_value_length = 64, _sliding_window = 128, ) # 1/4 heuristic: 24 // 4 = 6 global, 18 SWA 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) assert b._estimate_kv_cache_bytes(131072, "f16") == expected def test_gemma4_per_layer_swa_metadata(self): b = self._swa_backend( _n_layers = 30, _n_kv_heads = None, _n_kv_heads_by_layer = [8, 8, 8, 8, 8, 2] * 5, _n_heads = 16, _embedding_length = 2816, _kv_key_length = 512, _kv_value_length = 512, _sliding_window = 1024, _sliding_window_pattern = [True, True, True, True, True, False] * 5, _kv_key_length_swa = 256, _kv_value_length_swa = 256, ) full_layers = 5 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 return int(full + sliding) for ctx in (4096, 46500, 262144): assert b._estimate_kv_cache_bytes(ctx, "f16") == expected(ctx) def test_ctx_smaller_than_window(self): """When ctx < 2 * sliding_window, SWA cache caps at ctx.""" 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) assert b._estimate_kv_cache_bytes(ctx, "f16") == expected def test_odd_layer_count(self): b = self._swa_backend(_n_layers = 63) 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) assert b._estimate_kv_cache_bytes(1000, "f16") == expected # F. Path 4: Standard GQA Estimation class TestStandardGQAEstimation: """Standard GQA with explicit key_length/value_length.""" def _gqa_backend(self, **overrides): defaults = { "_n_layers": 28, "_n_kv_heads": 8, "_n_heads": 16, "_embedding_length": 1024, "_kv_key_length": 128, "_kv_value_length": 128, } defaults.update(overrides) b = LlamaCppBackend() for k, v in defaults.items(): setattr(b, k, v) return b def test_qwen3_06b(self): b = self._gqa_backend() expected = 28 * 40960 * 8 * (128 + 128) * 2 assert b._estimate_kv_cache_bytes(40960, "f16") == expected def test_asymmetric_kv_dims(self): """key_length != value_length (some architectures have this).""" b = self._gqa_backend(_kv_key_length = 192, _kv_value_length = 64) expected = 28 * 4096 * 8 * (192 + 64) * 2 assert b._estimate_kv_cache_bytes(4096, "f16") == expected def test_differs_from_legacy(self): """GQA path differs from legacy when key_length != embed//n_heads.""" b = self._gqa_backend() head_dim = 1024 // 16 # 64 gqa_result = b._estimate_kv_cache_bytes(4096, "f16") # Legacy: 2 * 8 * 64 * 28 * 4096 * 2 legacy_result = int(2 * 8 * head_dim * 28 * 4096 * 2) # GQA: 28 * 4096 * 8 * (128+128) * 2 -- uses actual key_length=128 assert gqa_result != legacy_result assert gqa_result > legacy_result # key_length (128) > head_dim (64) # G. Path 5: Legacy Fallback Estimation class TestLegacyEstimation: """Legacy: embed // n_heads, for old GGUFs without new fields.""" def _legacy_backend(self, **overrides): defaults = { "_n_layers": 32, "_n_kv_heads": 8, "_n_heads": 32, "_embedding_length": 4096, } defaults.update(overrides) b = LlamaCppBackend() for k, v in defaults.items(): setattr(b, k, v) return b def test_basic_legacy(self): b = self._legacy_backend() head_dim = 4096 // 32 # 128 expected = int(2 * 8 * 128 * 32 * 4096 * 2) assert b._estimate_kv_cache_bytes(4096, "f16") == expected def test_legacy_with_only_n_heads(self): """n_kv_heads is None, falls back to n_heads.""" b = self._legacy_backend(_n_kv_heads = None) head_dim = 4096 // 32 expected = int(2 * 32 * head_dim * 32 * 4096 * 2) assert b._estimate_kv_cache_bytes(4096, "f16") == expected def test_legacy_identical_to_old_formula(self): """Legacy path matches the pre-PR formula.""" b = self._legacy_backend() n_layers = 32 n_kv_heads = 8 head_dim = 4096 // 32 n_ctx = 8192 bpe = 2.0 old_formula = int(2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe) assert b._estimate_kv_cache_bytes(n_ctx, "f16") == old_formula # H. Path Priority (selection order) class TestPathPriority: """Confirm: MLA > Hybrid Mamba > SWA > GQA > Legacy.""" def test_mla_takes_priority_over_all(self): """If kv_lora_rank is set, MLA path wins even with other fields present.""" b = LlamaCppBackend() b._n_layers = 61 b._n_kv_heads = 1 b._n_heads = 128 b._embedding_length = 7168 b._kv_key_length = 576 b._kv_value_length = 512 b._kv_lora_rank = 512 b._ssm_inner_size = 4096 # Would trigger Hybrid 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) assert b._estimate_kv_cache_bytes(1000, "f16") == expected_mla def test_hybrid_over_swa(self): """Hybrid takes priority over SWA when both fields present.""" b = LlamaCppBackend() b._n_layers = 64 b._n_kv_heads = 4 b._n_heads = 24 b._embedding_length = 5120 b._kv_key_length = 256 b._kv_value_length = 256 b._ssm_inner_size = 6144 b._full_attention_interval = 4 b._sliding_window = 1024 # Would trigger SWA n_attn = 64 // 4 expected_hybrid = int(n_attn * 1000 * 4 * (256 + 256) * 2) assert b._estimate_kv_cache_bytes(1000, "f16") == expected_hybrid def test_all_paths_produce_different_values(self): """With chosen params, each path yields a distinct value.""" # embedding_length=768 so legacy head_dim (768//16=48) != key_length # (256), and MLA key_len (256) != legacy K+V (2*48=96). params = { "_n_layers": 40, "_n_kv_heads": 4, "_n_heads": 16, "_embedding_length": 768, "_kv_key_length": 256, "_kv_value_length": 256, } ctx = 4096 # Path 4: Standard GQA b_gqa = LlamaCppBackend() for k, v in params.items(): setattr(b_gqa, k, v) gqa_val = b_gqa._estimate_kv_cache_bytes(ctx, "f16") # Path 1: MLA b_mla = LlamaCppBackend() for k, v in params.items(): setattr(b_mla, k, v) b_mla._kv_lora_rank = 512 mla_val = b_mla._estimate_kv_cache_bytes(ctx, "f16") # Path 2: Hybrid Mamba b_hybrid = LlamaCppBackend() for k, v in params.items(): setattr(b_hybrid, k, v) b_hybrid._ssm_inner_size = 4096 b_hybrid._full_attention_interval = 4 hybrid_val = b_hybrid._estimate_kv_cache_bytes(ctx, "f16") # Path 3: SWA b_swa = LlamaCppBackend() for k, v in params.items(): setattr(b_swa, k, v) b_swa._sliding_window = 512 swa_val = b_swa._estimate_kv_cache_bytes(ctx, "f16") # Path 5: Legacy (no key_length/value_length) b_legacy = LlamaCppBackend() b_legacy._n_layers = 40 b_legacy._n_kv_heads = 4 b_legacy._n_heads = 16 b_legacy._embedding_length = 768 legacy_val = b_legacy._estimate_kv_cache_bytes(ctx, "f16") values = [mla_val, hybrid_val, swa_val, gqa_val, legacy_val] assert len(set(values)) == 5, f"Expected 5 distinct values, got {values}" # I. KV Cache Quantization class TestQuantization: """All supported cache_type_kv values scale correctly.""" @pytest.mark.parametrize( "cache_type,expected_bpe", [ ("f32", 4.0), ("f16", 2.0), ("bf16", 2.0), ("q8_0", 34 / 32), ("q5_1", 0.75), ("q5_0", 0.6875), ("q4_1", 0.625), ("q4_0", 0.5625), ("iq4_nl", 0.5625), (None, 2.0), # default is f16 ("unknown", 2.0), # unknown falls back to f16 ], ) def test_quantization_scaling(self, cache_type, expected_bpe): b = LlamaCppBackend() b._n_layers = 10 b._n_kv_heads = 1 b._n_heads = 8 b._embedding_length = 512 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) assert result == expected # J. Edge Cases class TestEdgeCases: """Boundary conditions and degenerate inputs.""" def test_zero_context(self): b = LlamaCppBackend() b._n_layers = 32 b._kv_key_length = 128 assert b._estimate_kv_cache_bytes(0, "f16") == 0 def test_negative_context(self): b = LlamaCppBackend() b._n_layers = 32 b._kv_key_length = 128 assert b._estimate_kv_cache_bytes(-1, "f16") == 0 def test_context_of_one(self): b = LlamaCppBackend() b._n_layers = 10 b._n_kv_heads = 1 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) def test_very_large_context(self): """1M context should not overflow or crash.""" b = LlamaCppBackend() b._n_layers = 10 b._n_kv_heads = 1 b._kv_key_length = 128 b._kv_value_length = 128 result = b._estimate_kv_cache_bytes(1_000_000, "f16") assert result > 0 assert isinstance(result, int) def test_n_kv_heads_none_falls_to_n_heads(self): b = LlamaCppBackend() b._n_layers = 10 b._n_kv_heads = None b._n_heads = 8 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) assert result == expected def test_both_heads_none_falls_to_one(self): b = LlamaCppBackend() b._n_layers = 10 b._n_kv_heads = None b._n_heads = None 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) assert result == expected # J2. Server-flag knobs (--swa-full, --kv-unified/--parallel, # --ctx-checkpoints, --kv-offload) class TestServerFlags: """Estimator should mirror llama-server CLI flags that change KV size.""" def _swa_backend(self, **overrides): defaults = { "_n_layers": 26, "_n_kv_heads": 4, "_n_heads": 8, "_embedding_length": 1152, "_kv_key_length": 256, "_kv_value_length": 256, "_sliding_window": 512, "_sliding_window_pattern": [True, True, True, True, True, False] * 4 + [True, True], } defaults.update(overrides) b = LlamaCppBackend() for k, v in defaults.items(): setattr(b, k, v) return b def _gqa_backend(self, **overrides): defaults = { "_n_layers": 28, "_n_kv_heads": 8, "_n_heads": 16, "_embedding_length": 1024, "_kv_key_length": 128, "_kv_value_length": 128, } defaults.update(overrides) b = LlamaCppBackend() for k, v in defaults.items(): setattr(b, k, v) return b # ── --swa-full ────────────────────────────────────────────────── def test_swa_full_collapses_pattern_path_to_full_ctx(self): b = self._swa_backend() ctx = 32_768 flagged = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True) # swa_full: every layer caches n_ctx -- equals path 4 sizing. kv_per_token = 4 * (256 + 256) * 2 # n_kv_heads * (k+v) * f16 expected = 26 * ctx * kv_per_token assert flagged == expected assert flagged > b._estimate_kv_cache_bytes(ctx, "f16") def test_swa_full_collapses_legacy_path_to_full_ctx(self): # No per-layer pattern -> 1/4-global heuristic; swa_full overrides. b = self._swa_backend(_sliding_window_pattern = None) ctx = 16_384 flagged = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True) n_global = max(1, 26 // 4) n_swa = 26 - n_global kv_per = 4 * (256 + 256) * 2 # swa_cells == n_ctx when swa_full=True expected = n_global * ctx * kv_per + n_swa * ctx * kv_per assert flagged == expected def test_swa_full_no_op_for_non_swa_model(self): b = self._gqa_backend() baseline = b._estimate_kv_cache_bytes(8192, "f16") flagged = b._estimate_kv_cache_bytes(8192, "f16", swa_full = True) assert flagged == baseline def test_swa_full_suppresses_checkpoint_term(self): b = self._swa_backend() with_cp = b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 8) with_cp_full = b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 8, swa_full = True) no_cp_full = b._estimate_kv_cache_bytes(8192, "f16", swa_full = True) # Checkpoints only matter when SWA layers don't already keep n_ctx. assert with_cp_full == no_cp_full assert with_cp > b._estimate_kv_cache_bytes(8192, "f16") # ── --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). def test_gqa_kv_constant_across_parallel(self): b = self._gqa_backend() baseline = b._estimate_kv_cache_bytes(4096, "f16") for slots in (1, 2, 4, 8): for unified in (True, False): assert ( b._estimate_kv_cache_bytes(4096, "f16", n_parallel = slots, kv_unified = unified) == baseline ) def test_zero_parallel_floors_at_one(self): b = self._gqa_backend() baseline = b._estimate_kv_cache_bytes(4096, "f16") for unified in (True, False): assert ( b._estimate_kv_cache_bytes(4096, "f16", n_parallel = 0, kv_unified = unified) == baseline ) def test_swa_path_scales_only_swa_portion(self): b = self._swa_backend() ctx = 8192 baseline = b._estimate_kv_cache_bytes(ctx, "f16") # Decompose baseline by walking the estimator's own loop. 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 global_bytes = sum( ctx * 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 ) # Sanity: parallel=1 reproduces baseline exactly assert global_bytes + swa_bytes_per_slot == baseline # Only the SWA portion scales by parallel 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 ) assert scaled == global_bytes + slots * swa_bps def test_mla_kv_constant_across_parallel(self): b = LlamaCppBackend() b._n_layers = 60 b._n_kv_heads = 1 b._kv_lora_rank = 512 b._key_length_mla = 64 b._kv_key_length = 576 baseline = b._estimate_kv_cache_bytes(8192, "f16") for slots in (1, 2, 4, 8): for unified in (True, False): assert ( b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = unified) == baseline ) # ── --ctx-checkpoints ────────────────────────────────────────── def test_ctx_checkpoints_zero_is_no_op(self): b = self._swa_backend() baseline = b._estimate_kv_cache_bytes(8192, "f16") assert b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 0) == baseline def test_ctx_checkpoints_no_op_for_non_swa(self): b = self._gqa_backend() baseline = b._estimate_kv_cache_bytes(8192, "f16") assert b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 32) == baseline def test_ctx_checkpoints_pattern_path_adds_known_bytes(self): b = self._swa_backend() ctx = 8192 baseline = b._estimate_kv_cache_bytes(ctx, "f16") flagged = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4) # 22 SWA layers * 4 cps * 512 cells * 4 heads * (256+256) * 2 bytes n_swa_layers = sum(1 for f in [True, True, True, True, True, False] * 4 + [True, True] if f) per_layer = 4 * 512 * 4 * (256 + 256) * 2 assert flagged == baseline + n_swa_layers * per_layer def test_ctx_checkpoints_legacy_path_adds_known_bytes(self): b = self._swa_backend(_sliding_window_pattern = None) ctx = 8192 baseline = b._estimate_kv_cache_bytes(ctx, "f16") flagged = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4) n_global = max(1, 26 // 4) n_swa = 26 - n_global kv_per = 4 * (256 + 256) * 2 extra = 4 * n_swa * 512 * kv_per # ctx_checkpoints * n_swa * sliding * kv_per assert flagged == baseline + extra def test_ctx_checkpoints_compose_with_n_parallel(self): # Only the SWA + checkpoint portion scales by n_parallel; the # global-layer portion is constant. b = self._swa_backend() 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 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) # ── --kv-offload (kv_on_gpu) ─────────────────────────────────── def test_fit_returns_requested_when_kv_off_gpu(self): b = self._gqa_backend() # Tiny VRAM budget -- would normally force a reduction. fitted = b._fit_context_to_vram( requested_ctx = 32_768, available_mib = 1, model_size_bytes = 100, cache_type_kv = "f16", kv_on_gpu = False, ) assert fitted == 32_768 def test_fit_reduces_when_kv_on_gpu(self): b = self._gqa_backend() fitted = b._fit_context_to_vram( requested_ctx = 32_768, available_mib = 64, model_size_bytes = 1024 * 1024, # 1 MiB cache_type_kv = "f16", kv_on_gpu = True, ) assert fitted < 32_768 def test_fit_mtp_engaged_returns_smaller_or_equal_context(self): # Flat MTP fallback budget is _CTX_FIT_VRAM_FRACTION - 0.05; non-MTP is # the full fraction. On a tight budget MTP must yield <= non-MTP. b = self._gqa_backend() common = dict( requested_ctx = 32_768, available_mib = 128, model_size_bytes = 8 * 1024 * 1024, cache_type_kv = "f16", ) baseline = b._fit_context_to_vram(**common) mtp = b._fit_context_to_vram(**common, mtp_engaged = True) assert mtp <= baseline def test_fit_mtp_engaged_unchanged_when_kv_off_gpu(self): # kv_on_gpu=False short-circuits the fit; mtp_engaged irrelevant. b = self._gqa_backend() fitted = b._fit_context_to_vram( requested_ctx = 32_768, available_mib = 1, model_size_bytes = 100, cache_type_kv = "f16", kv_on_gpu = False, mtp_engaged = True, ) assert fitted == 32_768 def test_fit_threads_swa_full_through_estimator(self): # SWA model, generous budget; both should fit but cache size differs. b = self._swa_backend() ctx = 8192 kv_default = b._estimate_kv_cache_bytes(ctx, "f16") kv_full = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True) assert kv_full > kv_default # Budget = model + kv_default (rounded up) -- swa_full must not fit. budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / _CTX_FIT_VRAM_FRACTION + 1 fitted_default = b._fit_context_to_vram( requested_ctx = ctx, available_mib = int(budget_mib), model_size_bytes = 1024 * 1024, cache_type_kv = "f16", ) fitted_full = b._fit_context_to_vram( requested_ctx = ctx, available_mib = int(budget_mib), model_size_bytes = 1024 * 1024, cache_type_kv = "f16", swa_full = True, ) assert fitted_default == ctx assert fitted_full < ctx # 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). 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. """ def _gqa_backend(self, **overrides): defaults = { "_n_layers": 28, "_n_kv_heads": 8, "_n_heads": 16, "_embedding_length": 1024, "_kv_key_length": 128, "_kv_value_length": 128, } defaults.update(overrides) b = LlamaCppBackend() for k, v in defaults.items(): setattr(b, k, v) return b def _swa_backend(self, **overrides): defaults = { "_n_layers": 18, "_n_kv_heads": 1, "_n_heads": 4, "_embedding_length": 1024, "_kv_key_length": 256, "_kv_value_length": 256, "_sliding_window": 512, # 15 SWA + 3 global, mirrors gemma-3-270m "_sliding_window_pattern": [t == "swa" for t in (["swa"] * 5 + ["global"]) * 3], } defaults.update(overrides) b = LlamaCppBackend() for k, v in defaults.items(): setattr(b, k, v) return b # ── non-SWA paths: constant ──────────────────────────────────── def test_pure_gqa_constant_across_parallel(self): b = self._gqa_backend() baseline = b._estimate_kv_cache_bytes(8192, "f16") for slots in (1, 2, 4, 8): for unified in (True, False): assert ( b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = unified) == baseline ) def test_mla_constant_across_parallel(self): b = LlamaCppBackend() b._n_layers = 60 b._n_kv_heads = 1 b._kv_lora_rank = 512 b._key_length_mla = 64 b._kv_key_length = 576 baseline = b._estimate_kv_cache_bytes(8192, "f16") for slots in (1, 2, 4, 8): assert b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots) == baseline def test_hybrid_constant_across_parallel(self): b = LlamaCppBackend() b._n_layers = 64 b._n_kv_heads = 16 b._n_heads = 32 b._embedding_length = 4096 b._kv_key_length = 128 b._kv_value_length = 128 b._ssm_inner_size = 4096 b._full_attention_interval = 4 baseline = b._estimate_kv_cache_bytes(8192, "f16") for slots in (1, 2, 4, 8): assert b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots) == baseline def test_legacy_constant_across_parallel(self): b = LlamaCppBackend() b._n_layers = 32 b._n_kv_heads = 8 b._n_heads = 8 b._embedding_length = 4096 baseline = b._estimate_kv_cache_bytes(8192, "f16") 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_swa_pattern_scales_only_swa_portion(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): got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = unified) assert got == global_bytes + slots * swa_bps def test_swa_fallback_scales_only_swa_portion(self): # No per-layer pattern -> 1/4-global heuristic. b = self._swa_backend(_sliding_window_pattern = None) ctx = 8192 swa = b._sliding_window n_layers = 18 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 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. 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 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. b = self._swa_backend() ctx = 8192 baseline = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True) for slots in (1, 2, 4, 8): assert ( b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True, n_parallel = slots) == baseline ) # ── kv_unified: no-op for memory math ────────────────────────── 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" # ── Empirical Gemma-3 270m formula ───────────────────────────── def test_matches_empirical_gemma3_270m_formula(self): """Exact match against the 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, sliding=512, f16. """ b = LlamaCppBackend() b._n_layers = 18 b._n_kv_heads = 1 b._n_heads = 4 b._embedding_length = 1024 b._kv_key_length = 256 b._kv_value_length = 256 b._sliding_window = 512 # Mirrors the bootstrap-resolved gemma3 pattern (period 6) on an # 18-layer model: 15 SWA, 3 global. b._sliding_window_pattern = [(i + 1) % 6 != 0 for i in range(18)] n_global = 3 n_swa = 15 # 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_mib = got_bytes / (1024 * 1024) assert ( got_mib == expected_mib ), f"slots={slots}: got {got_mib} MiB, expected {expected_mib} MiB" # J3. shared_kv_layers (Gemma 3n / Gemma 4) class TestSharedKVLayers: """``.attention.shared_kv_layers`` reduces the layer count that allocates KV. The trailing ``shared_kv_layers`` blocks reuse earlier caches (Gemma 3n: 35 layers, 15 shared -> 20 allocate; Gemma 4 same field). Unset on every other arch -> no behavioural change.""" def _gemma3n_backend(self, **overrides): # Mirrors google/gemma-3n-E4B-it: 35 layers, 15 shared, SWA window # 1024, period 5 (4 sliding + 1 full repeating). defaults = { "_n_layers": 35, "_n_kv_heads": 4, "_n_heads": 8, "_embedding_length": 2048, "_kv_key_length": 256, "_kv_value_length": 256, "_sliding_window": 1024, "_sliding_window_pattern": [ t == "sliding_attention" for t in (["sliding_attention"] * 4 + ["full_attention"]) * 7 ], "_shared_kv_layers": 15, } defaults.update(overrides) b = LlamaCppBackend() for k, v in defaults.items(): setattr(b, k, v) return b def _gqa_backend(self, **overrides): defaults = { "_n_layers": 28, "_n_kv_heads": 8, "_n_heads": 16, "_embedding_length": 1024, "_kv_key_length": 128, "_kv_value_length": 128, } defaults.update(overrides) b = LlamaCppBackend() for k, v in defaults.items(): setattr(b, k, v) return b def test_field_initialises_to_none(self): b = LlamaCppBackend() assert b._shared_kv_layers is None def test_unset_field_is_noop(self): b = self._gqa_backend() baseline = b._estimate_kv_cache_bytes(8192, "f16") b._shared_kv_layers = None assert b._estimate_kv_cache_bytes(8192, "f16") == baseline b._shared_kv_layers = 0 assert b._estimate_kv_cache_bytes(8192, "f16") == baseline def test_path4_drops_shared_layers(self): b = self._gqa_backend(_shared_kv_layers = 4) ctx = 4096 kv_per = 8 * (128 + 128) * 2 # 28 - 4 = 24 layers actually allocate assert b._estimate_kv_cache_bytes(ctx, "f16") == 24 * ctx * kv_per def test_path5_drops_shared_layers(self): b = LlamaCppBackend() b._n_layers = 32 b._n_kv_heads = 8 b._n_heads = 8 b._embedding_length = 4096 b._shared_kv_layers = 8 ctx = 4096 head_dim = 4096 // 8 # 512 # 32 - 8 = 24 layers expected = 2 * 8 * head_dim * 24 * ctx * 2 assert b._estimate_kv_cache_bytes(ctx, "f16") == expected def test_path1_mla_drops_shared_layers(self): b = LlamaCppBackend() b._n_layers = 60 b._n_kv_heads = 1 b._kv_lora_rank = 512 b._key_length_mla = 64 b._kv_key_length = 576 b._shared_kv_layers = 10 ctx = 8192 # 60 - 10 = 50 assert b._estimate_kv_cache_bytes(ctx, "f16") == 50 * ctx * 1 * 576 * 2 def test_path3_pattern_loops_only_unshared_layers(self): b = self._gemma3n_backend() ctx = 8192 # First 20 layers contribute; layers 20..34 skipped. Pattern # [s,s,s,s,F] repeated -> in layers 0..19: sliding 16, full 4. sliding_in_unshared = sum(b._sliding_window_pattern[:20]) full_in_unshared = 20 - sliding_in_unshared 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 assert b._estimate_kv_cache_bytes(ctx, "f16") == expected def test_shared_layers_reduces_estimate(self): b = self._gemma3n_backend() with_shared = b._estimate_kv_cache_bytes(8192, "f16") b._shared_kv_layers = 0 without_shared = b._estimate_kv_cache_bytes(8192, "f16") # 20/35 = 0.571 of the work; ~43% reduction. ratio = with_shared / without_shared assert 0.5 < ratio < 0.65 def test_path3_pattern_with_swa_full_and_shared(self): b = self._gemma3n_backend() ctx = 8192 flagged = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True) # Every unshared layer caches n_ctx -> path-4-style sizing over # only the 20 unshared layers. kv_per = 4 * (256 + 256) * 2 assert flagged == 20 * ctx * kv_per def test_path3_fallback_uses_unshared_count(self): # No per-layer pattern -> 1/4-global heuristic over n_layers_kv, # not n_layers. b = self._gemma3n_backend(_sliding_window_pattern = None) ctx = 8192 n_layers_kv = 35 - 15 # 20 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 assert b._estimate_kv_cache_bytes(ctx, "f16") == expected def test_shared_floors_at_one_layer(self): # Pathological: shared >= n_layers must not zero out the cache. b = self._gqa_backend(_shared_kv_layers = 99) ctx = 4096 kv_per = 8 * (128 + 128) * 2 assert b._estimate_kv_cache_bytes(ctx, "f16") == 1 * ctx * kv_per def test_composes_with_n_parallel(self): # Only the SWA portion of unshared layers scales by n_parallel; # the global portion is constant. b = self._gemma3n_backend() ctx = 8192 swa = b._sliding_window per_token = 4 * (256 + 256) * 2 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 flagged = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False) assert flagged == global_bytes + slots * swa_bytes_per_slot def test_composes_with_ctx_checkpoints(self): b = self._gemma3n_backend() ctx = 8192 baseline = b._estimate_kv_cache_bytes(ctx, "f16") with_cp = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4) # Checkpoints count only over UNSHARED SWA layers (16 of them). sliding_in_unshared = sum(b._sliding_window_pattern[:20]) per_cp_layer = 4 * 1024 * 4 * (256 + 256) * 2 # cps * swa * heads * (k+v) * bpe assert with_cp == baseline + sliding_in_unshared * per_cp_layer def test_unload_resets_shared_kv_layers(self): b = LlamaCppBackend() b._shared_kv_layers = 12 b.unload_model() assert b._shared_kv_layers is None # K. Lifecycle Tests class TestLifecycle: """Init, unload, and reparse field management.""" def test_init_fields_none(self): b = LlamaCppBackend() for attr in [ "_kv_key_length", "_kv_value_length", "_sliding_window", "_sliding_window_pattern", "_full_attention_interval", "_kv_lora_rank", "_key_length_mla", "_kv_key_length_swa", "_kv_value_length_swa", "_ssm_inner_size", "_ssm_state_size", "_shared_kv_layers", ]: assert getattr(b, attr) is None assert b._n_kv_heads_by_layer is None def test_unload_resets_fields(self): b = LlamaCppBackend() b._n_layers = 32 b._kv_key_length = 128 b._kv_lora_rank = 512 b._sliding_window = 1024 b._sliding_window_pattern = [True, False] b._n_kv_heads_by_layer = [8, 2] b._kv_key_length_swa = 64 b._kv_value_length_swa = 64 b._ssm_inner_size = 4096 b._full_attention_interval = 4 b._shared_kv_layers = 8 b.unload_model() for attr in [ "_kv_key_length", "_kv_value_length", "_sliding_window", "_sliding_window_pattern", "_full_attention_interval", "_kv_lora_rank", "_key_length_mla", "_kv_key_length_swa", "_kv_value_length_swa", "_ssm_inner_size", "_ssm_state_size", "_shared_kv_layers", ]: assert getattr(b, attr) is None assert b._n_kv_heads_by_layer is None def test_end_to_end_synthetic_mla(self): """Round-trip: write GGUF -> parse -> estimate.""" b = _backend_from_gguf( "deepseek2", { "context_length": 163840, "block_count": 61, "attention.head_count_kv": 1, "attention.head_count": 128, "embedding_length": 7168, "attention.key_length": 576, "attention.value_length": 512, "attention.kv_lora_rank": 512, "attention.key_length_mla": 192, }, ) assert b._can_estimate_kv() result = b._estimate_kv_cache_bytes(163840, "f16") expected = 61 * 163840 * 1 * 576 * 2 assert result == expected def test_end_to_end_synthetic_hybrid(self): b = _backend_from_gguf( "qwen35", { "context_length": 262144, "block_count": 64, "attention.head_count_kv": 4, "attention.head_count": 24, "embedding_length": 5120, "attention.key_length": 256, "attention.value_length": 256, "full_attention_interval": 4, "ssm.inner_size": 6144, "ssm.state_size": 128, }, ) assert b._can_estimate_kv() result = b._estimate_kv_cache_bytes(262144, "f16") n_attn = 64 // 4 expected = n_attn * 262144 * 4 * (256 + 256) * 2 assert result == expected def test_end_to_end_synthetic_swa(self): b = _backend_from_gguf( "gemma3", { "context_length": 131072, "block_count": 62, "attention.head_count_kv": 16, "attention.head_count": 32, "embedding_length": 5376, "attention.key_length": 128, "attention.value_length": 128, "attention.sliding_window": 1024, }, ) 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. period = 6 kv_per = 16 * 256 * 2 expected = 0 for i in range(62): is_swa = (i + 1) % period != 0 layer_ctx = min(131072, 2 * 1024) if is_swa else 131072 expected += layer_ctx * kv_per assert result == expected def test_end_to_end_synthetic_shared_kv_round_trip(self): # Mirrors gemma3n_text: 35 layers, 15 shared, sliding_window=1024. b = _backend_from_gguf( "gemma3n_text", { "context_length": 32768, "block_count": 35, "attention.head_count_kv": 4, "attention.head_count": 8, "embedding_length": 2048, "attention.key_length": 256, "attention.value_length": 256, "attention.sliding_window": 1024, "attention.shared_kv_layers": 15, }, ) assert b._can_estimate_kv() assert b._shared_kv_layers == 15 # Bootstrap for gemma3n_text -> period 5; resolver synthesises a # 35-entry bool array. Only the first 20 (n_layers - shared) # allocate KV. result = b._estimate_kv_cache_bytes(8192, "f16") assert result > 0 # Sanity: shared back to 0 -> strictly larger estimate (more # layers allocate). b._shared_kv_layers = 0 unshared = b._estimate_kv_cache_bytes(8192, "f16") assert unshared > result def test_end_to_end_synthetic_gqa(self): b = _backend_from_gguf( "qwen3", { "context_length": 40960, "block_count": 28, "attention.head_count_kv": 8, "attention.head_count": 16, "embedding_length": 1024, "attention.key_length": 128, "attention.value_length": 128, }, ) assert b._can_estimate_kv() result = b._estimate_kv_cache_bytes(40960, "f16") expected = 28 * 40960 * 8 * 256 * 2 assert result == expected def test_end_to_end_synthetic_legacy(self): b = _backend_from_gguf( "llama", { "context_length": 4096, "block_count": 32, "attention.head_count_kv": 8, "attention.head_count": 32, "embedding_length": 4096, }, ) assert b._can_estimate_kv() result = b._estimate_kv_cache_bytes(4096, "f16") head_dim = 4096 // 32 expected = int(2 * 8 * head_dim * 32 * 4096 * 2) assert result == expected