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

studio: thread llama-server KV flags through the estimator

Browse source 
Adds keyword-only knobs to _estimate_kv_cache_bytes and _fit_context_to_vram
that mirror the llama-server CLI options that change KV memory:

  --swa-full         (swa_full)        SWA layers cache the full n_ctx
                                       instead of 2 * sliding_window cells
  --parallel N       (n_parallel)      number of server slots
  --kv-unified       (kv_unified)      single shared KV buffer; when off,
                                       multiplies KV by n_parallel
  --ctx-checkpoints  (ctx_checkpoints) per-slot SWA snapshots, each one
                                       sliding-window of state per SWA layer
  --kv-offload       (kv_on_gpu)       when off, KV lives in CPU RAM and is
                                       not subtracted from the VRAM budget

Defaults preserve the previous behavior (swa_full=False, n_parallel=1,
kv_unified=True, ctx_checkpoints=0, kv_on_gpu=True) so existing call sites
are unaffected. All five paths (MLA, hybrid, SWA pattern, SWA fallback,
GQA, legacy) now apply the per-slot replication factor; the SWA paths
also honor swa_full and add the checkpoint term when applicable.

Tests: TestServerFlags (17 cases) covers every flag, the no-op cases, the
swa_full + ctx_checkpoints interaction, slot multiplication on each path,
and the kv_on_gpu shortcut in _fit_context_to_vram.
This commit is contained in:
Daniel Han 2026-05-03 02:56:36 +00:00
commit 29dcf96eb0
2 changed files with 311 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -994,7 +994,14 @@ class LlamaCppBackend:
return fallback
def _estimate_kv_cache_bytes(
self, n_ctx: int, cache_type_kv: Optional[str] = None
self,
n_ctx: int,
cache_type_kv: Optional[str] = None,
*,
swa_full: bool = False,
n_parallel: int = 1,
kv_unified: bool = True,
ctx_checkpoints: int = 0,
) -> int:
"""Estimate KV cache VRAM for a given context length.
@ -1005,6 +1012,19 @@ class LlamaCppBackend:
4. GQA -- standard full KV with explicit key/value dimensions
5. Legacy -- fallback using embed // n_heads
Server-flag knobs (mirror llama-server's CLI):
swa_full -- ``--swa-full``: force SWA layers to cache the
full ``n_ctx`` (collapses path 3 to path 4
sizing for the SWA layers).
n_parallel -- ``--parallel``: number of server slots.
kv_unified -- ``--kv-unified`` (default on): single shared
KV buffer across slots. When False, allocates
one buffer per slot (multiplies KV by
``n_parallel``).
ctx_checkpoints -- ``--ctx-checkpoints``: SWA snapshot count per
slot (PR #15293). Each snapshot stores one
sliding-window of state per SWA layer.
Returns 0 if metadata is insufficient for estimation.
"""
if not self._can_estimate_kv() or n_ctx <= 0:
@ -1026,6 +1046,9 @@ class LlamaCppBackend:
"iq4_nl": 0.5625,
}.get(cache_type_kv or "f16", 2.0)
# Per-slot replication when slots don't share the unified buffer.
slot_factor = max(1, n_parallel) if not kv_unified else 1
# Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5)
# MLA stores one compressed KV latent per token/layer (shared across heads).
# V is reconstructed from the latent on the fly -- no separate V cache.
@ -1036,7 +1059,7 @@ class LlamaCppBackend:
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)
return int(n_layers * n_ctx * n_kv_mla * key_len * bpe * slot_factor)
key_len = self._kv_key_length
val_len = self._kv_value_length
@ -1050,9 +1073,11 @@ class LlamaCppBackend:
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)
return int(
n_attn * n_ctx * n_kv * (key_len + val_len) * bpe * slot_factor
)
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)
return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe * slot_factor)
# Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...).
# Pattern is filled in by the resolver at parse time; if absent,
@ -1060,8 +1085,11 @@ class LlamaCppBackend:
# llama.cpp double-buffers the SWA cache (so it can keep the
# current and next windows during the shift) -- it allocates
# `2 * sliding_window` cells per SWA layer, capped at n_ctx.
# Verified via `llama_kv_cache_iswa: ... SWA KV cache,
# size = N cells` log line on Gemma-3 270m / 1b GGUFs.
# ``--swa-full`` forces full n_ctx for SWA layers instead.
# ``--ctx-checkpoints N`` adds N snapshots per SWA layer (one
# sliding-window of state each) per slot for context-shift
# recovery; only meaningful when SWA layers don't already cache
# n_ctx.
if (
self._sliding_window is not None
and self._sliding_window > 0
@ -1069,11 +1097,12 @@ class LlamaCppBackend:
and val_len is not None
):
swa = self._sliding_window
swa_cells = min(n_ctx, 2 * swa)
swa_cells = n_ctx if swa_full else min(n_ctx, 2 * swa)
key_len_swa = self._kv_key_length_swa or key_len
val_len_swa = self._kv_value_length_swa or val_len
if self._sliding_window_pattern is not None:
total = 0.0
checkpoint_extra = 0.0
for layer_idx in range(n_layers):
layer_n_kv = self._kv_heads_for_layer(layer_idx, n_kv)
is_swa = (
@ -1086,21 +1115,38 @@ class LlamaCppBackend:
total += (
layer_ctx * layer_n_kv * (layer_key_len + layer_val_len) * bpe
)
return int(total)
if is_swa and ctx_checkpoints > 0 and not swa_full:
checkpoint_extra += (
ctx_checkpoints
* swa
* layer_n_kv
* (layer_key_len + layer_val_len)
* bpe
)
return int((total + checkpoint_extra) * slot_factor)
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 * swa_cells * kv_per_token
kv_per_token_swa = n_kv * (key_len_swa + val_len_swa) * bpe
base = (
n_global * n_ctx * kv_per_token + n_swa * swa_cells * kv_per_token_swa
)
checkpoint_extra = (
ctx_checkpoints * n_swa * swa * kv_per_token_swa
if ctx_checkpoints > 0 and not swa_full
else 0.0
)
return int((base + checkpoint_extra) * slot_factor)
# 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)
return int(
n_layers * n_ctx * n_kv * (key_len + val_len) * bpe * slot_factor
)
# 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)
return int(2 * n_kv * head_dim * n_layers * n_ctx * bpe * slot_factor)
def _fit_context_to_vram(
self,
@ -1109,6 +1155,12 @@ class LlamaCppBackend:
model_size_bytes: int,
cache_type_kv: Optional[str] = None,
min_ctx: int = 4096,
*,
swa_full: bool = False,
n_parallel: int = 1,
kv_unified: bool = True,
ctx_checkpoints: int = 0,
kv_on_gpu: bool = True,
) -> int:
"""Return the largest context length that fits in GPU VRAM.
@ -1116,6 +1168,11 @@ class LlamaCppBackend:
threshold -- 10% reserved for compute buffers, CUDA context,
scratch space, flash-attn workspace, etc.).
If the model weights alone don't fit, returns min_ctx unchanged.
``kv_on_gpu`` mirrors ``--kv-offload`` (default on). When False
the KV cache lives in CPU RAM and doesn't compete with weights
for VRAM; the requested context is honored verbatim. The other
keyword args mirror ``_estimate_kv_cache_bytes``.
"""
if not self._can_estimate_kv():
logger.debug(
@ -1125,11 +1182,22 @@ class LlamaCppBackend:
)
return requested_ctx
# KV lives off-GPU: no VRAM accounting needed for the cache itself.
if not kv_on_gpu:
return requested_ctx
kv_kwargs = dict(
swa_full = swa_full,
n_parallel = n_parallel,
kv_unified = kv_unified,
ctx_checkpoints = ctx_checkpoints,
)
budget_bytes = available_mib * 1024 * 1024 * 0.90
model_footprint = model_size_bytes
# Check if requested context already fits
kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv)
kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs)
if model_footprint + kv <= budget_bytes:
return requested_ctx
@ -1151,7 +1219,7 @@ class LlamaCppBackend:
best = effective_min
while lo <= hi:
mid = (lo + hi) // 2
kv = self._estimate_kv_cache_bytes(mid, cache_type_kv)
kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs)
if kv <= remaining:
best = mid
lo = mid + 1

229
studio/backend/tests/test_kv_cache_estimation.py
View file

@ -1302,6 +1302,235 @@ class TestEdgeCases:
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)
# With 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 ──────────────────────────────────
def test_unified_kv_ignores_n_parallel(self):
b = self._gqa_backend()
baseline = b._estimate_kv_cache_bytes(4096, "f16")
for slots in (1, 2, 4, 8):
assert (
b._estimate_kv_cache_bytes(
4096, "f16", n_parallel = slots, kv_unified = True
)
== baseline
)
def test_non_unified_multiplies_kv_by_n_parallel(self):
b = self._gqa_backend()
baseline = b._estimate_kv_cache_bytes(4096, "f16")
for slots in (1, 2, 4, 8):
scaled = b._estimate_kv_cache_bytes(
4096, "f16", n_parallel = slots, kv_unified = False
)
assert scaled == baseline * slots
def test_non_unified_with_zero_parallel_floors_at_one(self):
b = self._gqa_backend()
baseline = b._estimate_kv_cache_bytes(4096, "f16")
scaled = b._estimate_kv_cache_bytes(
4096, "f16", n_parallel = 0, kv_unified = False
)
assert scaled == baseline
def test_non_unified_scales_swa_path(self):
b = self._swa_backend()
baseline = b._estimate_kv_cache_bytes(8192, "f16")
scaled = b._estimate_kv_cache_bytes(
8192, "f16", n_parallel = 3, kv_unified = False
)
assert scaled == baseline * 3
def test_non_unified_scales_mla_path(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")
scaled = b._estimate_kv_cache_bytes(
8192, "f16", n_parallel = 4, kv_unified = False
)
assert scaled == baseline * 4
# ── --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 checkpoints * 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):
b = self._swa_backend()
ctx = 8192
single = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4)
triple = b._estimate_kv_cache_bytes(
ctx, "f16", ctx_checkpoints = 4, n_parallel = 3, kv_unified = False
)
assert triple == single * 3
# ── --kv-offload (kv_on_gpu) ───────────────────────────────────
def test_fit_returns_requested_when_kv_off_gpu(self):
b = self._gqa_backend()
# Tiny VRAM budget -- normally would 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_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 should not fit.
budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / 0.90 + 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
# ---------------------------------------------------------------------------
# K. Lifecycle Tests
# ---------------------------------------------------------------------------