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

feat(studio): architecture-aware KV cache VRAM estimation (#4757)

Browse source 
* 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.

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

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

* 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.

* 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.

* 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.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-04-01 06:04:12 -07:00 • committed by GitHub
commit 77e1a9edc9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 107 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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,10 +356,17 @@ class LlamaCppBackend:
def _can_estimate_kv(self) -> bool:
"""True if we have enough GGUF metadata to estimate KV cache size."""
return (
self._n_layers is not None
and self._embedding_length is not None
and (self._n_kv_heads is not None or self._n_heads is not None)
if self._n_layers is None:
return False
# 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
)
def _estimate_kv_cache_bytes(
@ -358,14 +374,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 +402,60 @@ 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 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_mla * 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 # 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]
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).
if (
self._sliding_window is not None
and self._sliding_window > 0
and key_len is not None
and val_len is not None
):
swa = self._sliding_window
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(
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 +659,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 +701,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 +1513,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: