From 6b13cab74651528b4e57e2b86e0803cf2a0e21c6 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Tue, 5 May 2026 16:36:46 +0530 Subject: [PATCH] fix KVCache estimates for gemma4 style sliding window models (#5225) * fix KVCache estimates for gemma4 style sliding window models Signed-off-by: Datta Nimmaturi * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: add per-arch SWA pattern fallback + n_kv_heads mirror for PR #5225 The pattern-aware SWA estimator added in this PR only fires when the GGUF carries `.attention.sliding_window_pattern`. Today's Gemma-2 / Gemma-3 / Gemma-3n / gpt-oss / Phi-3 GGUFs ship `attention.sliding_window` but not the pattern field (llama.cpp's converter strips it), so the new branch is bypassed and we fall back to the legacy 1/4-global heuristic on the most popular SWA arches in our catalogue (gemma3 alone has 6+ variants in the unsloth/* top 30 by downloads, plus gpt-oss-20b/120b). Two additions on top of this PR: 1. `_SWA_PATTERN_DEFAULTS_BY_ARCH` table keyed by GGUF arch name. When the GGUF reports a sliding window but no pattern, we synthesise the pattern from the architecture's canonical period (gemma2=2, gemma3=6, gemma3n=5, gpt_oss=2, phi3=1, cohere2=4). Periods sourced from a survey of the top 150 unsloth/* HF configs against `text_config.layer_types` and `Gemma*Config.sliding_window_pattern`. 2. Mirror `_n_kv_heads_by_layer` into the scalar `_n_kv_heads` (using max as a conservative upper bound) when the head_count_kv array is read. Without this, any non-SWA estimator path (GQA, legacy) on a Gemma-4-style model falls through to `n_heads`, which can be many times larger than the real per-layer KV head count. Also let `_can_estimate_kv` accept the array directly as belt-and-suspenders. End-to-end check on `unsloth/gemma-3-270m-it-Q4_K_M.gguf` (18 layers, sliding_window=512, no pattern field): the parser now resolves the pattern to period=6 (3 global, 15 SWA), matching the actual Gemma3TextConfig default. KV estimate at 32k context drops from 141 MB (legacy 1/4) to 108 MB (per-layer), a 23% reduction that directly translates into more headroom for `_fit_context_to_vram` and fewer cases where the slider lands on the 4096 floor. Tests: extended `test_kv_cache_estimation.py` with `TestArchSwaPatternDefaults` covering the six tabled arches, an unknown-arch negative, explicit-pattern precedence, and a no-sliding-window negative; updated `test_array_fields_parsed` to reflect the new mirror semantics; updated `test_end_to_end_synthetic_swa` to use the period=6 expectation. All 102 tests in the kv-cache / context-fit / max-context suites pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: tighten arch SWA table to verified non-regressing entries Audit of every unsloth/* HF model (1334 repos, all config.json fetched in scripts/survey_all_unsloth.py) plus end-to-end checks against five real GGUFs (gemma-3-270m, gemma-3-1b, qwen2.5-0.5b, phi-3.5-mini, falcon-h1-0.5b, granite-4.1-8b) confirms: * Pure-GQA arches (llama, qwen3, mistral3, glm4, llama4, ...) and the qwen2 family with use_sliding_window=False all reach Path 4 GQA cleanly. The llama.cpp converter strips `attention.sliding_window` for qwen2/qwen2_vl/qwen2_5_vl when use_sliding_window=False, so the SWA path never fires for them. Verified on Qwen2.5-0.5B-GGUF: no sliding_window field in metadata. * MLA arches (deepseek_v3/v32/v4, glm4_moe_lite, glm_moe_dsa, kimi_k25) emit `kv_lora_rank` -> Path 1 fires correctly. * Hybrid Mamba/Attn arches that emit both ssm.* and full_attention_interval (qwen3_5, qwen3_5_moe, qwen3_next) -> Path 2 fires correctly. Two table changes: 1. Drop the `phi3` entry. Phi-3 GGUFs emit `phi3.attention.sliding_window=262144` but never emit `attention.key_length`/`value_length`, so the SWA path is gated off and the estimator falls to the legacy formula. The huge sliding_window also means SWA layers and global layers cache identical numbers of tokens at any practical context, so a fallback would be a no-op anyway. The previous `phi3: 1` entry was also semantically wrong: period=1 with the (i+1)%N!=0 rule produces all-global, not the all-SWA you'd want for Phi-3. 2. Document the audit findings in the table comment, including the two arches we deliberately skip (phi3, qwen2*) and the one architecture family that is not a regression vs. main but is also not yet optimal (mistral v0.1/v0.2 all-SWA every-layer cannot be expressed with the period sentinel). Tests: added `test_non_swa_arch_uses_full_attention_path` parametrized over llama / qwen2 / qwen3 / mistral / mistral3 / glm4 / llama4 to pin the invariant that pure-GQA arches never receive a synthetic SWA pattern. Removed phi3 from the parametrize list of `test_arch_default_pattern_applied`. All 108 tests pass. Known separately tracked (not addressed here): falcon-h1 GGUFs ship ssm.* + key_length but no full_attention_interval, so the hybrid path 2 cannot fire and the estimator falls to GQA path 4, which counts every block as an attention layer. Affects 8 unsloth/* repos (~500 downloads). Same gap exists for granitemoehybrid-class GGUFs. Worth a follow-up that adds either a HYBRID_ATTENTION_INTERVAL_BY_ARCH table or a tensor-name probe. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: make SWA pattern resolver dynamic so new models work without code changes The static `_SWA_PATTERN_DEFAULTS_BY_ARCH` dict only covered architectures we knew about at PR-merge time. New SWA models would need a code change here every time a new arch shipped, which doesn't scale. Replaced with a 4-tier resolver so any newly-released model that lands on Hugging Face with a normal `config.json` is covered automatically: Tier 0 (parser) -- explicit GGUF metadata if the converter emits it (BOOL array or scalar period). Already supported. Tier 1 (cache) -- $UNSLOTH_STUDIO_HOME/swa_cache.json. Populated by previous Tier 3 fetches. Survives restarts. Tier 2 (bootstrap) -- `_BOOTSTRAP_SWA_DEFAULTS` shipped with Studio. Same five entries as the old static table (gemma2/3/3n/gpt_oss/cohere2). Lets fully-offline installs keep working for popular SWA arches with zero network. Tier 3 (HF fetch) -- pulls `config.json` from the GGUF's source HF repo and reads `sliding_window_pattern` (int) or `text_config.layer_types` (string array). Result is cached to Tier 1 so subsequent loads are offline-fast. Disabled by `UNSLOTH_STUDIO_OFFLINE=1`. Network errors and missing repos fall through silently. Tier 4 (caller) -- legacy 1/4-global SWA estimate (unchanged). The GGUF parser now also extracts a handful of `general.*` keys (`source.huggingface.repository`, `source.url`, `source.repo_url`, `base_model.0.repo_url`, `base_model.0.organization` + `.name`, `organization` + `basename`) so the resolver has source-repo candidates to try. End-to-end smoke against a brand-new arch (`never_seen_before_arch`, not in the bootstrap dict) pointing at `google/gemma-3-1b-it`: resolver fetched the HF config, derived period=6, materialised the 26-layer mask with 4 global layers (indices 5/11/17/23), and wrote `{"never_seen_before_arch": 6}` to the on-disk cache. Next load hits Tier 1 with no network. Tests: added `TestDynamicSwaResolver` with 10 tests covering each tier (period derivation, aperiodic mask handling, URL parsing, bootstrap precedence, cache precedence, HF fetch + persistence, candidate fallback, offline env knob, network failure). All 118 tests in the kv-cache / context-fit / max-context suites pass. The `_SWA_PATTERN_DEFAULTS_BY_ARCH` name was retired in favour of `_BOOTSTRAP_SWA_DEFAULTS` to make the tier semantics explicit. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: add Tier 2.5 transformers introspection to SWA resolver Slots a new tier between bootstrap and HF fetch that asks the locally-installed `transformers` package directly. Two strategies, in order, both offline-friendly: a. Default-instantiate the matching `Config` class via `CONFIG_MAPPING[arch]()` and read `sliding_window_pattern` / `text_config.layer_types`. Drills into `text_config` for multimodal wrappers. b. `inspect.getsource(cfg_class)` regex-parse for `sliding_window_pattern: int = N` defaults. Catches configs whose constructor raises (missing required args), or where the default is bound only in the __init__ signature. Walks `cfg_class.sub_configs["text_config"]` too so multimodal wrappers that delegate to a TextConfig still get inspected. Resolver chain is now 5 tiers: GGUF metadata, on-disk cache, bootstrap defaults, transformers introspection, HF Hub fetch, legacy fallback. Tier 2.5 results are persisted to the same on-disk cache as Tier 3 so subsequent loads skip the import overhead. `_arch_aliases` normalises hyphen vs underscore variants (`falcon-h1` vs `falcon_h1`) since GGUF and HF disagree for a handful of arches. Cross-version verification (probe at `temp/swa_probe/`): ``` arch transformers 4.57.6 transformers 5.7.0 gemma3 6 6 gemma2 2 2 cohere2 4 4 gpt_oss 2 2 gemma3n 5 5 gemma4 ARCH-MISSING 6 <- new arch picked up automatically falcon_h1 None [True]*32 <- per-layer mask used verbatim phi3 None None mistral None None qwen2 1 (all-global) 1 llama None None deepseek_v3 None None ``` The `gemma4` and `falcon_h1` rows are the headline: a brand-new arch that lands in transformers (gemma4 is 5.x-only) is supported by the resolver the moment a user upgrades the package, with zero edits to this file. Same applies to any future arch with a `Config` class. Tests: added `TestTransformersIntrospection` with 6 cases covering arch-alias normalisation, real-arch resolution against the live transformers, inspect.getsource fallback when default-init raises, graceful behaviour when transformers is unavailable, unknown-arch returns None, and Tier-2.5-before-Tier-3 ordering. Also adjusted the existing Tier 3 failure test to mock Tier 2.5 out so it specifically exercises the network-failure path. All 124 tests in the kv-cache / context-fit / max-context suites pass. Updated the module-level resolver comment from "4-tier" to "5-tier" to document the new tier. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: consolidate verbose comments and docstrings in SWA resolver Net -345 lines: -210 in llama_cpp.py, -357/+111 in test_kv_cache_estimation.py. No behaviour change; only comments, docstrings, and one tests-only `_SWA_FIELDS` helper to remove the copy-pasted GGUF metadata dict from each resolver test. Code changes only delete or shorten: * 5-tier resolver header collapsed from a 38-line block diagram to a 9-line summary; the rest is the function bodies. * Bootstrap dict per-arch comments collapsed to one-line `Config` references. * `_swa_cache_path`, `_save_swa_cache`, `_period_from_layer_types`, `_arch_aliases`, `_swa_entry_from_config_obj`, `_resolve_swa_pattern`, `_resolve_swa_entry_from_transformers` docstrings stripped to one line or removed when the body is self-evident. * Tier-by-tier inline comments inside `_resolve_swa_pattern` removed (function body reads top-to-bottom in tier order). * Path-3 SWA estimator comment shortened from a 12-line tier breakdown to 3 lines. * Parser fallback comment block (originally explained the resolver in-line) trimmed to two lines pointing at the resolver. * `_can_estimate_kv` legacy-clause comment shortened to one line. * GGUF `general.*` WANTED block comment shortened to one line. Test changes: * Per-test docstrings dropped where the test name and body already explain intent. * Class-level docstrings reduced to one line. * Common GGUF field dict factored to module-level `_SWA_FIELDS`. * Multi-line URL/list assertions collapsed to one-liners. All 124 tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: account for SWA cache double-buffering in path 3 estimate Cross-check against llama.cpp ground truth (running llama-server with --parallel 1 and reading the `llama_kv_cache: size = X MiB ( N cells, M layers, ... )` log lines) showed the SWA path under-counted by ~20% on Gemma-3-shaped models: GGUF pred MiB actual MiB ratio gemma-3-270m-it-Q4_K_M 31.50 39.00 0.81 gemma-3-1b-it-Q2_K 43.00 54.00 0.80 Root cause: llama.cpp double-buffers the SWA cache so it can keep the current and next windows during the shift, allocating `2 * sliding_window` cells per SWA layer (capped at n_ctx). My formula was using `min(n_ctx, sliding_window)` instead of `min(n_ctx, 2 * sliding_window)`. Verified directly: llama_kv_cache_iswa: creating SWA KV cache, size = 1024 cells llama_kv_cache: size = 15.00 MiB (1024 cells, 15 layers, ...) with `gemma3.attention.sliding_window = 512` -> 2 * 512 = 1024 cells. Fix: introduce `swa_cells = min(n_ctx, 2 * swa)` in path 3 and use that for both the per-layer-pattern branch and the legacy 1/4-global fallback. Re-run after the fix: GGUF pred MiB actual MiB ratio gemma-3-270m-it-Q4_K_M 39.00 39.00 1.000 gemma-3-1b-it-Q2_K 54.00 54.00 1.000 qwen2.5-0.5b-instruct-q4_k_m 96.00 96.00 1.000 Phi-3.5-mini-instruct-Q4_K_M 3072.00 3072.00 1.000 Falcon-H1-0.5B-Instruct-Q4_K_M 144.00 144.00 1.000 granite-4.1-8b-Q3_K_M 1280.00 1280.00 1.000 All 5 paths now match llama.cpp's actual allocation exactly under single-sequence inference (Studio's default). Tests: updated `test_gemma3`, `test_gpt_oss`, `test_gemma4_per_layer_swa_metadata`, `test_ctx_smaller_than_window`, `test_odd_layer_count`, and `test_end_to_end_synthetic_swa` to use the doubled SWA cell count. All 124 tests in the kv-cache / context-fit / max-context suites pass. * studio: tolerate truncated GGUF input so resolver fallback still runs Wraps each iteration of the GGUF KV-pair loop in a try/except that breaks out cleanly on `struct.error` or `UnicodeDecodeError`, instead of letting the outer try eat the exception and skip the SWA resolver fallback at the end. The motivating use case is reading the GGUF metadata via an HF Hub HTTP byte-range fetch. The first ~128 KiB of a typical GGUF contains all the metadata we need (arch, block_count, attention.*, sliding window, ssm, MLA fields, plus the tokenizer config) -- but for models with large tokenizer vocabs (Gemma 3 has 262144 tokens) the tokenizer arrays spill past the 128 KiB boundary. The truncation used to bubble out as `unpack requires a buffer of 8 bytes`, abandoning the resolver fallback and leaving us with no SWA pattern (so the SWA path fell through to the legacy 1/4 estimate). Verified end to end against `unsloth/gemma-3-1b-it-GGUF`: Range-fetch first 128 KiB of `gemma-3-1b-it-Q2_K.gguf` over HTTP (HTTP 206 Partial Content), parse: arch = gemma3 block_count = 26 attention.sliding_window = 512 sliding_window_pattern = set (4 global) <- via Tier 2 bootstrap KV @ ctx=8192 = 54.00 MiB <- matches llama.cpp ground truth This means Studio can preview KV-cache requirements (and therefore auto-context fit) for any HF GGUF without downloading the weights. All 124 tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: thread llama-server KV flags through the estimator 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: account for shared_kv_layers (Gemma 3n / Gemma 4) Gemma 3n and Gemma 4 set .attention.shared_kv_layers in the GGUF metadata (convert_hf_to_gguf.py lines 7578 and 7712). The trailing N layers of the model reuse KV from earlier layers and don't allocate their own cache, so n_layers in the per-layer formulas overcounted by exactly that many blocks. For google/gemma-3n-E4B-it (35 layers, 15 shared), this puts ~43% of the KV estimate back on the table. Changes: - _read_gguf_metadata parses .attention.shared_kv_layers into self._shared_kv_layers; init / unload / reparse all reset it. - _estimate_kv_cache_bytes computes n_layers_kv = max(1, n_layers - shared_kv_layers) and substitutes it for n_layers in: Path 1 (MLA), Path 3 (SWA pattern loop bound and the no-pattern fallback), Path 4 (GQA), Path 5 (legacy). Path 2 (hybrid) keeps n_layers since hybrid + shared_kv combined isn't a thing today and the semantics would need to specify which attention layers are shared. - max(1, ...) floor protects against pathological GGUFs where shared >= n_layers. - Composes naturally with --swa-full, --kv-unified / --parallel, --ctx-checkpoints, and the per-layer SWA pattern from the dynamic resolver. When the field is unset (every other arch) the math is byte-identical to before. Tests: TestSharedKVLayers (13 cases) covers each path's drop, the no-op-when-unset case, the floor at one layer, composition with the server-flag knobs, and lifecycle reset. test_end_to_end_synthetic_shared_kv_round_trip exercises the full GGUF parse -> estimate path on a synthetic gemma3n_text blob. Existing TestLifecycle tests extended to cover the new field. Full suite: 132 passing. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: only stub httpx in tests when the real lib is missing The unit suite stubs httpx unconditionally so tests can run on a minimal Python install. Surfaced during a fresh-venv simulation: when the stub is installed via setdefault on a system that DOES have httpx, huggingface_hub.errors fails to import HTTPError / Response at module load time, which the transformers introspection tier swallows via its bare except. Result: TestTransformersIntrospection passes in venvs where httpx happened to be imported first (workspace) and silently fails in venvs where it doesn't (fresh uv venv). Switch to "only stub when real lib unavailable", and round out the stub with HTTPError, RequestError, and Response so any test environment without httpx still gets a complete enough surface for huggingface_hub to import. * studio: per-layer-type --parallel N memory accounting for SWA Empirical verification against llama-server (see workspace_5/temp/sim_pr5225/probe_parallel_full_matrix.py and verify_parallel_matches_server.py) showed the prior whole-cache slot_factor multiplication in _estimate_kv_cache_bytes was wrong for n_parallel > 1. The actual rule, verified bit-exact across the full (parallel x ctx) grid for both SWA and pure-GQA models: * non-SWA layers: total cells = n_ctx, partitioned across slots (per-slot ctx = n_ctx / parallel). Total memory is CONSTANT in n_parallel. * SWA layers: per-slot cells = 2 * sliding_window (clamped at n_ctx and at per_slot_ctx when ctx is split among many slots). Total memory grows LINEARLY in n_parallel. * --kv-unified: no measurable difference to total memory; both modes yield the same byte total in measured cases. Retained as accepted-but-ignored kwarg for API forward-compat. Closed form (Path 3 with per-layer pattern): total_kv = sum_global_layers(n_ctx * n_kv * (k+v) * bpe) + parallel * sum_swa_layers( min(2*sliding_window, n_ctx, n_ctx//parallel) * n_kv_layer * (k_swa + v_swa) * bpe ) + parallel * checkpoint_extra_per_slot (when ctx_checkpoints > 0) Changes to _estimate_kv_cache_bytes: - Path 3 (SWA pattern): accumulate global_bytes and swa_bytes_per_slot separately; final result = global_bytes + slots * (swa_bps + cp_bps). - Path 3 (no-pattern fallback): same split using the 1/4-global heuristic. - Paths 1 / 2 / 4 / 5: drop the slot_factor multiplication. Non-SWA caches don't scale with --parallel. - swa_full=True: SWA cells = per_slot_ctx (was n_ctx), so slots cancels out and total stays constant. Matches llama-server's --swa-full --parallel N output exactly. Production wiring fix in start(): - Seven internal calls to _estimate_kv_cache_bytes / _fit_context_to_vram used the default n_parallel=1, even though load_model accepts the caller's n_parallel value (forwarded to llama-server via --parallel on the command line). Pass n_parallel through all seven so VRAM budgeting is correct when an operator sets parallel slots above 1. Studio's default ships at 1 so production today is unaffected; this completes the wiring for operators who tune it. Tests: - TestParallelSWAScaling (10 new cases): closed-form invariants per path, swa_full + parallel collapse, kv_unified no-op proof, per-slot SWA cell clamping, and the empirical Gemma-3 270m formula (24 + parallel * 15 MiB at ctx=8192) baked from the verifier. - TestServerFlags: rewrote 4 assertions and renamed 2 to reflect the per-layer rule; non-SWA paths now correctly assert constancy. - TestSharedKVLayers::test_composes_with_n_parallel: rewrote to assert only the SWA portion of the unshared layers scales. Backward compatibility: at n_parallel=1 the output is bit-identical to before this change (verified across 120,960 sweep combinations and the 141-test suite in both workspace and fresh-uv-venv environments). Verifier output at --parallel in {1,2,4,8} x ctx in {4096,8192,16384} shows ratio 1.000 against llama-server for both SWA and pure-GQA models (24/24 cells exact match). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: accept new estimator kwargs in load-time test stubs `test_llama_cpp_context_fit.py` and `test_llama_cpp_max_context_threshold.py` patch `_estimate_kv_cache_bytes` with constant per-token stubs and then call the real `_fit_context_to_vram`. After 29dcf96e threaded the llama-server flag kwargs (`swa_full`, `n_parallel`, `kv_unified`, `ctx_checkpoints`) through `_fit_context_to_vram`, the production method forwards them to the stubbed estimator and the old positional-only stubs raise `TypeError`. These two suites exercise the load-time fit decision and the max-context threshold property with a constant per-token KV cost; SWA / parallel-slot accounting is intentionally out of scope, so the stubs absorb the new kwargs and ignore them. No production change. Restores both files to fully passing: 15/15 in `test_llama_cpp_context_fit` and 8/8 in `test_llama_cpp_max_context_threshold`. Combined with the existing 141/141 in `test_kv_cache_estimation`, the three KV-cache test modules are 164/164 green. --------- Signed-off-by: Datta Nimmaturi Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: Daniel Han --- studio/backend/core/inference/llama_cpp.py | 627 +++++++- .../backend/tests/test_kv_cache_estimation.py | 1340 ++++++++++++++++- .../tests/test_llama_cpp_context_fit.py | 6 +- .../test_llama_cpp_max_context_threshold.py | 6 +- 4 files changed, 1871 insertions(+), 108 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index bddd71c301..f768764c22 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -11,6 +11,7 @@ through its OpenAI-compatible /v1/chat/completions endpoint. import atexit import contextlib import json +import os import re import struct import structlog @@ -76,6 +77,238 @@ _SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$") _SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$") +# ── Sliding-window-pattern resolver ─────────────────────────── +# Resolves the per-layer SWA mask when a GGUF reports a sliding window +# but no `sliding_window_pattern` field. Tier order in +# `_resolve_swa_pattern`: GGUF metadata, on-disk cache, bootstrap dict +# below, transformers introspection, HF Hub config.json, legacy 1/4 +# fallback. Period N means layer i is SWA iff `(i + 1) % N != 0`, +# matching transformers. Skipped on purpose: phi3 (no key/val length +# in GGUF, window >= ctx anyway), qwen2 family (converter strips +# sliding_window when use_sliding_window=False), mistral v0.1/v0.2 +# (all-SWA can't be expressed as a period). +_BOOTSTRAP_SWA_DEFAULTS: dict[str, int] = { + "gemma2": 2, # Gemma2Config.sliding_window_pattern + "gemma3": 6, # Gemma3TextConfig.sliding_window_pattern + "gemma3n": 5, # text_config.layer_types: SWA*4 + FULL + "gpt_oss": 2, # text_config.layer_types: alternating + "cohere2": 4, # Cohere2Config.sliding_window_pattern +} + +# Process-wide cache backed by JSON on disk. Values are int period or +# list[bool] mask. Lazy-loaded. +_SWA_CACHE: Optional[dict] = None +_SWA_CACHE_LOCK = threading.Lock() + + +def _swa_cache_path() -> Path: + home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") + base = Path(home) if home else Path.home() / ".unsloth" / "studio" + return base / "swa_cache.json" + + +def _load_swa_cache() -> dict: + global _SWA_CACHE + with _SWA_CACHE_LOCK: + if _SWA_CACHE is not None: + return _SWA_CACHE + try: + with open(_swa_cache_path()) as f: + _SWA_CACHE = json.load(f) + if not isinstance(_SWA_CACHE, dict): + _SWA_CACHE = {} + except (FileNotFoundError, json.JSONDecodeError, OSError): + _SWA_CACHE = {} + return _SWA_CACHE + + +def _save_swa_cache(cache: dict) -> None: + try: + path = _swa_cache_path() + path.parent.mkdir(parents = True, exist_ok = True) + tmp = path.with_suffix(".json.tmp") + with open(tmp, "w") as f: + json.dump(cache, f, indent = 2, sort_keys = True) + tmp.replace(path) + except OSError: + pass + + +def _period_from_layer_types(layer_types: list) -> Optional[int]: + """Smallest period N where `(i+1) % N != 0` matches the SWA mask, + or None if no fixed period fits.""" + if not layer_types: + return None + is_swa = ["full" not in str(t).lower() for t in layer_types] + n = len(is_swa) + for N in range(1, n + 1): + if all(((i + 1) % N != 0) == is_swa[i] for i in range(n)): + return N + return None + + +def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]: + try: + from huggingface_hub import hf_hub_download + + cfg_path = hf_hub_download(repo_id, "config.json", repo_type = "model") + with open(cfg_path) as f: + cfg = json.load(f) + except Exception: + return None + + src = cfg.get("text_config") if isinstance(cfg.get("text_config"), dict) else cfg + period = src.get("sliding_window_pattern") + if isinstance(period, int) and period > 0: + return period + lt = src.get("layer_types") + if isinstance(lt, list) and lt: + return _period_from_layer_types(lt) or [ + "full" not in str(t).lower() for t in lt + ] + return None + + +def _arch_aliases(arch: str) -> tuple: + # GGUF emits `falcon-h1`; HF model_type is `falcon_h1`. Normalise both ways. + seen = [] + for a in (arch, arch.replace("-", "_"), arch.replace("_", "-")): + if a and a not in seen: + seen.append(a) + return tuple(seen) + + +def _swa_entry_from_config_obj(cfg) -> Optional[object]: + src = getattr(cfg, "text_config", None) or cfg + period = getattr(src, "sliding_window_pattern", None) + if isinstance(period, int) and period > 0: + return period + lt = getattr(src, "layer_types", None) + if isinstance(lt, list) and lt: + return _period_from_layer_types(lt) or [ + "full" not in str(t).lower() for t in lt + ] + return None + + +_SWA_PATTERN_SOURCE_RE = re.compile( + r"sliding_window_pattern\s*(?::\s*[\w\[\], ]*)?\s*=\s*(\d+)" +) + + +def _resolve_swa_entry_from_transformers(arch: str) -> Optional[object]: + """Default-instantiate the matching Config; on failure, regex-parse + its source for `sliding_window_pattern = N`.""" + try: + from transformers.models.auto.configuration_auto import ( + CONFIG_MAPPING, + CONFIG_MAPPING_NAMES, + ) + except Exception: + return None + + cfg_class = None + for alias in _arch_aliases(arch): + if alias in CONFIG_MAPPING_NAMES: + try: + cfg_class = CONFIG_MAPPING[alias] + break + except Exception: + cfg_class = None + if cfg_class is None: + return None + + try: + if (entry := _swa_entry_from_config_obj(cfg_class())) is not None: + return entry + except Exception: + pass + + import inspect + + candidates = [cfg_class] + text_cfg_class = getattr(cfg_class, "sub_configs", {}).get("text_config") + if text_cfg_class is not None: + candidates.append(text_cfg_class) + for cls in candidates: + try: + src = inspect.getsource(cls) + except (OSError, TypeError): + continue + if m := _SWA_PATTERN_SOURCE_RE.search(src): + period = int(m.group(1)) + if period > 0: + return period + return None + + +def _resolve_swa_pattern( + arch: Optional[str], + n_layers: Optional[int], + source_repo_candidates: tuple = (), + *, + allow_network: Optional[bool] = None, +) -> Optional[list]: + if not arch or not n_layers: + return None + if allow_network is None: + allow_network = os.environ.get("UNSLOTH_STUDIO_OFFLINE", "0") not in ( + "1", + "true", + "True", + "yes", + ) + + cache = _load_swa_cache() + + def _entry_to_mask(entry): + if isinstance(entry, int) and entry > 0: + return [(i + 1) % entry != 0 for i in range(n_layers)] + if isinstance(entry, list) and entry: + return [bool(entry[i % len(entry)]) for i in range(n_layers)] + return None + + def _persist(entry): + with _SWA_CACHE_LOCK: + cache[arch] = entry + _save_swa_cache(cache) + + if (entry := cache.get(arch)) is not None: + if (mask := _entry_to_mask(entry)) is not None: + return mask + + if (entry := _BOOTSTRAP_SWA_DEFAULTS.get(arch)) is not None: + return _entry_to_mask(entry) + + entry = _resolve_swa_entry_from_transformers(arch) + if entry is not None: + _persist(entry) + return _entry_to_mask(entry) + + # Tier 3: live HF fetch (with persistent caching of the result) + if allow_network: + for repo_id in source_repo_candidates: + if not repo_id: + continue + entry = _fetch_swa_entry_from_hf(repo_id) + if entry is not None: + _persist(entry) + return _entry_to_mask(entry) + + return None + + +def _hf_repo_from_url(url: Optional[str]) -> Optional[str]: + """Strip `https://huggingface.co/owner/name(/...)` to `owner/name`.""" + if not url or "huggingface.co/" not in url: + return None + tail = url.split("huggingface.co/", 1)[1].rstrip("/") + parts = tail.split("/") + if len(parts) < 2: + return None + return f"{parts[0]}/{parts[1]}" + + # Model size extraction — lazy import to avoid pulling in transformers # at module level. See PR description for the full explanation. def _extract_model_size_b(model_id: str): @@ -215,17 +448,24 @@ class LlamaCppBackend: # KV-cache estimation fields (populated by _read_gguf_metadata) self._n_layers: Optional[int] = None self._n_kv_heads: Optional[int] = None + self._n_kv_heads_by_layer: Optional[list[int]] = None self._n_heads: Optional[int] = None self._embedding_length: Optional[int] = None - # Architecture-aware KV fields (8 new fields for 5-path estimation) + # Architecture-aware KV fields for 5-path estimation self._kv_key_length: Optional[int] = None self._kv_value_length: Optional[int] = None self._sliding_window: Optional[int] = None + self._sliding_window_pattern: Optional[list[bool]] = None self._full_attention_interval: Optional[int] = None self._kv_lora_rank: Optional[int] = None self._key_length_mla: Optional[int] = None + self._kv_key_length_swa: Optional[int] = None + self._kv_value_length_swa: Optional[int] = None self._ssm_inner_size: Optional[int] = None self._ssm_state_size: Optional[int] = None + # Last N layers reuse KV from earlier layers and don't allocate + # their own cache (Gemma 3n / Gemma 4: .attention.shared_kv_layers). + self._shared_kv_layers: Optional[int] = None self._lock = threading.Lock() self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None @@ -744,13 +984,29 @@ class LlamaCppBackend: # New-style: need both explicit key AND value dimensions if self._kv_key_length is not None and self._kv_value_length is not None: return True - # Legacy: need embedding_length + head count + # Legacy: need embedding_length + a head count (scalar or per-layer). return self._embedding_length is not None and ( - self._n_kv_heads is not None or self._n_heads is not None + self._n_kv_heads is not None + or self._n_heads is not None + or self._n_kv_heads_by_layer is not None ) + def _kv_heads_for_layer(self, layer_idx: int, fallback: int) -> int: + if self._n_kv_heads_by_layer is not None and layer_idx < len( + self._n_kv_heads_by_layer + ): + return self._n_kv_heads_by_layer[layer_idx] + 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. @@ -761,12 +1017,34 @@ 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. + Verified empirically against llama-server: + non-SWA layers stay constant (cells split + across slots), SWA layers scale linearly + (per-slot window). + kv_unified -- ``--kv-unified`` (default on): retained for + API forward-compat. Currently a no-op for + memory math because the unified buffer total + matches per-slot buffers in measured cases. + 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: return 0 n_layers = self._n_layers # type: ignore[assignment] + # Gemma 3n / Gemma 4 reuse KV from earlier layers in the last + # ``shared_kv_layers`` blocks -- those don't allocate their own + # cache. Floor at 1 so a misconfigured GGUF can't zero out KV. + shared = self._shared_kv_layers or 0 + n_layers_kv = max(1, n_layers - shared) n_kv = self._n_kv_heads or self._n_heads or 1 # type: ignore[assignment] # Bytes per element depends on KV cache quantization @@ -782,6 +1060,8 @@ class LlamaCppBackend: "iq4_nl": 0.5625, }.get(cache_type_kv or "f16", 2.0) + slots = max(1, n_parallel) + # 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. @@ -792,7 +1072,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_kv * n_ctx * n_kv_mla * key_len * bpe) key_len = self._kv_key_length val_len = self._kv_value_length @@ -810,11 +1090,19 @@ class LlamaCppBackend: head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe) - # Path 3: Sliding Window (Gemma-3, gpt-oss) - # SWA layers only cache min(ctx, window) tokens; global layers cache full ctx. - # Most SWA architectures use few global layers (e.g., Gemma-3 uses 1 in 6). - # Without an explicit field, we conservatively assume 1/4 of layers are global - # which is still far more accurate than the legacy formula (which ignores SWA). + # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). + # Pattern is filled in by the resolver at parse time; if absent, + # falls through to the legacy 1/4-global heuristic below. + # Per-layer-type ``--parallel N`` accounting (verified empirically + # against ``llama-server``): + # * non-SWA layers: total cells = n_ctx, partitioned across + # slots -> total memory CONSTANT in slots. + # * SWA layers: per-slot cells = 2 * sliding_window + # (capped at n_ctx and at per_slot_ctx + # when ctx is split among many slots) -> + # total memory grows LINEARLY in slots. + # ``--swa-full`` forces full n_ctx for SWA layers instead. + # ``--ctx-checkpoints N`` adds N snapshots per SWA layer per slot. if ( self._sliding_window is not None and self._sliding_window > 0 @@ -822,20 +1110,72 @@ class LlamaCppBackend: and val_len is not None ): swa = self._sliding_window - n_global = max(1, n_layers // 4) - n_swa = n_layers - n_global + per_slot_ctx = max(1, n_ctx // slots) + # ``--swa-full`` makes SWA layers cache the full context just + # like non-SWA: cells get partitioned across slots, so per-slot + # cells = per_slot_ctx and the slots*per-slot product collapses + # back to the constant ``n_ctx`` total. Otherwise SWA caches + # 2*sliding_window per slot, clamped at the per-slot ctx. + swa_cells_per_slot = ( + per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx) + ) + 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: + global_bytes = 0.0 # constant across slots + swa_bytes_per_slot = 0.0 # multiplied by slots + checkpoint_extra_per_slot = 0.0 + # Iterate only over layers that allocate their own KV; + # the trailing ``shared`` layers reuse earlier caches. + for layer_idx in range(n_layers_kv): + layer_n_kv = self._kv_heads_for_layer(layer_idx, n_kv) + is_swa = ( + layer_idx < len(self._sliding_window_pattern) + and self._sliding_window_pattern[layer_idx] + ) + if is_swa: + swa_bytes_per_slot += ( + swa_cells_per_slot + * layer_n_kv + * (key_len_swa + val_len_swa) + * bpe + ) + if ctx_checkpoints > 0 and not swa_full: + checkpoint_extra_per_slot += ( + ctx_checkpoints + * swa + * layer_n_kv + * (key_len_swa + val_len_swa) + * bpe + ) + else: + global_bytes += n_ctx * layer_n_kv * (key_len + val_len) * bpe + return int( + global_bytes + + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot) + ) + n_global = max(1, n_layers_kv // 4) + n_swa = n_layers_kv - n_global kv_per_token = n_kv * (key_len + val_len) * bpe + kv_per_token_swa = n_kv * (key_len_swa + val_len_swa) * bpe + global_bytes = n_global * n_ctx * kv_per_token + swa_bytes_per_slot = n_swa * swa_cells_per_slot * kv_per_token_swa + checkpoint_extra_per_slot = ( + ctx_checkpoints * n_swa * swa * kv_per_token_swa + if ctx_checkpoints > 0 and not swa_full + else 0.0 + ) return int( - n_global * n_ctx * kv_per_token + n_swa * min(n_ctx, swa) * kv_per_token + global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot) ) # 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_kv * n_ctx * n_kv * (key_len + val_len) * bpe) # Path 5: Legacy fallback (old GGUFs without explicit dimensions) head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] - return int(2 * n_kv * head_dim * n_layers * n_ctx * bpe) + return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe) def _fit_context_to_vram( self, @@ -844,6 +1184,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. @@ -851,6 +1197,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( @@ -860,11 +1211,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 @@ -886,7 +1248,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 @@ -1019,6 +1381,19 @@ class LlamaCppBackend: for _ in range(alen): LlamaCppBackend._gguf_skip_value(f, atype) + @staticmethod + def _gguf_read_array_value(f, atype: int, alen: int) -> Optional[list]: + if atype == 4: # UINT32 + return [struct.unpack(" None: """Read context_length, architecture params, and chat_template from a GGUF header. @@ -1037,23 +1412,44 @@ class LlamaCppBackend: self._supports_tools = False self._n_layers = None self._n_kv_heads = None + self._n_kv_heads_by_layer = None self._n_heads = None self._embedding_length = None self._kv_key_length = None self._kv_value_length = None self._sliding_window = None + self._sliding_window_pattern = None self._full_attention_interval = None self._kv_lora_rank = None self._key_length_mla = None + self._kv_key_length_swa = None + self._kv_value_length_swa = None self._ssm_inner_size = None self._ssm_state_size = None + self._shared_kv_layers = None try: - WANTED = {"general.architecture", "tokenizer.chat_template"} + WANTED = { + "general.architecture", + "tokenizer.chat_template", + # Source-repo hints for the SWA resolver's HF fallback. + "general.source.huggingface.repository", + "general.source.url", + "general.source.repo_url", + "general.base_model.0.repo_url", + "general.base_model.0.organization", + "general.base_model.0.name", + "general.basename", + "general.organization", + "general.size_label", + "general.finetune", + } # Additional arch-specific keys are added dynamically once # we know the architecture name. arch_keys: dict[str, str] = {} # gguf_key -> attribute name arch = None + sliding_window_pattern_period: Optional[int] = None + general: dict[str, str] = {} with open(gguf_path, "rb") as f: magic = struct.unpack(" {effective_ctx} " f"(model: {model_size / (1024**3):.1f} GB, " @@ -1586,7 +2082,7 @@ class LlamaCppBackend: ) kv_cache_bytes = self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv + effective_ctx, cache_type_kv, n_parallel = n_parallel ) logger.info( f"GGUF size: {model_size / (1024**3):.1f} GB, " @@ -2015,16 +2511,21 @@ class LlamaCppBackend: self._speculative_type = None self._n_layers = None self._n_kv_heads = None + self._n_kv_heads_by_layer = None self._n_heads = None self._embedding_length = None self._kv_key_length = None self._kv_value_length = None self._sliding_window = None + self._sliding_window_pattern = None self._full_attention_interval = None self._kv_lora_rank = None self._key_length_mla = None + self._kv_key_length_swa = None + self._kv_value_length_swa = None self._ssm_inner_size = None self._ssm_state_size = None + self._shared_kv_layers = None # Clean up temp chat template file if hasattr(self, "_chat_template_file") and self._chat_template_file: try: diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py index 2640ded90d..29d87804ff 100644 --- a/studio/backend/tests/test_kv_cache_estimation.py +++ b/studio/backend/tests/test_kv_cache_estimation.py @@ -12,6 +12,7 @@ Cross-platform: Linux, macOS, Windows, WSL. """ import io +import json import struct import sys import types as _types @@ -37,35 +38,43 @@ sys.modules.setdefault("loggers", _loggers_stub) _structlog_stub = _types.ModuleType("structlog") sys.modules.setdefault("structlog", _structlog_stub) -# httpx -_httpx_stub = _types.ModuleType("httpx") -for _exc_name in ( - "ConnectError", - "TimeoutException", - "ReadTimeout", - "ReadError", - "RemoteProtocolError", - "CloseError", -): - setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) +# httpx -- only stub when the real library isn't installed. Stubbing +# unconditionally would shadow ``HTTPError`` / ``Response`` etc. that +# ``huggingface_hub.errors`` imports at module load time, which causes +# the transformers introspection tier to silently return None inside +# the test process. +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 -class _FakeTimeout: - def __init__(self, *a, **kw): - pass - - -_httpx_stub.Timeout = _FakeTimeout -_httpx_stub.Client = type( - "Client", - (), - { - "__init__": lambda self, **kw: None, - "__enter__": lambda self: self, - "__exit__": lambda self, *a: None, - }, -) -sys.modules.setdefault("httpx", _httpx_stub) + _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 LlamaCppBackend @@ -77,8 +86,7 @@ from core.inference.llama_cpp import LlamaCppBackend def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes: """Build a minimal GGUF v3 binary blob with the given KV metadata. - Only supports UINT32 (type 4), UINT64 (type 10), and STRING (type 8) - values, which is all the metadata parser reads. + Supports the scalar and simple array metadata used by the parser. """ buf = io.BytesIO() # Header: magic, version, tensor_count, kv_count @@ -96,6 +104,17 @@ def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes: val_bytes = val.encode("utf-8") buf.write(struct.pack(" bytes: return buf.getvalue() -def _backend_from_gguf(arch: str, fields: dict) -> LlamaCppBackend: - """Create a LlamaCppBackend with parsed GGUF metadata from given fields.""" +def _backend_from_gguf( + arch: str, fields: dict, general: dict | None = None +) -> LlamaCppBackend: + """Create a LlamaCppBackend with parsed GGUF metadata from given fields. + + `general` lets a test inject extra `general.*` metadata (used 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 @@ -133,7 +161,7 @@ def _backend_from_gguf(arch: str, fields: dict) -> LlamaCppBackend: class TestGGUFParserNewFields: - """Verify that the 8 new architecture-aware fields are correctly parsed.""" + """Verify that architecture-aware fields are correctly parsed.""" @pytest.mark.parametrize( "field,gguf_key,value", @@ -158,15 +186,189 @@ class TestGGUFParserNewFields: "_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_all_13_fields_parsed_together(self): + 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 estimator paths and any caller 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): + # Period=6 is the gemma3 default; 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: GGUF has 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 the 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, @@ -176,9 +378,12 @@ class TestGGUFParserNewFields: "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, } @@ -191,13 +396,294 @@ class TestGGUFParserNewFields: 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 per 6), gpt-oss (alternating), gemma3n (1 per 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) + # Override bootstrap=6 with a cached period=3. + 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 gets 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 the failure into the Tier 3 path; 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 + ): + # With 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: """Verify that fields are properly reset between parses.""" @@ -209,11 +695,19 @@ class TestGGUFParserReset: "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 should be None @@ -230,6 +724,10 @@ class TestGGUFParserReset: 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 @@ -455,7 +953,9 @@ class TestSlidingWindowEstimation: n_global = max(1, 62 // 4) # 15 n_swa = 62 - n_global # 47 kv_per = 16 * (128 + 128) * 2 - expected = int(n_global * 131072 * kv_per + n_swa * min(131072, 1024) * kv_per) + # 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): @@ -472,27 +972,52 @@ class TestSlidingWindowEstimation: n_global = max(1, 24 // 4) # 6 n_swa = 24 - n_global # 18 kv_per = 8 * (64 + 64) * 2 - expected = int(n_global * 131072 * kv_per + n_swa * min(131072, 128) * kv_per) + 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 context < sliding_window, SWA layers use full context anyway.""" + """When context < 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, 8192) * kv_per) - # min(4096, 8192) = 4096, so both pools use full ctx + 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): - """Odd layer count: n_global = max(1, n//4), n_swa = n - n_global.""" b = self._swa_backend(_n_layers = 63) n_global = max(1, 63 // 4) # 15 n_swa = 63 - n_global # 48 kv_per = 16 * (128 + 128) * 2 - expected = int(n_global * 1000 * kv_per + n_swa * min(1000, 1024) * kv_per) + expected = int(n_global * 1000 * kv_per + n_swa * min(1000, 2 * 1024) * kv_per) assert b._estimate_kv_cache_bytes(1000, "f16") == expected @@ -785,6 +1310,686 @@ 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 ────────────────────────────────── + # Empirically verified against llama-server: non-SWA caches partition + # n_ctx across slots (total memory constant); SWA layers are the only + # portion that scales with --parallel. --kv-unified is currently 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 same loop the estimator does. + 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 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 get clamped 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 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): + # Only the SWA + checkpoint portion scales by n_parallel; the + # global-layer portion stays 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 -- 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 + + +# --------------------------------------------------------------------------- +# J2.5. --parallel N memory accounting (per-layer-type scaling rule) +# --------------------------------------------------------------------------- + + +class TestParallelSWAScaling: + """Verifies the per-layer-type scaling rule against 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 yield 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 should 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; result is the all-global + # GQA-style total, which is 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): + # Both unified=True and unified=False must produce the same + # total bytes for every backend type and every 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 + # 5-period [swa,swa,swa,swa,full] * 3 + [swa,swa,swa]: mirrors the + # bootstrap-resolved pattern for gemma3 (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: + """``.attention.shared_kv_layers`` reduces the layer count that + actually 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 are 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; expect ~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; equals 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 should 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 the unshared layers scales by n_parallel; + # the global portion stays 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 only count 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 # --------------------------------------------------------------------------- @@ -799,13 +2004,18 @@ class TestLifecycle: "_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() @@ -813,20 +2023,30 @@ class TestLifecycle: 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): """Full round-trip: write GGUF -> parse -> estimate.""" @@ -887,12 +2107,46 @@ class TestLifecycle: ) assert b._can_estimate_kv() result = b._estimate_kv_cache_bytes(131072, "f16") - n_global = max(1, 62 // 4) # 15 - n_swa = 62 - n_global # 47 + # gemma3 -> period 6 from the bootstrap table, SWA cache + # double-buffered to 2 * sliding_window cells. + period = 6 kv_per = 16 * 256 * 2 - expected = int(n_global * 131072 * kv_per + n_swa * 1024 * kv_per) + 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 table for gemma3n_text -> period 5; the resolver + # synthesises a 35-entry bool array. The first 20 entries + # (n_layers - shared) are the only ones that allocate KV. + result = b._estimate_kv_cache_bytes(8192, "f16") + assert result > 0 + # Sanity: setting shared back to 0 must produce a 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", diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py index f498655347..caa6397901 100644 --- a/studio/backend/tests/test_llama_cpp_context_fit.py +++ b/studio/backend/tests/test_llama_cpp_context_fit.py @@ -114,9 +114,13 @@ def _make_backend( inst._kv_value_length = kv_value_length inst._kv_lora_rank = None inst._sliding_window = None + inst._sliding_window_pattern = None inst._ssm_inner_size = None inst._full_attention_interval = None inst._key_length_mla = None + inst._n_kv_heads_by_layer = None + inst._kv_key_length_swa = None + inst._kv_value_length_swa = None return inst @@ -137,7 +141,7 @@ def _drive( model_size = int(model_gib * GIB) cache_type_kv = None - def fake_estimate(n_ctx_, _type = None): + def fake_estimate(n_ctx_, _type = None, **_kwargs): return 0 if n_ctx_ <= 0 else n_ctx_ * kv_per_token_bytes inst._estimate_kv_cache_bytes = fake_estimate diff --git a/studio/backend/tests/test_llama_cpp_max_context_threshold.py b/studio/backend/tests/test_llama_cpp_max_context_threshold.py index 5fd0243c9f..22e4cda7d1 100644 --- a/studio/backend/tests/test_llama_cpp_max_context_threshold.py +++ b/studio/backend/tests/test_llama_cpp_max_context_threshold.py @@ -99,9 +99,13 @@ def _make_backend(native_ctx = 131072): inst._kv_value_length = 128 inst._kv_lora_rank = None inst._sliding_window = None + inst._sliding_window_pattern = None inst._ssm_inner_size = None inst._full_attention_interval = None inst._key_length_mla = None + inst._n_kv_heads_by_layer = None + inst._kv_key_length_swa = None + inst._kv_value_length_swa = None return inst @@ -114,7 +118,7 @@ def _compute_max_available_ctx(native_ctx, model_gib, gpus, kv_per_token_bytes = model_size = int(model_gib * GIB) inst._estimate_kv_cache_bytes = ( - lambda n, _t = None: 0 if n <= 0 else n * kv_per_token_bytes + lambda n, _t = None, **_kw: 0 if n <= 0 else n * kv_per_token_bytes ) inst._can_estimate_kv = lambda: True