[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

@@ -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 (
+        if self._n_layers is None:
-            self._n_layers is not None
+            return False
-            and self._embedding_length is not None
+        # MLA: kv_lora_rank is sufficient (K-only cache)
-            and (self._n_kv_heads is not None or self._n_heads is not None)
+        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]
+        n_kv = self._n_kv_heads or self._n_heads or 1  # type: ignore[assignment]
-        head_dim = self._embedding_length // self._n_heads if self._n_heads else 128  # type: ignore[operator]
 
         # 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
+        # Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5)
-        return int(2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe)
+        # 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:

studio/backend/tests/test_kv_cache_estimation.py

@@ -0,0 +1,929 @@
+# 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("<I", 0x46554747))  # GGUF magic
+    buf.write(struct.pack("<I", 3))  # version 3
+    buf.write(struct.pack("<Q", 0))  # tensor_count
+    buf.write(struct.pack("<Q", len(kv_pairs)))
+
+    for key, val in kv_pairs.items():
+        key_bytes = key.encode("utf-8")
+        buf.write(struct.pack("<Q", len(key_bytes)))
+        buf.write(key_bytes)
+        if isinstance(val, str):
+            buf.write(struct.pack("<I", 8))  # STRING
+            val_bytes = val.encode("utf-8")
+            buf.write(struct.pack("<Q", len(val_bytes)))
+            buf.write(val_bytes)
+        elif isinstance(val, int):
+            if val <= 0xFFFFFFFF:
+                buf.write(struct.pack("<I", 4))  # UINT32
+                buf.write(struct.pack("<I", val))
+            else:
+                buf.write(struct.pack("<I", 10))  # UINT64
+                buf.write(struct.pack("<Q", val))
+        else:
+            raise TypeError(f"Unsupported value type: {type(val)}")
+    return buf.getvalue()
+
+
+def _backend_from_gguf(arch: str, fields: dict) -> 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