From bac2cfac7cab5629cee376dec92e5a1f3ca6503e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Apr 2026 12:29:20 +0000 Subject: [PATCH 1/5] 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/5] [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/5] 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/5] 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/5] 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]