wehub-resource-sync
e93507a09c
Lockfile supply-chain audit / lockfile supply-chain audit (push) Has been cancelled
Windows Studio GGUF CI / GPU prebuilt resolves without Visual Studio (push) Has been cancelled
Windows Studio GGUF CI / setup.ps1 unit tests (VS 2026 / CMake guard) (push) Has been cancelled
Windows Studio GGUF CI / real-VS detection (VS 2022) (push) Has been cancelled
Windows Studio GGUF CI / real-VS detection (VS 2026) (push) Has been cancelled
Windows Studio GGUF CI / VC++ runtime detect + install round-trip (windows-2025-vs2026) (push) Has been cancelled
Windows Studio GGUF CI / VC++ runtime detect + install round-trip (windows-latest) (push) Has been cancelled
Windows Studio Update CI / Studio Updating Tests (push) Has been cancelled
Wheel CI / Wheel build + content sanity + import smoke (push) Has been cancelled
Lint CI / Source lint (Python + shell + YAML + JSON + safety nets) (push) Has been cancelled
MLX CI on Mac M1 / dispatch (push) Has been cancelled
Security audit / advisory audit (pip + npm + cargo) (push) Has been cancelled
Security audit / pip scan-packages :: extras (push) Has been cancelled
Security audit / pip scan-packages :: studio (push) Has been cancelled
Security audit / pip scan-packages :: hf-stack (push) Has been cancelled
Security audit / npm scan-packages (Studio frontend tarballs) (push) Has been cancelled
Security audit / workflow-trigger lint (pull_request_target / cache-poisoning) (push) Has been cancelled
Security audit / pytest tests/security (push) Has been cancelled
Security audit / npm provenance + new install-script diff (push) Has been cancelled
Studio API CI / Studio API & Auth Tests (push) Has been cancelled
Backend CI / (Python 3.10) (push) Has been cancelled
Backend CI / (Python 3.11) (push) Has been cancelled
Backend CI / (Python 3.12) (push) Has been cancelled
Backend CI / (Python 3.13) (push) Has been cancelled
Backend CI / Repo tests (CPU) (push) Has been cancelled
Frontend CI / Frontend build + bundle sanity (push) Has been cancelled
Studio GGUF CI / OpenAI, Anthropic API tests (push) Has been cancelled
Studio GGUF CI / Tool calling Tests (push) Has been cancelled
Studio GGUF CI / JSON, images (push) Has been cancelled
Mac Studio GGUF CI / OpenAI, Anthropic API tests (push) Has been cancelled
Mac Studio GGUF CI / Tool calling Tests (push) Has been cancelled
Mac Studio GGUF CI / JSON, images (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-14) (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-15) (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-26) (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-15-intel) (push) Has been cancelled
Mac Studio API CI / Studio API & Auth Tests (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-26-intel) (push) Has been cancelled
Mac Studio UI CI / Chat UI Tests (push) Has been cancelled
Studio Tauri CI / Tauri Linux debug build (no codesign) (push) Has been cancelled
Mac Studio Update CI / Studio Updating Tests (push) Has been cancelled
Studio UI CI / Chat UI Tests (push) Has been cancelled
Windows Studio API CI / Studio API & Auth Tests (push) Has been cancelled
Windows Studio UI CI / Chat UI Tests (push) Has been cancelled
Studio Update CI / Studio Updating Tests (push) Has been cancelled
Core / Core (HF=default + TRL=default) (push) Has been cancelled
Core / Core (HF=4.57.6 + TRL<1) (push) Has been cancelled
Core / Core (HF=latest + TRL=latest) (push) Has been cancelled
Core / llama.cpp build + smoke (push) Has been cancelled
Windows Studio GGUF CI / OpenAI, Anthropic API tests (push) Has been cancelled
Windows Studio GGUF CI / Tool calling Tests (push) Has been cancelled
Windows Studio GGUF CI / JSON, images (push) Has been cancelled
Windows Studio GGUF CI / Studio install + inference without Visual Studio (push) Has been cancelled
Studio export capability / capability (macos-latest) (push) Has been cancelled
Studio export capability / capability (ubuntu-latest) (push) Has been cancelled
Studio export capability / capability (windows-latest) (push) Has been cancelled
Cross-platform parity / parity (macos-latest) (push) Has been cancelled
Cross-platform parity / parity (windows-latest) (push) Has been cancelled
Scorecard supply-chain security / Scorecard analysis (push) Has been cancelled
Studio load-orchestrator CI / test (push) Has been cancelled
2113 lines
77 KiB
Python
2113 lines
77 KiB
Python
# 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).
|
|
|
|
No GPU, network, or libraries beyond pytest. Cross-platform.
|
|
"""
|
|
|
|
import io
|
|
import json
|
|
import struct
|
|
import sys
|
|
import types as _types
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
# Stub heavy / unavailable deps 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 -- only stub when the real library is missing. Unconditional stubbing
|
|
# shadows HTTPError/Response that huggingface_hub.errors imports at load time,
|
|
# silently breaking the transformers introspection tier.
|
|
try:
|
|
import httpx as _httpx_real # noqa: F401
|
|
except ImportError:
|
|
_httpx_stub = _types.ModuleType("httpx")
|
|
for _exc_name in (
|
|
"ConnectError",
|
|
"TimeoutException",
|
|
"ReadTimeout",
|
|
"ReadError",
|
|
"RemoteProtocolError",
|
|
"CloseError",
|
|
"HTTPError",
|
|
"RequestError",
|
|
):
|
|
setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
|
|
|
|
class _FakeTimeout:
|
|
def __init__(self, *a, **kw):
|
|
pass
|
|
|
|
_httpx_stub.Timeout = _FakeTimeout
|
|
_httpx_stub.Response = type("Response", (), {})
|
|
_httpx_stub.Client = type(
|
|
"Client",
|
|
(),
|
|
{
|
|
"__init__": lambda self, **kw: None,
|
|
"__enter__": lambda self: self,
|
|
"__exit__": lambda self, *a: None,
|
|
},
|
|
)
|
|
sys.modules["httpx"] = _httpx_stub
|
|
|
|
from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend
|
|
|
|
# Helpers
|
|
|
|
|
|
def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes:
|
|
"""Build a minimal GGUF v3 blob with the given KV metadata.
|
|
|
|
Supports the scalar and simple array metadata the parser uses.
|
|
"""
|
|
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, list):
|
|
buf.write(struct.pack("<I", 9)) # ARRAY
|
|
is_bool_array = all(isinstance(x, bool) for x in val)
|
|
buf.write(struct.pack("<I", 7 if is_bool_array else 5))
|
|
buf.write(struct.pack("<Q", len(val)))
|
|
if is_bool_array:
|
|
for item in val:
|
|
buf.write(struct.pack("<?", item))
|
|
else:
|
|
for item in val:
|
|
buf.write(struct.pack("<i", item))
|
|
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,
|
|
general: dict | None = None,
|
|
) -> LlamaCppBackend:
|
|
"""Create a LlamaCppBackend with parsed GGUF metadata from given fields.
|
|
|
|
`general` injects extra `general.*` metadata, to verify the dynamic
|
|
SWA resolver picks up source-repo hints from GGUFs that ship them.
|
|
"""
|
|
kv = {"general.architecture": arch}
|
|
for k, v in (general or {}).items():
|
|
kv[k] = v
|
|
for k, v in fields.items():
|
|
kv[f"{arch}.{k}"] = v
|
|
import tempfile, os
|
|
|
|
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:
|
|
"""Architecture-aware fields are parsed correctly."""
|
|
|
|
@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",
|
|
"_sliding_window_pattern",
|
|
"_full_attention_interval",
|
|
"_kv_lora_rank",
|
|
"_key_length_mla",
|
|
"_kv_key_length_swa",
|
|
"_kv_value_length_swa",
|
|
"_ssm_inner_size",
|
|
"_ssm_state_size",
|
|
]:
|
|
assert getattr(b, attr) is None
|
|
|
|
def test_array_fields_parsed(self):
|
|
b = _backend_from_gguf(
|
|
"gemma4",
|
|
{
|
|
"block_count": 6,
|
|
"attention.head_count_kv": [8, 8, 8, 8, 8, 2],
|
|
"attention.sliding_window_pattern": [
|
|
True,
|
|
True,
|
|
True,
|
|
True,
|
|
True,
|
|
False,
|
|
],
|
|
},
|
|
)
|
|
# Per-layer KV head count is preserved exactly...
|
|
assert b._n_kv_heads_by_layer == [8, 8, 8, 8, 8, 2]
|
|
# ...and mirrored into the scalar field as a conservative max, so
|
|
# non-SWA paths and callers using `n_kv = self._n_kv_heads or ...`
|
|
# get a safe upper bound.
|
|
assert b._n_kv_heads == 8
|
|
assert b._sliding_window_pattern == [True, True, True, True, True, False]
|
|
|
|
|
|
class TestArchSwaPatternDefaults:
|
|
"""Bootstrap arch table fires when GGUF reports `sliding_window` but
|
|
no per-layer pattern (true for every Gemma 2/3/3n/gpt-oss GGUF today)."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"arch,n_layers,expected_period",
|
|
[
|
|
("gemma2", 26, 2),
|
|
("gemma3", 18, 6),
|
|
("gemma3n", 35, 5),
|
|
("gpt_oss", 24, 2),
|
|
("cohere2", 32, 4),
|
|
],
|
|
)
|
|
def test_arch_default_pattern_applied(self, arch, n_layers, expected_period):
|
|
b = _backend_from_gguf(
|
|
arch,
|
|
{
|
|
"block_count": n_layers,
|
|
"attention.head_count": 4,
|
|
"attention.head_count_kv": 1,
|
|
"attention.key_length": 256,
|
|
"attention.value_length": 256,
|
|
"attention.sliding_window": 512,
|
|
},
|
|
)
|
|
expected_pattern = [(i + 1) % expected_period != 0 for i in range(n_layers)]
|
|
assert (
|
|
b._sliding_window_pattern == expected_pattern
|
|
), f"{arch} should expand to period={expected_period}"
|
|
|
|
def test_unknown_arch_no_default(self):
|
|
b = _backend_from_gguf(
|
|
"totallymadeupv7",
|
|
{
|
|
"block_count": 24,
|
|
"attention.head_count": 4,
|
|
"attention.head_count_kv": 1,
|
|
"attention.key_length": 128,
|
|
"attention.value_length": 128,
|
|
"attention.sliding_window": 1024,
|
|
},
|
|
)
|
|
assert b._sliding_window_pattern is None
|
|
|
|
def test_explicit_pattern_overrides_arch_default(self):
|
|
# gemma3 default is period=6; the explicit array must win.
|
|
b = _backend_from_gguf(
|
|
"gemma3",
|
|
{
|
|
"block_count": 6,
|
|
"attention.head_count": 4,
|
|
"attention.head_count_kv": 1,
|
|
"attention.key_length": 256,
|
|
"attention.value_length": 256,
|
|
"attention.sliding_window": 512,
|
|
"attention.sliding_window_pattern": [
|
|
True,
|
|
False,
|
|
True,
|
|
False,
|
|
True,
|
|
False,
|
|
],
|
|
},
|
|
)
|
|
assert b._sliding_window_pattern == [True, False, True, False, True, False]
|
|
|
|
def test_no_sliding_window_no_pattern(self):
|
|
b = _backend_from_gguf(
|
|
"gemma3",
|
|
{
|
|
"block_count": 18,
|
|
"attention.head_count": 4,
|
|
"attention.head_count_kv": 1,
|
|
"attention.key_length": 256,
|
|
"attention.value_length": 256,
|
|
# no sliding_window key
|
|
},
|
|
)
|
|
assert b._sliding_window_pattern is None
|
|
|
|
@pytest.mark.parametrize(
|
|
"arch", ["llama", "qwen2", "qwen3", "mistral", "mistral3", "glm4", "llama4"]
|
|
)
|
|
def test_non_swa_arch_uses_full_attention_path(self, arch):
|
|
# Pure-GQA arches: no sliding_window, no synthetic pattern,
|
|
# estimator hits Path 4.
|
|
b = _backend_from_gguf(
|
|
arch,
|
|
{
|
|
"block_count": 32,
|
|
"attention.head_count": 32,
|
|
"attention.head_count_kv": 8,
|
|
"attention.key_length": 128,
|
|
"attention.value_length": 128,
|
|
"embedding_length": 4096,
|
|
},
|
|
)
|
|
assert b._sliding_window_pattern is None
|
|
assert b._sliding_window is None
|
|
kv = b._estimate_kv_cache_bytes(8192, "f16")
|
|
gqa_expected = 32 * 8192 * 8 * (128 + 128) * 2
|
|
assert kv == gqa_expected
|
|
|
|
def test_arch_default_reduces_kv_estimate_vs_legacy(self):
|
|
common = {
|
|
"block_count": 62,
|
|
"attention.head_count": 32,
|
|
"attention.head_count_kv": 16,
|
|
"attention.key_length": 128,
|
|
"attention.value_length": 128,
|
|
"attention.sliding_window": 1024,
|
|
"embedding_length": 5376,
|
|
}
|
|
with_default = _backend_from_gguf("gemma3", common)
|
|
# Arch not in table -> legacy 1/4 path.
|
|
without_default = _backend_from_gguf("totallymadeupv7", common)
|
|
|
|
kv_default = with_default._estimate_kv_cache_bytes(131072, "f16")
|
|
kv_legacy = without_default._estimate_kv_cache_bytes(131072, "f16")
|
|
assert kv_default > 0
|
|
assert kv_legacy > 0
|
|
assert kv_default < kv_legacy, (
|
|
f"arch fallback should under-shoot legacy estimate: " f"{kv_default} >= {kv_legacy}"
|
|
)
|
|
|
|
def test_scalar_sliding_window_pattern_expanded(self):
|
|
block_count = 8
|
|
b = _backend_from_gguf(
|
|
"gemma3",
|
|
{
|
|
"attention.sliding_window_pattern": 4,
|
|
"block_count": block_count,
|
|
"attention.head_count_kv": 4,
|
|
"attention.key_length": 256,
|
|
"attention.value_length": 256,
|
|
"attention.sliding_window": 1024,
|
|
},
|
|
)
|
|
expected = [(i + 1) % 4 != 0 for i in range(block_count)]
|
|
assert isinstance(b._sliding_window_pattern, list)
|
|
assert b._sliding_window_pattern == expected
|
|
assert b._estimate_kv_cache_bytes(4096, "f16") > 0
|
|
|
|
def test_all_fields_parsed_together(self):
|
|
fields = {
|
|
"context_length": 131072,
|
|
"block_count": 62,
|
|
"attention.head_count_kv": 16,
|
|
"attention.head_count": 32,
|
|
"embedding_length": 5376,
|
|
"attention.key_length": 128,
|
|
"attention.value_length": 128,
|
|
"attention.sliding_window": 1024,
|
|
"attention.sliding_window_pattern": [True, False],
|
|
"full_attention_interval": 6,
|
|
"attention.kv_lora_rank": 512,
|
|
"attention.key_length_mla": 256,
|
|
"attention.key_length_swa": 64,
|
|
"attention.value_length_swa": 64,
|
|
"ssm.inner_size": 4096,
|
|
"ssm.state_size": 128,
|
|
}
|
|
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._sliding_window_pattern == [True, False]
|
|
assert b._full_attention_interval == 6
|
|
assert b._kv_lora_rank == 512
|
|
assert b._key_length_mla == 256
|
|
assert b._kv_key_length_swa == 64
|
|
assert b._kv_value_length_swa == 64
|
|
assert b._ssm_inner_size == 4096
|
|
assert b._ssm_state_size == 128
|
|
|
|
|
|
_SWA_FIELDS = {
|
|
"block_count": 12,
|
|
"attention.head_count": 4,
|
|
"attention.head_count_kv": 1,
|
|
"attention.key_length": 256,
|
|
"attention.value_length": 256,
|
|
"attention.sliding_window": 512,
|
|
}
|
|
|
|
|
|
class TestDynamicSwaResolver:
|
|
"""4-tier resolver: GGUF metadata, on-disk cache, bootstrap, HF fetch."""
|
|
|
|
def _isolate_cache(self, monkeypatch, tmp_path):
|
|
from core.inference import llama_cpp as lc
|
|
|
|
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
|
|
monkeypatch.setattr(lc, "_SWA_CACHE", None)
|
|
return tmp_path
|
|
|
|
def test_period_from_layer_types_finds_smallest_period(self):
|
|
from core.inference.llama_cpp import _period_from_layer_types
|
|
|
|
# gemma3 (1 global/6), gpt-oss (alternating), gemma3n (1/5).
|
|
assert _period_from_layer_types((["sliding_attention"] * 5 + ["full_attention"]) * 4) == 6
|
|
assert _period_from_layer_types(["sliding_attention", "full_attention"] * 12) == 2
|
|
assert _period_from_layer_types((["sliding_attention"] * 4 + ["full_attention"]) * 7) == 5
|
|
|
|
def test_period_from_layer_types_returns_none_for_aperiodic(self):
|
|
from core.inference.llama_cpp import _period_from_layer_types
|
|
lt = [
|
|
"sliding_attention",
|
|
"full_attention",
|
|
"sliding_attention",
|
|
"sliding_attention",
|
|
"full_attention",
|
|
"sliding_attention",
|
|
"sliding_attention",
|
|
"sliding_attention",
|
|
]
|
|
assert _period_from_layer_types(lt) is None
|
|
|
|
def test_hf_repo_from_url(self):
|
|
from core.inference.llama_cpp import _hf_repo_from_url
|
|
|
|
assert (
|
|
_hf_repo_from_url("https://huggingface.co/google/gemma-3-1b-it")
|
|
== "google/gemma-3-1b-it"
|
|
)
|
|
assert (
|
|
_hf_repo_from_url("https://huggingface.co/google/gemma-3-1b-it/blob/main/config.json")
|
|
== "google/gemma-3-1b-it"
|
|
)
|
|
for bad in [
|
|
"https://huggingface.co/google",
|
|
"https://example.com/foo/bar",
|
|
None,
|
|
"",
|
|
]:
|
|
assert _hf_repo_from_url(bad) is None
|
|
|
|
def test_bootstrap_tier_used_when_no_cache(self, monkeypatch, tmp_path):
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
from core.inference import llama_cpp as lc
|
|
|
|
def boom(*a, **kw):
|
|
raise AssertionError("HF fetch must not run when bootstrap covers the arch")
|
|
|
|
monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", boom)
|
|
b = _backend_from_gguf("gemma3", dict(_SWA_FIELDS, block_count = 18))
|
|
assert b._sliding_window_pattern == [(i + 1) % 6 != 0 for i in range(18)]
|
|
|
|
def test_disk_cache_takes_precedence_over_bootstrap(self, monkeypatch, tmp_path):
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
# Cached period=3 overrides bootstrap=6.
|
|
with open(tmp_path / "swa_cache.json", "w") as f:
|
|
json.dump({"gemma3": 3}, f)
|
|
b = _backend_from_gguf("gemma3", dict(_SWA_FIELDS, block_count = 18))
|
|
assert b._sliding_window_pattern == [(i + 1) % 3 != 0 for i in range(18)]
|
|
|
|
def test_disk_cache_supports_array_entries(self, monkeypatch, tmp_path):
|
|
# Aperiodic mask is tiled across n_layers.
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
mask = [True, False, True, True, False, True, False, False]
|
|
with open(tmp_path / "swa_cache.json", "w") as f:
|
|
json.dump({"customarch": mask}, f)
|
|
b = _backend_from_gguf("customarch", dict(_SWA_FIELDS, block_count = 16))
|
|
assert b._sliding_window_pattern == [bool(mask[i % 8]) for i in range(16)]
|
|
|
|
def test_hf_fetch_populates_cache(self, monkeypatch, tmp_path):
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
from core.inference import llama_cpp as lc
|
|
|
|
calls = []
|
|
|
|
def fake_fetch(repo_id):
|
|
calls.append(repo_id)
|
|
return 4 if repo_id == "vendor/newmodel-1b-instruct" else None
|
|
|
|
monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", fake_fetch)
|
|
b = _backend_from_gguf(
|
|
"newmodel",
|
|
_SWA_FIELDS,
|
|
general = {"general.source.huggingface.repository": "vendor/newmodel-1b-instruct"},
|
|
)
|
|
assert b._sliding_window_pattern == [(i + 1) % 4 != 0 for i in range(12)]
|
|
assert calls == ["vendor/newmodel-1b-instruct"]
|
|
with open(tmp_path / "swa_cache.json") as f:
|
|
assert json.load(f) == {"newmodel": 4}
|
|
|
|
def test_hf_fetch_falls_back_to_other_candidates(self, monkeypatch, tmp_path):
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
from core.inference import llama_cpp as lc
|
|
|
|
monkeypatch.setattr(
|
|
lc,
|
|
"_fetch_swa_entry_from_hf",
|
|
lambda r: 6 if r == "vendor/newmodel-base" else None,
|
|
)
|
|
b = _backend_from_gguf(
|
|
"newmodel",
|
|
_SWA_FIELDS,
|
|
general = {
|
|
"general.base_model.0.repo_url": "https://huggingface.co/vendor/newmodel-base"
|
|
},
|
|
)
|
|
assert b._sliding_window_pattern == [(i + 1) % 6 != 0 for i in range(12)]
|
|
|
|
def test_offline_env_skips_network(self, monkeypatch, tmp_path):
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
monkeypatch.setenv("UNSLOTH_STUDIO_OFFLINE", "1")
|
|
from core.inference import llama_cpp as lc
|
|
|
|
def boom(*a, **kw):
|
|
raise AssertionError("HF fetch must not run when offline=1")
|
|
|
|
monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", boom)
|
|
b = _backend_from_gguf(
|
|
"newmodel",
|
|
_SWA_FIELDS,
|
|
general = {"general.source.huggingface.repository": "vendor/newmodel"},
|
|
)
|
|
assert b._sliding_window_pattern is None
|
|
|
|
def test_hf_fetch_failure_falls_through_silently(self, monkeypatch, tmp_path):
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
from core.inference import llama_cpp as lc
|
|
|
|
monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", lambda repo_id: None)
|
|
# Force failure into Tier 3; bypass Tier 2.5.
|
|
monkeypatch.setattr(lc, "_resolve_swa_entry_from_transformers", lambda arch: None)
|
|
b = _backend_from_gguf(
|
|
"newmodel",
|
|
_SWA_FIELDS,
|
|
general = {"general.source.huggingface.repository": "vendor/does-not-exist"},
|
|
)
|
|
assert b._sliding_window_pattern is None
|
|
assert not (tmp_path / "swa_cache.json").exists()
|
|
|
|
|
|
class TestTransformersIntrospection:
|
|
"""Tier 2.5: default-init the matching Config; on failure, parse via inspect."""
|
|
|
|
def _isolate_cache(self, monkeypatch, tmp_path):
|
|
from core.inference import llama_cpp as lc
|
|
|
|
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
|
|
monkeypatch.setattr(lc, "_SWA_CACHE", None)
|
|
return tmp_path
|
|
|
|
def test_arch_aliases_normalises_hyphen_underscore(self):
|
|
from core.inference.llama_cpp import _arch_aliases
|
|
|
|
aliases = _arch_aliases("falcon-h1")
|
|
assert aliases[0] == "falcon-h1" and "falcon_h1" in aliases
|
|
assert _arch_aliases("gemma3") == ("gemma3",)
|
|
assert _arch_aliases("") == ()
|
|
|
|
def test_resolves_real_transformers_arches(self):
|
|
from core.inference.llama_cpp import _resolve_swa_entry_from_transformers
|
|
|
|
assert _resolve_swa_entry_from_transformers("gemma3") == 6
|
|
assert _resolve_swa_entry_from_transformers("gemma2") == 2
|
|
assert _resolve_swa_entry_from_transformers("cohere2") == 4
|
|
|
|
def test_falls_back_to_inspect_when_default_init_raises(self, monkeypatch):
|
|
from core.inference import llama_cpp as lc
|
|
|
|
class _FakeBrokenConfig:
|
|
"""Class with sliding_window_pattern: int = 7 in its docstring."""
|
|
|
|
def __init__(self, required_arg):
|
|
raise TypeError("requires an argument")
|
|
|
|
class _FakeLazyMapping(dict):
|
|
def __getitem__(self, k):
|
|
return _FakeBrokenConfig if k == "brokenarch" else super().__getitem__(k)
|
|
|
|
import sys, types as _types
|
|
|
|
fake_auto = _types.ModuleType("transformers.models.auto.configuration_auto")
|
|
fake_auto.CONFIG_MAPPING_NAMES = {"brokenarch": "FakeBroken"}
|
|
fake_auto.CONFIG_MAPPING = _FakeLazyMapping({"brokenarch": "FakeBroken"})
|
|
monkeypatch.setitem(sys.modules, "transformers.models.auto.configuration_auto", fake_auto)
|
|
assert lc._resolve_swa_entry_from_transformers("brokenarch") == 7
|
|
|
|
def test_returns_none_when_transformers_unavailable(self, monkeypatch):
|
|
from core.inference import llama_cpp as lc
|
|
import sys
|
|
|
|
orig_import = (
|
|
__builtins__["__import__"]
|
|
if isinstance(__builtins__, dict)
|
|
else __builtins__.__import__
|
|
)
|
|
|
|
def fake_import(name, *a, **kw):
|
|
if name.startswith("transformers"):
|
|
raise ImportError("transformers not installed")
|
|
return orig_import(name, *a, **kw)
|
|
|
|
monkeypatch.setattr("builtins.__import__", fake_import)
|
|
for k in list(sys.modules):
|
|
if k.startswith("transformers"):
|
|
monkeypatch.delitem(sys.modules, k, raising = False)
|
|
assert lc._resolve_swa_entry_from_transformers("gemma3") is None
|
|
|
|
def test_returns_none_for_arch_unknown_to_transformers(self):
|
|
from core.inference.llama_cpp import _resolve_swa_entry_from_transformers
|
|
assert _resolve_swa_entry_from_transformers("totally-fake-arch-xyz") is None
|
|
|
|
def test_full_resolver_uses_transformers_before_hf_fetch(self, monkeypatch, tmp_path):
|
|
# Bootstrap empty: Tier 2.5 must answer before Tier 3 fires.
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
from core.inference import llama_cpp as lc
|
|
|
|
monkeypatch.setattr(lc, "_BOOTSTRAP_SWA_DEFAULTS", {})
|
|
|
|
def boom(repo_id):
|
|
raise AssertionError("Tier 3 must not run when Tier 2.5 has the answer")
|
|
|
|
monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", boom)
|
|
b = _backend_from_gguf(
|
|
"gemma3",
|
|
dict(_SWA_FIELDS, block_count = 18),
|
|
general = {"general.source.huggingface.repository": "google/gemma-3-1b-it"},
|
|
)
|
|
assert b._sliding_window_pattern == [(i + 1) % 6 != 0 for i in range(18)]
|
|
with open(tmp_path / "swa_cache.json") as f:
|
|
assert json.load(f) == {"gemma3": 6}
|
|
|
|
|
|
class TestGGUFParserReset:
|
|
"""Fields are reset between parses."""
|
|
|
|
def test_reset_between_parses(self):
|
|
# First parse: all fields set
|
|
b = _backend_from_gguf(
|
|
"arch1",
|
|
{
|
|
"block_count": 32,
|
|
"attention.key_length": 128,
|
|
"attention.kv_lora_rank": 512,
|
|
"attention.head_count_kv": [8, 2],
|
|
"attention.sliding_window_pattern": [True, False],
|
|
"attention.key_length_swa": 64,
|
|
"attention.value_length_swa": 64,
|
|
"ssm.inner_size": 4096,
|
|
},
|
|
)
|
|
assert b._kv_key_length == 128
|
|
assert b._kv_lora_rank == 512
|
|
assert b._n_kv_heads_by_layer == [8, 2]
|
|
assert b._sliding_window_pattern == [True, False]
|
|
assert b._kv_key_length_swa == 64
|
|
assert b._kv_value_length_swa == 64
|
|
assert b._ssm_inner_size == 4096
|
|
|
|
# Second parse without those fields -- they must 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._n_kv_heads_by_layer is None
|
|
assert b._sliding_window_pattern is None
|
|
assert b._kv_key_length_swa is None
|
|
assert b._kv_value_length_swa is None
|
|
assert b._ssm_inner_size is None
|
|
assert b._n_layers == 64
|
|
|
|
|
|
# B. _can_estimate_kv Gate Tests
|
|
|
|
|
|
class TestCanEstimateKV:
|
|
"""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 is NOT 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 must NOT add value_length -- V is reconstructed from the latent."""
|
|
b = self._mla_backend()
|
|
result = b._estimate_kv_cache_bytes(1000, "f16")
|
|
# 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):
|
|
"""No key_length: fall back to kv_lora_rank + key_length_mla."""
|
|
b = self._mla_backend(_kv_key_length = None)
|
|
# default _key_length_mla=192, 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):
|
|
"""No key_length and no key_length_mla: fall back 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 uses 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")
|
|
# Uses 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):
|
|
"""Fall back 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 must not 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
|
|
# SWA cache is double-buffered: 2 * sliding_window cells, capped at n_ctx.
|
|
swa_cells = min(131072, 2 * 1024)
|
|
expected = int(n_global * 131072 * kv_per + n_swa * swa_cells * kv_per)
|
|
assert b._estimate_kv_cache_bytes(131072, "f16") == expected
|
|
|
|
def test_gpt_oss(self):
|
|
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
|
|
swa_cells = min(131072, 2 * 128)
|
|
expected = int(n_global * 131072 * kv_per + n_swa * swa_cells * kv_per)
|
|
assert b._estimate_kv_cache_bytes(131072, "f16") == expected
|
|
|
|
def test_gemma4_per_layer_swa_metadata(self):
|
|
b = self._swa_backend(
|
|
_n_layers = 30,
|
|
_n_kv_heads = None,
|
|
_n_kv_heads_by_layer = [8, 8, 8, 8, 8, 2] * 5,
|
|
_n_heads = 16,
|
|
_embedding_length = 2816,
|
|
_kv_key_length = 512,
|
|
_kv_value_length = 512,
|
|
_sliding_window = 1024,
|
|
_sliding_window_pattern = [True, True, True, True, True, False] * 5,
|
|
_kv_key_length_swa = 256,
|
|
_kv_value_length_swa = 256,
|
|
)
|
|
|
|
full_layers = 5
|
|
sliding_layers = 25
|
|
|
|
def expected(ctx):
|
|
full = full_layers * ctx * 2 * (512 + 512) * 2
|
|
sliding = sliding_layers * min(ctx, 2 * 1024) * 8 * (256 + 256) * 2
|
|
return int(full + sliding)
|
|
|
|
for ctx in (4096, 46500, 262144):
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected(ctx)
|
|
|
|
def test_ctx_smaller_than_window(self):
|
|
"""When ctx < 2 * sliding_window, SWA cache caps at ctx."""
|
|
b = self._swa_backend(_sliding_window = 8192)
|
|
n_global = max(1, 62 // 4) # 15
|
|
n_swa = 62 - n_global # 47
|
|
kv_per = 16 * (128 + 128) * 2
|
|
ctx = 4096
|
|
expected = int(n_global * ctx * kv_per + n_swa * min(ctx, 2 * 8192) * kv_per)
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
|
|
|
|
def test_odd_layer_count(self):
|
|
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, 2 * 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 differs 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: 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):
|
|
"""Legacy path matches 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 wins even with other fields 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 chosen params, each path yields a distinct value."""
|
|
# embedding_length=768 so legacy head_dim (768//16=48) != 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:
|
|
"""All supported cache_type_kv values scale correctly."""
|
|
|
|
@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
|
|
|
|
|
|
# J2. Server-flag knobs (--swa-full, --kv-unified/--parallel,
|
|
# --ctx-checkpoints, --kv-offload)
|
|
|
|
|
|
class TestServerFlags:
|
|
"""Estimator should mirror llama-server CLI flags that change KV size."""
|
|
|
|
def _swa_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 26,
|
|
"_n_kv_heads": 4,
|
|
"_n_heads": 8,
|
|
"_embedding_length": 1152,
|
|
"_kv_key_length": 256,
|
|
"_kv_value_length": 256,
|
|
"_sliding_window": 512,
|
|
"_sliding_window_pattern": [True, True, True, True, True, False] * 4 + [True, True],
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
def _gqa_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 28,
|
|
"_n_kv_heads": 8,
|
|
"_n_heads": 16,
|
|
"_embedding_length": 1024,
|
|
"_kv_key_length": 128,
|
|
"_kv_value_length": 128,
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
# ── --swa-full ──────────────────────────────────────────────────
|
|
|
|
def test_swa_full_collapses_pattern_path_to_full_ctx(self):
|
|
b = self._swa_backend()
|
|
ctx = 32_768
|
|
flagged = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
|
|
# swa_full: every layer caches n_ctx -- equals path 4 sizing.
|
|
kv_per_token = 4 * (256 + 256) * 2 # n_kv_heads * (k+v) * f16
|
|
expected = 26 * ctx * kv_per_token
|
|
assert flagged == expected
|
|
assert flagged > b._estimate_kv_cache_bytes(ctx, "f16")
|
|
|
|
def test_swa_full_collapses_legacy_path_to_full_ctx(self):
|
|
# No per-layer pattern -> 1/4-global heuristic; swa_full overrides.
|
|
b = self._swa_backend(_sliding_window_pattern = None)
|
|
ctx = 16_384
|
|
flagged = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
|
|
n_global = max(1, 26 // 4)
|
|
n_swa = 26 - n_global
|
|
kv_per = 4 * (256 + 256) * 2
|
|
# swa_cells == n_ctx when swa_full=True
|
|
expected = n_global * ctx * kv_per + n_swa * ctx * kv_per
|
|
assert flagged == expected
|
|
|
|
def test_swa_full_no_op_for_non_swa_model(self):
|
|
b = self._gqa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
flagged = b._estimate_kv_cache_bytes(8192, "f16", swa_full = True)
|
|
assert flagged == baseline
|
|
|
|
def test_swa_full_suppresses_checkpoint_term(self):
|
|
b = self._swa_backend()
|
|
with_cp = b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 8)
|
|
with_cp_full = b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 8, swa_full = True)
|
|
no_cp_full = b._estimate_kv_cache_bytes(8192, "f16", swa_full = True)
|
|
# Checkpoints only matter when SWA layers don't already keep n_ctx.
|
|
assert with_cp_full == no_cp_full
|
|
assert with_cp > b._estimate_kv_cache_bytes(8192, "f16")
|
|
|
|
# ── --parallel + --kv-unified ──────────────────────────────────
|
|
# Verified against llama-server: non-SWA caches partition n_ctx across
|
|
# slots (total memory constant); only SWA layers scale with --parallel.
|
|
# --kv-unified is a no-op for memory math (kept for API forward-compat).
|
|
|
|
def test_gqa_kv_constant_across_parallel(self):
|
|
b = self._gqa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(4096, "f16")
|
|
for slots in (1, 2, 4, 8):
|
|
for unified in (True, False):
|
|
assert (
|
|
b._estimate_kv_cache_bytes(4096, "f16", n_parallel = slots, kv_unified = unified)
|
|
== baseline
|
|
)
|
|
|
|
def test_zero_parallel_floors_at_one(self):
|
|
b = self._gqa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(4096, "f16")
|
|
for unified in (True, False):
|
|
assert (
|
|
b._estimate_kv_cache_bytes(4096, "f16", n_parallel = 0, kv_unified = unified)
|
|
== baseline
|
|
)
|
|
|
|
def test_swa_path_scales_only_swa_portion(self):
|
|
b = self._swa_backend()
|
|
ctx = 8192
|
|
baseline = b._estimate_kv_cache_bytes(ctx, "f16")
|
|
# Decompose baseline by walking the estimator's own loop.
|
|
swa = b._sliding_window
|
|
per_token_global = 4 * (256 + 256) * 2 # n_kv * (k+v) * f16
|
|
per_token_swa = 4 * (256 + 256) * 2 # k_swa/val_swa fall back
|
|
per_slot_swa_cells = min(ctx, 2 * swa) # not clamped at parallel=1
|
|
global_bytes = sum(
|
|
ctx * per_token_global for f in b._sliding_window_pattern[: b._n_layers] if not f
|
|
)
|
|
swa_bytes_per_slot = sum(
|
|
per_slot_swa_cells * per_token_swa
|
|
for f in b._sliding_window_pattern[: b._n_layers]
|
|
if f
|
|
)
|
|
# Sanity: parallel=1 reproduces baseline exactly
|
|
assert global_bytes + swa_bytes_per_slot == baseline
|
|
# Only the SWA portion scales by parallel
|
|
for slots in (1, 2, 3, 4):
|
|
scaled = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False)
|
|
# SWA cells clamp to per_slot_ctx when ctx/slots < 2*swa
|
|
per_slot_ctx = max(1, ctx // slots)
|
|
cells = min(ctx, 2 * swa, per_slot_ctx)
|
|
swa_bps = sum(
|
|
cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f
|
|
)
|
|
assert scaled == global_bytes + slots * swa_bps
|
|
|
|
def test_mla_kv_constant_across_parallel(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 60
|
|
b._n_kv_heads = 1
|
|
b._kv_lora_rank = 512
|
|
b._key_length_mla = 64
|
|
b._kv_key_length = 576
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
for slots in (1, 2, 4, 8):
|
|
for unified in (True, False):
|
|
assert (
|
|
b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = unified)
|
|
== baseline
|
|
)
|
|
|
|
# ── --ctx-checkpoints ──────────────────────────────────────────
|
|
|
|
def test_ctx_checkpoints_zero_is_no_op(self):
|
|
b = self._swa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
assert b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 0) == baseline
|
|
|
|
def test_ctx_checkpoints_no_op_for_non_swa(self):
|
|
b = self._gqa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
assert b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 32) == baseline
|
|
|
|
def test_ctx_checkpoints_pattern_path_adds_known_bytes(self):
|
|
b = self._swa_backend()
|
|
ctx = 8192
|
|
baseline = b._estimate_kv_cache_bytes(ctx, "f16")
|
|
flagged = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4)
|
|
# 22 SWA layers * 4 cps * 512 cells * 4 heads * (256+256) * 2 bytes
|
|
n_swa_layers = sum(1 for f in [True, True, True, True, True, False] * 4 + [True, True] if f)
|
|
per_layer = 4 * 512 * 4 * (256 + 256) * 2
|
|
assert flagged == baseline + n_swa_layers * per_layer
|
|
|
|
def test_ctx_checkpoints_legacy_path_adds_known_bytes(self):
|
|
b = self._swa_backend(_sliding_window_pattern = None)
|
|
ctx = 8192
|
|
baseline = b._estimate_kv_cache_bytes(ctx, "f16")
|
|
flagged = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4)
|
|
n_global = max(1, 26 // 4)
|
|
n_swa = 26 - n_global
|
|
kv_per = 4 * (256 + 256) * 2
|
|
extra = 4 * n_swa * 512 * kv_per # ctx_checkpoints * n_swa * sliding * kv_per
|
|
assert flagged == baseline + extra
|
|
|
|
def test_ctx_checkpoints_compose_with_n_parallel(self):
|
|
# Only the SWA + checkpoint portion scales by n_parallel; the
|
|
# global-layer portion is constant.
|
|
b = self._swa_backend()
|
|
ctx = 8192
|
|
swa = b._sliding_window
|
|
per_token = 4 * (256 + 256) * 2
|
|
global_bytes = sum(
|
|
ctx * per_token for f in b._sliding_window_pattern[: b._n_layers] if not f
|
|
)
|
|
n_swa_layers = sum(1 for f in b._sliding_window_pattern[: b._n_layers] if f)
|
|
slots = 3
|
|
per_slot_ctx = max(1, ctx // slots)
|
|
swa_cells = min(ctx, 2 * swa, per_slot_ctx)
|
|
swa_bytes_per_slot = n_swa_layers * swa_cells * per_token
|
|
cp_extra_per_slot = n_swa_layers * 4 * swa * per_token # 4 checkpoints
|
|
flagged = b._estimate_kv_cache_bytes(
|
|
ctx, "f16", ctx_checkpoints = 4, n_parallel = slots, kv_unified = False
|
|
)
|
|
assert flagged == global_bytes + slots * (swa_bytes_per_slot + cp_extra_per_slot)
|
|
|
|
# ── --kv-offload (kv_on_gpu) ───────────────────────────────────
|
|
|
|
def test_fit_returns_requested_when_kv_off_gpu(self):
|
|
b = self._gqa_backend()
|
|
# Tiny VRAM budget -- would normally force a reduction.
|
|
fitted = b._fit_context_to_vram(
|
|
requested_ctx = 32_768,
|
|
available_mib = 1,
|
|
model_size_bytes = 100,
|
|
cache_type_kv = "f16",
|
|
kv_on_gpu = False,
|
|
)
|
|
assert fitted == 32_768
|
|
|
|
def test_fit_reduces_when_kv_on_gpu(self):
|
|
b = self._gqa_backend()
|
|
fitted = b._fit_context_to_vram(
|
|
requested_ctx = 32_768,
|
|
available_mib = 64,
|
|
model_size_bytes = 1024 * 1024, # 1 MiB
|
|
cache_type_kv = "f16",
|
|
kv_on_gpu = True,
|
|
)
|
|
assert fitted < 32_768
|
|
|
|
def test_fit_mtp_engaged_returns_smaller_or_equal_context(self):
|
|
# Flat MTP fallback budget is _CTX_FIT_VRAM_FRACTION - 0.05; non-MTP is
|
|
# the full fraction. On a tight budget MTP must yield <= non-MTP.
|
|
b = self._gqa_backend()
|
|
common = dict(
|
|
requested_ctx = 32_768,
|
|
available_mib = 128,
|
|
model_size_bytes = 8 * 1024 * 1024,
|
|
cache_type_kv = "f16",
|
|
)
|
|
baseline = b._fit_context_to_vram(**common)
|
|
mtp = b._fit_context_to_vram(**common, mtp_engaged = True)
|
|
assert mtp <= baseline
|
|
|
|
def test_fit_mtp_engaged_unchanged_when_kv_off_gpu(self):
|
|
# kv_on_gpu=False short-circuits the fit; mtp_engaged irrelevant.
|
|
b = self._gqa_backend()
|
|
fitted = b._fit_context_to_vram(
|
|
requested_ctx = 32_768,
|
|
available_mib = 1,
|
|
model_size_bytes = 100,
|
|
cache_type_kv = "f16",
|
|
kv_on_gpu = False,
|
|
mtp_engaged = True,
|
|
)
|
|
assert fitted == 32_768
|
|
|
|
def test_fit_threads_swa_full_through_estimator(self):
|
|
# SWA model, generous budget; both should fit but cache size differs.
|
|
b = self._swa_backend()
|
|
ctx = 8192
|
|
kv_default = b._estimate_kv_cache_bytes(ctx, "f16")
|
|
kv_full = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
|
|
assert kv_full > kv_default
|
|
# Budget = model + kv_default (rounded up) -- swa_full must not fit.
|
|
budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / _CTX_FIT_VRAM_FRACTION + 1
|
|
fitted_default = b._fit_context_to_vram(
|
|
requested_ctx = ctx,
|
|
available_mib = int(budget_mib),
|
|
model_size_bytes = 1024 * 1024,
|
|
cache_type_kv = "f16",
|
|
)
|
|
fitted_full = b._fit_context_to_vram(
|
|
requested_ctx = ctx,
|
|
available_mib = int(budget_mib),
|
|
model_size_bytes = 1024 * 1024,
|
|
cache_type_kv = "f16",
|
|
swa_full = True,
|
|
)
|
|
assert fitted_default == ctx
|
|
assert fitted_full < ctx
|
|
|
|
|
|
# J2.5. --parallel N memory accounting (per-layer-type scaling rule)
|
|
|
|
|
|
class TestParallelSWAScaling:
|
|
"""Per-layer-type scaling rule vs the closed form measured from
|
|
llama-server. Empirical formula on Gemma-3 270m at ctx=8192:
|
|
total_kv = 24 + parallel * 15 (MiB).
|
|
|
|
Rule (verified vs ``llama-server`` log on real GGUFs):
|
|
* non-SWA layers: total cells = n_ctx, partitioned across slots,
|
|
memory CONSTANT in n_parallel.
|
|
* SWA layers: per-slot cells = 2 * sliding_window (clamped at
|
|
n_ctx and at per_slot_ctx); memory LINEAR in n_parallel.
|
|
* --kv-unified is a no-op for memory math; both modes give the
|
|
same total in measured cases.
|
|
"""
|
|
|
|
def _gqa_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 28,
|
|
"_n_kv_heads": 8,
|
|
"_n_heads": 16,
|
|
"_embedding_length": 1024,
|
|
"_kv_key_length": 128,
|
|
"_kv_value_length": 128,
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
def _swa_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 18,
|
|
"_n_kv_heads": 1,
|
|
"_n_heads": 4,
|
|
"_embedding_length": 1024,
|
|
"_kv_key_length": 256,
|
|
"_kv_value_length": 256,
|
|
"_sliding_window": 512,
|
|
# 15 SWA + 3 global, mirrors gemma-3-270m
|
|
"_sliding_window_pattern": [t == "swa" for t in (["swa"] * 5 + ["global"]) * 3],
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
# ── non-SWA paths: constant ────────────────────────────────────
|
|
|
|
def test_pure_gqa_constant_across_parallel(self):
|
|
b = self._gqa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
for slots in (1, 2, 4, 8):
|
|
for unified in (True, False):
|
|
assert (
|
|
b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = unified)
|
|
== baseline
|
|
)
|
|
|
|
def test_mla_constant_across_parallel(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 60
|
|
b._n_kv_heads = 1
|
|
b._kv_lora_rank = 512
|
|
b._key_length_mla = 64
|
|
b._kv_key_length = 576
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
for slots in (1, 2, 4, 8):
|
|
assert b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots) == baseline
|
|
|
|
def test_hybrid_constant_across_parallel(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 64
|
|
b._n_kv_heads = 16
|
|
b._n_heads = 32
|
|
b._embedding_length = 4096
|
|
b._kv_key_length = 128
|
|
b._kv_value_length = 128
|
|
b._ssm_inner_size = 4096
|
|
b._full_attention_interval = 4
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
for slots in (1, 2, 4, 8):
|
|
assert b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots) == baseline
|
|
|
|
def test_legacy_constant_across_parallel(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 32
|
|
b._n_kv_heads = 8
|
|
b._n_heads = 8
|
|
b._embedding_length = 4096
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
for slots in (1, 2, 4, 8):
|
|
assert b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots) == baseline
|
|
|
|
# ── SWA paths: scale only the SWA portion ──────────────────────
|
|
|
|
def test_swa_pattern_scales_only_swa_portion(self):
|
|
b = self._swa_backend()
|
|
ctx = 8192
|
|
swa = b._sliding_window
|
|
per_token = 1 * (256 + 256) * 2 # n_kv * (k+v) * f16
|
|
n_global = sum(1 for f in b._sliding_window_pattern if not f)
|
|
n_swa = sum(1 for f in b._sliding_window_pattern if f)
|
|
global_bytes = n_global * ctx * per_token
|
|
for slots in (1, 2, 4, 8):
|
|
per_slot_ctx = max(1, ctx // slots)
|
|
cells = min(ctx, 2 * swa, per_slot_ctx)
|
|
swa_bps = n_swa * cells * per_token
|
|
for unified in (True, False):
|
|
got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = unified)
|
|
assert got == global_bytes + slots * swa_bps
|
|
|
|
def test_swa_fallback_scales_only_swa_portion(self):
|
|
# No per-layer pattern -> 1/4-global heuristic.
|
|
b = self._swa_backend(_sliding_window_pattern = None)
|
|
ctx = 8192
|
|
swa = b._sliding_window
|
|
n_layers = 18
|
|
n_global = max(1, n_layers // 4)
|
|
n_swa = n_layers - n_global
|
|
per_token = 1 * (256 + 256) * 2
|
|
global_bytes = n_global * ctx * per_token
|
|
for slots in (1, 2, 4, 8):
|
|
per_slot_ctx = max(1, ctx // slots)
|
|
cells = min(ctx, 2 * swa, per_slot_ctx)
|
|
swa_bps = n_swa * cells * per_token
|
|
got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots)
|
|
assert got == global_bytes + slots * swa_bps
|
|
|
|
def test_swa_per_slot_clamped_when_ctx_lt_slots_x_2window(self):
|
|
# ctx=4096 / slots=8 -> per_slot_ctx=512, but 2*sliding=1024.
|
|
# SWA cells clamp at per_slot_ctx (512), not 2*sliding.
|
|
b = self._swa_backend()
|
|
ctx = 4096
|
|
per_slot_ctx_at_8 = ctx // 8
|
|
assert per_slot_ctx_at_8 < 2 * b._sliding_window
|
|
# Build expected with the clamped formula
|
|
n_swa = sum(1 for f in b._sliding_window_pattern if f)
|
|
n_global = sum(1 for f in b._sliding_window_pattern if not f)
|
|
per_token = 1 * (256 + 256) * 2
|
|
global_bytes = n_global * ctx * per_token
|
|
cells = min(ctx, 2 * b._sliding_window, per_slot_ctx_at_8)
|
|
assert cells == per_slot_ctx_at_8
|
|
expected = global_bytes + 8 * (n_swa * cells * per_token)
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = 8) == expected
|
|
|
|
def test_swa_full_does_not_scale_under_parallel(self):
|
|
# swa_full forces every layer to n_ctx -> all-global GQA-style
|
|
# total, constant in parallel.
|
|
b = self._swa_backend()
|
|
ctx = 8192
|
|
baseline = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
|
|
for slots in (1, 2, 4, 8):
|
|
assert (
|
|
b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True, n_parallel = slots) == baseline
|
|
)
|
|
|
|
# ── kv_unified: no-op for memory math ──────────────────────────
|
|
|
|
def test_kv_unified_is_no_op_for_memory_math(self):
|
|
# unified=True and unified=False must give the same total bytes
|
|
# for every backend type and parallel value.
|
|
backends = [
|
|
("gqa", self._gqa_backend()),
|
|
("swa", self._swa_backend()),
|
|
]
|
|
for label, b in backends:
|
|
for slots in (1, 2, 4, 8):
|
|
u = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = True)
|
|
nu = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = False)
|
|
assert u == nu, f"{label} parallel={slots} unified-mismatch"
|
|
|
|
# ── Empirical Gemma-3 270m formula ─────────────────────────────
|
|
|
|
def test_matches_empirical_gemma3_270m_formula(self):
|
|
"""Exact match against the formula measured from llama-server:
|
|
total_kv = 24 + parallel * 15 (MiB) at ctx=8192.
|
|
|
|
Geometry: 18 layers (3 global + 15 SWA), n_kv=1, head_dim=256,
|
|
sliding=512, f16.
|
|
"""
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 18
|
|
b._n_kv_heads = 1
|
|
b._n_heads = 4
|
|
b._embedding_length = 1024
|
|
b._kv_key_length = 256
|
|
b._kv_value_length = 256
|
|
b._sliding_window = 512
|
|
# Mirrors the bootstrap-resolved gemma3 pattern (period 6) on an
|
|
# 18-layer model: 15 SWA, 3 global.
|
|
b._sliding_window_pattern = [(i + 1) % 6 != 0 for i in range(18)]
|
|
n_global = 3
|
|
n_swa = 15
|
|
# Confirm pattern shape
|
|
assert sum(b._sliding_window_pattern) == n_swa
|
|
for slots, expected_mib in [(1, 39), (2, 54), (4, 84)]:
|
|
got_bytes = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots)
|
|
got_mib = got_bytes / (1024 * 1024)
|
|
assert (
|
|
got_mib == expected_mib
|
|
), f"slots={slots}: got {got_mib} MiB, expected {expected_mib} MiB"
|
|
|
|
|
|
# J3. shared_kv_layers (Gemma 3n / Gemma 4)
|
|
|
|
|
|
class TestSharedKVLayers:
|
|
"""``<arch>.attention.shared_kv_layers`` reduces the layer count that
|
|
allocates KV. The trailing ``shared_kv_layers`` blocks reuse earlier
|
|
caches (Gemma 3n: 35 layers, 15 shared -> 20 allocate; Gemma 4 same
|
|
field). Unset on every other arch -> no behavioural change."""
|
|
|
|
def _gemma3n_backend(self, **overrides):
|
|
# Mirrors google/gemma-3n-E4B-it: 35 layers, 15 shared, SWA window
|
|
# 1024, period 5 (4 sliding + 1 full repeating).
|
|
defaults = {
|
|
"_n_layers": 35,
|
|
"_n_kv_heads": 4,
|
|
"_n_heads": 8,
|
|
"_embedding_length": 2048,
|
|
"_kv_key_length": 256,
|
|
"_kv_value_length": 256,
|
|
"_sliding_window": 1024,
|
|
"_sliding_window_pattern": [
|
|
t == "sliding_attention"
|
|
for t in (["sliding_attention"] * 4 + ["full_attention"]) * 7
|
|
],
|
|
"_shared_kv_layers": 15,
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
def _gqa_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 28,
|
|
"_n_kv_heads": 8,
|
|
"_n_heads": 16,
|
|
"_embedding_length": 1024,
|
|
"_kv_key_length": 128,
|
|
"_kv_value_length": 128,
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
def test_field_initialises_to_none(self):
|
|
b = LlamaCppBackend()
|
|
assert b._shared_kv_layers is None
|
|
|
|
def test_unset_field_is_noop(self):
|
|
b = self._gqa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
b._shared_kv_layers = None
|
|
assert b._estimate_kv_cache_bytes(8192, "f16") == baseline
|
|
b._shared_kv_layers = 0
|
|
assert b._estimate_kv_cache_bytes(8192, "f16") == baseline
|
|
|
|
def test_path4_drops_shared_layers(self):
|
|
b = self._gqa_backend(_shared_kv_layers = 4)
|
|
ctx = 4096
|
|
kv_per = 8 * (128 + 128) * 2
|
|
# 28 - 4 = 24 layers actually allocate
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == 24 * ctx * kv_per
|
|
|
|
def test_path5_drops_shared_layers(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 32
|
|
b._n_kv_heads = 8
|
|
b._n_heads = 8
|
|
b._embedding_length = 4096
|
|
b._shared_kv_layers = 8
|
|
ctx = 4096
|
|
head_dim = 4096 // 8 # 512
|
|
# 32 - 8 = 24 layers
|
|
expected = 2 * 8 * head_dim * 24 * ctx * 2
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
|
|
|
|
def test_path1_mla_drops_shared_layers(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 60
|
|
b._n_kv_heads = 1
|
|
b._kv_lora_rank = 512
|
|
b._key_length_mla = 64
|
|
b._kv_key_length = 576
|
|
b._shared_kv_layers = 10
|
|
ctx = 8192
|
|
# 60 - 10 = 50
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == 50 * ctx * 1 * 576 * 2
|
|
|
|
def test_path3_pattern_loops_only_unshared_layers(self):
|
|
b = self._gemma3n_backend()
|
|
ctx = 8192
|
|
# First 20 layers contribute; layers 20..34 skipped. Pattern
|
|
# [s,s,s,s,F] repeated -> in layers 0..19: sliding 16, full 4.
|
|
sliding_in_unshared = sum(b._sliding_window_pattern[:20])
|
|
full_in_unshared = 20 - sliding_in_unshared
|
|
assert sliding_in_unshared == 16
|
|
assert full_in_unshared == 4
|
|
kv_per = 4 * (256 + 256) * 2
|
|
swa_cells = min(ctx, 2 * 1024)
|
|
expected = full_in_unshared * ctx * kv_per + sliding_in_unshared * swa_cells * kv_per
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
|
|
|
|
def test_shared_layers_reduces_estimate(self):
|
|
b = self._gemma3n_backend()
|
|
with_shared = b._estimate_kv_cache_bytes(8192, "f16")
|
|
b._shared_kv_layers = 0
|
|
without_shared = b._estimate_kv_cache_bytes(8192, "f16")
|
|
# 20/35 = 0.571 of the work; ~43% reduction.
|
|
ratio = with_shared / without_shared
|
|
assert 0.5 < ratio < 0.65
|
|
|
|
def test_path3_pattern_with_swa_full_and_shared(self):
|
|
b = self._gemma3n_backend()
|
|
ctx = 8192
|
|
flagged = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
|
|
# Every unshared layer caches n_ctx -> path-4-style sizing over
|
|
# only the 20 unshared layers.
|
|
kv_per = 4 * (256 + 256) * 2
|
|
assert flagged == 20 * ctx * kv_per
|
|
|
|
def test_path3_fallback_uses_unshared_count(self):
|
|
# No per-layer pattern -> 1/4-global heuristic over n_layers_kv,
|
|
# not n_layers.
|
|
b = self._gemma3n_backend(_sliding_window_pattern = None)
|
|
ctx = 8192
|
|
n_layers_kv = 35 - 15 # 20
|
|
n_global = max(1, n_layers_kv // 4) # 5
|
|
n_swa = n_layers_kv - n_global # 15
|
|
kv_per = 4 * (256 + 256) * 2
|
|
swa_cells = min(ctx, 2 * 1024)
|
|
expected = n_global * ctx * kv_per + n_swa * swa_cells * kv_per
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
|
|
|
|
def test_shared_floors_at_one_layer(self):
|
|
# Pathological: shared >= n_layers must not zero out the cache.
|
|
b = self._gqa_backend(_shared_kv_layers = 99)
|
|
ctx = 4096
|
|
kv_per = 8 * (128 + 128) * 2
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == 1 * ctx * kv_per
|
|
|
|
def test_composes_with_n_parallel(self):
|
|
# Only the SWA portion of unshared layers scales by n_parallel;
|
|
# the global portion is constant.
|
|
b = self._gemma3n_backend()
|
|
ctx = 8192
|
|
swa = b._sliding_window
|
|
per_token = 4 * (256 + 256) * 2
|
|
unshared_pattern = b._sliding_window_pattern[:20] # 35 - 15 shared
|
|
sliding_in_unshared = sum(unshared_pattern)
|
|
global_in_unshared = len(unshared_pattern) - sliding_in_unshared
|
|
global_bytes = global_in_unshared * ctx * per_token
|
|
slots = 3
|
|
per_slot_ctx = max(1, ctx // slots)
|
|
swa_cells = min(ctx, 2 * swa, per_slot_ctx)
|
|
swa_bytes_per_slot = sliding_in_unshared * swa_cells * per_token
|
|
flagged = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False)
|
|
assert flagged == global_bytes + slots * swa_bytes_per_slot
|
|
|
|
def test_composes_with_ctx_checkpoints(self):
|
|
b = self._gemma3n_backend()
|
|
ctx = 8192
|
|
baseline = b._estimate_kv_cache_bytes(ctx, "f16")
|
|
with_cp = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4)
|
|
# Checkpoints count only over UNSHARED SWA layers (16 of them).
|
|
sliding_in_unshared = sum(b._sliding_window_pattern[:20])
|
|
per_cp_layer = 4 * 1024 * 4 * (256 + 256) * 2 # cps * swa * heads * (k+v) * bpe
|
|
assert with_cp == baseline + sliding_in_unshared * per_cp_layer
|
|
|
|
def test_unload_resets_shared_kv_layers(self):
|
|
b = LlamaCppBackend()
|
|
b._shared_kv_layers = 12
|
|
b.unload_model()
|
|
assert b._shared_kv_layers is None
|
|
|
|
|
|
# K. Lifecycle Tests
|
|
|
|
|
|
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",
|
|
"_sliding_window_pattern",
|
|
"_full_attention_interval",
|
|
"_kv_lora_rank",
|
|
"_key_length_mla",
|
|
"_kv_key_length_swa",
|
|
"_kv_value_length_swa",
|
|
"_ssm_inner_size",
|
|
"_ssm_state_size",
|
|
"_shared_kv_layers",
|
|
]:
|
|
assert getattr(b, attr) is None
|
|
assert b._n_kv_heads_by_layer is None
|
|
|
|
def test_unload_resets_fields(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 32
|
|
b._kv_key_length = 128
|
|
b._kv_lora_rank = 512
|
|
b._sliding_window = 1024
|
|
b._sliding_window_pattern = [True, False]
|
|
b._n_kv_heads_by_layer = [8, 2]
|
|
b._kv_key_length_swa = 64
|
|
b._kv_value_length_swa = 64
|
|
b._ssm_inner_size = 4096
|
|
b._full_attention_interval = 4
|
|
b._shared_kv_layers = 8
|
|
b.unload_model()
|
|
for attr in [
|
|
"_kv_key_length",
|
|
"_kv_value_length",
|
|
"_sliding_window",
|
|
"_sliding_window_pattern",
|
|
"_full_attention_interval",
|
|
"_kv_lora_rank",
|
|
"_key_length_mla",
|
|
"_kv_key_length_swa",
|
|
"_kv_value_length_swa",
|
|
"_ssm_inner_size",
|
|
"_ssm_state_size",
|
|
"_shared_kv_layers",
|
|
]:
|
|
assert getattr(b, attr) is None
|
|
assert b._n_kv_heads_by_layer is None
|
|
|
|
def test_end_to_end_synthetic_mla(self):
|
|
"""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")
|
|
# gemma3 -> period 6 from bootstrap; SWA cache double-buffered to
|
|
# 2 * sliding_window cells.
|
|
period = 6
|
|
kv_per = 16 * 256 * 2
|
|
expected = 0
|
|
for i in range(62):
|
|
is_swa = (i + 1) % period != 0
|
|
layer_ctx = min(131072, 2 * 1024) if is_swa else 131072
|
|
expected += layer_ctx * kv_per
|
|
assert result == expected
|
|
|
|
def test_end_to_end_synthetic_shared_kv_round_trip(self):
|
|
# Mirrors gemma3n_text: 35 layers, 15 shared, sliding_window=1024.
|
|
b = _backend_from_gguf(
|
|
"gemma3n_text",
|
|
{
|
|
"context_length": 32768,
|
|
"block_count": 35,
|
|
"attention.head_count_kv": 4,
|
|
"attention.head_count": 8,
|
|
"embedding_length": 2048,
|
|
"attention.key_length": 256,
|
|
"attention.value_length": 256,
|
|
"attention.sliding_window": 1024,
|
|
"attention.shared_kv_layers": 15,
|
|
},
|
|
)
|
|
assert b._can_estimate_kv()
|
|
assert b._shared_kv_layers == 15
|
|
# Bootstrap for gemma3n_text -> period 5; resolver synthesises a
|
|
# 35-entry bool array. Only the first 20 (n_layers - shared)
|
|
# allocate KV.
|
|
result = b._estimate_kv_cache_bytes(8192, "f16")
|
|
assert result > 0
|
|
# Sanity: shared back to 0 -> strictly larger estimate (more
|
|
# layers allocate).
|
|
b._shared_kv_layers = 0
|
|
unshared = b._estimate_kv_cache_bytes(8192, "f16")
|
|
assert unshared > result
|
|
|
|
def test_end_to_end_synthetic_gqa(self):
|
|
b = _backend_from_gguf(
|
|
"qwen3",
|
|
{
|
|
"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
|