From bac2cfac7cab5629cee376dec92e5a1f3ca6503e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Apr 2026 12:29:20 +0000 Subject: [PATCH 1/7] feat(studio): architecture-aware KV cache VRAM estimation Replace the single legacy formula (2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe) with 5-path estimation that reads 8 additional GGUF metadata fields: 1. MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5) -- K-only cache using compressed KV latent + RoPE; no separate V allocation 2. Hybrid Mamba (Qwen3.5-27B, Qwen3.5-35B-A3B) -- only attention layers (1 in N) carry KV; Mamba layers have none 3. Sliding Window (Gemma-3, gpt-oss) -- SWA layers cache min(ctx, window) tokens instead of the full context 4. Standard GQA -- uses explicit key_length/value_length from GGUF instead of embed // n_heads (which is wrong for many models) 5. Legacy fallback -- identical to old formula for old GGUFs New GGUF fields parsed: attention.key_length, attention.value_length, attention.sliding_window, full_attention_interval, attention.kv_lora_rank, attention.key_length_mla, ssm.inner_size, ssm.state_size. Validated against 9 real GGUF files (72/72 field checks pass). The legacy formula was off by +682% for Gemma-3 and -81% for DeepSeek-V3.1. --- studio/backend/core/inference/llama_cpp.py | 99 ++++++++++++++++++++-- 1 file changed, 93 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index eb3776e603..87e570e94c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -61,6 +61,15 @@ class LlamaCppBackend: self._n_kv_heads: Optional[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) + self._kv_key_length: Optional[int] = None + self._kv_value_length: Optional[int] = None + self._sliding_window: Optional[int] = None + self._full_attention_interval: Optional[int] = None + self._kv_lora_rank: Optional[int] = None + self._key_length_mla: Optional[int] = None + self._ssm_inner_size: Optional[int] = None + self._ssm_state_size: Optional[int] = None self._lock = threading.Lock() self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None @@ -347,9 +356,17 @@ class LlamaCppBackend: def _can_estimate_kv(self) -> bool: """True if we have enough GGUF metadata to estimate KV cache size.""" + if self._n_layers is None: + return False + # New-style: explicit key/value dimensions from GGUF + if self._kv_key_length is not None: + return True + # MLA: kv_lora_rank is sufficient + if self._kv_lora_rank is not None: + return True + # Legacy: need embedding_length + head count return ( - self._n_layers is not None - and self._embedding_length is not None + self._embedding_length is not None and (self._n_kv_heads is not None or self._n_heads is not None) ) @@ -358,14 +375,20 @@ class LlamaCppBackend: ) -> int: """Estimate KV cache VRAM for a given context length. + Uses 5-path architecture-aware estimation: + 1. MLA -- compressed KV latent + RoPE, K-only (no separate V) + 2. Hybrid -- only attention layers need KV (Mamba layers don't) + 3. SWA -- sliding-window layers cache min(ctx, window) tokens + 4. GQA -- standard full KV with explicit key/value dimensions + 5. Legacy -- fallback using embed // n_heads + 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] - n_kv_heads = self._n_kv_heads or self._n_heads # type: ignore[assignment] - head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] + n_kv = self._n_kv_heads or self._n_heads or 1 # type: ignore[assignment] # Bytes per element depends on KV cache quantization bpe = { @@ -380,8 +403,47 @@ class LlamaCppBackend: "iq4_nl": 0.5625, }.get(cache_type_kv or "f16", 2.0) - # K + V caches: 2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe - return int(2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe) + # Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5) + # MLA stores only the compressed KV latent + RoPE in the K cache. + # V is reconstructed from the latent on the fly -- no separate V cache. + # key_length = kv_lora_rank + rope_dim (the full compressed representation). + if self._kv_lora_rank is not None: + key_len = self._kv_key_length or (self._kv_lora_rank + 64) + return int(n_layers * n_ctx * n_kv * key_len * bpe) + + key_len = self._kv_key_length + val_len = self._kv_value_length + + # Path 2: Hybrid Mamba/Attention (Qwen3.5-27B, Qwen3.5-35B-A3B) + # Only 1 in N layers is attention; the rest are Mamba (no KV cache). + if self._ssm_inner_size is not None and self._full_attention_interval is not None: + fai = self._full_attention_interval + n_attn = n_layers // fai if fai > 0 else n_layers + if key_len is not None and val_len is not None: + return int(n_attn * n_ctx * n_kv * (key_len + val_len) * bpe) + 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. + # Conservative: assume half layers are global, half are SWA. + if self._sliding_window is not None and key_len is not None and val_len is not None: + swa = self._sliding_window + n_global = n_layers // 2 + n_swa = n_layers - n_global + kv_per_token = n_kv * (key_len + val_len) * bpe + return int( + n_global * n_ctx * kv_per_token + + n_swa * min(n_ctx, swa) * kv_per_token + ) + + # 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) + + # 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) def _fit_context_to_vram( self, @@ -585,6 +647,14 @@ class LlamaCppBackend: self._n_kv_heads = None self._n_heads = None self._embedding_length = None + self._kv_key_length = None + self._kv_value_length = None + self._sliding_window = None + self._full_attention_interval = None + self._kv_lora_rank = None + self._key_length_mla = None + self._ssm_inner_size = None + self._ssm_state_size = None try: WANTED = {"general.architecture", "tokenizer.chat_template"} @@ -619,6 +689,15 @@ class LlamaCppBackend: f"{arch}.attention.head_count_kv": "n_kv_heads", f"{arch}.attention.head_count": "n_heads", f"{arch}.embedding_length": "embedding_length", + # Architecture-aware KV cache fields + f"{arch}.attention.key_length": "kv_key_length", + f"{arch}.attention.value_length": "kv_value_length", + f"{arch}.attention.sliding_window": "sliding_window", + f"{arch}.full_attention_interval": "full_attention_interval", + f"{arch}.attention.kv_lora_rank": "kv_lora_rank", + f"{arch}.attention.key_length_mla": "key_length_mla", + f"{arch}.ssm.inner_size": "ssm_inner_size", + f"{arch}.ssm.state_size": "ssm_state_size", } elif key == "tokenizer.chat_template": self._chat_template = val_s @@ -1422,6 +1501,14 @@ class LlamaCppBackend: self._n_kv_heads = None self._n_heads = None self._embedding_length = None + self._kv_key_length = None + self._kv_value_length = None + self._sliding_window = None + self._full_attention_interval = None + self._kv_lora_rank = None + self._key_length_mla = None + self._ssm_inner_size = None + self._ssm_state_size = None # Clean up temp chat template file if hasattr(self, "_chat_template_file") and self._chat_template_file: try: From 41198342d9028c7196ce25d5a1e77c1eeb67de73 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 1 Apr 2026 12:29:57 +0000 Subject: [PATCH 2/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/llama_cpp.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 87e570e94c..40c80c2c92 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -365,9 +365,8 @@ class LlamaCppBackend: if self._kv_lora_rank is not None: return True # Legacy: need embedding_length + head count - return ( - self._embedding_length is not None - and (self._n_kv_heads is not None or self._n_heads is not None) + return self._embedding_length is not None and ( + self._n_kv_heads is not None or self._n_heads is not None ) def _estimate_kv_cache_bytes( @@ -416,7 +415,10 @@ class LlamaCppBackend: # Path 2: Hybrid Mamba/Attention (Qwen3.5-27B, Qwen3.5-35B-A3B) # Only 1 in N layers is attention; the rest are Mamba (no KV cache). - if self._ssm_inner_size is not None and self._full_attention_interval is not None: + if ( + self._ssm_inner_size is not None + and self._full_attention_interval is not None + ): fai = self._full_attention_interval n_attn = n_layers // fai if fai > 0 else n_layers if key_len is not None and val_len is not None: @@ -427,14 +429,17 @@ class LlamaCppBackend: # Path 3: Sliding Window (Gemma-3, gpt-oss) # SWA layers only cache min(ctx, window) tokens; global layers cache full ctx. # Conservative: assume half layers are global, half are SWA. - if self._sliding_window is not None and key_len is not None and val_len is not None: + if ( + self._sliding_window is not None + and key_len is not None + and val_len is not None + ): swa = self._sliding_window n_global = n_layers // 2 n_swa = n_layers - n_global kv_per_token = n_kv * (key_len + val_len) * bpe return int( - n_global * n_ctx * kv_per_token - + n_swa * min(n_ctx, swa) * kv_per_token + n_global * n_ctx * kv_per_token + n_swa * min(n_ctx, swa) * kv_per_token ) # Path 4: Standard GQA with explicit key/value dimensions From ae6fb93b6f0b5df09b6a8ee5f34ac0a375724da7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Apr 2026 12:42:44 +0000 Subject: [PATCH 3/7] Fix MLA fallback and SWA global/local ratio heuristic Two fixes based on review findings: 1. MLA fallback now uses key_length_mla from GGUF metadata instead of hardcoded rope_dim=64. Falls back to 64 only when key_length_mla is absent. This ensures correct estimates for MLA variants that use rope dimensions other than 64. 2. SWA global/local layer ratio changed from 50/50 to 1/4 (25% global, 75% SWA). Most sliding window architectures have predominantly local layers (Gemma-3 uses ~17% global, gpt-oss uses ~50%). The 1/4 heuristic is closer to the common case and still a large improvement over the legacy formula which ignores SWA entirely. --- studio/backend/core/inference/llama_cpp.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 40c80c2c92..a1b872c8e6 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -407,7 +407,8 @@ class LlamaCppBackend: # V is reconstructed from the latent on the fly -- no separate V cache. # key_length = kv_lora_rank + rope_dim (the full compressed representation). if self._kv_lora_rank is not None: - key_len = self._kv_key_length or (self._kv_lora_rank + 64) + 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 * key_len * bpe) key_len = self._kv_key_length @@ -428,14 +429,16 @@ class LlamaCppBackend: # Path 3: Sliding Window (Gemma-3, gpt-oss) # SWA layers only cache min(ctx, window) tokens; global layers cache full ctx. - # Conservative: assume half layers are global, half are SWA. + # 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). if ( self._sliding_window is not None and key_len is not None and val_len is not None ): swa = self._sliding_window - n_global = n_layers // 2 + n_global = max(1, n_layers // 4) n_swa = n_layers - n_global kv_per_token = n_kv * (key_len + val_len) * bpe return int( From 434dee96185a6d1018e5593f4829bd4ba001476e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Apr 2026 12:49:41 +0000 Subject: [PATCH 4/7] Tighten _can_estimate_kv gate and treat sliding_window=0 as disabled Two additional fixes from review round 1 (5/8 and 4/8 reviewer consensus): 1. _can_estimate_kv now requires BOTH key_length AND value_length for the explicit-dims path. Previously key_length alone was enough, which could cause silent fallthrough to the legacy formula with fabricated defaults (n_kv=1, head_dim=128) when value_length was absent from the GGUF. 2. SWA path now requires sliding_window > 0. Some GGUFs use 0 as a disabled sentinel. Without this guard, min(ctx, 0) would zero out all SWA layer contributions, severely underestimating KV cache. --- studio/backend/core/inference/llama_cpp.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index a1b872c8e6..6635764e01 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -358,12 +358,12 @@ class LlamaCppBackend: """True if we have enough GGUF metadata to estimate KV cache size.""" if self._n_layers is None: return False - # New-style: explicit key/value dimensions from GGUF - if self._kv_key_length is not None: - return True - # MLA: kv_lora_rank is sufficient + # MLA: kv_lora_rank is sufficient (K-only cache) if self._kv_lora_rank is not None: return True + # 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 return self._embedding_length is not None and ( self._n_kv_heads is not None or self._n_heads is not None @@ -434,6 +434,7 @@ class LlamaCppBackend: # which is still far more accurate than the legacy formula (which ignores SWA). if ( self._sliding_window is not None + and self._sliding_window > 0 and key_len is not None and val_len is not None ): From 87e5385b44470c65851ce2d73e6281b585575ab6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Apr 2026 12:51:43 +0000 Subject: [PATCH 5/7] Fix MLA n_kv safety and use ceiling division for hybrid path Addresses Gemini Code Assist review findings: 1. MLA path now uses n_kv_mla = n_kv_heads or 1 (not n_heads). This prevents a 128x overestimate for DeepSeek-V3 if head_count_kv is absent from the GGUF (n_heads=128 would have been used instead). 2. Hybrid path now uses ceiling division for attention layer count. This prevents undercounting by 1 when n_layers is not perfectly divisible by full_attention_interval. --- studio/backend/core/inference/llama_cpp.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 6635764e01..873b72bba1 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -403,13 +403,16 @@ class LlamaCppBackend: }.get(cache_type_kv or "f16", 2.0) # Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5) - # MLA stores only the compressed KV latent + RoPE in the K cache. + # 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. # key_length = kv_lora_rank + rope_dim (the full compressed representation). + # MLA GGUFs set head_count_kv=1; default to 1 if absent to avoid + # falling back to n_heads (e.g., 128 for DeepSeek-V3) which would 128x. if self._kv_lora_rank is not None: + 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 * key_len * bpe) + return int(n_layers * n_ctx * n_kv_mla * key_len * bpe) key_len = self._kv_key_length val_len = self._kv_value_length @@ -421,7 +424,7 @@ class LlamaCppBackend: and self._full_attention_interval is not None ): fai = self._full_attention_interval - n_attn = n_layers // fai if fai > 0 else n_layers + n_attn = -(-n_layers // fai) if fai > 0 else n_layers # ceiling division if key_len is not None and val_len is not None: return int(n_attn * n_ctx * n_kv * (key_len + val_len) * bpe) head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] From f4efbba2a04ef3bd2ab3496f68df5ed59c76d803 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Apr 2026 12:51:55 +0000 Subject: [PATCH 6/7] test: add MLA n_kv default safety test (66 tests total) Added test_mla_defaults_n_kv_to_1_when_heads_absent to verify MLA path uses n_kv=1 (not n_heads) when head_count_kv is absent. --- .../backend/tests/test_kv_cache_estimation.py | 858 ++++++++++++++++++ 1 file changed, 858 insertions(+) create mode 100644 studio/backend/tests/test_kv_cache_estimation.py diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py new file mode 100644 index 0000000000..26d384c915 --- /dev/null +++ b/studio/backend/tests/test_kv_cache_estimation.py @@ -0,0 +1,858 @@ +# 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). + +Requires no GPU, network, or external libraries beyond pytest. +Cross-platform: Linux, macOS, Windows, WSL. +""" + +import io +import struct +import sys +import types as _types +from pathlib import Path + +import pytest + +# --------------------------------------------------------------------------- +# Stub heavy / unavailable external dependencies 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 +_httpx_stub = _types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", "TimeoutException", "ReadTimeout", + "ReadError", "RemoteProtocolError", "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + +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) + +from core.inference.llama_cpp import LlamaCppBackend + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +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. + """ + 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.""" + kv = {"general.architecture": arch} + 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: + """Verify that the 8 new architecture-aware fields are correctly parsed.""" + + @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", + "_full_attention_interval", "_kv_lora_rank", "_key_length_mla", + "_ssm_inner_size", "_ssm_state_size", + ]: + assert getattr(b, attr) is None + + def test_all_13_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, + "full_attention_interval": 6, + "attention.kv_lora_rank": 512, + "attention.key_length_mla": 256, + "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._full_attention_interval == 6 + assert b._kv_lora_rank == 512 + assert b._key_length_mla == 256 + assert b._ssm_inner_size == 4096 + assert b._ssm_state_size == 128 + + +class TestGGUFParserReset: + """Verify that fields are properly reset between parses.""" + + def test_reset_between_parses(self): + # First parse with all fields + b = _backend_from_gguf("arch1", { + "block_count": 32, + "attention.key_length": 128, + "attention.kv_lora_rank": 512, + "ssm.inner_size": 4096, + }) + assert b._kv_key_length == 128 + assert b._kv_lora_rank == 512 + assert b._ssm_inner_size == 4096 + + # Second parse without those fields -- they should 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._ssm_inner_size is None + assert b._n_layers == 64 + + +# --------------------------------------------------------------------------- +# B. _can_estimate_kv Gate Tests +# --------------------------------------------------------------------------- + +class TestCanEstimateKV: + """Verify 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 should NOT be 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 should NOT add value_length -- V is reconstructed from the latent.""" + b = self._mla_backend() + result = b._estimate_kv_cache_bytes(1000, "f16") + # Should be 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): + """If key_length is missing, fallback to kv_lora_rank + key_length_mla.""" + b = self._mla_backend(_kv_key_length=None) + # _key_length_mla=192 in default, 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): + """If both key_length and key_length_mla are missing, fallback 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 should use 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") + # Should use 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): + """Fallback 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 should not cause 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 + expected = int(n_global * 131072 * kv_per + n_swa * min(131072, 1024) * 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 + expected = int(n_global * 131072 * kv_per + n_swa * min(131072, 128) * kv_per) + assert b._estimate_kv_cache_bytes(131072, "f16") == expected + + def test_ctx_smaller_than_window(self): + """When context < sliding_window, SWA layers use full context anyway.""" + 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 + 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) + 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 should differ 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 would use: 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): + """Confirm legacy path produces the same result as 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 is used even if other fields are 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 carefully chosen params, each path should yield a distinct value.""" + # Use embedding_length=768 so legacy head_dim (768//16=48) differs from + # 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: + """Verify all supported cache_type_kv values produce correct scaling.""" + + @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 + + +# --------------------------------------------------------------------------- +# 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", + "_full_attention_interval", "_kv_lora_rank", "_key_length_mla", + "_ssm_inner_size", "_ssm_state_size", + ]: + assert getattr(b, attr) 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._ssm_inner_size = 4096 + b._full_attention_interval = 4 + b.unload_model() + for attr in [ + "_kv_key_length", "_kv_value_length", "_sliding_window", + "_full_attention_interval", "_kv_lora_rank", "_key_length_mla", + "_ssm_inner_size", "_ssm_state_size", + ]: + assert getattr(b, attr) is None + + def test_end_to_end_synthetic_mla(self): + """Full 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") + n_global = max(1, 62 // 4) # 15 + n_swa = 62 - n_global # 47 + kv_per = 16 * 256 * 2 + expected = int(n_global * 131072 * kv_per + n_swa * 1024 * kv_per) + assert result == expected + + 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 From a9d6c26bb5fb5f477b32ee9958d9242aa9b4491b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 1 Apr 2026 12:52:54 +0000 Subject: [PATCH 7/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../backend/tests/test_kv_cache_estimation.py | 301 +++++++++++------- 1 file changed, 186 insertions(+), 115 deletions(-) diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py index 26d384c915..2640ded90d 100644 --- a/studio/backend/tests/test_kv_cache_estimation.py +++ b/studio/backend/tests/test_kv_cache_estimation.py @@ -40,18 +40,26 @@ sys.modules.setdefault("structlog", _structlog_stub) # httpx _httpx_stub = _types.ModuleType("httpx") for _exc_name in ( - "ConnectError", "TimeoutException", "ReadTimeout", - "ReadError", "RemoteProtocolError", "CloseError", + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", ): setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + class _FakeTimeout: def __init__(self, *a, **kw): pass + _httpx_stub.Timeout = _FakeTimeout _httpx_stub.Client = type( - "Client", (), { + "Client", + (), + { "__init__": lambda self, **kw: None, "__enter__": lambda self: self, "__exit__": lambda self, *a: None, @@ -65,6 +73,7 @@ from core.inference.llama_cpp import LlamaCppBackend # Helpers # --------------------------------------------------------------------------- + def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes: """Build a minimal GGUF v3 binary blob with the given KV metadata. @@ -74,8 +83,8 @@ def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes: buf = io.BytesIO() # Header: magic, version, tensor_count, kv_count buf.write(struct.pack(" LlamaCppBackend: 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") + fd, path = tempfile.mkstemp(suffix = ".gguf") try: os.write(fd, data) os.close(fd) @@ -121,19 +131,23 @@ def _backend_from_gguf(arch: str, fields: dict) -> LlamaCppBackend: # A. GGUF Parser Tests # --------------------------------------------------------------------------- + class TestGGUFParserNewFields: """Verify that the 8 new architecture-aware fields are correctly parsed.""" - @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), - ]) + @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 @@ -141,9 +155,14 @@ class TestGGUFParserNewFields: 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", - "_full_attention_interval", "_kv_lora_rank", "_key_length_mla", - "_ssm_inner_size", "_ssm_state_size", + "_kv_key_length", + "_kv_value_length", + "_sliding_window", + "_full_attention_interval", + "_kv_lora_rank", + "_key_length_mla", + "_ssm_inner_size", + "_ssm_state_size", ]: assert getattr(b, attr) is None @@ -184,12 +203,15 @@ class TestGGUFParserReset: def test_reset_between_parses(self): # First parse with all fields - b = _backend_from_gguf("arch1", { - "block_count": 32, - "attention.key_length": 128, - "attention.kv_lora_rank": 512, - "ssm.inner_size": 4096, - }) + b = _backend_from_gguf( + "arch1", + { + "block_count": 32, + "attention.key_length": 128, + "attention.kv_lora_rank": 512, + "ssm.inner_size": 4096, + }, + ) assert b._kv_key_length == 128 assert b._kv_lora_rank == 512 assert b._ssm_inner_size == 4096 @@ -197,8 +219,9 @@ class TestGGUFParserReset: # Second parse without those fields -- they should 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") + fd, path = tempfile.mkstemp(suffix = ".gguf") os.write(fd, data) os.close(fd) try: @@ -215,6 +238,7 @@ class TestGGUFParserReset: # B. _can_estimate_kv Gate Tests # --------------------------------------------------------------------------- + class TestCanEstimateKV: """Verify gate logic for all field combinations.""" @@ -274,6 +298,7 @@ class TestCanEstimateKV: # C. Path 1: MLA Estimation # --------------------------------------------------------------------------- + class TestMLAEstimation: """MLA: K-only cache using compressed KV latent + RoPE.""" @@ -310,7 +335,7 @@ class TestMLAEstimation: def test_mla_fallback_when_no_key_length(self): """If key_length is missing, fallback to kv_lora_rank + key_length_mla.""" - b = self._mla_backend(_kv_key_length=None) + b = self._mla_backend(_kv_key_length = None) # _key_length_mla=192 in default, so rope_dim=192 result = b._estimate_kv_cache_bytes(1000, "f16") expected = 61 * 1000 * 1 * (512 + 192) * 2 # 704 @@ -318,14 +343,14 @@ class TestMLAEstimation: def test_mla_fallback_no_key_length_mla(self): """If both key_length and key_length_mla are missing, fallback to +64.""" - b = self._mla_backend(_kv_key_length=None, _key_length_mla=None) + 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 should use 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 + b = self._mla_backend(_n_kv_heads = None) # n_heads=128 still set result = b._estimate_kv_cache_bytes(1000, "f16") # Should use n_kv_mla=1, NOT n_heads=128 expected = 61 * 1000 * 1 * 576 * 2 @@ -344,6 +369,7 @@ class TestMLAEstimation: # D. Path 2: Hybrid Mamba Estimation # --------------------------------------------------------------------------- + class TestHybridMambaEstimation: """Hybrid Mamba: only attention layers (1 in N) need KV cache.""" @@ -373,8 +399,11 @@ class TestHybridMambaEstimation: 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_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 @@ -382,14 +411,14 @@ class TestHybridMambaEstimation: def test_hybrid_without_explicit_dims(self): """Fallback to head_dim when key_length/value_length are missing.""" - b = self._hybrid_backend(_kv_key_length=None, _kv_value_length=None) + 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 should not cause ZeroDivisionError.""" - b = self._hybrid_backend(_full_attention_interval=0) + 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 @@ -400,6 +429,7 @@ class TestHybridMambaEstimation: # E. Path 3: Sliding Window Estimation # --------------------------------------------------------------------------- + class TestSlidingWindowEstimation: """SWA: half global (full ctx) + half sliding window.""" @@ -423,29 +453,33 @@ class TestSlidingWindowEstimation: 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 + 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) 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, + _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 + 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) assert b._estimate_kv_cache_bytes(131072, "f16") == expected def test_ctx_smaller_than_window(self): """When context < sliding_window, SWA layers use full context anyway.""" - b = self._swa_backend(_sliding_window=8192) + b = self._swa_backend(_sliding_window = 8192) n_global = max(1, 62 // 4) # 15 - n_swa = 62 - n_global # 47 + 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) @@ -454,9 +488,9 @@ class TestSlidingWindowEstimation: 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) + b = self._swa_backend(_n_layers = 63) n_global = max(1, 63 // 4) # 15 - n_swa = 63 - n_global # 48 + 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) assert b._estimate_kv_cache_bytes(1000, "f16") == expected @@ -466,6 +500,7 @@ class TestSlidingWindowEstimation: # F. Path 4: Standard GQA Estimation # --------------------------------------------------------------------------- + class TestStandardGQAEstimation: """Standard GQA with explicit key_length/value_length.""" @@ -491,7 +526,7 @@ class TestStandardGQAEstimation: 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) + 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 @@ -511,6 +546,7 @@ class TestStandardGQAEstimation: # G. Path 5: Legacy Fallback Estimation # --------------------------------------------------------------------------- + class TestLegacyEstimation: """Legacy: embed // n_heads, for old GGUFs without new fields.""" @@ -535,7 +571,7 @@ class TestLegacyEstimation: 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) + 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 @@ -556,6 +592,7 @@ class TestLegacyEstimation: # H. Path Priority (selection order) # --------------------------------------------------------------------------- + class TestPathPriority: """Confirm: MLA > Hybrid Mamba > SWA > GQA > Legacy.""" @@ -599,8 +636,11 @@ class TestPathPriority: # Use embedding_length=768 so legacy head_dim (768//16=48) differs from # 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, + "_n_layers": 40, + "_n_kv_heads": 4, + "_n_heads": 16, + "_embedding_length": 768, + "_kv_key_length": 256, "_kv_value_length": 256, } ctx = 4096 @@ -649,22 +689,26 @@ class TestPathPriority: # I. KV Cache Quantization # --------------------------------------------------------------------------- + class TestQuantization: """Verify all supported cache_type_kv values produce correct scaling.""" - @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 - ]) + @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 @@ -682,6 +726,7 @@ class TestQuantization: # J. Edge Cases # --------------------------------------------------------------------------- + class TestEdgeCases: """Boundary conditions and degenerate inputs.""" @@ -744,15 +789,21 @@ class TestEdgeCases: # 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", - "_full_attention_interval", "_kv_lora_rank", "_key_length_mla", - "_ssm_inner_size", "_ssm_state_size", + "_kv_key_length", + "_kv_value_length", + "_sliding_window", + "_full_attention_interval", + "_kv_lora_rank", + "_key_length_mla", + "_ssm_inner_size", + "_ssm_state_size", ]: assert getattr(b, attr) is None @@ -766,43 +817,54 @@ class TestLifecycle: b._full_attention_interval = 4 b.unload_model() for attr in [ - "_kv_key_length", "_kv_value_length", "_sliding_window", - "_full_attention_interval", "_kv_lora_rank", "_key_length_mla", - "_ssm_inner_size", "_ssm_state_size", + "_kv_key_length", + "_kv_value_length", + "_sliding_window", + "_full_attention_interval", + "_kv_lora_rank", + "_key_length_mla", + "_ssm_inner_size", + "_ssm_state_size", ]: assert getattr(b, attr) is None def test_end_to_end_synthetic_mla(self): """Full 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, - }) + 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, - }) + 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 @@ -810,47 +872,56 @@ class TestLifecycle: 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, - }) + 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") n_global = max(1, 62 // 4) # 15 - n_swa = 62 - n_global # 47 + n_swa = 62 - n_global # 47 kv_per = 16 * 256 * 2 expected = int(n_global * 131072 * kv_per + n_swa * 1024 * kv_per) assert result == expected 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, - }) + 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, - }) + 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