fix KVCache estimates for gemma4 style sliding window models (#5225)
* fix KVCache estimates for gemma4 style sliding window models
Signed-off-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>
* [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 `<arch>.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 <arch>.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 <arch>.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 <venkatadattasainimmaturi@gmail.com>
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 <danielhanchen@gmail.com>
This commit is contained in:
Datta Nimmaturi
GitHub
parent
b39f4b282a
commit
6b13cab746
4 changed files with 1874 additions and 111 deletions
|
|
@ -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: <arch>.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("<I", f.read(4))[0] for _ in range(alen)]
|
||||
if atype == 5: # INT32
|
||||
return [struct.unpack("<i", f.read(4))[0] for _ in range(alen)]
|
||||
if atype == 7: # BOOL
|
||||
return [struct.unpack("<?", f.read(1))[0] for _ in range(alen)]
|
||||
|
||||
for _ in range(alen):
|
||||
LlamaCppBackend._gguf_skip_value(f, atype)
|
||||
return None
|
||||
|
||||
def _read_gguf_metadata(self, gguf_path: str) -> 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("<I", f.read(4))[0]
|
||||
|
|
@ -1063,49 +1459,141 @@ class LlamaCppBackend:
|
|||
_tensor_count, kv_count = struct.unpack("<QQ", f.read(16))
|
||||
|
||||
for _ in range(kv_count):
|
||||
key_len = struct.unpack("<Q", f.read(8))[0]
|
||||
key = f.read(key_len).decode("utf-8")
|
||||
vtype = struct.unpack("<I", f.read(4))[0]
|
||||
# Tolerate truncated input (e.g., a partial header
|
||||
# fetched via HTTP byte-range): bail out gracefully
|
||||
# so the resolver fallback still runs on whatever
|
||||
# we did manage to parse.
|
||||
try:
|
||||
key_len_bytes = f.read(8)
|
||||
if len(key_len_bytes) < 8:
|
||||
break
|
||||
key_len = struct.unpack("<Q", key_len_bytes)[0]
|
||||
key_bytes = f.read(key_len)
|
||||
if len(key_bytes) < key_len:
|
||||
break
|
||||
key = key_bytes.decode("utf-8")
|
||||
vtype_bytes = f.read(4)
|
||||
if len(vtype_bytes) < 4:
|
||||
break
|
||||
vtype = struct.unpack("<I", vtype_bytes)[0]
|
||||
except (struct.error, UnicodeDecodeError):
|
||||
break
|
||||
|
||||
if key in WANTED or key in arch_keys:
|
||||
# Read this value
|
||||
if vtype == 8: # STRING
|
||||
slen = struct.unpack("<Q", f.read(8))[0]
|
||||
val_s = f.read(slen).decode("utf-8")
|
||||
if key == "general.architecture":
|
||||
arch = val_s
|
||||
# Register arch-specific keys to look for
|
||||
arch_keys = {
|
||||
f"{arch}.context_length": "context_length",
|
||||
f"{arch}.block_count": "n_layers",
|
||||
f"{arch}.attention.head_count_kv": "n_kv_heads",
|
||||
f"{arch}.attention.head_count": "n_heads",
|
||||
f"{arch}.embedding_length": "embedding_length",
|
||||
# Architecture-aware KV cache fields
|
||||
f"{arch}.attention.key_length": "kv_key_length",
|
||||
f"{arch}.attention.value_length": "kv_value_length",
|
||||
f"{arch}.attention.sliding_window": "sliding_window",
|
||||
f"{arch}.full_attention_interval": "full_attention_interval",
|
||||
f"{arch}.attention.kv_lora_rank": "kv_lora_rank",
|
||||
f"{arch}.attention.key_length_mla": "key_length_mla",
|
||||
f"{arch}.ssm.inner_size": "ssm_inner_size",
|
||||
f"{arch}.ssm.state_size": "ssm_state_size",
|
||||
}
|
||||
elif key == "tokenizer.chat_template":
|
||||
self._chat_template = val_s
|
||||
elif vtype in (4, 10): # UINT32 or UINT64
|
||||
val_i = (
|
||||
struct.unpack("<I", f.read(4))[0]
|
||||
if vtype == 4
|
||||
else struct.unpack("<Q", f.read(8))[0]
|
||||
)
|
||||
attr = arch_keys.get(key)
|
||||
if attr:
|
||||
setattr(self, f"_{attr}", val_i)
|
||||
try:
|
||||
if key in WANTED or key in arch_keys:
|
||||
if vtype == 8: # STRING
|
||||
slen = struct.unpack("<Q", f.read(8))[0]
|
||||
val_s = f.read(slen).decode("utf-8")
|
||||
if (
|
||||
key.startswith("general.")
|
||||
and key != "general.architecture"
|
||||
):
|
||||
general[key] = val_s
|
||||
if key == "general.architecture":
|
||||
arch = val_s
|
||||
arch_keys = {
|
||||
f"{arch}.context_length": "context_length",
|
||||
f"{arch}.block_count": "n_layers",
|
||||
f"{arch}.attention.head_count_kv": "n_kv_heads",
|
||||
f"{arch}.attention.head_count": "n_heads",
|
||||
f"{arch}.embedding_length": "embedding_length",
|
||||
f"{arch}.attention.key_length": "kv_key_length",
|
||||
f"{arch}.attention.value_length": "kv_value_length",
|
||||
f"{arch}.attention.sliding_window": "sliding_window",
|
||||
f"{arch}.attention.sliding_window_pattern": "sliding_window_pattern",
|
||||
f"{arch}.full_attention_interval": "full_attention_interval",
|
||||
f"{arch}.attention.kv_lora_rank": "kv_lora_rank",
|
||||
f"{arch}.attention.key_length_mla": "key_length_mla",
|
||||
f"{arch}.attention.key_length_swa": "kv_key_length_swa",
|
||||
f"{arch}.attention.value_length_swa": "kv_value_length_swa",
|
||||
f"{arch}.attention.shared_kv_layers": "shared_kv_layers",
|
||||
f"{arch}.ssm.inner_size": "ssm_inner_size",
|
||||
f"{arch}.ssm.state_size": "ssm_state_size",
|
||||
}
|
||||
elif key == "tokenizer.chat_template":
|
||||
self._chat_template = val_s
|
||||
elif vtype in (4, 10): # UINT32 or UINT64
|
||||
val_i = (
|
||||
struct.unpack("<I", f.read(4))[0]
|
||||
if vtype == 4
|
||||
else struct.unpack("<Q", f.read(8))[0]
|
||||
)
|
||||
attr = arch_keys.get(key)
|
||||
if attr:
|
||||
if attr == "sliding_window_pattern":
|
||||
sliding_window_pattern_period = val_i
|
||||
else:
|
||||
setattr(self, f"_{attr}", val_i)
|
||||
elif vtype == 9: # ARRAY
|
||||
atype = struct.unpack("<I", f.read(4))[0]
|
||||
alen = struct.unpack("<Q", f.read(8))[0]
|
||||
val_a = self._gguf_read_array_value(f, atype, alen)
|
||||
attr = arch_keys.get(key)
|
||||
if attr == "n_kv_heads" and val_a is not None:
|
||||
self._n_kv_heads_by_layer = [int(x) for x in val_a]
|
||||
if self._n_kv_heads is None and val_a:
|
||||
self._n_kv_heads = max(int(x) for x in val_a)
|
||||
elif (
|
||||
attr == "sliding_window_pattern"
|
||||
and val_a is not None
|
||||
):
|
||||
self._sliding_window_pattern = [
|
||||
bool(x) for x in val_a
|
||||
]
|
||||
sliding_window_pattern_period = None
|
||||
else:
|
||||
self._gguf_skip_value(f, vtype)
|
||||
else:
|
||||
self._gguf_skip_value(f, vtype)
|
||||
else:
|
||||
self._gguf_skip_value(f, vtype)
|
||||
except (struct.error, UnicodeDecodeError):
|
||||
# Truncated input (e.g., HTTP byte-range fetch
|
||||
# of just the GGUF header); break so the
|
||||
# resolver fallback still runs on what we have.
|
||||
break
|
||||
|
||||
# Expand a scalar period straight from the GGUF first.
|
||||
if (
|
||||
self._sliding_window_pattern is None
|
||||
and sliding_window_pattern_period
|
||||
and self._n_layers
|
||||
):
|
||||
self._sliding_window_pattern = [
|
||||
(i + 1) % sliding_window_pattern_period != 0
|
||||
for i in range(self._n_layers)
|
||||
]
|
||||
|
||||
# Otherwise hand off to the resolver (cache / bootstrap /
|
||||
# transformers / HF). See `_resolve_swa_pattern`.
|
||||
if (
|
||||
self._sliding_window_pattern is None
|
||||
and self._sliding_window
|
||||
and self._n_layers
|
||||
):
|
||||
hf_repo_candidates = (
|
||||
general.get("general.source.huggingface.repository"),
|
||||
_hf_repo_from_url(general.get("general.source.url")),
|
||||
_hf_repo_from_url(general.get("general.source.repo_url")),
|
||||
_hf_repo_from_url(general.get("general.base_model.0.repo_url")),
|
||||
(
|
||||
f"{general['general.base_model.0.organization']}/"
|
||||
f"{general['general.base_model.0.name']}".replace(" ", "-")
|
||||
if general.get("general.base_model.0.organization")
|
||||
and general.get("general.base_model.0.name")
|
||||
else None
|
||||
),
|
||||
(
|
||||
f"{general['general.organization']}/"
|
||||
f"{general['general.basename']}".replace(" ", "-")
|
||||
if general.get("general.organization")
|
||||
and general.get("general.basename")
|
||||
else None
|
||||
),
|
||||
)
|
||||
self._sliding_window_pattern = _resolve_swa_pattern(
|
||||
arch,
|
||||
self._n_layers,
|
||||
hf_repo_candidates,
|
||||
)
|
||||
|
||||
if self._context_length:
|
||||
logger.info(f"GGUF metadata: context_length={self._context_length}")
|
||||
|
|
@ -1504,8 +1992,11 @@ class LlamaCppBackend:
|
|||
pool_mib,
|
||||
model_size,
|
||||
cache_type_kv,
|
||||
n_parallel = n_parallel,
|
||||
)
|
||||
kv = self._estimate_kv_cache_bytes(
|
||||
capped, cache_type_kv, n_parallel = n_parallel
|
||||
)
|
||||
kv = self._estimate_kv_cache_bytes(capped, cache_type_kv)
|
||||
total_mib = (model_size + kv) / (1024 * 1024)
|
||||
if total_mib <= pool_mib * 0.90:
|
||||
best_cap = max(best_cap, capped)
|
||||
|
|
@ -1529,7 +2020,7 @@ class LlamaCppBackend:
|
|||
# have surfaced the "might be slower" warning before
|
||||
# the user submitted a ctx above the fit ceiling.
|
||||
requested_total = model_size + self._estimate_kv_cache_bytes(
|
||||
effective_ctx, cache_type_kv
|
||||
effective_ctx, cache_type_kv, n_parallel = n_parallel
|
||||
)
|
||||
gpu_indices, use_fit = self._select_gpus(requested_total, gpus)
|
||||
# No silent shrink: effective_ctx stays == n_ctx.
|
||||
|
|
@ -1544,8 +2035,11 @@ class LlamaCppBackend:
|
|||
pool_mib,
|
||||
model_size,
|
||||
cache_type_kv,
|
||||
n_parallel = n_parallel,
|
||||
)
|
||||
kv = self._estimate_kv_cache_bytes(
|
||||
capped, cache_type_kv, n_parallel = n_parallel
|
||||
)
|
||||
kv = self._estimate_kv_cache_bytes(capped, cache_type_kv)
|
||||
total_mib = (model_size + kv) / (1024 * 1024)
|
||||
if total_mib <= pool_mib * 0.90:
|
||||
effective_ctx = capped
|
||||
|
|
@ -1578,7 +2072,9 @@ class LlamaCppBackend:
|
|||
)
|
||||
|
||||
if effective_ctx < original_ctx:
|
||||
kv_est = self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv)
|
||||
kv_est = self._estimate_kv_cache_bytes(
|
||||
effective_ctx, cache_type_kv, n_parallel = n_parallel
|
||||
)
|
||||
logger.info(
|
||||
f"Context auto-reduced: {original_ctx} -> {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:
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue