diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index d577175bca..b959005d36 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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 diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py index 5c16365c9f..b7c037cd6a 100644 --- a/studio/backend/tests/test_kv_cache_estimation.py +++ b/studio/backend/tests/test_kv_cache_estimation.py @@ -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 # ---------------------------------------------------------------------------