diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index bddd71c301..f768764c22 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -11,6 +11,7 @@ through its OpenAI-compatible /v1/chat/completions endpoint. import atexit import contextlib import json +import os import re import struct import structlog @@ -76,6 +77,238 @@ _SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$") _SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$") +# ── Sliding-window-pattern resolver ─────────────────────────── +# Resolves the per-layer SWA mask when a GGUF reports a sliding window +# but no `sliding_window_pattern` field. Tier order in +# `_resolve_swa_pattern`: GGUF metadata, on-disk cache, bootstrap dict +# below, transformers introspection, HF Hub config.json, legacy 1/4 +# fallback. Period N means layer i is SWA iff `(i + 1) % N != 0`, +# matching transformers. Skipped on purpose: phi3 (no key/val length +# in GGUF, window >= ctx anyway), qwen2 family (converter strips +# sliding_window when use_sliding_window=False), mistral v0.1/v0.2 +# (all-SWA can't be expressed as a period). +_BOOTSTRAP_SWA_DEFAULTS: dict[str, int] = { + "gemma2": 2, # Gemma2Config.sliding_window_pattern + "gemma3": 6, # Gemma3TextConfig.sliding_window_pattern + "gemma3n": 5, # text_config.layer_types: SWA*4 + FULL + "gpt_oss": 2, # text_config.layer_types: alternating + "cohere2": 4, # Cohere2Config.sliding_window_pattern +} + +# Process-wide cache backed by JSON on disk. Values are int period or +# list[bool] mask. Lazy-loaded. +_SWA_CACHE: Optional[dict] = None +_SWA_CACHE_LOCK = threading.Lock() + + +def _swa_cache_path() -> Path: + home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") + base = Path(home) if home else Path.home() / ".unsloth" / "studio" + return base / "swa_cache.json" + + +def _load_swa_cache() -> dict: + global _SWA_CACHE + with _SWA_CACHE_LOCK: + if _SWA_CACHE is not None: + return _SWA_CACHE + try: + with open(_swa_cache_path()) as f: + _SWA_CACHE = json.load(f) + if not isinstance(_SWA_CACHE, dict): + _SWA_CACHE = {} + except (FileNotFoundError, json.JSONDecodeError, OSError): + _SWA_CACHE = {} + return _SWA_CACHE + + +def _save_swa_cache(cache: dict) -> None: + try: + path = _swa_cache_path() + path.parent.mkdir(parents = True, exist_ok = True) + tmp = path.with_suffix(".json.tmp") + with open(tmp, "w") as f: + json.dump(cache, f, indent = 2, sort_keys = True) + tmp.replace(path) + except OSError: + pass + + +def _period_from_layer_types(layer_types: list) -> Optional[int]: + """Smallest period N where `(i+1) % N != 0` matches the SWA mask, + or None if no fixed period fits.""" + if not layer_types: + return None + is_swa = ["full" not in str(t).lower() for t in layer_types] + n = len(is_swa) + for N in range(1, n + 1): + if all(((i + 1) % N != 0) == is_swa[i] for i in range(n)): + return N + return None + + +def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]: + try: + from huggingface_hub import hf_hub_download + + cfg_path = hf_hub_download(repo_id, "config.json", repo_type = "model") + with open(cfg_path) as f: + cfg = json.load(f) + except Exception: + return None + + src = cfg.get("text_config") if isinstance(cfg.get("text_config"), dict) else cfg + period = src.get("sliding_window_pattern") + if isinstance(period, int) and period > 0: + return period + lt = src.get("layer_types") + if isinstance(lt, list) and lt: + return _period_from_layer_types(lt) or [ + "full" not in str(t).lower() for t in lt + ] + return None + + +def _arch_aliases(arch: str) -> tuple: + # GGUF emits `falcon-h1`; HF model_type is `falcon_h1`. Normalise both ways. + seen = [] + for a in (arch, arch.replace("-", "_"), arch.replace("_", "-")): + if a and a not in seen: + seen.append(a) + return tuple(seen) + + +def _swa_entry_from_config_obj(cfg) -> Optional[object]: + src = getattr(cfg, "text_config", None) or cfg + period = getattr(src, "sliding_window_pattern", None) + if isinstance(period, int) and period > 0: + return period + lt = getattr(src, "layer_types", None) + if isinstance(lt, list) and lt: + return _period_from_layer_types(lt) or [ + "full" not in str(t).lower() for t in lt + ] + return None + + +_SWA_PATTERN_SOURCE_RE = re.compile( + r"sliding_window_pattern\s*(?::\s*[\w\[\], ]*)?\s*=\s*(\d+)" +) + + +def _resolve_swa_entry_from_transformers(arch: str) -> Optional[object]: + """Default-instantiate the matching Config; on failure, regex-parse + its source for `sliding_window_pattern = N`.""" + try: + from transformers.models.auto.configuration_auto import ( + CONFIG_MAPPING, + CONFIG_MAPPING_NAMES, + ) + except Exception: + return None + + cfg_class = None + for alias in _arch_aliases(arch): + if alias in CONFIG_MAPPING_NAMES: + try: + cfg_class = CONFIG_MAPPING[alias] + break + except Exception: + cfg_class = None + if cfg_class is None: + return None + + try: + if (entry := _swa_entry_from_config_obj(cfg_class())) is not None: + return entry + except Exception: + pass + + import inspect + + candidates = [cfg_class] + text_cfg_class = getattr(cfg_class, "sub_configs", {}).get("text_config") + if text_cfg_class is not None: + candidates.append(text_cfg_class) + for cls in candidates: + try: + src = inspect.getsource(cls) + except (OSError, TypeError): + continue + if m := _SWA_PATTERN_SOURCE_RE.search(src): + period = int(m.group(1)) + if period > 0: + return period + return None + + +def _resolve_swa_pattern( + arch: Optional[str], + n_layers: Optional[int], + source_repo_candidates: tuple = (), + *, + allow_network: Optional[bool] = None, +) -> Optional[list]: + if not arch or not n_layers: + return None + if allow_network is None: + allow_network = os.environ.get("UNSLOTH_STUDIO_OFFLINE", "0") not in ( + "1", + "true", + "True", + "yes", + ) + + cache = _load_swa_cache() + + def _entry_to_mask(entry): + if isinstance(entry, int) and entry > 0: + return [(i + 1) % entry != 0 for i in range(n_layers)] + if isinstance(entry, list) and entry: + return [bool(entry[i % len(entry)]) for i in range(n_layers)] + return None + + def _persist(entry): + with _SWA_CACHE_LOCK: + cache[arch] = entry + _save_swa_cache(cache) + + if (entry := cache.get(arch)) is not None: + if (mask := _entry_to_mask(entry)) is not None: + return mask + + if (entry := _BOOTSTRAP_SWA_DEFAULTS.get(arch)) is not None: + return _entry_to_mask(entry) + + entry = _resolve_swa_entry_from_transformers(arch) + if entry is not None: + _persist(entry) + return _entry_to_mask(entry) + + # Tier 3: live HF fetch (with persistent caching of the result) + if allow_network: + for repo_id in source_repo_candidates: + if not repo_id: + continue + entry = _fetch_swa_entry_from_hf(repo_id) + if entry is not None: + _persist(entry) + return _entry_to_mask(entry) + + return None + + +def _hf_repo_from_url(url: Optional[str]) -> Optional[str]: + """Strip `https://huggingface.co/owner/name(/...)` to `owner/name`.""" + if not url or "huggingface.co/" not in url: + return None + tail = url.split("huggingface.co/", 1)[1].rstrip("/") + parts = tail.split("/") + if len(parts) < 2: + return None + return f"{parts[0]}/{parts[1]}" + + # Model size extraction — lazy import to avoid pulling in transformers # at module level. See PR description for the full explanation. def _extract_model_size_b(model_id: str): @@ -215,17 +448,24 @@ class LlamaCppBackend: # KV-cache estimation fields (populated by _read_gguf_metadata) self._n_layers: Optional[int] = None self._n_kv_heads: Optional[int] = None + self._n_kv_heads_by_layer: Optional[list[int]] = None self._n_heads: Optional[int] = None self._embedding_length: Optional[int] = None - # Architecture-aware KV fields (8 new fields for 5-path estimation) + # Architecture-aware KV fields for 5-path estimation self._kv_key_length: Optional[int] = None self._kv_value_length: Optional[int] = None self._sliding_window: Optional[int] = None + self._sliding_window_pattern: Optional[list[bool]] = None self._full_attention_interval: Optional[int] = None self._kv_lora_rank: Optional[int] = None self._key_length_mla: Optional[int] = None + self._kv_key_length_swa: Optional[int] = None + self._kv_value_length_swa: Optional[int] = None self._ssm_inner_size: Optional[int] = None self._ssm_state_size: Optional[int] = None + # Last N layers reuse KV from earlier layers and don't allocate + # their own cache (Gemma 3n / Gemma 4: .attention.shared_kv_layers). + self._shared_kv_layers: Optional[int] = None self._lock = threading.Lock() self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None @@ -744,13 +984,29 @@ class LlamaCppBackend: # New-style: need both explicit key AND value dimensions if self._kv_key_length is not None and self._kv_value_length is not None: return True - # Legacy: need embedding_length + head count + # Legacy: need embedding_length + a head count (scalar or per-layer). return self._embedding_length is not None and ( - self._n_kv_heads is not None or self._n_heads is not None + self._n_kv_heads is not None + or self._n_heads is not None + or self._n_kv_heads_by_layer is not None ) + def _kv_heads_for_layer(self, layer_idx: int, fallback: int) -> int: + if self._n_kv_heads_by_layer is not None and layer_idx < len( + self._n_kv_heads_by_layer + ): + return self._n_kv_heads_by_layer[layer_idx] + return fallback + def _estimate_kv_cache_bytes( - self, n_ctx: int, cache_type_kv: Optional[str] = None + self, + n_ctx: int, + cache_type_kv: Optional[str] = None, + *, + swa_full: bool = False, + n_parallel: int = 1, + kv_unified: bool = True, + ctx_checkpoints: int = 0, ) -> int: """Estimate KV cache VRAM for a given context length. @@ -761,12 +1017,34 @@ class LlamaCppBackend: 4. GQA -- standard full KV with explicit key/value dimensions 5. Legacy -- fallback using embed // n_heads + Server-flag knobs (mirror llama-server's CLI): + swa_full -- ``--swa-full``: force SWA layers to cache the + full ``n_ctx`` (collapses path 3 to path 4 + sizing for the SWA layers). + n_parallel -- ``--parallel``: number of server slots. + Verified empirically against llama-server: + non-SWA layers stay constant (cells split + across slots), SWA layers scale linearly + (per-slot window). + kv_unified -- ``--kv-unified`` (default on): retained for + API forward-compat. Currently a no-op for + memory math because the unified buffer total + matches per-slot buffers in measured cases. + ctx_checkpoints -- ``--ctx-checkpoints``: SWA snapshot count per + slot (PR #15293). Each snapshot stores one + sliding-window of state per SWA layer. + Returns 0 if metadata is insufficient for estimation. """ if not self._can_estimate_kv() or n_ctx <= 0: return 0 n_layers = self._n_layers # type: ignore[assignment] + # Gemma 3n / Gemma 4 reuse KV from earlier layers in the last + # ``shared_kv_layers`` blocks -- those don't allocate their own + # cache. Floor at 1 so a misconfigured GGUF can't zero out KV. + shared = self._shared_kv_layers or 0 + n_layers_kv = max(1, n_layers - shared) n_kv = self._n_kv_heads or self._n_heads or 1 # type: ignore[assignment] # Bytes per element depends on KV cache quantization @@ -782,6 +1060,8 @@ class LlamaCppBackend: "iq4_nl": 0.5625, }.get(cache_type_kv or "f16", 2.0) + slots = max(1, n_parallel) + # Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5) # MLA stores one compressed KV latent per token/layer (shared across heads). # V is reconstructed from the latent on the fly -- no separate V cache. @@ -792,7 +1072,7 @@ class LlamaCppBackend: n_kv_mla = self._n_kv_heads or 1 rope_dim = self._key_length_mla or 64 key_len = self._kv_key_length or (self._kv_lora_rank + rope_dim) - return int(n_layers * n_ctx * n_kv_mla * key_len * bpe) + return int(n_layers_kv * n_ctx * n_kv_mla * key_len * bpe) key_len = self._kv_key_length val_len = self._kv_value_length @@ -810,11 +1090,19 @@ class LlamaCppBackend: head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe) - # Path 3: Sliding Window (Gemma-3, gpt-oss) - # SWA layers only cache min(ctx, window) tokens; global layers cache full ctx. - # Most SWA architectures use few global layers (e.g., Gemma-3 uses 1 in 6). - # Without an explicit field, we conservatively assume 1/4 of layers are global - # which is still far more accurate than the legacy formula (which ignores SWA). + # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). + # Pattern is filled in by the resolver at parse time; if absent, + # falls through to the legacy 1/4-global heuristic below. + # Per-layer-type ``--parallel N`` accounting (verified empirically + # against ``llama-server``): + # * non-SWA layers: total cells = n_ctx, partitioned across + # slots -> total memory CONSTANT in slots. + # * SWA layers: per-slot cells = 2 * sliding_window + # (capped at n_ctx and at per_slot_ctx + # when ctx is split among many slots) -> + # total memory grows LINEARLY in slots. + # ``--swa-full`` forces full n_ctx for SWA layers instead. + # ``--ctx-checkpoints N`` adds N snapshots per SWA layer per slot. if ( self._sliding_window is not None and self._sliding_window > 0 @@ -822,20 +1110,72 @@ class LlamaCppBackend: and val_len is not None ): swa = self._sliding_window - n_global = max(1, n_layers // 4) - n_swa = n_layers - n_global + per_slot_ctx = max(1, n_ctx // slots) + # ``--swa-full`` makes SWA layers cache the full context just + # like non-SWA: cells get partitioned across slots, so per-slot + # cells = per_slot_ctx and the slots*per-slot product collapses + # back to the constant ``n_ctx`` total. Otherwise SWA caches + # 2*sliding_window per slot, clamped at the per-slot ctx. + swa_cells_per_slot = ( + per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx) + ) + key_len_swa = self._kv_key_length_swa or key_len + val_len_swa = self._kv_value_length_swa or val_len + if self._sliding_window_pattern is not None: + global_bytes = 0.0 # constant across slots + swa_bytes_per_slot = 0.0 # multiplied by slots + checkpoint_extra_per_slot = 0.0 + # Iterate only over layers that allocate their own KV; + # the trailing ``shared`` layers reuse earlier caches. + for layer_idx in range(n_layers_kv): + layer_n_kv = self._kv_heads_for_layer(layer_idx, n_kv) + is_swa = ( + layer_idx < len(self._sliding_window_pattern) + and self._sliding_window_pattern[layer_idx] + ) + if is_swa: + swa_bytes_per_slot += ( + swa_cells_per_slot + * layer_n_kv + * (key_len_swa + val_len_swa) + * bpe + ) + if ctx_checkpoints > 0 and not swa_full: + checkpoint_extra_per_slot += ( + ctx_checkpoints + * swa + * layer_n_kv + * (key_len_swa + val_len_swa) + * bpe + ) + else: + global_bytes += n_ctx * layer_n_kv * (key_len + val_len) * bpe + return int( + global_bytes + + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot) + ) + n_global = max(1, n_layers_kv // 4) + n_swa = n_layers_kv - n_global kv_per_token = n_kv * (key_len + val_len) * bpe + kv_per_token_swa = n_kv * (key_len_swa + val_len_swa) * bpe + global_bytes = n_global * n_ctx * kv_per_token + swa_bytes_per_slot = n_swa * swa_cells_per_slot * kv_per_token_swa + checkpoint_extra_per_slot = ( + ctx_checkpoints * n_swa * swa * kv_per_token_swa + if ctx_checkpoints > 0 and not swa_full + else 0.0 + ) return int( - n_global * n_ctx * kv_per_token + n_swa * min(n_ctx, swa) * kv_per_token + global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot) ) # Path 4: Standard GQA with explicit key/value dimensions if key_len is not None and val_len is not None: - return int(n_layers * n_ctx * n_kv * (key_len + val_len) * bpe) + return int(n_layers_kv * n_ctx * n_kv * (key_len + val_len) * bpe) # Path 5: Legacy fallback (old GGUFs without explicit dimensions) head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] - return int(2 * n_kv * head_dim * n_layers * n_ctx * bpe) + return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe) def _fit_context_to_vram( self, @@ -844,6 +1184,12 @@ class LlamaCppBackend: model_size_bytes: int, cache_type_kv: Optional[str] = None, min_ctx: int = 4096, + *, + swa_full: bool = False, + n_parallel: int = 1, + kv_unified: bool = True, + ctx_checkpoints: int = 0, + kv_on_gpu: bool = True, ) -> int: """Return the largest context length that fits in GPU VRAM. @@ -851,6 +1197,11 @@ class LlamaCppBackend: threshold -- 10% reserved for compute buffers, CUDA context, scratch space, flash-attn workspace, etc.). If the model weights alone don't fit, returns min_ctx unchanged. + + ``kv_on_gpu`` mirrors ``--kv-offload`` (default on). When False + the KV cache lives in CPU RAM and doesn't compete with weights + for VRAM; the requested context is honored verbatim. The other + keyword args mirror ``_estimate_kv_cache_bytes``. """ if not self._can_estimate_kv(): logger.debug( @@ -860,11 +1211,22 @@ class LlamaCppBackend: ) return requested_ctx + # KV lives off-GPU: no VRAM accounting needed for the cache itself. + if not kv_on_gpu: + return requested_ctx + + kv_kwargs = dict( + swa_full = swa_full, + n_parallel = n_parallel, + kv_unified = kv_unified, + ctx_checkpoints = ctx_checkpoints, + ) + budget_bytes = available_mib * 1024 * 1024 * 0.90 model_footprint = model_size_bytes # Check if requested context already fits - kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv) + kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs) if model_footprint + kv <= budget_bytes: return requested_ctx @@ -886,7 +1248,7 @@ class LlamaCppBackend: best = effective_min while lo <= hi: mid = (lo + hi) // 2 - kv = self._estimate_kv_cache_bytes(mid, cache_type_kv) + kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs) if kv <= remaining: best = mid lo = mid + 1 @@ -1019,6 +1381,19 @@ class LlamaCppBackend: for _ in range(alen): LlamaCppBackend._gguf_skip_value(f, atype) + @staticmethod + def _gguf_read_array_value(f, atype: int, alen: int) -> Optional[list]: + if atype == 4: # UINT32 + return [struct.unpack(" None: """Read context_length, architecture params, and chat_template from a GGUF header. @@ -1037,23 +1412,44 @@ class LlamaCppBackend: self._supports_tools = False self._n_layers = None self._n_kv_heads = None + self._n_kv_heads_by_layer = None self._n_heads = None self._embedding_length = None self._kv_key_length = None self._kv_value_length = None self._sliding_window = None + self._sliding_window_pattern = None self._full_attention_interval = None self._kv_lora_rank = None self._key_length_mla = None + self._kv_key_length_swa = None + self._kv_value_length_swa = None self._ssm_inner_size = None self._ssm_state_size = None + self._shared_kv_layers = None try: - WANTED = {"general.architecture", "tokenizer.chat_template"} + WANTED = { + "general.architecture", + "tokenizer.chat_template", + # Source-repo hints for the SWA resolver's HF fallback. + "general.source.huggingface.repository", + "general.source.url", + "general.source.repo_url", + "general.base_model.0.repo_url", + "general.base_model.0.organization", + "general.base_model.0.name", + "general.basename", + "general.organization", + "general.size_label", + "general.finetune", + } # Additional arch-specific keys are added dynamically once # we know the architecture name. arch_keys: dict[str, str] = {} # gguf_key -> attribute name arch = None + sliding_window_pattern_period: Optional[int] = None + general: dict[str, str] = {} with open(gguf_path, "rb") as f: magic = struct.unpack(" {effective_ctx} " f"(model: {model_size / (1024**3):.1f} GB, " @@ -1586,7 +2082,7 @@ class LlamaCppBackend: ) kv_cache_bytes = self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv + effective_ctx, cache_type_kv, n_parallel = n_parallel ) logger.info( f"GGUF size: {model_size / (1024**3):.1f} GB, " @@ -2015,16 +2511,21 @@ class LlamaCppBackend: self._speculative_type = None self._n_layers = None self._n_kv_heads = None + self._n_kv_heads_by_layer = None self._n_heads = None self._embedding_length = None self._kv_key_length = None self._kv_value_length = None self._sliding_window = None + self._sliding_window_pattern = None self._full_attention_interval = None self._kv_lora_rank = None self._key_length_mla = None + self._kv_key_length_swa = None + self._kv_value_length_swa = None self._ssm_inner_size = None self._ssm_state_size = None + self._shared_kv_layers = None # Clean up temp chat template file if hasattr(self, "_chat_template_file") and self._chat_template_file: try: diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py index 2640ded90d..29d87804ff 100644 --- a/studio/backend/tests/test_kv_cache_estimation.py +++ b/studio/backend/tests/test_kv_cache_estimation.py @@ -12,6 +12,7 @@ Cross-platform: Linux, macOS, Windows, WSL. """ import io +import json import struct import sys import types as _types @@ -37,35 +38,43 @@ sys.modules.setdefault("loggers", _loggers_stub) _structlog_stub = _types.ModuleType("structlog") sys.modules.setdefault("structlog", _structlog_stub) -# httpx -_httpx_stub = _types.ModuleType("httpx") -for _exc_name in ( - "ConnectError", - "TimeoutException", - "ReadTimeout", - "ReadError", - "RemoteProtocolError", - "CloseError", -): - setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) +# httpx -- only stub when the real library isn't installed. Stubbing +# unconditionally would shadow ``HTTPError`` / ``Response`` etc. that +# ``huggingface_hub.errors`` imports at module load time, which causes +# the transformers introspection tier to silently return None inside +# the test process. +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 -class _FakeTimeout: - def __init__(self, *a, **kw): - pass - - -_httpx_stub.Timeout = _FakeTimeout -_httpx_stub.Client = type( - "Client", - (), - { - "__init__": lambda self, **kw: None, - "__enter__": lambda self: self, - "__exit__": lambda self, *a: None, - }, -) -sys.modules.setdefault("httpx", _httpx_stub) + _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 LlamaCppBackend @@ -77,8 +86,7 @@ from core.inference.llama_cpp import LlamaCppBackend def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes: """Build a minimal GGUF v3 binary blob with the given KV metadata. - Only supports UINT32 (type 4), UINT64 (type 10), and STRING (type 8) - values, which is all the metadata parser reads. + Supports the scalar and simple array metadata used by the parser. """ buf = io.BytesIO() # Header: magic, version, tensor_count, kv_count @@ -96,6 +104,17 @@ def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes: val_bytes = val.encode("utf-8") buf.write(struct.pack(" bytes: return buf.getvalue() -def _backend_from_gguf(arch: str, fields: dict) -> LlamaCppBackend: - """Create a LlamaCppBackend with parsed GGUF metadata from given fields.""" +def _backend_from_gguf( + arch: str, fields: dict, general: dict | None = None +) -> LlamaCppBackend: + """Create a LlamaCppBackend with parsed GGUF metadata from given fields. + + `general` lets a test inject extra `general.*` metadata (used 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 @@ -133,7 +161,7 @@ def _backend_from_gguf(arch: str, fields: dict) -> LlamaCppBackend: class TestGGUFParserNewFields: - """Verify that the 8 new architecture-aware fields are correctly parsed.""" + """Verify that architecture-aware fields are correctly parsed.""" @pytest.mark.parametrize( "field,gguf_key,value", @@ -158,15 +186,189 @@ class TestGGUFParserNewFields: "_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_all_13_fields_parsed_together(self): + 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 estimator paths and any caller 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): + # Period=6 is the gemma3 default; 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: GGUF has 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 the 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, @@ -176,9 +378,12 @@ class TestGGUFParserNewFields: "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, } @@ -191,13 +396,294 @@ class TestGGUFParserNewFields: 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 per 6), gpt-oss (alternating), gemma3n (1 per 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) + # Override bootstrap=6 with a cached period=3. + 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 gets 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 the failure into the Tier 3 path; 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 + ): + # With 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: """Verify that fields are properly reset between parses.""" @@ -209,11 +695,19 @@ class TestGGUFParserReset: "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 should be None @@ -230,6 +724,10 @@ class TestGGUFParserReset: 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 @@ -455,7 +953,9 @@ class TestSlidingWindowEstimation: n_global = max(1, 62 // 4) # 15 n_swa = 62 - n_global # 47 kv_per = 16 * (128 + 128) * 2 - expected = int(n_global * 131072 * kv_per + n_swa * min(131072, 1024) * kv_per) + # 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): @@ -472,27 +972,52 @@ class TestSlidingWindowEstimation: n_global = max(1, 24 // 4) # 6 n_swa = 24 - n_global # 18 kv_per = 8 * (64 + 64) * 2 - expected = int(n_global * 131072 * kv_per + n_swa * min(131072, 128) * kv_per) + 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 context < sliding_window, SWA layers use full context anyway.""" + """When context < 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, 8192) * kv_per) - # min(4096, 8192) = 4096, so both pools use full ctx + 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): - """Odd layer count: n_global = max(1, n//4), n_swa = n - n_global.""" 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, 1024) * kv_per) + expected = int(n_global * 1000 * kv_per + n_swa * min(1000, 2 * 1024) * kv_per) assert b._estimate_kv_cache_bytes(1000, "f16") == expected @@ -785,6 +1310,686 @@ class TestEdgeCases: 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) + # With 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 ────────────────────────────────── + # Empirically verified against llama-server: non-SWA caches partition + # n_ctx across slots (total memory constant); SWA layers are the only + # portion that scales with --parallel. --kv-unified is currently 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 same loop the estimator does. + 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 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 get clamped 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 checkpoints * 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 stays 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 -- normally would 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_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 should not fit. + budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / 0.90 + 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: + """Verifies the per-layer-type scaling rule against 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 yield 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 should 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; result is the all-global + # GQA-style total, which is 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): + # Both unified=True and unified=False must produce the same + # total bytes for every backend type and every 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 + # 5-period [swa,swa,swa,swa,full] * 3 + [swa,swa,swa]: mirrors the + # bootstrap-resolved pattern for gemma3 (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 + actually 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 are 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; expect ~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; equals 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 should 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 the unshared layers scales by n_parallel; + # the global portion stays 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 only count 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 # --------------------------------------------------------------------------- @@ -799,13 +2004,18 @@ class TestLifecycle: "_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() @@ -813,20 +2023,30 @@ class TestLifecycle: 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): """Full round-trip: write GGUF -> parse -> estimate.""" @@ -887,12 +2107,46 @@ class TestLifecycle: ) assert b._can_estimate_kv() result = b._estimate_kv_cache_bytes(131072, "f16") - n_global = max(1, 62 // 4) # 15 - n_swa = 62 - n_global # 47 + # gemma3 -> period 6 from the bootstrap table, SWA cache + # double-buffered to 2 * sliding_window cells. + period = 6 kv_per = 16 * 256 * 2 - expected = int(n_global * 131072 * kv_per + n_swa * 1024 * kv_per) + 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 table for gemma3n_text -> period 5; the resolver + # synthesises a 35-entry bool array. The first 20 entries + # (n_layers - shared) are the only ones that allocate KV. + result = b._estimate_kv_cache_bytes(8192, "f16") + assert result > 0 + # Sanity: setting shared back to 0 must produce a 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", diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py index f498655347..caa6397901 100644 --- a/studio/backend/tests/test_llama_cpp_context_fit.py +++ b/studio/backend/tests/test_llama_cpp_context_fit.py @@ -114,9 +114,13 @@ def _make_backend( inst._kv_value_length = kv_value_length inst._kv_lora_rank = None inst._sliding_window = None + inst._sliding_window_pattern = None inst._ssm_inner_size = None inst._full_attention_interval = None inst._key_length_mla = None + inst._n_kv_heads_by_layer = None + inst._kv_key_length_swa = None + inst._kv_value_length_swa = None return inst @@ -137,7 +141,7 @@ def _drive( model_size = int(model_gib * GIB) cache_type_kv = None - def fake_estimate(n_ctx_, _type = None): + def fake_estimate(n_ctx_, _type = None, **_kwargs): return 0 if n_ctx_ <= 0 else n_ctx_ * kv_per_token_bytes inst._estimate_kv_cache_bytes = fake_estimate diff --git a/studio/backend/tests/test_llama_cpp_max_context_threshold.py b/studio/backend/tests/test_llama_cpp_max_context_threshold.py index 5fd0243c9f..22e4cda7d1 100644 --- a/studio/backend/tests/test_llama_cpp_max_context_threshold.py +++ b/studio/backend/tests/test_llama_cpp_max_context_threshold.py @@ -99,9 +99,13 @@ def _make_backend(native_ctx = 131072): inst._kv_value_length = 128 inst._kv_lora_rank = None inst._sliding_window = None + inst._sliding_window_pattern = None inst._ssm_inner_size = None inst._full_attention_interval = None inst._key_length_mla = None + inst._n_kv_heads_by_layer = None + inst._kv_key_length_swa = None + inst._kv_value_length_swa = None return inst @@ -114,7 +118,7 @@ def _compute_max_available_ctx(native_ctx, model_gib, gpus, kv_per_token_bytes = model_size = int(model_gib * GIB) inst._estimate_kv_cache_bytes = ( - lambda n, _t = None: 0 if n <= 0 else n * kv_per_token_bytes + lambda n, _t = None, **_kw: 0 if n <= 0 else n * kv_per_token_bytes ) inst._can_estimate_kv = lambda: True