0e69614878* studio: reserve MTP draft VRAM in GGUF auto-fit Auto-fit advertised a context (for example ~110k for the Qwen3.6-27B MTP GGUF) that fit on paper but OOMed mid-generation or during tool calls once MTP speculative decoding was active. The MTP draft path's VRAM was reserved as a flat 5% of total VRAM, which tracks neither of the two real costs: the MTP head keeps its own attention KV cache that grows with context, and the speculative verification buffer grows with --spec-draft-n-max. On the hybrid Mamba/attention Qwen3.6 models the main KV is small, so auto-fit happily kept a near-native context while the draft path pushed the load over budget at runtime. Replace the flat fraction with a byte-accurate, context- and n_max-aware reserve sized from GGUF dims: draft KV from nextn_predict_layers and the attention dims at f16 (llama.cpp's MTP draft context uses f16 KV regardless of the main cache type), plus a verify buffer per embedding-unit per draft token. The reserve is evaluated per candidate context inside the fit binary search and added to every pin/fit check, including the tensor-parallel planner and its even-split decision. Coefficients were calibrated against llama-server VRAM measurements on the Qwen3.6-27B MTP GGUF (RMS 14 MiB). The flat fraction remains as a fallback when GGUF dims are unavailable, so non-MTP loads are unchanged. The budget now also engages when the user wires MTP through extra args (--spec-type draft-mtp, including chains), reads the effective draft depth from --spec-draft-n-max or the legacy --draft-max with extras taking precedence over the first-class field, reserves a separate drafter's weights when supplied via --model-draft/--spec-draft-model/-md, and mirrors _build_speculative_flags so it never reserves for MTP the launch resolver will not emit (needs a head/drafter and a binary that supports --spec-type mtp). Adds tests/test_mtp_vram_budget.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: total-based VRAM budget + deterministic compute-graph buffer Build on the byte-accurate MTP reserve with three changes that make the GGUF auto-fit budget deterministic across architectures and recover usable context, especially for MTP models on a single tight card. 1. Total-based budget. Cap GPU occupancy at a fraction of TOTAL VRAM rather than a fraction of FREE VRAM, and raise the fraction from 0.90 to 0.95: budget = free - (1 - 0.95) * total (per GPU, summed for a pool) The reserve is now absolute (a fixed slice of the card) instead of shrinking as the GPU fills, so a partly-used GPU keeps a constant cushion for compute/CUDA/verify buffers instead of over-promising context and spilling to CPU at runtime. _get_gpu_memory() reads memory.total alongside memory.free; _fit_context_to_vram, _select_gpus and the load_model pool loops thread the totals through. Multi-GPU layer-split pools sum(free_i - 0.05*total_i); tensor mode reserves per device. 2. Deterministic compute-graph buffer. Replace the flat 5 GB/device tensor reserve (a magic constant that over-reserved about 8x on a 27B model) with _estimate_compute_buffer_bytes, sized from GGUF dims and the launch flags: out = n_vocab * n_ubatch * 4 # vocab-width output buffer act = 4 * n_embd * n_ubatch * 4 # activation scratch pipeline_per_device = act + out * (n_parallel - 1) tensor_per_device = 2*act + out * n_parallel The buffer is context-independent and scales with --parallel (serving slots), not with how the model is split across GPUs. It is now reserved in BOTH multi-GPU paths (layer split folds one buffer into the pooled footprint; tensor mode reserves it per device). The flat 5 GB stays only as a fallback when vocab/embedding dims are unavailable. Calibrated against llama-server measurements (parallel 1/2/4/8 give 36/492/1388/3220 MiB on a single GPU; about 600 MiB/device tensor); the estimate is a small upper bound. 3. GGUF parsing. Read vocab size (tokenizer tokens array length) and feed_forward_length for the compute-buffer estimate. Effect on the Qwen3.6-27B MTP Q6_K case (MTP on): a single 32 GB card at about 31 GB free advertises f16 23k to 64k, q8_0 44k to 115k, q4_0 82k to 200k; 2x 24 GB tensor mode recovers the full 262k window for f16 (was about 134k). Validated on hardware: 1x 32 GB f16 at 64768 loads at 29.3 GB / 120 t/s; 2x 23 GB tensor f16 at 262144 loads at 22.2 GB/device / 98 t/s; both within 0.4% of the estimate. Adds test_compute_buffer.py and updates the KV/context-fit/MTP-budget tests for the 0.95 constant and the new budget. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: tighten comments in the VRAM auto-fit changes Condense the docstrings and inline comments added by this PR (internal backend helpers): drop restated-signature docstrings, fold multi-line block comments to one or two lines, and remove notes that just repeat the code. No behavior change (AST-verified comment/docstring-only via comment_tools.py); the backend test suite is unchanged and green. * studio: address review findings in the VRAM auto-fit budget Five fixes from a parallel-reviewer pass on this PR; all confirmed against the real functions and covered by new tests. - Tensor mode now honors the total-based VRAM cap. _plan_tensor_parallel took total_by_idx and budgets each GPU at free - (1-frac)*total, mirroring the layer-split paths; previously it fit against raw free and could spend the 5% safety cushion on a partly-used multi-GPU box (reproduced ~3.3 GB over). - Draft K and V cache types are parsed and accounted independently. A one-sided override (e.g. --cache-type-k-draft q4_0, V left f16) no longer applies the small quant to both axes and under-reserves the f16 axis. The embedded-head formula sizes per axis; the separate-drafter path uses the heavier type so it never under-reserves. - The compute-graph buffer honors a user --ubatch / --ubatch-size / -ub override (parsed and threaded into every _estimate_compute_buffer_bytes call and the tensor planner); it previously always assumed the 512 default, under-reserving up to ~8x at --ubatch 4096. - GPU ranking uses the usable budget (free - (1-frac)*total) instead of raw free in _select_gpus and both auto-context subset loops, so a more-used large card no longer outranks a less-used small card that has more usable room. Adds regression tests for each (tensor total cap, ubatch reserve scaling, split K/V no-under-reserve, --ubatch parser, usable-ranking GPU selection). Full targeted backend suite green (321 passed). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: gate tensor-parallel admission on usable VRAM budget The tensor-parallel GPU admission filters still used raw free VRAM after the total-based budget landed, an asymmetric fix: a partly-used large card can clear the per-device compute-buffer reserve on raw free while its usable budget (free - (1-frac)*total) does not, so the planner admitted it and the even split could emit a near-zero weight slice for a GPU that should have been excluded. - _plan_tensor_parallel: admit GPUs by usable budget, not raw free (move the _usable helper above the filter). - load_model: admit the tensor set by _gpu_usable, and downgrade to layer split when the pooled usable budget cannot hold weights plus per-device compute buffers (the planner can only floor the context, not stop an overcommitted launch). Adds regression tests: planner drops a GPU whose usable budget is below the reserve, and a source-level check that load_model admits on the usable budget and carries the pooled-weight downgrade. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: size the MTP reserve for the user's overriding drafter A user --model-draft passed in extra_args is appended last and wins at the llama-server launch, but the VRAM budget preferred Studio's auto-detected drafter (mtp_draft_path or extras), so a larger custom drafter was under-reserved. Flip the precedence to extras-first, matching the draft-depth (n_max) resolution two lines above. Adds a source-level regression test. * studio: account for MTP reserve in tensor gate, restore 2-col GPU probe Two issues found by re-review of the prior fix: - The tensor-parallel capacity gate only checked the model weights against the pooled budget, not the MTP reserve. A separate-drafter MTP load whose weights fit but weights + drafter do not could still launch overcommitted in tensor mode. Add the non-shrinkable MTP reserve (drafter weights + floor draft KV, or the flat 2 GiB fallback when dims are unavailable) to the gate. - The nvidia-smi probe was switched to a three-column query (index,free,total) for the total-based budget but required exactly three columns, so a driver or mock returning the legacy two-column "index,free" was dropped and the probe fell through to the real GPUs. Accept two columns (total 0) and treat an unknown total as the legacy free*fraction in _select_gpus. Tests: tensor gate asserts the MTP term is included; _get_gpu_memory parses both two- and three-column output; the existing two-column GPU-detection mocks pass again. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: keep the VRAM cushion in tensor planning when GPU totals are unknown _plan_tensor_parallel fell back to raw free VRAM when a GPU's total was unavailable (a two-column nvidia-smi probe reporting total 0), while _select_gpus and the load_model ranking both fall back to free*fraction. That let tensor planning spend the 5% cushion the rest of the fit preserves and over-advertise context in exactly that path. Align the fallback to free*_CTX_FIT_VRAM_FRACTION. Updates the no-totals planner test expectations (now free*frac) and adds a regression test that total 0 keeps the cushion. * studio: honor LLAMA_ARG_* env overrides and HF draft flags in the VRAM budget The budget parsed llama-server flags only from the request's extra_args, but the child process inherits Studio's full environment (child_env_without_native_path_secret copies os.environ), and llama-server honors LLAMA_ARG_* env vars for the same options. So a service-level override the child acts on was invisible to the fit, which could then advertise a context/GPU set that OOMs at load. - _extra_args_n_ubatch: fall back to LLAMA_ARG_UBATCH (drives the compute buffer; an unseen 4096 vs the 512 default under-reserves ~8x). - _extra_args_mtp_draft_path: also recognize the HF draft-repo flags (--spec-draft-hf/-hfd/-hfrd/--hf-repo-draft) and fall back to LLAMA_ARG_SPEC_DRAFT_MODEL / LLAMA_ARG_SPEC_DRAFT_HF_REPO. An HF repo isn't a local file so it can't be sized, but recognizing it routes to the flat reserve instead of mis-sizing Studio's auto/embedded drafter. - _extra_args_draft_cache_types: fall back to LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K/_V per axis. CLI extra_args win over env (they are appended last at launch). Each parser takes an injectable env for deterministic tests. Adds env-fallback and HF-flag tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: review polish - drop non-flag --ubatch, harden GPU probe, document buffer Non-blocking items from a second review pass; no behavior change in the common path: - _extra_args_n_ubatch: drop --ubatch; the binary only accepts --ubatch-size/-ub, so parsing --ubatch implied support it does not have (it would over-reserve for a launch that fails on the unknown flag). - _get_gpu_memory: skip a malformed nvidia-smi line instead of letting one bad line raise and drop the whole NVIDIA probe to the torch fallback. - _estimate_compute_buffer_bytes: document that the per-slot output-buffer model assumes a small n_outputs_max (chat decode); it would under-count for embeddings / --logits-all / reranking, which Studio does not run on this path. * studio: honor LLAMA_ARG_SPEC_TYPE when deciding the MTP reserve _extra_args_requests_mtp only checked extra_args, but the child inherits Studio's env and llama-server honors LLAMA_ARG_SPEC_TYPE. So a service-level LLAMA_ARG_SPEC_TYPE=draft-mtp would run MTP while the fit skipped the draft reserve and could advertise a context/GPU set that OOMs at load. Recognize the env value (CLI still wins). Completes the env-override coverage alongside ubatch, draft model, and draft cache types. Adds an env regression test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: reserve VRAM for non-MTP model-based draft modes too The draft reserve only engaged for MTP. A user passing a non-MTP model-based draft mode (--spec-type draft-simple / draft-eagle3) with a --model-draft loads a separate draft model whose weights + KV consume GPU memory, but the fit reserved nothing and could OOM at load. Engage the existing drafter reserve for those modes when extras (or LLAMA_ARG_SPEC_TYPE) name a drafter; ngram-* load no model and are unaffected. Purely additive (reserves where there was none). Adds parser + gate tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: floor quantized embedded MTP draft KV at f16; fix two test issues Address PR review feedback (three findings): 1. Quantized embedded MTP draft KV was underpriced. The embedded head is a single draft layer, so llama.cpp cannot amortize quantized-KV overhead over many layers the way the main model does: a quantized draft KV (e.g. --spec-draft-type-k q4_0) actually fits LESS context than f16, not more (ggml-org/llama.cpp#24102, where a collaborator recommends f16 for the draft KV). Pricing q4_0 at 0.5625 of an element (~28% of f16) under-reserved, so a quantized override could advertise a context that shrinks or OOMs at load. Floor the embedded draft KV bytes-per-element at f16 (quantized types priced as f16, f32 still its full 4 bytes). The separate multi-layer drafter, where quantization does amortize, keeps the user's real type. 2. test_load_model_reserves_for_non_mtp_draft_modes asserted an exact one-line source substring that pre-commit black wrapped across lines, breaking CI. Strip whitespace before matching so the check survives any line-wrapping. 3. test_compute_buffer.py installed a partial httpx stub via setdefault that, if collected before test_kv_cache_estimation.py, leaked into sys.modules without HTTPError/Response and could break the transformers introspection tier by collection order. Adopt the sister file's pattern: only stub when real httpx is absent, and include the full symbol set. Updates the affected draft-KV tests to assert the f16 floor. * studio: guard httpx stub in test_mtp_vram_budget too test_mtp_vram_budget.py installed a partial httpx stub via setdefault that, like test_compute_buffer.py before it, lacked HTTPError/Response and could leak into sys.modules ahead of tests that need huggingface_hub/transformers, breaking the introspection tier by collection order. Apply the same guard used by test_kv_cache_estimation.py: only stub when real httpx is absent, with the full symbol set. * studio: per-device layer-split reserve, effective spec-type, drafter weights, KV restore Address PR review feedback (four findings in the auto-fit budget): A. Reserve the per-device layer-split overhead. A layer (pipeline) split allocates a fixed per-device overhead (CUDA context + per-device compute scratch) on every participating GPU, beyond the slot-scaling compute buffer that is conserved across the split. Measured ~0.9 GB/device on the Qwen3.6-27B GGUF (b9625), independent of --parallel: layer-split TOTAL VRAM grew +894 MiB (parallel=8) / +946 MiB (parallel=1) per extra GPU, ~linear to +2.6 GB at 4 GPUs. The fit folded a single compute buffer for all subset sizes, so a k-GPU layer split was short by ~(k-1)*0.9 GB and could pin a context that fits the pool on paper but OOMs a device. Reserve (k-1) * _PIPELINE_PER_DEVICE_OVERHEAD_MIB per subset in the layer-split fit; k=1 adds nothing, so single-GPU sizing (and the validated benchmark rows) is unchanged. B. Track the effective --spec-type. _extra_args_requests_mtp returned true on the first MTP-ish --spec-type and consulted LLAMA_ARG_SPEC_TYPE even when a CLI --spec-type was present, contrary to llama.cpp (last CLI value wins; a CLI flag overrides the env). So `--spec-type draft-mtp --spec-type ngram-mod` or a non-MTP CLI value with a stale MTP env over-reserved a drafter the launch won't load (shrinking context / selecting extra GPUs). Route both detectors through a new _effective_spec_type helper. C. Keep known drafter weights in the fallback reserve. When a separate drafter's KV metadata can't be sized, _estimate_mtp_overhead_bytes returned None and discarded the drafter's known weight bytes, falling back to the flat 5% reserve; a drafter larger than that cushion could launch over budget and OOM. Reserve the known weights even when KV sizing fails (None only when nothing is known). D. Restore quantized KV on tensor->layer-split downgrade. The tensor attempt drops a quantized KV cache (tensor mode aborts on it). When the GPU-count or capacity gate then downgrades to layer split -- which supports quantized KV -- the dropped type was lost and the launch used f16, using more VRAM and shrinking context. Remember the dropped type and restore it on downgrade (the launch re-emits it from the var). Adds regression tests for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: per-device overhead in GPU pin, skip CPU draft, gate env spec-type Address PR review feedback (three follow-up findings): F1. Reserve the per-device layer-split overhead in the pin path too. The earlier per-device reserve was added to the auto-context fit loops but not to _select_gpus, which the explicit-ctx and file-size-only paths use to PIN GPUs with -ngl -1 (no --fit fallback). A 2+ GPU pin within ~1 GiB/extra-GPU of the budget could OOM a device at load. Add a per_device_overhead_bytes arg to _select_gpus so a k-GPU pin must hold model + (k-1)*overhead; pass the pipeline overhead at both pin call sites. Single-GPU pins are unchanged. F2. Don't charge a CPU-offloaded drafter against the GPU budget. A user passing --spec-draft-ngl 0 or --spec-draft-device none/cpu keeps the separate draft model's weights + KV on CPU, but the budget still charged the full drafter GGUF size, auto-reducing context or downgrading GPU selection. Detect the CPU-offload flags and drop the separate drafter (and its flat fallback) from the budget; an embedded head follows the main -ngl and is unaffected. F3. Consult LLAMA_ARG_SPEC_TYPE only when it can reach the child. llama-server's CLI args override env, and _build_speculative_flags emits a --spec-type/--spec-default for every UI mode except "off". So a stale MTP env on a non-MTP model (auto mode) made the fit reserve MTP that the emitted --spec-default disables, shrinking context / picking extra GPUs. Gate the env consult on "no user --spec-type and UI mode off"; the MTP-model auto path still engages via Studio's own detection. Adds regression tests for each. * studio: drafter budget precedence and --spec-default in effective spec-type Two spec-precedence fixes surfaced by an independent multi-reviewer pass: R3. Size the drafter the launch actually loads. _mtp_draft_for_budget consulted LLAMA_ARG_SPEC_DRAFT_MODEL (via _extra_args_mtp_draft_path's env fallback) before Studio's resolved mtp_draft_path, but _build_speculative_flags emits --model-draft mtp_draft_path, which overrides the env at launch. With a stale (smaller) env drafter, the budget under-reserved and could OOM. Order the budget by what actually launches: CLI extras --model-draft (appended last, wins), then Studio's emitted mtp_draft_path (when MTP engages and the user doesn't own --spec-type), then the env drafter. R4. Treat --spec-default as a CLI spec override in _effective_spec_type. It only recognized --spec-type, so extras=["--spec-default"] with LLAMA_ARG_SPEC_TYPE= draft-mtp fell through to the env and over-reserved MTP, even though the CLI --spec-default overrides the env to a non-MTP default. Recognize it as a CLI spec flag (resolves to "default", non-MTP) that suppresses the env fallback. Adds regression tests for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: refine MTP draft reserve (parallel slots, last-wins, KV cushion, ranking) Address PR review feedback (five follow-up findings, all edges of this session's earlier MTP/auto-fit changes): G1. Price the separate drafter's KV per --parallel slot. _mtp_draft_kv_bytes called the drafter's _estimate_kv_cache_bytes with the default n_parallel=1, but the drafter is served under the main model's slot count; a sliding-window drafter (Gemma) grows KV per slot and was under-reserved. Thread n_parallel through the draft KV / overhead estimate and the fit closure. G2. Honor last-wins for the draft-offload flags. _extra_args_draft_offloaded_to_cpu returned True on the first CPU value, so --spec-draft-ngl 0 --spec-draft-ngl -1 (final = GPU) wrongly dropped the drafter reserve while the server kept it on GPU -> OOM. Decide on the final value of each flag only. G3. Keep the flat cushion when only the drafter weights could be sized. The weights fallback installs mtp_overhead_fn, which made callers drop the flat MTP reserve, leaving the still-unsized draft KV with no cushion. Keep the flat fraction on in that weights-only case, on top of the byte-accurate weights. G4. Rank auto/cap GPU subsets by the active budget fraction. The ranking used a hard-coded 0.95 while the fit tests _pin_fraction (lowered by the flat MTP reserve); on mixed-total GPUs that could order subsets differently and pick a worse plan. Rank with the same fraction the fit uses. G5. Keep the embedded-head flat reserve under a draft CPU-offload flag. F2's not-_draft_on_cpu guard also dropped the reserve for an embedded MTP head, which is part of the main model and stays on GPU regardless of --spec-draft-ngl. Only suppress the flat reserve for a CPU-offloaded separate drafter (no embedded head). Adds regression tests for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: keep GPU on non-integer total, keep tensor flat reserve for weights-only Two review findings: - _get_gpu_memory dropped a whole GPU when nvidia-smi reported a non-integer memory.total ("N/A" on some drivers / MIG / vGPU): index, free and total were parsed in one try/except that skipped the line on any ValueError, so the GPU vanished from the probe and the load could silently spill to CPU. Parse index and free (required) first, then total separately, defaulting to 0 (the fit then uses the free*frac path for that GPU). Adds N/A and bad-free test cases. - Tensor planning skipped the flat MTP reserve for a weights-only drafter (file size known, KV unsizable): the capacity gate used the byte floor whenever mtp_overhead_fn was set, so it reserved only the drafter weights and no draft KV. Tensor mode has no --fit valve, so that could overcommit and OOM. Keep the flat reserve (never below the byte floor) in the weights-only case too, mirroring the layer-split _mtp_kv_unsized handling. Adds a regression test. (A third suggestion -- fold --batch-size into the compute-buffer reserve -- was checked on hardware and declined: -b 8192 -ub 512 used identical VRAM to the default at -c 64000, so the logical batch does not size the graph buffer; the estimate correctly uses the physical micro-batch.) * studio: budget the main KV from LLAMA_ARG_CACHE_TYPE env when Studio emits none The child inherits LLAMA_ARG_CACHE_TYPE_K / LLAMA_ARG_CACHE_TYPE_V, but Studio emits --cache-type-k/-v only when the param or extras set the type. When neither does, a heavier env type (f32) reaches the child while the auto-fit budget assumed the f16 default, under-reserving the main KV and risking OOM at the advertised context. This is the one main-KV axis that lacked the env-aware handling the other axes already have (spec-type, draft model, draft cache type, ubatch). load_model now adopts the heavier of the two env types when it exceeds f16 (only f32 does), and the launch re-emits it so child and budget stay byte-consistent. Quantized env types are <= f16 and remain safely over-reserved by the default, so they are left untouched (no change). A single value is used because the budget's KV estimate has one cache_type_kv knob, matching parse_cache_override's existing key/value collapse. Adds _env_main_cache_type_for_budget plus regression tests covering f32 adoption, the K/V heavier-of collapse, quantized/unknown no-ops, and the load_model source precedence. * studio: budget tensor parallel when LLAMA_ARG_SPLIT_MODE env selects it Studio emits --split-mode tensor only on its tensor branch; the default layer-split path emits nothing and resolve_tensor_parallel consults only extras. The child inherits LLAMA_ARG_SPLIT_MODE, so a tensor env on a layer-split plan silently runs the child tensor-parallel (heavier per-device compute buffer) while the budget reserved only the layer-split per-device overhead, under- reserving on multi-GPU. load_model now flips the plan to tensor when extras do not set a split mode and the env selects tensor, so Studio plans, reserves, and emits tensor consistently. The flip is one-directional (guarded on not tensor_parallel and no extras split-mode) so an existing tensor plan is never downgraded and extras keep precedence. Other env modes (layer/row/none) are not a runtime-heavier surprise and are left untouched. Adds _env_split_mode_is_tensor plus unit and load_model source-level tests. * studio: reconcile inherited llama.cpp env with the budgeted launch decision Addresses a review pass over the VRAM auto-fit work. The budget now sizes the right amount, but the child process inherits LLAMA_ARG_* env (see child_env_without_native_path_secret), and a few axes could still run the child in a mode Studio neither chose nor budgeted. Mixed known/unknown GPU totals over-advertised the pooled layer-split budget. _pool_budget_mib pooled free and total separately, so an unknown-total GPU (MIG/vGPU/N/A) contributed its full free with no cushion when mixed with known-total GPUs (~(1-frac)*free over-advertise, about 500 MiB in a two-GPU case). It now sums each GPU's own usable budget, and the layer-split fit calls take that as an absolute budget (budget_frac=1.0, total_mib=None) so the fit and the footprint check agree. All-known-total pools are unchanged. LLAMA_ARG_SPLIT_MODE=tensor survived a tensor-to-layer downgrade. The downgrade only stripped CLI extras, so the inherited env still ran the child tensor while Studio budgeted layer split. When the final decision is layer split, a non-layer inherited split mode (and any paired LLAMA_ARG_TENSOR_SPLIT) is now cleared from the child env. Inherited quantized LLAMA_ARG_CACHE_TYPE_K/_V crashed tensor mode. Tensor mode aborts on a quantized KV cache; Studio drops a quantized cache_type_kv for the tensor attempt but the inherited env reached the child anyway. When the final decision is tensor split, a quantized cache-type env is now cleared so the child uses the tensor-safe default that was budgeted. Env-derived cache budget no longer mutates the emitted launch flags. An env-only main KV type now informs the budget only; it is not re-emitted, so an asymmetric K=f32,V=f16 env reaches the child as set instead of being rewritten to symmetric --cache-type-k/-v f32. Adds source-level regression tests for all four and confirms the documented single-GPU/tensor/pipeline numbers are byte-identical before and after. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: tighten comments in the VRAM auto-fit code Compress the verbose docstrings and inline comments added by this work to succinct 2-4 line versions, drop restated/obvious ones, and cut duplicated rationale across the two tensor-downgrade branches. Keeps the non-obvious intent (env-inheritance precedence, the #24102 embedded-draft floor, the per-device overhead and pool-budget rationale) while removing roughly 120 lines of comment text from llama_cpp.py. Also trims the few longest test-comment blocks; concise per-test scenario notes are left intact. No logic change: verified with comment_tools.py check --strip-docstrings (code- only signature unchanged vs the prior commit) and the full backend suite still passes (824). * studio: lock in env-drafter engagement for the separate-draft reserve A review suggested an env-provided LLAMA_ARG_SPEC_DRAFT_MODEL would skip the draft reserve and OOM. It does not: the gate's _extra_args_mtp_draft_path(extra_args) call defaults env=None, which consults os.environ, so an env-only drafter still sets _user_draft_via_extras and is sized via _env_draft_for_budget. Add a source guard that the gate keeps the env-inclusive form (not extras-only env={}) and a behavioral test mirroring the reviewed scenario, so a future cleanup can't regress it. No production change. * studio: carry the unsized MTP reserve and env split/offload into tensor planning Addresses a review pass over the multi-GPU and env-inheritance paths. Tensor planner dropped the unsized draft-KV cushion. When a separate drafter has known weights but unreadable KV metadata, _plan_tensor_parallel receives a non-None weights-only mtp_overhead_fn and applied the flat 2 GiB reserve only for the no-fn case, so its binary search spent the unsized-KV cushion on context and over-advertised. Add mtp_flat_reserve_bytes (subtracted from the pooled budget and the even-split check), and pass it from load_model whenever _mtp_kv_unsized. The layer path and the tensor pre-gate already kept this cushion. Stale LLAMA_ARG_TENSOR_SPLIT survived in tensor mode. When the planner picks an even split it emits no --tensor-split, so an inherited tensor-split env reached the child and overrode the budgeted split. The layer downgrade branch cleared it; the tensor branch now does too. Env-only draft CPU offload was ignored. _extra_args_draft_offloaded_to_cpu checked extras but not LLAMA_ARG_N_GPU_LAYERS_DRAFT, so an env-offloaded drafter was still charged GPU budget and under-advertised context. It now consults that env (the device flag has no env), called with env=os.environ. Layer-split compute buffer had no fallback when GGUF dims are missing. The estimate returns 0 then, so the layer path folded no buffer while the tensor path falls back to the flat reserve. Use the flat reserve for the layer path too (a safe upper bound, since the tensor buffer >= the layer one). All four are gated on conditions the documented benchmarks don't hit; the single-GPU/tensor/pipeline reconfirm numbers are byte-identical, and the full backend suite passes (830) with regression tests for each fix. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: share the env-aware tensor decision across load and dedup matchers A review pass found the inherited-LLAMA_ARG_SPLIT_MODE=tensor flip lived only in load_model, so the two duplicate-load matchers disagreed with it. Consolidate the decision into _effective_tensor_parallel (extras + toggle, then flip on when extras set no split mode and the child inherits a tensor split env). load_model, the backend matcher (_already_in_target_state) and the route matcher (_request_matches_loaded_settings) now all call it. Before, an env-driven tensor server compared against resolve_tensor_parallel (env-blind) in both matchers, so a follow-up load that should dedup was seen as a mismatch and the healthy server was needlessly killed and reloaded. Also finish the tensor cache-type handling: when the tensor attempt drops a quantized KV it now re-adopts a heavier inherited env cache type (f32) for the budget, mirroring the initial adoption; and the two layer-split downgrades clear _cache_type_from_env so the restored quantized type is actually re-emitted rather than left to a stale inherited env. All gated on inherited env the documented benchmarks don't set; the single-GPU, tensor and pipeline reconfirm numbers are byte-identical, and the full backend suite passes (832) with unit + source regression tests for the shared helper and the route matcher. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: complete the env-aware tensor/spec handling across all paths A second review pass found the env-aware tensor/MTP handling was applied asymmetrically: some paths inherited LLAMA_ARG_* env, others didn't. Three real follow-ups, plus a small consolidation so the env semantics live in one place. 1. Tensor fallback ignored the inherited tensor env. load_with_tensor_fallback computed its retry gate with the env-blind resolve_tensor_parallel, so an env-only tensor load (toggle off, no --split-mode extra) that crashed on a tensor-incompatible GGUF re-raised instead of retrying layer split. It now uses the env-aware decision; and since the inherited env would otherwise re-engage tensor on the retry (CLI args persist, the env does too), the retry forces --split-mode layer (CLI wins over env) so it can't re-crash. 2. Duplicate-load matchers looped reloads after a tensor->layer downgrade. Both matchers compared the env-expanded tensor decision against the loaded server, but load_model may downgrade tensor to layer (capacity/buffer) and scrub the child env. The still-set parent env then made every identical request look like a mismatch, killing and reloading a healthy layer server. Add _tensor_parallel_matches_loaded, which only lets an inherited tensor env raise a match against a server that actually launched tensor; a downgraded server matches the same request (an identical load would downgrade the same way). 3. MTP binary-capability fallback leaked an inherited LLAMA_ARG_SPEC_TYPE. When the binary lacks MTP, _emit_mtp degraded but emitted no spec flag, so an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp still reached the child and attempted MTP the gate had budgeted off. It now emits --spec-default (CLI wins over env) like the sibling no-head / non-MTP fallbacks. Consolidation: moved _env_split_mode_is_tensor / _effective_tensor_parallel into llama_server_args.py (with the new _tensor_parallel_matches_loaded) so the lightweight tensor_fallback module can share them without importing llama_cpp; llama_cpp re-exports them for back-compat. All gated on inherited env the documented benchmarks don't set; the single-GPU, tensor and pipeline reconfirm numbers are byte-identical, and the full backend suite passes (883) with regression tests for each fix. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: budget the heavier axis of asymmetric --cache-type-k/-v extras A review pass found the explicit-extras counterpart of the env cache-type fix. load_model adopts the heavier inherited LLAMA_ARG_CACHE_TYPE_K/_V env for the reserve, but the explicit-extras path used resolve_cache_type_kv, which collapses both axes to one last-wins value. So extras such as --cache-type-k f32 --cache-type-v f16 (lighter axis last) budgeted f16 for both axes while the child allocates f32 on K, over-advertising context and re-opening the OOM path this PR closes. Add parse_cache_override_per_axis (keeps the K/V last-wins values apart) and _extra_args_main_cache_type_for_budget (the heavier of the two by bytes/elem), and budget from it. The user's extras are appended last and win per axis at the child, so this only raises the reserve; the emitted command and the asymmetric child cache are unchanged, and the common single-axis / symmetric cases resolve to the same type as before. Reconfirm numbers (single-GPU table, tensor, pipeline) are byte-identical, and the full backend suite passes (892) with per-axis parser and heavier-axis budget regression tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: fix tensor-safety masking and strip inherited HF drafter selectors A review pass found two more env/extras edge cases on the speculative and tensor cache paths. Tensor-safety could miss a quantized axis. The previous change budgets the heavier-by-bytes cache type, but that masks a quantized axis paired with a heavier one: --cache-type-k f16 --cache-type-v q4_0 resolves to f16, so the tensor-safety block did not fire and the q4_0 axis survived into tensor mode, which aborts on quantized KV. Test each explicit --cache-type-k/-v axis (not just the budget type) so any quantized axis drops the cache for the tensor attempt. Inherited HF drafter selectors were not stripped. _extra_args_mtp_draft_path treats --spec-draft-hf / -hfd / -hfrd / --hf-repo-draft as drafter selectors, but _SPEC_FLAGS only stripped the local --model-draft selectors, so on an inherited- extras Apply a stale HF drafter survived and last-wins-overrode Studio's re-derived spec choice. Add the HF aliases to _SPEC_FLAGS. The per-drafter tuning knobs (--spec-draft-type-*, -ngld, --spec-draft-device) are intentionally left in place: the VRAM budget reads them via the same parsers the child honors, so they stay consistent on inherit, and stripping them would silently move a CPU-offloaded drafter back onto the GPU. A third flagged item -- that the HF draft env var should be LLAMA_ARG_HFD_REPO -- was a false positive from a stale manpage; the bundled binary's common/arg.cpp sets LLAMA_ARG_SPEC_DRAFT_HF_REPO for --spec-draft-hf, which the code already uses, so it is left unchanged. Reconfirm numbers (single-GPU table, tensor, pipeline) are byte-identical, and the full backend suite passes (899) with regression tests for both fixes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: preserve asymmetric cache on tensor downgrade and skip CPU-drafter reserve A review pass found two more tensor-path edges, one a regression from the per-axis cache change. Tensor-to-layer downgrade collapsed asymmetric cache extras. The per-axis tensor-safety check strips an asymmetric --cache-type-k/-v (tensor rejects quantized KV), but the downgrade restored only the scalar heavier type, so a layer fallback silently rewrote --cache-type-k q4_0 --cache-type-v f16 to symmetric f16/f16 even though layer split supports the original. Save the original extras before the tensor strip and restore them verbatim (minus the user --split-mode) on both downgrade points; the budget still uses the heavier scalar, the child gets the real asymmetric cache. Before the per-axis change this case happened to survive (last-wins was f16, untouched), so this restores that. Tensor mode reserved GPU VRAM for a CPU-offloaded drafter. The layer path drops the flat MTP reserve when the only drafter is a separate CPU one with no embedded head, but the tensor capacity gate and planner still charged it, under-advertising context. Gate the tensor reserve on the same condition via _mtp_reserves_gpu. Reconfirm numbers (single-GPU table, tensor, pipeline) are byte-identical (both fixes are gated on conditions the benchmarks don't hit), and the full backend suite passes (901) with regression tests for each. * studio: drop now-unused llama_server_args imports from llama_cpp The refactor re-pointed load_model and the matchers off resolve_tensor_parallel / resolve_cache_type_kv and moved the env split-mode helper into llama_server_args, leaving those three names imported but unused in llama_cpp. The repo's import-hoist safety-net lint blocks that, so drop them; the env split-mode test now imports _env_split_mode_is_tensor from its real home (llama_server_args). --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2112 lines
77 KiB
Python
2112 lines
77 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
|
|
|
"""Tests for 5-path architecture-aware KV cache VRAM estimation.
|
|
|
|
Covers the GGUF metadata parser, _can_estimate_kv gate, all 5 estimation
|
|
paths (MLA, Hybrid Mamba, Sliding Window, Standard GQA, Legacy), KV cache
|
|
quantization, edge cases, and lifecycle (init/unload/reparse).
|
|
|
|
No GPU, network, or libraries beyond pytest. Cross-platform.
|
|
"""
|
|
|
|
import io
|
|
import json
|
|
import struct
|
|
import sys
|
|
import types as _types
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
# Stub heavy / unavailable deps before importing the module under test.
|
|
# Same pattern as test_native_context_length.py.
|
|
|
|
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
|
if _BACKEND_DIR not in sys.path:
|
|
sys.path.insert(0, _BACKEND_DIR)
|
|
|
|
# loggers
|
|
_loggers_stub = _types.ModuleType("loggers")
|
|
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
|
sys.modules.setdefault("loggers", _loggers_stub)
|
|
|
|
# structlog
|
|
_structlog_stub = _types.ModuleType("structlog")
|
|
sys.modules.setdefault("structlog", _structlog_stub)
|
|
|
|
# httpx -- only stub when the real library is missing. Unconditional stubbing
|
|
# shadows HTTPError/Response that huggingface_hub.errors imports at load time,
|
|
# silently breaking the transformers introspection tier.
|
|
try:
|
|
import httpx as _httpx_real # noqa: F401
|
|
except ImportError:
|
|
_httpx_stub = _types.ModuleType("httpx")
|
|
for _exc_name in (
|
|
"ConnectError",
|
|
"TimeoutException",
|
|
"ReadTimeout",
|
|
"ReadError",
|
|
"RemoteProtocolError",
|
|
"CloseError",
|
|
"HTTPError",
|
|
"RequestError",
|
|
):
|
|
setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
|
|
|
|
class _FakeTimeout:
|
|
def __init__(self, *a, **kw):
|
|
pass
|
|
|
|
_httpx_stub.Timeout = _FakeTimeout
|
|
_httpx_stub.Response = type("Response", (), {})
|
|
_httpx_stub.Client = type(
|
|
"Client",
|
|
(),
|
|
{
|
|
"__init__": lambda self, **kw: None,
|
|
"__enter__": lambda self: self,
|
|
"__exit__": lambda self, *a: None,
|
|
},
|
|
)
|
|
sys.modules["httpx"] = _httpx_stub
|
|
|
|
from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend
|
|
|
|
# Helpers
|
|
|
|
|
|
def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes:
|
|
"""Build a minimal GGUF v3 blob with the given KV metadata.
|
|
|
|
Supports the scalar and simple array metadata the parser uses.
|
|
"""
|
|
buf = io.BytesIO()
|
|
# Header: magic, version, tensor_count, kv_count
|
|
buf.write(struct.pack("<I", 0x46554747)) # GGUF magic
|
|
buf.write(struct.pack("<I", 3)) # version 3
|
|
buf.write(struct.pack("<Q", 0)) # tensor_count
|
|
buf.write(struct.pack("<Q", len(kv_pairs)))
|
|
|
|
for key, val in kv_pairs.items():
|
|
key_bytes = key.encode("utf-8")
|
|
buf.write(struct.pack("<Q", len(key_bytes)))
|
|
buf.write(key_bytes)
|
|
if isinstance(val, str):
|
|
buf.write(struct.pack("<I", 8)) # STRING
|
|
val_bytes = val.encode("utf-8")
|
|
buf.write(struct.pack("<Q", len(val_bytes)))
|
|
buf.write(val_bytes)
|
|
elif isinstance(val, list):
|
|
buf.write(struct.pack("<I", 9)) # ARRAY
|
|
is_bool_array = all(isinstance(x, bool) for x in val)
|
|
buf.write(struct.pack("<I", 7 if is_bool_array else 5))
|
|
buf.write(struct.pack("<Q", len(val)))
|
|
if is_bool_array:
|
|
for item in val:
|
|
buf.write(struct.pack("<?", item))
|
|
else:
|
|
for item in val:
|
|
buf.write(struct.pack("<i", item))
|
|
elif isinstance(val, int):
|
|
if val <= 0xFFFFFFFF:
|
|
buf.write(struct.pack("<I", 4)) # UINT32
|
|
buf.write(struct.pack("<I", val))
|
|
else:
|
|
buf.write(struct.pack("<I", 10)) # UINT64
|
|
buf.write(struct.pack("<Q", val))
|
|
else:
|
|
raise TypeError(f"Unsupported value type: {type(val)}")
|
|
return buf.getvalue()
|
|
|
|
|
|
def _backend_from_gguf(
|
|
arch: str,
|
|
fields: dict,
|
|
general: dict | None = None,
|
|
) -> LlamaCppBackend:
|
|
"""Create a LlamaCppBackend with parsed GGUF metadata from given fields.
|
|
|
|
`general` injects extra `general.*` metadata, to verify the dynamic
|
|
SWA resolver picks up source-repo hints from GGUFs that ship them.
|
|
"""
|
|
kv = {"general.architecture": arch}
|
|
for k, v in (general or {}).items():
|
|
kv[k] = v
|
|
for k, v in fields.items():
|
|
kv[f"{arch}.{k}"] = v
|
|
import tempfile, os
|
|
|
|
data = _make_gguf_bytes(arch, kv)
|
|
fd, path = tempfile.mkstemp(suffix = ".gguf")
|
|
try:
|
|
os.write(fd, data)
|
|
os.close(fd)
|
|
b = LlamaCppBackend()
|
|
b._read_gguf_metadata(path)
|
|
return b
|
|
finally:
|
|
os.unlink(path)
|
|
|
|
|
|
# A. GGUF Parser Tests
|
|
|
|
|
|
class TestGGUFParserNewFields:
|
|
"""Architecture-aware fields are parsed correctly."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"field,gguf_key,value",
|
|
[
|
|
("_kv_key_length", "attention.key_length", 128),
|
|
("_kv_value_length", "attention.value_length", 128),
|
|
("_sliding_window", "attention.sliding_window", 1024),
|
|
("_full_attention_interval", "full_attention_interval", 4),
|
|
("_kv_lora_rank", "attention.kv_lora_rank", 512),
|
|
("_key_length_mla", "attention.key_length_mla", 256),
|
|
("_ssm_inner_size", "ssm.inner_size", 6144),
|
|
("_ssm_state_size", "ssm.state_size", 128),
|
|
],
|
|
)
|
|
def test_field_parsed(self, field, gguf_key, value):
|
|
b = _backend_from_gguf("testarch", {gguf_key: value})
|
|
assert getattr(b, field) == value
|
|
|
|
def test_missing_fields_are_none(self):
|
|
b = _backend_from_gguf("testarch", {"block_count": 10})
|
|
for attr in [
|
|
"_kv_key_length",
|
|
"_kv_value_length",
|
|
"_sliding_window",
|
|
"_sliding_window_pattern",
|
|
"_full_attention_interval",
|
|
"_kv_lora_rank",
|
|
"_key_length_mla",
|
|
"_kv_key_length_swa",
|
|
"_kv_value_length_swa",
|
|
"_ssm_inner_size",
|
|
"_ssm_state_size",
|
|
]:
|
|
assert getattr(b, attr) is None
|
|
|
|
def test_array_fields_parsed(self):
|
|
b = _backend_from_gguf(
|
|
"gemma4",
|
|
{
|
|
"block_count": 6,
|
|
"attention.head_count_kv": [8, 8, 8, 8, 8, 2],
|
|
"attention.sliding_window_pattern": [
|
|
True,
|
|
True,
|
|
True,
|
|
True,
|
|
True,
|
|
False,
|
|
],
|
|
},
|
|
)
|
|
# Per-layer KV head count is preserved exactly...
|
|
assert b._n_kv_heads_by_layer == [8, 8, 8, 8, 8, 2]
|
|
# ...and mirrored into the scalar field as a conservative max, so
|
|
# non-SWA paths and callers using `n_kv = self._n_kv_heads or ...`
|
|
# get a safe upper bound.
|
|
assert b._n_kv_heads == 8
|
|
assert b._sliding_window_pattern == [True, True, True, True, True, False]
|
|
|
|
|
|
class TestArchSwaPatternDefaults:
|
|
"""Bootstrap arch table fires when GGUF reports `sliding_window` but
|
|
no per-layer pattern (true for every Gemma 2/3/3n/gpt-oss GGUF today)."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"arch,n_layers,expected_period",
|
|
[
|
|
("gemma2", 26, 2),
|
|
("gemma3", 18, 6),
|
|
("gemma3n", 35, 5),
|
|
("gpt_oss", 24, 2),
|
|
("cohere2", 32, 4),
|
|
],
|
|
)
|
|
def test_arch_default_pattern_applied(self, arch, n_layers, expected_period):
|
|
b = _backend_from_gguf(
|
|
arch,
|
|
{
|
|
"block_count": n_layers,
|
|
"attention.head_count": 4,
|
|
"attention.head_count_kv": 1,
|
|
"attention.key_length": 256,
|
|
"attention.value_length": 256,
|
|
"attention.sliding_window": 512,
|
|
},
|
|
)
|
|
expected_pattern = [(i + 1) % expected_period != 0 for i in range(n_layers)]
|
|
assert (
|
|
b._sliding_window_pattern == expected_pattern
|
|
), f"{arch} should expand to period={expected_period}"
|
|
|
|
def test_unknown_arch_no_default(self):
|
|
b = _backend_from_gguf(
|
|
"totallymadeupv7",
|
|
{
|
|
"block_count": 24,
|
|
"attention.head_count": 4,
|
|
"attention.head_count_kv": 1,
|
|
"attention.key_length": 128,
|
|
"attention.value_length": 128,
|
|
"attention.sliding_window": 1024,
|
|
},
|
|
)
|
|
assert b._sliding_window_pattern is None
|
|
|
|
def test_explicit_pattern_overrides_arch_default(self):
|
|
# gemma3 default is period=6; the explicit array must win.
|
|
b = _backend_from_gguf(
|
|
"gemma3",
|
|
{
|
|
"block_count": 6,
|
|
"attention.head_count": 4,
|
|
"attention.head_count_kv": 1,
|
|
"attention.key_length": 256,
|
|
"attention.value_length": 256,
|
|
"attention.sliding_window": 512,
|
|
"attention.sliding_window_pattern": [
|
|
True,
|
|
False,
|
|
True,
|
|
False,
|
|
True,
|
|
False,
|
|
],
|
|
},
|
|
)
|
|
assert b._sliding_window_pattern == [True, False, True, False, True, False]
|
|
|
|
def test_no_sliding_window_no_pattern(self):
|
|
b = _backend_from_gguf(
|
|
"gemma3",
|
|
{
|
|
"block_count": 18,
|
|
"attention.head_count": 4,
|
|
"attention.head_count_kv": 1,
|
|
"attention.key_length": 256,
|
|
"attention.value_length": 256,
|
|
# no sliding_window key
|
|
},
|
|
)
|
|
assert b._sliding_window_pattern is None
|
|
|
|
@pytest.mark.parametrize(
|
|
"arch", ["llama", "qwen2", "qwen3", "mistral", "mistral3", "glm4", "llama4"]
|
|
)
|
|
def test_non_swa_arch_uses_full_attention_path(self, arch):
|
|
# Pure-GQA arches: no sliding_window, no synthetic pattern,
|
|
# estimator hits Path 4.
|
|
b = _backend_from_gguf(
|
|
arch,
|
|
{
|
|
"block_count": 32,
|
|
"attention.head_count": 32,
|
|
"attention.head_count_kv": 8,
|
|
"attention.key_length": 128,
|
|
"attention.value_length": 128,
|
|
"embedding_length": 4096,
|
|
},
|
|
)
|
|
assert b._sliding_window_pattern is None
|
|
assert b._sliding_window is None
|
|
kv = b._estimate_kv_cache_bytes(8192, "f16")
|
|
gqa_expected = 32 * 8192 * 8 * (128 + 128) * 2
|
|
assert kv == gqa_expected
|
|
|
|
def test_arch_default_reduces_kv_estimate_vs_legacy(self):
|
|
common = {
|
|
"block_count": 62,
|
|
"attention.head_count": 32,
|
|
"attention.head_count_kv": 16,
|
|
"attention.key_length": 128,
|
|
"attention.value_length": 128,
|
|
"attention.sliding_window": 1024,
|
|
"embedding_length": 5376,
|
|
}
|
|
with_default = _backend_from_gguf("gemma3", common)
|
|
# Arch not in table -> legacy 1/4 path.
|
|
without_default = _backend_from_gguf("totallymadeupv7", common)
|
|
|
|
kv_default = with_default._estimate_kv_cache_bytes(131072, "f16")
|
|
kv_legacy = without_default._estimate_kv_cache_bytes(131072, "f16")
|
|
assert kv_default > 0
|
|
assert kv_legacy > 0
|
|
assert kv_default < kv_legacy, (
|
|
f"arch fallback should under-shoot legacy estimate: " f"{kv_default} >= {kv_legacy}"
|
|
)
|
|
|
|
def test_scalar_sliding_window_pattern_expanded(self):
|
|
block_count = 8
|
|
b = _backend_from_gguf(
|
|
"gemma3",
|
|
{
|
|
"attention.sliding_window_pattern": 4,
|
|
"block_count": block_count,
|
|
"attention.head_count_kv": 4,
|
|
"attention.key_length": 256,
|
|
"attention.value_length": 256,
|
|
"attention.sliding_window": 1024,
|
|
},
|
|
)
|
|
expected = [(i + 1) % 4 != 0 for i in range(block_count)]
|
|
assert isinstance(b._sliding_window_pattern, list)
|
|
assert b._sliding_window_pattern == expected
|
|
assert b._estimate_kv_cache_bytes(4096, "f16") > 0
|
|
|
|
def test_all_fields_parsed_together(self):
|
|
fields = {
|
|
"context_length": 131072,
|
|
"block_count": 62,
|
|
"attention.head_count_kv": 16,
|
|
"attention.head_count": 32,
|
|
"embedding_length": 5376,
|
|
"attention.key_length": 128,
|
|
"attention.value_length": 128,
|
|
"attention.sliding_window": 1024,
|
|
"attention.sliding_window_pattern": [True, False],
|
|
"full_attention_interval": 6,
|
|
"attention.kv_lora_rank": 512,
|
|
"attention.key_length_mla": 256,
|
|
"attention.key_length_swa": 64,
|
|
"attention.value_length_swa": 64,
|
|
"ssm.inner_size": 4096,
|
|
"ssm.state_size": 128,
|
|
}
|
|
b = _backend_from_gguf("testarch", fields)
|
|
assert b._context_length == 131072
|
|
assert b._n_layers == 62
|
|
assert b._n_kv_heads == 16
|
|
assert b._n_heads == 32
|
|
assert b._embedding_length == 5376
|
|
assert b._kv_key_length == 128
|
|
assert b._kv_value_length == 128
|
|
assert b._sliding_window == 1024
|
|
assert b._sliding_window_pattern == [True, False]
|
|
assert b._full_attention_interval == 6
|
|
assert b._kv_lora_rank == 512
|
|
assert b._key_length_mla == 256
|
|
assert b._kv_key_length_swa == 64
|
|
assert b._kv_value_length_swa == 64
|
|
assert b._ssm_inner_size == 4096
|
|
assert b._ssm_state_size == 128
|
|
|
|
|
|
_SWA_FIELDS = {
|
|
"block_count": 12,
|
|
"attention.head_count": 4,
|
|
"attention.head_count_kv": 1,
|
|
"attention.key_length": 256,
|
|
"attention.value_length": 256,
|
|
"attention.sliding_window": 512,
|
|
}
|
|
|
|
|
|
class TestDynamicSwaResolver:
|
|
"""4-tier resolver: GGUF metadata, on-disk cache, bootstrap, HF fetch."""
|
|
|
|
def _isolate_cache(self, monkeypatch, tmp_path):
|
|
from core.inference import llama_cpp as lc
|
|
|
|
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
|
|
monkeypatch.setattr(lc, "_SWA_CACHE", None)
|
|
return tmp_path
|
|
|
|
def test_period_from_layer_types_finds_smallest_period(self):
|
|
from core.inference.llama_cpp import _period_from_layer_types
|
|
|
|
# gemma3 (1 global/6), gpt-oss (alternating), gemma3n (1/5).
|
|
assert _period_from_layer_types((["sliding_attention"] * 5 + ["full_attention"]) * 4) == 6
|
|
assert _period_from_layer_types(["sliding_attention", "full_attention"] * 12) == 2
|
|
assert _period_from_layer_types((["sliding_attention"] * 4 + ["full_attention"]) * 7) == 5
|
|
|
|
def test_period_from_layer_types_returns_none_for_aperiodic(self):
|
|
from core.inference.llama_cpp import _period_from_layer_types
|
|
lt = [
|
|
"sliding_attention",
|
|
"full_attention",
|
|
"sliding_attention",
|
|
"sliding_attention",
|
|
"full_attention",
|
|
"sliding_attention",
|
|
"sliding_attention",
|
|
"sliding_attention",
|
|
]
|
|
assert _period_from_layer_types(lt) is None
|
|
|
|
def test_hf_repo_from_url(self):
|
|
from core.inference.llama_cpp import _hf_repo_from_url
|
|
|
|
assert (
|
|
_hf_repo_from_url("https://huggingface.co/google/gemma-3-1b-it")
|
|
== "google/gemma-3-1b-it"
|
|
)
|
|
assert (
|
|
_hf_repo_from_url("https://huggingface.co/google/gemma-3-1b-it/blob/main/config.json")
|
|
== "google/gemma-3-1b-it"
|
|
)
|
|
for bad in [
|
|
"https://huggingface.co/google",
|
|
"https://example.com/foo/bar",
|
|
None,
|
|
"",
|
|
]:
|
|
assert _hf_repo_from_url(bad) is None
|
|
|
|
def test_bootstrap_tier_used_when_no_cache(self, monkeypatch, tmp_path):
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
from core.inference import llama_cpp as lc
|
|
|
|
def boom(*a, **kw):
|
|
raise AssertionError("HF fetch must not run when bootstrap covers the arch")
|
|
|
|
monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", boom)
|
|
b = _backend_from_gguf("gemma3", dict(_SWA_FIELDS, block_count = 18))
|
|
assert b._sliding_window_pattern == [(i + 1) % 6 != 0 for i in range(18)]
|
|
|
|
def test_disk_cache_takes_precedence_over_bootstrap(self, monkeypatch, tmp_path):
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
# Cached period=3 overrides bootstrap=6.
|
|
with open(tmp_path / "swa_cache.json", "w") as f:
|
|
json.dump({"gemma3": 3}, f)
|
|
b = _backend_from_gguf("gemma3", dict(_SWA_FIELDS, block_count = 18))
|
|
assert b._sliding_window_pattern == [(i + 1) % 3 != 0 for i in range(18)]
|
|
|
|
def test_disk_cache_supports_array_entries(self, monkeypatch, tmp_path):
|
|
# Aperiodic mask is tiled across n_layers.
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
mask = [True, False, True, True, False, True, False, False]
|
|
with open(tmp_path / "swa_cache.json", "w") as f:
|
|
json.dump({"customarch": mask}, f)
|
|
b = _backend_from_gguf("customarch", dict(_SWA_FIELDS, block_count = 16))
|
|
assert b._sliding_window_pattern == [bool(mask[i % 8]) for i in range(16)]
|
|
|
|
def test_hf_fetch_populates_cache(self, monkeypatch, tmp_path):
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
from core.inference import llama_cpp as lc
|
|
|
|
calls = []
|
|
|
|
def fake_fetch(repo_id):
|
|
calls.append(repo_id)
|
|
return 4 if repo_id == "vendor/newmodel-1b-instruct" else None
|
|
|
|
monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", fake_fetch)
|
|
b = _backend_from_gguf(
|
|
"newmodel",
|
|
_SWA_FIELDS,
|
|
general = {"general.source.huggingface.repository": "vendor/newmodel-1b-instruct"},
|
|
)
|
|
assert b._sliding_window_pattern == [(i + 1) % 4 != 0 for i in range(12)]
|
|
assert calls == ["vendor/newmodel-1b-instruct"]
|
|
with open(tmp_path / "swa_cache.json") as f:
|
|
assert json.load(f) == {"newmodel": 4}
|
|
|
|
def test_hf_fetch_falls_back_to_other_candidates(self, monkeypatch, tmp_path):
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
from core.inference import llama_cpp as lc
|
|
|
|
monkeypatch.setattr(
|
|
lc,
|
|
"_fetch_swa_entry_from_hf",
|
|
lambda r: 6 if r == "vendor/newmodel-base" else None,
|
|
)
|
|
b = _backend_from_gguf(
|
|
"newmodel",
|
|
_SWA_FIELDS,
|
|
general = {
|
|
"general.base_model.0.repo_url": "https://huggingface.co/vendor/newmodel-base"
|
|
},
|
|
)
|
|
assert b._sliding_window_pattern == [(i + 1) % 6 != 0 for i in range(12)]
|
|
|
|
def test_offline_env_skips_network(self, monkeypatch, tmp_path):
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
monkeypatch.setenv("UNSLOTH_STUDIO_OFFLINE", "1")
|
|
from core.inference import llama_cpp as lc
|
|
|
|
def boom(*a, **kw):
|
|
raise AssertionError("HF fetch must not run when offline=1")
|
|
|
|
monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", boom)
|
|
b = _backend_from_gguf(
|
|
"newmodel",
|
|
_SWA_FIELDS,
|
|
general = {"general.source.huggingface.repository": "vendor/newmodel"},
|
|
)
|
|
assert b._sliding_window_pattern is None
|
|
|
|
def test_hf_fetch_failure_falls_through_silently(self, monkeypatch, tmp_path):
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
from core.inference import llama_cpp as lc
|
|
|
|
monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", lambda repo_id: None)
|
|
# Force failure into Tier 3; bypass Tier 2.5.
|
|
monkeypatch.setattr(lc, "_resolve_swa_entry_from_transformers", lambda arch: None)
|
|
b = _backend_from_gguf(
|
|
"newmodel",
|
|
_SWA_FIELDS,
|
|
general = {"general.source.huggingface.repository": "vendor/does-not-exist"},
|
|
)
|
|
assert b._sliding_window_pattern is None
|
|
assert not (tmp_path / "swa_cache.json").exists()
|
|
|
|
|
|
class TestTransformersIntrospection:
|
|
"""Tier 2.5: default-init the matching Config; on failure, parse via inspect."""
|
|
|
|
def _isolate_cache(self, monkeypatch, tmp_path):
|
|
from core.inference import llama_cpp as lc
|
|
|
|
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
|
|
monkeypatch.setattr(lc, "_SWA_CACHE", None)
|
|
return tmp_path
|
|
|
|
def test_arch_aliases_normalises_hyphen_underscore(self):
|
|
from core.inference.llama_cpp import _arch_aliases
|
|
|
|
aliases = _arch_aliases("falcon-h1")
|
|
assert aliases[0] == "falcon-h1" and "falcon_h1" in aliases
|
|
assert _arch_aliases("gemma3") == ("gemma3",)
|
|
assert _arch_aliases("") == ()
|
|
|
|
def test_resolves_real_transformers_arches(self):
|
|
from core.inference.llama_cpp import _resolve_swa_entry_from_transformers
|
|
|
|
assert _resolve_swa_entry_from_transformers("gemma3") == 6
|
|
assert _resolve_swa_entry_from_transformers("gemma2") == 2
|
|
assert _resolve_swa_entry_from_transformers("cohere2") == 4
|
|
|
|
def test_falls_back_to_inspect_when_default_init_raises(self, monkeypatch):
|
|
from core.inference import llama_cpp as lc
|
|
|
|
class _FakeBrokenConfig:
|
|
"""Class with sliding_window_pattern: int = 7 in its docstring."""
|
|
|
|
def __init__(self, required_arg):
|
|
raise TypeError("requires an argument")
|
|
|
|
class _FakeLazyMapping(dict):
|
|
def __getitem__(self, k):
|
|
return _FakeBrokenConfig if k == "brokenarch" else super().__getitem__(k)
|
|
|
|
import sys, types as _types
|
|
|
|
fake_auto = _types.ModuleType("transformers.models.auto.configuration_auto")
|
|
fake_auto.CONFIG_MAPPING_NAMES = {"brokenarch": "FakeBroken"}
|
|
fake_auto.CONFIG_MAPPING = _FakeLazyMapping({"brokenarch": "FakeBroken"})
|
|
monkeypatch.setitem(sys.modules, "transformers.models.auto.configuration_auto", fake_auto)
|
|
assert lc._resolve_swa_entry_from_transformers("brokenarch") == 7
|
|
|
|
def test_returns_none_when_transformers_unavailable(self, monkeypatch):
|
|
from core.inference import llama_cpp as lc
|
|
import sys
|
|
|
|
orig_import = (
|
|
__builtins__["__import__"]
|
|
if isinstance(__builtins__, dict)
|
|
else __builtins__.__import__
|
|
)
|
|
|
|
def fake_import(name, *a, **kw):
|
|
if name.startswith("transformers"):
|
|
raise ImportError("transformers not installed")
|
|
return orig_import(name, *a, **kw)
|
|
|
|
monkeypatch.setattr("builtins.__import__", fake_import)
|
|
for k in list(sys.modules):
|
|
if k.startswith("transformers"):
|
|
monkeypatch.delitem(sys.modules, k, raising = False)
|
|
assert lc._resolve_swa_entry_from_transformers("gemma3") is None
|
|
|
|
def test_returns_none_for_arch_unknown_to_transformers(self):
|
|
from core.inference.llama_cpp import _resolve_swa_entry_from_transformers
|
|
assert _resolve_swa_entry_from_transformers("totally-fake-arch-xyz") is None
|
|
|
|
def test_full_resolver_uses_transformers_before_hf_fetch(self, monkeypatch, tmp_path):
|
|
# Bootstrap empty: Tier 2.5 must answer before Tier 3 fires.
|
|
self._isolate_cache(monkeypatch, tmp_path)
|
|
from core.inference import llama_cpp as lc
|
|
|
|
monkeypatch.setattr(lc, "_BOOTSTRAP_SWA_DEFAULTS", {})
|
|
|
|
def boom(repo_id):
|
|
raise AssertionError("Tier 3 must not run when Tier 2.5 has the answer")
|
|
|
|
monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", boom)
|
|
b = _backend_from_gguf(
|
|
"gemma3",
|
|
dict(_SWA_FIELDS, block_count = 18),
|
|
general = {"general.source.huggingface.repository": "google/gemma-3-1b-it"},
|
|
)
|
|
assert b._sliding_window_pattern == [(i + 1) % 6 != 0 for i in range(18)]
|
|
with open(tmp_path / "swa_cache.json") as f:
|
|
assert json.load(f) == {"gemma3": 6}
|
|
|
|
|
|
class TestGGUFParserReset:
|
|
"""Fields are reset between parses."""
|
|
|
|
def test_reset_between_parses(self):
|
|
# First parse: all fields set
|
|
b = _backend_from_gguf(
|
|
"arch1",
|
|
{
|
|
"block_count": 32,
|
|
"attention.key_length": 128,
|
|
"attention.kv_lora_rank": 512,
|
|
"attention.head_count_kv": [8, 2],
|
|
"attention.sliding_window_pattern": [True, False],
|
|
"attention.key_length_swa": 64,
|
|
"attention.value_length_swa": 64,
|
|
"ssm.inner_size": 4096,
|
|
},
|
|
)
|
|
assert b._kv_key_length == 128
|
|
assert b._kv_lora_rank == 512
|
|
assert b._n_kv_heads_by_layer == [8, 2]
|
|
assert b._sliding_window_pattern == [True, False]
|
|
assert b._kv_key_length_swa == 64
|
|
assert b._kv_value_length_swa == 64
|
|
assert b._ssm_inner_size == 4096
|
|
|
|
# Second parse without those fields -- they must be None
|
|
kv = {"general.architecture": "arch2", "arch2.block_count": 64}
|
|
import tempfile, os
|
|
|
|
data = _make_gguf_bytes("arch2", kv)
|
|
fd, path = tempfile.mkstemp(suffix = ".gguf")
|
|
os.write(fd, data)
|
|
os.close(fd)
|
|
try:
|
|
b._read_gguf_metadata(path)
|
|
finally:
|
|
os.unlink(path)
|
|
assert b._kv_key_length is None
|
|
assert b._kv_lora_rank is None
|
|
assert b._n_kv_heads_by_layer is None
|
|
assert b._sliding_window_pattern is None
|
|
assert b._kv_key_length_swa is None
|
|
assert b._kv_value_length_swa is None
|
|
assert b._ssm_inner_size is None
|
|
assert b._n_layers == 64
|
|
|
|
|
|
# B. _can_estimate_kv Gate Tests
|
|
|
|
|
|
class TestCanEstimateKV:
|
|
"""Gate logic for all field combinations."""
|
|
|
|
def test_no_layers_returns_false(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = None
|
|
b._kv_key_length = 128
|
|
assert not b._can_estimate_kv()
|
|
|
|
def test_explicit_both_dims_sufficient(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 32
|
|
b._kv_key_length = 128
|
|
b._kv_value_length = 128
|
|
assert b._can_estimate_kv()
|
|
|
|
def test_key_length_alone_insufficient(self):
|
|
"""key_length without value_length is NOT enough."""
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 32
|
|
b._kv_key_length = 128
|
|
assert not b._can_estimate_kv()
|
|
|
|
def test_kv_lora_rank_sufficient(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 61
|
|
b._kv_lora_rank = 512
|
|
assert b._can_estimate_kv()
|
|
|
|
def test_legacy_embed_plus_heads(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 28
|
|
b._embedding_length = 1024
|
|
b._n_heads = 16
|
|
assert b._can_estimate_kv()
|
|
|
|
def test_legacy_embed_plus_kv_heads(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 28
|
|
b._embedding_length = 1024
|
|
b._n_kv_heads = 8
|
|
assert b._can_estimate_kv()
|
|
|
|
def test_legacy_no_embed_returns_false(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 28
|
|
b._n_heads = 16
|
|
# No embedding_length, no new-style fields
|
|
assert not b._can_estimate_kv()
|
|
|
|
def test_fresh_backend_returns_false(self):
|
|
b = LlamaCppBackend()
|
|
assert not b._can_estimate_kv()
|
|
|
|
|
|
# C. Path 1: MLA Estimation
|
|
|
|
|
|
class TestMLAEstimation:
|
|
"""MLA: K-only cache using compressed KV latent + RoPE."""
|
|
|
|
def _mla_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 61,
|
|
"_n_kv_heads": 1,
|
|
"_n_heads": 128,
|
|
"_embedding_length": 7168,
|
|
"_kv_key_length": 576,
|
|
"_kv_value_length": 512,
|
|
"_kv_lora_rank": 512,
|
|
"_key_length_mla": 192,
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
def test_deepseek_v3_f16(self):
|
|
b = self._mla_backend()
|
|
# 61 layers * 163840 ctx * 1 head * 576 key_len * 2 bpe
|
|
expected = 61 * 163840 * 1 * 576 * 2
|
|
assert b._estimate_kv_cache_bytes(163840, "f16") == expected
|
|
|
|
def test_mla_ignores_value_length(self):
|
|
"""MLA must NOT add value_length -- V is reconstructed from the latent."""
|
|
b = self._mla_backend()
|
|
result = b._estimate_kv_cache_bytes(1000, "f16")
|
|
# n_layers * ctx * 1 * key_len(576) * 2
|
|
expected = 61 * 1000 * 1 * 576 * 2
|
|
assert result == expected
|
|
|
|
def test_mla_fallback_when_no_key_length(self):
|
|
"""No key_length: fall back to kv_lora_rank + key_length_mla."""
|
|
b = self._mla_backend(_kv_key_length = None)
|
|
# default _key_length_mla=192, so rope_dim=192
|
|
result = b._estimate_kv_cache_bytes(1000, "f16")
|
|
expected = 61 * 1000 * 1 * (512 + 192) * 2 # 704
|
|
assert result == expected
|
|
|
|
def test_mla_fallback_no_key_length_mla(self):
|
|
"""No key_length and no key_length_mla: fall back to +64."""
|
|
b = self._mla_backend(_kv_key_length = None, _key_length_mla = None)
|
|
result = b._estimate_kv_cache_bytes(1000, "f16")
|
|
expected = 61 * 1000 * 1 * (512 + 64) * 2 # 576
|
|
assert result == expected
|
|
|
|
def test_mla_defaults_n_kv_to_1_when_heads_absent(self):
|
|
"""MLA uses n_kv=1 even if n_kv_heads is None (not n_heads)."""
|
|
b = self._mla_backend(_n_kv_heads = None) # n_heads=128 still set
|
|
result = b._estimate_kv_cache_bytes(1000, "f16")
|
|
# Uses n_kv_mla=1, NOT n_heads=128
|
|
expected = 61 * 1000 * 1 * 576 * 2
|
|
assert result == expected
|
|
|
|
def test_mla_q4_quantization(self):
|
|
b = self._mla_backend()
|
|
result_f16 = b._estimate_kv_cache_bytes(1000, "f16")
|
|
result_q4 = b._estimate_kv_cache_bytes(1000, "q4_0")
|
|
assert result_q4 < result_f16
|
|
# q4_0 bpe = 0.5625, f16 bpe = 2.0
|
|
assert result_q4 == int(61 * 1000 * 1 * 576 * 0.5625)
|
|
|
|
|
|
# D. Path 2: Hybrid Mamba Estimation
|
|
|
|
|
|
class TestHybridMambaEstimation:
|
|
"""Hybrid Mamba: only attention layers (1 in N) need KV cache."""
|
|
|
|
def _hybrid_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 64,
|
|
"_n_kv_heads": 4,
|
|
"_n_heads": 24,
|
|
"_embedding_length": 5120,
|
|
"_kv_key_length": 256,
|
|
"_kv_value_length": 256,
|
|
"_full_attention_interval": 4,
|
|
"_ssm_inner_size": 6144,
|
|
"_ssm_state_size": 128,
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
def test_qwen35_27b(self):
|
|
b = self._hybrid_backend()
|
|
# n_attn = 64 // 4 = 16
|
|
expected = 16 * 262144 * 4 * (256 + 256) * 2
|
|
assert b._estimate_kv_cache_bytes(262144, "f16") == expected
|
|
|
|
def test_qwen35_35b_a3b(self):
|
|
b = self._hybrid_backend(
|
|
_n_layers = 40,
|
|
_n_kv_heads = 2,
|
|
_n_heads = 16,
|
|
_embedding_length = 2048,
|
|
_ssm_inner_size = 4096,
|
|
)
|
|
# n_attn = 40 // 4 = 10
|
|
expected = 10 * 262144 * 2 * (256 + 256) * 2
|
|
assert b._estimate_kv_cache_bytes(262144, "f16") == expected
|
|
|
|
def test_hybrid_without_explicit_dims(self):
|
|
"""Fall back to head_dim when key_length/value_length are missing."""
|
|
b = self._hybrid_backend(_kv_key_length = None, _kv_value_length = None)
|
|
head_dim = 5120 // 24 # 213
|
|
expected = 16 * 4096 * 4 * 2 * head_dim * 2
|
|
assert b._estimate_kv_cache_bytes(4096, "f16") == expected
|
|
|
|
def test_fai_zero_safety(self):
|
|
"""full_attention_interval=0 must not ZeroDivisionError."""
|
|
b = self._hybrid_backend(_full_attention_interval = 0)
|
|
result = b._estimate_kv_cache_bytes(4096, "f16")
|
|
# fai=0 -> n_attn = n_layers (all layers)
|
|
expected = 64 * 4096 * 4 * (256 + 256) * 2
|
|
assert result == expected
|
|
|
|
|
|
# E. Path 3: Sliding Window Estimation
|
|
|
|
|
|
class TestSlidingWindowEstimation:
|
|
"""SWA: half global (full ctx) + half sliding window."""
|
|
|
|
def _swa_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 62,
|
|
"_n_kv_heads": 16,
|
|
"_n_heads": 32,
|
|
"_embedding_length": 5376,
|
|
"_kv_key_length": 128,
|
|
"_kv_value_length": 128,
|
|
"_sliding_window": 1024,
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
def test_gemma3(self):
|
|
b = self._swa_backend()
|
|
# 1/4 heuristic: 62 // 4 = 15 global, 47 SWA
|
|
n_global = max(1, 62 // 4) # 15
|
|
n_swa = 62 - n_global # 47
|
|
kv_per = 16 * (128 + 128) * 2
|
|
# SWA cache is double-buffered: 2 * sliding_window cells, capped at n_ctx.
|
|
swa_cells = min(131072, 2 * 1024)
|
|
expected = int(n_global * 131072 * kv_per + n_swa * swa_cells * kv_per)
|
|
assert b._estimate_kv_cache_bytes(131072, "f16") == expected
|
|
|
|
def test_gpt_oss(self):
|
|
b = self._swa_backend(
|
|
_n_layers = 24,
|
|
_n_kv_heads = 8,
|
|
_n_heads = 64,
|
|
_embedding_length = 2880,
|
|
_kv_key_length = 64,
|
|
_kv_value_length = 64,
|
|
_sliding_window = 128,
|
|
)
|
|
# 1/4 heuristic: 24 // 4 = 6 global, 18 SWA
|
|
n_global = max(1, 24 // 4) # 6
|
|
n_swa = 24 - n_global # 18
|
|
kv_per = 8 * (64 + 64) * 2
|
|
swa_cells = min(131072, 2 * 128)
|
|
expected = int(n_global * 131072 * kv_per + n_swa * swa_cells * kv_per)
|
|
assert b._estimate_kv_cache_bytes(131072, "f16") == expected
|
|
|
|
def test_gemma4_per_layer_swa_metadata(self):
|
|
b = self._swa_backend(
|
|
_n_layers = 30,
|
|
_n_kv_heads = None,
|
|
_n_kv_heads_by_layer = [8, 8, 8, 8, 8, 2] * 5,
|
|
_n_heads = 16,
|
|
_embedding_length = 2816,
|
|
_kv_key_length = 512,
|
|
_kv_value_length = 512,
|
|
_sliding_window = 1024,
|
|
_sliding_window_pattern = [True, True, True, True, True, False] * 5,
|
|
_kv_key_length_swa = 256,
|
|
_kv_value_length_swa = 256,
|
|
)
|
|
|
|
full_layers = 5
|
|
sliding_layers = 25
|
|
|
|
def expected(ctx):
|
|
full = full_layers * ctx * 2 * (512 + 512) * 2
|
|
sliding = sliding_layers * min(ctx, 2 * 1024) * 8 * (256 + 256) * 2
|
|
return int(full + sliding)
|
|
|
|
for ctx in (4096, 46500, 262144):
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected(ctx)
|
|
|
|
def test_ctx_smaller_than_window(self):
|
|
"""When ctx < 2 * sliding_window, SWA cache caps at ctx."""
|
|
b = self._swa_backend(_sliding_window = 8192)
|
|
n_global = max(1, 62 // 4) # 15
|
|
n_swa = 62 - n_global # 47
|
|
kv_per = 16 * (128 + 128) * 2
|
|
ctx = 4096
|
|
expected = int(n_global * ctx * kv_per + n_swa * min(ctx, 2 * 8192) * kv_per)
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
|
|
|
|
def test_odd_layer_count(self):
|
|
b = self._swa_backend(_n_layers = 63)
|
|
n_global = max(1, 63 // 4) # 15
|
|
n_swa = 63 - n_global # 48
|
|
kv_per = 16 * (128 + 128) * 2
|
|
expected = int(n_global * 1000 * kv_per + n_swa * min(1000, 2 * 1024) * kv_per)
|
|
assert b._estimate_kv_cache_bytes(1000, "f16") == expected
|
|
|
|
|
|
# F. Path 4: Standard GQA Estimation
|
|
|
|
|
|
class TestStandardGQAEstimation:
|
|
"""Standard GQA with explicit key_length/value_length."""
|
|
|
|
def _gqa_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 28,
|
|
"_n_kv_heads": 8,
|
|
"_n_heads": 16,
|
|
"_embedding_length": 1024,
|
|
"_kv_key_length": 128,
|
|
"_kv_value_length": 128,
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
def test_qwen3_06b(self):
|
|
b = self._gqa_backend()
|
|
expected = 28 * 40960 * 8 * (128 + 128) * 2
|
|
assert b._estimate_kv_cache_bytes(40960, "f16") == expected
|
|
|
|
def test_asymmetric_kv_dims(self):
|
|
"""key_length != value_length (some architectures have this)."""
|
|
b = self._gqa_backend(_kv_key_length = 192, _kv_value_length = 64)
|
|
expected = 28 * 4096 * 8 * (192 + 64) * 2
|
|
assert b._estimate_kv_cache_bytes(4096, "f16") == expected
|
|
|
|
def test_differs_from_legacy(self):
|
|
"""GQA path differs from legacy when key_length != embed//n_heads."""
|
|
b = self._gqa_backend()
|
|
head_dim = 1024 // 16 # 64
|
|
gqa_result = b._estimate_kv_cache_bytes(4096, "f16")
|
|
# Legacy: 2 * 8 * 64 * 28 * 4096 * 2
|
|
legacy_result = int(2 * 8 * head_dim * 28 * 4096 * 2)
|
|
# GQA: 28 * 4096 * 8 * (128+128) * 2 -- uses actual key_length=128
|
|
assert gqa_result != legacy_result
|
|
assert gqa_result > legacy_result # key_length (128) > head_dim (64)
|
|
|
|
|
|
# G. Path 5: Legacy Fallback Estimation
|
|
|
|
|
|
class TestLegacyEstimation:
|
|
"""Legacy: embed // n_heads, for old GGUFs without new fields."""
|
|
|
|
def _legacy_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 32,
|
|
"_n_kv_heads": 8,
|
|
"_n_heads": 32,
|
|
"_embedding_length": 4096,
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
def test_basic_legacy(self):
|
|
b = self._legacy_backend()
|
|
head_dim = 4096 // 32 # 128
|
|
expected = int(2 * 8 * 128 * 32 * 4096 * 2)
|
|
assert b._estimate_kv_cache_bytes(4096, "f16") == expected
|
|
|
|
def test_legacy_with_only_n_heads(self):
|
|
"""n_kv_heads is None, falls back to n_heads."""
|
|
b = self._legacy_backend(_n_kv_heads = None)
|
|
head_dim = 4096 // 32
|
|
expected = int(2 * 32 * head_dim * 32 * 4096 * 2)
|
|
assert b._estimate_kv_cache_bytes(4096, "f16") == expected
|
|
|
|
def test_legacy_identical_to_old_formula(self):
|
|
"""Legacy path matches the pre-PR formula."""
|
|
b = self._legacy_backend()
|
|
n_layers = 32
|
|
n_kv_heads = 8
|
|
head_dim = 4096 // 32
|
|
n_ctx = 8192
|
|
bpe = 2.0
|
|
old_formula = int(2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe)
|
|
assert b._estimate_kv_cache_bytes(n_ctx, "f16") == old_formula
|
|
|
|
|
|
# H. Path Priority (selection order)
|
|
|
|
|
|
class TestPathPriority:
|
|
"""Confirm: MLA > Hybrid Mamba > SWA > GQA > Legacy."""
|
|
|
|
def test_mla_takes_priority_over_all(self):
|
|
"""If kv_lora_rank is set, MLA path wins even with other fields present."""
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 61
|
|
b._n_kv_heads = 1
|
|
b._n_heads = 128
|
|
b._embedding_length = 7168
|
|
b._kv_key_length = 576
|
|
b._kv_value_length = 512
|
|
b._kv_lora_rank = 512
|
|
b._ssm_inner_size = 4096 # Would trigger Hybrid
|
|
b._full_attention_interval = 4
|
|
b._sliding_window = 1024 # Would trigger SWA
|
|
|
|
# MLA: 61 * 1000 * 1 * 576 * 2
|
|
expected_mla = int(61 * 1000 * 1 * 576 * 2)
|
|
assert b._estimate_kv_cache_bytes(1000, "f16") == expected_mla
|
|
|
|
def test_hybrid_over_swa(self):
|
|
"""Hybrid takes priority over SWA when both fields present."""
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 64
|
|
b._n_kv_heads = 4
|
|
b._n_heads = 24
|
|
b._embedding_length = 5120
|
|
b._kv_key_length = 256
|
|
b._kv_value_length = 256
|
|
b._ssm_inner_size = 6144
|
|
b._full_attention_interval = 4
|
|
b._sliding_window = 1024 # Would trigger SWA
|
|
|
|
n_attn = 64 // 4
|
|
expected_hybrid = int(n_attn * 1000 * 4 * (256 + 256) * 2)
|
|
assert b._estimate_kv_cache_bytes(1000, "f16") == expected_hybrid
|
|
|
|
def test_all_paths_produce_different_values(self):
|
|
"""With chosen params, each path yields a distinct value."""
|
|
# embedding_length=768 so legacy head_dim (768//16=48) != key_length
|
|
# (256), and MLA key_len (256) != legacy K+V (2*48=96).
|
|
params = {
|
|
"_n_layers": 40,
|
|
"_n_kv_heads": 4,
|
|
"_n_heads": 16,
|
|
"_embedding_length": 768,
|
|
"_kv_key_length": 256,
|
|
"_kv_value_length": 256,
|
|
}
|
|
ctx = 4096
|
|
|
|
# Path 4: Standard GQA
|
|
b_gqa = LlamaCppBackend()
|
|
for k, v in params.items():
|
|
setattr(b_gqa, k, v)
|
|
gqa_val = b_gqa._estimate_kv_cache_bytes(ctx, "f16")
|
|
|
|
# Path 1: MLA
|
|
b_mla = LlamaCppBackend()
|
|
for k, v in params.items():
|
|
setattr(b_mla, k, v)
|
|
b_mla._kv_lora_rank = 512
|
|
mla_val = b_mla._estimate_kv_cache_bytes(ctx, "f16")
|
|
|
|
# Path 2: Hybrid Mamba
|
|
b_hybrid = LlamaCppBackend()
|
|
for k, v in params.items():
|
|
setattr(b_hybrid, k, v)
|
|
b_hybrid._ssm_inner_size = 4096
|
|
b_hybrid._full_attention_interval = 4
|
|
hybrid_val = b_hybrid._estimate_kv_cache_bytes(ctx, "f16")
|
|
|
|
# Path 3: SWA
|
|
b_swa = LlamaCppBackend()
|
|
for k, v in params.items():
|
|
setattr(b_swa, k, v)
|
|
b_swa._sliding_window = 512
|
|
swa_val = b_swa._estimate_kv_cache_bytes(ctx, "f16")
|
|
|
|
# Path 5: Legacy (no key_length/value_length)
|
|
b_legacy = LlamaCppBackend()
|
|
b_legacy._n_layers = 40
|
|
b_legacy._n_kv_heads = 4
|
|
b_legacy._n_heads = 16
|
|
b_legacy._embedding_length = 768
|
|
legacy_val = b_legacy._estimate_kv_cache_bytes(ctx, "f16")
|
|
|
|
values = [mla_val, hybrid_val, swa_val, gqa_val, legacy_val]
|
|
assert len(set(values)) == 5, f"Expected 5 distinct values, got {values}"
|
|
|
|
|
|
# I. KV Cache Quantization
|
|
|
|
|
|
class TestQuantization:
|
|
"""All supported cache_type_kv values scale correctly."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"cache_type,expected_bpe",
|
|
[
|
|
("f32", 4.0),
|
|
("f16", 2.0),
|
|
("bf16", 2.0),
|
|
("q8_0", 34 / 32),
|
|
("q5_1", 0.75),
|
|
("q5_0", 0.6875),
|
|
("q4_1", 0.625),
|
|
("q4_0", 0.5625),
|
|
("iq4_nl", 0.5625),
|
|
(None, 2.0), # default is f16
|
|
("unknown", 2.0), # unknown falls back to f16
|
|
],
|
|
)
|
|
def test_quantization_scaling(self, cache_type, expected_bpe):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 10
|
|
b._n_kv_heads = 1
|
|
b._n_heads = 8
|
|
b._embedding_length = 512
|
|
b._kv_key_length = 64
|
|
b._kv_value_length = 64
|
|
result = b._estimate_kv_cache_bytes(1000, cache_type)
|
|
expected = int(10 * 1000 * 1 * (64 + 64) * expected_bpe)
|
|
assert result == expected
|
|
|
|
|
|
# J. Edge Cases
|
|
|
|
|
|
class TestEdgeCases:
|
|
"""Boundary conditions and degenerate inputs."""
|
|
|
|
def test_zero_context(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 32
|
|
b._kv_key_length = 128
|
|
assert b._estimate_kv_cache_bytes(0, "f16") == 0
|
|
|
|
def test_negative_context(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 32
|
|
b._kv_key_length = 128
|
|
assert b._estimate_kv_cache_bytes(-1, "f16") == 0
|
|
|
|
def test_context_of_one(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 10
|
|
b._n_kv_heads = 1
|
|
b._kv_key_length = 64
|
|
b._kv_value_length = 64
|
|
result = b._estimate_kv_cache_bytes(1, "f16")
|
|
assert result == int(10 * 1 * 1 * (64 + 64) * 2)
|
|
|
|
def test_very_large_context(self):
|
|
"""1M context should not overflow or crash."""
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 10
|
|
b._n_kv_heads = 1
|
|
b._kv_key_length = 128
|
|
b._kv_value_length = 128
|
|
result = b._estimate_kv_cache_bytes(1_000_000, "f16")
|
|
assert result > 0
|
|
assert isinstance(result, int)
|
|
|
|
def test_n_kv_heads_none_falls_to_n_heads(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 10
|
|
b._n_kv_heads = None
|
|
b._n_heads = 8
|
|
b._kv_key_length = 64
|
|
b._kv_value_length = 64
|
|
result = b._estimate_kv_cache_bytes(100, "f16")
|
|
expected = int(10 * 100 * 8 * (64 + 64) * 2)
|
|
assert result == expected
|
|
|
|
def test_both_heads_none_falls_to_one(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 10
|
|
b._n_kv_heads = None
|
|
b._n_heads = None
|
|
b._kv_key_length = 64
|
|
b._kv_value_length = 64
|
|
result = b._estimate_kv_cache_bytes(100, "f16")
|
|
expected = int(10 * 100 * 1 * (64 + 64) * 2)
|
|
assert result == expected
|
|
|
|
|
|
# J2. Server-flag knobs (--swa-full, --kv-unified/--parallel,
|
|
# --ctx-checkpoints, --kv-offload)
|
|
|
|
|
|
class TestServerFlags:
|
|
"""Estimator should mirror llama-server CLI flags that change KV size."""
|
|
|
|
def _swa_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 26,
|
|
"_n_kv_heads": 4,
|
|
"_n_heads": 8,
|
|
"_embedding_length": 1152,
|
|
"_kv_key_length": 256,
|
|
"_kv_value_length": 256,
|
|
"_sliding_window": 512,
|
|
"_sliding_window_pattern": [True, True, True, True, True, False] * 4 + [True, True],
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
def _gqa_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 28,
|
|
"_n_kv_heads": 8,
|
|
"_n_heads": 16,
|
|
"_embedding_length": 1024,
|
|
"_kv_key_length": 128,
|
|
"_kv_value_length": 128,
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
# ── --swa-full ──────────────────────────────────────────────────
|
|
|
|
def test_swa_full_collapses_pattern_path_to_full_ctx(self):
|
|
b = self._swa_backend()
|
|
ctx = 32_768
|
|
flagged = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
|
|
# swa_full: every layer caches n_ctx -- equals path 4 sizing.
|
|
kv_per_token = 4 * (256 + 256) * 2 # n_kv_heads * (k+v) * f16
|
|
expected = 26 * ctx * kv_per_token
|
|
assert flagged == expected
|
|
assert flagged > b._estimate_kv_cache_bytes(ctx, "f16")
|
|
|
|
def test_swa_full_collapses_legacy_path_to_full_ctx(self):
|
|
# No per-layer pattern -> 1/4-global heuristic; swa_full overrides.
|
|
b = self._swa_backend(_sliding_window_pattern = None)
|
|
ctx = 16_384
|
|
flagged = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
|
|
n_global = max(1, 26 // 4)
|
|
n_swa = 26 - n_global
|
|
kv_per = 4 * (256 + 256) * 2
|
|
# swa_cells == n_ctx when swa_full=True
|
|
expected = n_global * ctx * kv_per + n_swa * ctx * kv_per
|
|
assert flagged == expected
|
|
|
|
def test_swa_full_no_op_for_non_swa_model(self):
|
|
b = self._gqa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
flagged = b._estimate_kv_cache_bytes(8192, "f16", swa_full = True)
|
|
assert flagged == baseline
|
|
|
|
def test_swa_full_suppresses_checkpoint_term(self):
|
|
b = self._swa_backend()
|
|
with_cp = b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 8)
|
|
with_cp_full = b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 8, swa_full = True)
|
|
no_cp_full = b._estimate_kv_cache_bytes(8192, "f16", swa_full = True)
|
|
# Checkpoints only matter when SWA layers don't already keep n_ctx.
|
|
assert with_cp_full == no_cp_full
|
|
assert with_cp > b._estimate_kv_cache_bytes(8192, "f16")
|
|
|
|
# ── --parallel + --kv-unified ──────────────────────────────────
|
|
# Verified against llama-server: non-SWA caches partition n_ctx across
|
|
# slots (total memory constant); only SWA layers scale with --parallel.
|
|
# --kv-unified is a no-op for memory math (kept for API forward-compat).
|
|
|
|
def test_gqa_kv_constant_across_parallel(self):
|
|
b = self._gqa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(4096, "f16")
|
|
for slots in (1, 2, 4, 8):
|
|
for unified in (True, False):
|
|
assert (
|
|
b._estimate_kv_cache_bytes(4096, "f16", n_parallel = slots, kv_unified = unified)
|
|
== baseline
|
|
)
|
|
|
|
def test_zero_parallel_floors_at_one(self):
|
|
b = self._gqa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(4096, "f16")
|
|
for unified in (True, False):
|
|
assert (
|
|
b._estimate_kv_cache_bytes(4096, "f16", n_parallel = 0, kv_unified = unified)
|
|
== baseline
|
|
)
|
|
|
|
def test_swa_path_scales_only_swa_portion(self):
|
|
b = self._swa_backend()
|
|
ctx = 8192
|
|
baseline = b._estimate_kv_cache_bytes(ctx, "f16")
|
|
# Decompose baseline by walking the estimator's own loop.
|
|
swa = b._sliding_window
|
|
per_token_global = 4 * (256 + 256) * 2 # n_kv * (k+v) * f16
|
|
per_token_swa = 4 * (256 + 256) * 2 # k_swa/val_swa fall back
|
|
per_slot_swa_cells = min(ctx, 2 * swa) # not clamped at parallel=1
|
|
global_bytes = sum(
|
|
ctx * per_token_global for f in b._sliding_window_pattern[: b._n_layers] if not f
|
|
)
|
|
swa_bytes_per_slot = sum(
|
|
per_slot_swa_cells * per_token_swa
|
|
for f in b._sliding_window_pattern[: b._n_layers]
|
|
if f
|
|
)
|
|
# Sanity: parallel=1 reproduces baseline exactly
|
|
assert global_bytes + swa_bytes_per_slot == baseline
|
|
# Only the SWA portion scales by parallel
|
|
for slots in (1, 2, 3, 4):
|
|
scaled = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False)
|
|
# SWA cells clamp to per_slot_ctx when ctx/slots < 2*swa
|
|
per_slot_ctx = max(1, ctx // slots)
|
|
cells = min(ctx, 2 * swa, per_slot_ctx)
|
|
swa_bps = sum(
|
|
cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f
|
|
)
|
|
assert scaled == global_bytes + slots * swa_bps
|
|
|
|
def test_mla_kv_constant_across_parallel(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 60
|
|
b._n_kv_heads = 1
|
|
b._kv_lora_rank = 512
|
|
b._key_length_mla = 64
|
|
b._kv_key_length = 576
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
for slots in (1, 2, 4, 8):
|
|
for unified in (True, False):
|
|
assert (
|
|
b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = unified)
|
|
== baseline
|
|
)
|
|
|
|
# ── --ctx-checkpoints ──────────────────────────────────────────
|
|
|
|
def test_ctx_checkpoints_zero_is_no_op(self):
|
|
b = self._swa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
assert b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 0) == baseline
|
|
|
|
def test_ctx_checkpoints_no_op_for_non_swa(self):
|
|
b = self._gqa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
assert b._estimate_kv_cache_bytes(8192, "f16", ctx_checkpoints = 32) == baseline
|
|
|
|
def test_ctx_checkpoints_pattern_path_adds_known_bytes(self):
|
|
b = self._swa_backend()
|
|
ctx = 8192
|
|
baseline = b._estimate_kv_cache_bytes(ctx, "f16")
|
|
flagged = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4)
|
|
# 22 SWA layers * 4 cps * 512 cells * 4 heads * (256+256) * 2 bytes
|
|
n_swa_layers = sum(1 for f in [True, True, True, True, True, False] * 4 + [True, True] if f)
|
|
per_layer = 4 * 512 * 4 * (256 + 256) * 2
|
|
assert flagged == baseline + n_swa_layers * per_layer
|
|
|
|
def test_ctx_checkpoints_legacy_path_adds_known_bytes(self):
|
|
b = self._swa_backend(_sliding_window_pattern = None)
|
|
ctx = 8192
|
|
baseline = b._estimate_kv_cache_bytes(ctx, "f16")
|
|
flagged = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4)
|
|
n_global = max(1, 26 // 4)
|
|
n_swa = 26 - n_global
|
|
kv_per = 4 * (256 + 256) * 2
|
|
extra = 4 * n_swa * 512 * kv_per # ctx_checkpoints * n_swa * sliding * kv_per
|
|
assert flagged == baseline + extra
|
|
|
|
def test_ctx_checkpoints_compose_with_n_parallel(self):
|
|
# Only the SWA + checkpoint portion scales by n_parallel; the
|
|
# global-layer portion is constant.
|
|
b = self._swa_backend()
|
|
ctx = 8192
|
|
swa = b._sliding_window
|
|
per_token = 4 * (256 + 256) * 2
|
|
global_bytes = sum(
|
|
ctx * per_token for f in b._sliding_window_pattern[: b._n_layers] if not f
|
|
)
|
|
n_swa_layers = sum(1 for f in b._sliding_window_pattern[: b._n_layers] if f)
|
|
slots = 3
|
|
per_slot_ctx = max(1, ctx // slots)
|
|
swa_cells = min(ctx, 2 * swa, per_slot_ctx)
|
|
swa_bytes_per_slot = n_swa_layers * swa_cells * per_token
|
|
cp_extra_per_slot = n_swa_layers * 4 * swa * per_token # 4 checkpoints
|
|
flagged = b._estimate_kv_cache_bytes(
|
|
ctx, "f16", ctx_checkpoints = 4, n_parallel = slots, kv_unified = False
|
|
)
|
|
assert flagged == global_bytes + slots * (swa_bytes_per_slot + cp_extra_per_slot)
|
|
|
|
# ── --kv-offload (kv_on_gpu) ───────────────────────────────────
|
|
|
|
def test_fit_returns_requested_when_kv_off_gpu(self):
|
|
b = self._gqa_backend()
|
|
# Tiny VRAM budget -- would normally force a reduction.
|
|
fitted = b._fit_context_to_vram(
|
|
requested_ctx = 32_768,
|
|
available_mib = 1,
|
|
model_size_bytes = 100,
|
|
cache_type_kv = "f16",
|
|
kv_on_gpu = False,
|
|
)
|
|
assert fitted == 32_768
|
|
|
|
def test_fit_reduces_when_kv_on_gpu(self):
|
|
b = self._gqa_backend()
|
|
fitted = b._fit_context_to_vram(
|
|
requested_ctx = 32_768,
|
|
available_mib = 64,
|
|
model_size_bytes = 1024 * 1024, # 1 MiB
|
|
cache_type_kv = "f16",
|
|
kv_on_gpu = True,
|
|
)
|
|
assert fitted < 32_768
|
|
|
|
def test_fit_mtp_engaged_returns_smaller_or_equal_context(self):
|
|
# Flat MTP fallback budget is _CTX_FIT_VRAM_FRACTION - 0.05; non-MTP is
|
|
# the full fraction. On a tight budget MTP must yield <= non-MTP.
|
|
b = self._gqa_backend()
|
|
common = dict(
|
|
requested_ctx = 32_768,
|
|
available_mib = 128,
|
|
model_size_bytes = 8 * 1024 * 1024,
|
|
cache_type_kv = "f16",
|
|
)
|
|
baseline = b._fit_context_to_vram(**common)
|
|
mtp = b._fit_context_to_vram(**common, mtp_engaged = True)
|
|
assert mtp <= baseline
|
|
|
|
def test_fit_mtp_engaged_unchanged_when_kv_off_gpu(self):
|
|
# kv_on_gpu=False short-circuits the fit; mtp_engaged irrelevant.
|
|
b = self._gqa_backend()
|
|
fitted = b._fit_context_to_vram(
|
|
requested_ctx = 32_768,
|
|
available_mib = 1,
|
|
model_size_bytes = 100,
|
|
cache_type_kv = "f16",
|
|
kv_on_gpu = False,
|
|
mtp_engaged = True,
|
|
)
|
|
assert fitted == 32_768
|
|
|
|
def test_fit_threads_swa_full_through_estimator(self):
|
|
# SWA model, generous budget; both should fit but cache size differs.
|
|
b = self._swa_backend()
|
|
ctx = 8192
|
|
kv_default = b._estimate_kv_cache_bytes(ctx, "f16")
|
|
kv_full = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
|
|
assert kv_full > kv_default
|
|
# Budget = model + kv_default (rounded up) -- swa_full must not fit.
|
|
budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / _CTX_FIT_VRAM_FRACTION + 1
|
|
fitted_default = b._fit_context_to_vram(
|
|
requested_ctx = ctx,
|
|
available_mib = int(budget_mib),
|
|
model_size_bytes = 1024 * 1024,
|
|
cache_type_kv = "f16",
|
|
)
|
|
fitted_full = b._fit_context_to_vram(
|
|
requested_ctx = ctx,
|
|
available_mib = int(budget_mib),
|
|
model_size_bytes = 1024 * 1024,
|
|
cache_type_kv = "f16",
|
|
swa_full = True,
|
|
)
|
|
assert fitted_default == ctx
|
|
assert fitted_full < ctx
|
|
|
|
|
|
# J2.5. --parallel N memory accounting (per-layer-type scaling rule)
|
|
|
|
|
|
class TestParallelSWAScaling:
|
|
"""Per-layer-type scaling rule vs the closed form measured from
|
|
llama-server. Empirical formula on Gemma-3 270m at ctx=8192:
|
|
total_kv = 24 + parallel * 15 (MiB).
|
|
|
|
Rule (verified vs ``llama-server`` log on real GGUFs):
|
|
* non-SWA layers: total cells = n_ctx, partitioned across slots,
|
|
memory CONSTANT in n_parallel.
|
|
* SWA layers: per-slot cells = 2 * sliding_window (clamped at
|
|
n_ctx and at per_slot_ctx); memory LINEAR in n_parallel.
|
|
* --kv-unified is a no-op for memory math; both modes give the
|
|
same total in measured cases.
|
|
"""
|
|
|
|
def _gqa_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 28,
|
|
"_n_kv_heads": 8,
|
|
"_n_heads": 16,
|
|
"_embedding_length": 1024,
|
|
"_kv_key_length": 128,
|
|
"_kv_value_length": 128,
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
def _swa_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 18,
|
|
"_n_kv_heads": 1,
|
|
"_n_heads": 4,
|
|
"_embedding_length": 1024,
|
|
"_kv_key_length": 256,
|
|
"_kv_value_length": 256,
|
|
"_sliding_window": 512,
|
|
# 15 SWA + 3 global, mirrors gemma-3-270m
|
|
"_sliding_window_pattern": [t == "swa" for t in (["swa"] * 5 + ["global"]) * 3],
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
# ── non-SWA paths: constant ────────────────────────────────────
|
|
|
|
def test_pure_gqa_constant_across_parallel(self):
|
|
b = self._gqa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
for slots in (1, 2, 4, 8):
|
|
for unified in (True, False):
|
|
assert (
|
|
b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = unified)
|
|
== baseline
|
|
)
|
|
|
|
def test_mla_constant_across_parallel(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 60
|
|
b._n_kv_heads = 1
|
|
b._kv_lora_rank = 512
|
|
b._key_length_mla = 64
|
|
b._kv_key_length = 576
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
for slots in (1, 2, 4, 8):
|
|
assert b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots) == baseline
|
|
|
|
def test_hybrid_constant_across_parallel(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 64
|
|
b._n_kv_heads = 16
|
|
b._n_heads = 32
|
|
b._embedding_length = 4096
|
|
b._kv_key_length = 128
|
|
b._kv_value_length = 128
|
|
b._ssm_inner_size = 4096
|
|
b._full_attention_interval = 4
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
for slots in (1, 2, 4, 8):
|
|
assert b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots) == baseline
|
|
|
|
def test_legacy_constant_across_parallel(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 32
|
|
b._n_kv_heads = 8
|
|
b._n_heads = 8
|
|
b._embedding_length = 4096
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
for slots in (1, 2, 4, 8):
|
|
assert b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots) == baseline
|
|
|
|
# ── SWA paths: scale only the SWA portion ──────────────────────
|
|
|
|
def test_swa_pattern_scales_only_swa_portion(self):
|
|
b = self._swa_backend()
|
|
ctx = 8192
|
|
swa = b._sliding_window
|
|
per_token = 1 * (256 + 256) * 2 # n_kv * (k+v) * f16
|
|
n_global = sum(1 for f in b._sliding_window_pattern if not f)
|
|
n_swa = sum(1 for f in b._sliding_window_pattern if f)
|
|
global_bytes = n_global * ctx * per_token
|
|
for slots in (1, 2, 4, 8):
|
|
per_slot_ctx = max(1, ctx // slots)
|
|
cells = min(ctx, 2 * swa, per_slot_ctx)
|
|
swa_bps = n_swa * cells * per_token
|
|
for unified in (True, False):
|
|
got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = unified)
|
|
assert got == global_bytes + slots * swa_bps
|
|
|
|
def test_swa_fallback_scales_only_swa_portion(self):
|
|
# No per-layer pattern -> 1/4-global heuristic.
|
|
b = self._swa_backend(_sliding_window_pattern = None)
|
|
ctx = 8192
|
|
swa = b._sliding_window
|
|
n_layers = 18
|
|
n_global = max(1, n_layers // 4)
|
|
n_swa = n_layers - n_global
|
|
per_token = 1 * (256 + 256) * 2
|
|
global_bytes = n_global * ctx * per_token
|
|
for slots in (1, 2, 4, 8):
|
|
per_slot_ctx = max(1, ctx // slots)
|
|
cells = min(ctx, 2 * swa, per_slot_ctx)
|
|
swa_bps = n_swa * cells * per_token
|
|
got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots)
|
|
assert got == global_bytes + slots * swa_bps
|
|
|
|
def test_swa_per_slot_clamped_when_ctx_lt_slots_x_2window(self):
|
|
# ctx=4096 / slots=8 -> per_slot_ctx=512, but 2*sliding=1024.
|
|
# SWA cells clamp at per_slot_ctx (512), not 2*sliding.
|
|
b = self._swa_backend()
|
|
ctx = 4096
|
|
per_slot_ctx_at_8 = ctx // 8
|
|
assert per_slot_ctx_at_8 < 2 * b._sliding_window
|
|
# Build expected with the clamped formula
|
|
n_swa = sum(1 for f in b._sliding_window_pattern if f)
|
|
n_global = sum(1 for f in b._sliding_window_pattern if not f)
|
|
per_token = 1 * (256 + 256) * 2
|
|
global_bytes = n_global * ctx * per_token
|
|
cells = min(ctx, 2 * b._sliding_window, per_slot_ctx_at_8)
|
|
assert cells == per_slot_ctx_at_8
|
|
expected = global_bytes + 8 * (n_swa * cells * per_token)
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = 8) == expected
|
|
|
|
def test_swa_full_does_not_scale_under_parallel(self):
|
|
# swa_full forces every layer to n_ctx -> all-global GQA-style
|
|
# total, constant in parallel.
|
|
b = self._swa_backend()
|
|
ctx = 8192
|
|
baseline = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
|
|
for slots in (1, 2, 4, 8):
|
|
assert (
|
|
b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True, n_parallel = slots) == baseline
|
|
)
|
|
|
|
# ── kv_unified: no-op for memory math ──────────────────────────
|
|
|
|
def test_kv_unified_is_no_op_for_memory_math(self):
|
|
# unified=True and unified=False must give the same total bytes
|
|
# for every backend type and parallel value.
|
|
backends = [
|
|
("gqa", self._gqa_backend()),
|
|
("swa", self._swa_backend()),
|
|
]
|
|
for label, b in backends:
|
|
for slots in (1, 2, 4, 8):
|
|
u = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = True)
|
|
nu = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = False)
|
|
assert u == nu, f"{label} parallel={slots} unified-mismatch"
|
|
|
|
# ── Empirical Gemma-3 270m formula ─────────────────────────────
|
|
|
|
def test_matches_empirical_gemma3_270m_formula(self):
|
|
"""Exact match against the formula measured from llama-server:
|
|
total_kv = 24 + parallel * 15 (MiB) at ctx=8192.
|
|
|
|
Geometry: 18 layers (3 global + 15 SWA), n_kv=1, head_dim=256,
|
|
sliding=512, f16.
|
|
"""
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 18
|
|
b._n_kv_heads = 1
|
|
b._n_heads = 4
|
|
b._embedding_length = 1024
|
|
b._kv_key_length = 256
|
|
b._kv_value_length = 256
|
|
b._sliding_window = 512
|
|
# Mirrors the bootstrap-resolved gemma3 pattern (period 6) on an
|
|
# 18-layer model: 15 SWA, 3 global.
|
|
b._sliding_window_pattern = [(i + 1) % 6 != 0 for i in range(18)]
|
|
n_global = 3
|
|
n_swa = 15
|
|
# Confirm pattern shape
|
|
assert sum(b._sliding_window_pattern) == n_swa
|
|
for slots, expected_mib in [(1, 39), (2, 54), (4, 84)]:
|
|
got_bytes = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots)
|
|
got_mib = got_bytes / (1024 * 1024)
|
|
assert (
|
|
got_mib == expected_mib
|
|
), f"slots={slots}: got {got_mib} MiB, expected {expected_mib} MiB"
|
|
|
|
|
|
# J3. shared_kv_layers (Gemma 3n / Gemma 4)
|
|
|
|
|
|
class TestSharedKVLayers:
|
|
"""``<arch>.attention.shared_kv_layers`` reduces the layer count that
|
|
allocates KV. The trailing ``shared_kv_layers`` blocks reuse earlier
|
|
caches (Gemma 3n: 35 layers, 15 shared -> 20 allocate; Gemma 4 same
|
|
field). Unset on every other arch -> no behavioural change."""
|
|
|
|
def _gemma3n_backend(self, **overrides):
|
|
# Mirrors google/gemma-3n-E4B-it: 35 layers, 15 shared, SWA window
|
|
# 1024, period 5 (4 sliding + 1 full repeating).
|
|
defaults = {
|
|
"_n_layers": 35,
|
|
"_n_kv_heads": 4,
|
|
"_n_heads": 8,
|
|
"_embedding_length": 2048,
|
|
"_kv_key_length": 256,
|
|
"_kv_value_length": 256,
|
|
"_sliding_window": 1024,
|
|
"_sliding_window_pattern": [
|
|
t == "sliding_attention"
|
|
for t in (["sliding_attention"] * 4 + ["full_attention"]) * 7
|
|
],
|
|
"_shared_kv_layers": 15,
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
def _gqa_backend(self, **overrides):
|
|
defaults = {
|
|
"_n_layers": 28,
|
|
"_n_kv_heads": 8,
|
|
"_n_heads": 16,
|
|
"_embedding_length": 1024,
|
|
"_kv_key_length": 128,
|
|
"_kv_value_length": 128,
|
|
}
|
|
defaults.update(overrides)
|
|
b = LlamaCppBackend()
|
|
for k, v in defaults.items():
|
|
setattr(b, k, v)
|
|
return b
|
|
|
|
def test_field_initialises_to_none(self):
|
|
b = LlamaCppBackend()
|
|
assert b._shared_kv_layers is None
|
|
|
|
def test_unset_field_is_noop(self):
|
|
b = self._gqa_backend()
|
|
baseline = b._estimate_kv_cache_bytes(8192, "f16")
|
|
b._shared_kv_layers = None
|
|
assert b._estimate_kv_cache_bytes(8192, "f16") == baseline
|
|
b._shared_kv_layers = 0
|
|
assert b._estimate_kv_cache_bytes(8192, "f16") == baseline
|
|
|
|
def test_path4_drops_shared_layers(self):
|
|
b = self._gqa_backend(_shared_kv_layers = 4)
|
|
ctx = 4096
|
|
kv_per = 8 * (128 + 128) * 2
|
|
# 28 - 4 = 24 layers actually allocate
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == 24 * ctx * kv_per
|
|
|
|
def test_path5_drops_shared_layers(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 32
|
|
b._n_kv_heads = 8
|
|
b._n_heads = 8
|
|
b._embedding_length = 4096
|
|
b._shared_kv_layers = 8
|
|
ctx = 4096
|
|
head_dim = 4096 // 8 # 512
|
|
# 32 - 8 = 24 layers
|
|
expected = 2 * 8 * head_dim * 24 * ctx * 2
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
|
|
|
|
def test_path1_mla_drops_shared_layers(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 60
|
|
b._n_kv_heads = 1
|
|
b._kv_lora_rank = 512
|
|
b._key_length_mla = 64
|
|
b._kv_key_length = 576
|
|
b._shared_kv_layers = 10
|
|
ctx = 8192
|
|
# 60 - 10 = 50
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == 50 * ctx * 1 * 576 * 2
|
|
|
|
def test_path3_pattern_loops_only_unshared_layers(self):
|
|
b = self._gemma3n_backend()
|
|
ctx = 8192
|
|
# First 20 layers contribute; layers 20..34 skipped. Pattern
|
|
# [s,s,s,s,F] repeated -> in layers 0..19: sliding 16, full 4.
|
|
sliding_in_unshared = sum(b._sliding_window_pattern[:20])
|
|
full_in_unshared = 20 - sliding_in_unshared
|
|
assert sliding_in_unshared == 16
|
|
assert full_in_unshared == 4
|
|
kv_per = 4 * (256 + 256) * 2
|
|
swa_cells = min(ctx, 2 * 1024)
|
|
expected = full_in_unshared * ctx * kv_per + sliding_in_unshared * swa_cells * kv_per
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
|
|
|
|
def test_shared_layers_reduces_estimate(self):
|
|
b = self._gemma3n_backend()
|
|
with_shared = b._estimate_kv_cache_bytes(8192, "f16")
|
|
b._shared_kv_layers = 0
|
|
without_shared = b._estimate_kv_cache_bytes(8192, "f16")
|
|
# 20/35 = 0.571 of the work; ~43% reduction.
|
|
ratio = with_shared / without_shared
|
|
assert 0.5 < ratio < 0.65
|
|
|
|
def test_path3_pattern_with_swa_full_and_shared(self):
|
|
b = self._gemma3n_backend()
|
|
ctx = 8192
|
|
flagged = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
|
|
# Every unshared layer caches n_ctx -> path-4-style sizing over
|
|
# only the 20 unshared layers.
|
|
kv_per = 4 * (256 + 256) * 2
|
|
assert flagged == 20 * ctx * kv_per
|
|
|
|
def test_path3_fallback_uses_unshared_count(self):
|
|
# No per-layer pattern -> 1/4-global heuristic over n_layers_kv,
|
|
# not n_layers.
|
|
b = self._gemma3n_backend(_sliding_window_pattern = None)
|
|
ctx = 8192
|
|
n_layers_kv = 35 - 15 # 20
|
|
n_global = max(1, n_layers_kv // 4) # 5
|
|
n_swa = n_layers_kv - n_global # 15
|
|
kv_per = 4 * (256 + 256) * 2
|
|
swa_cells = min(ctx, 2 * 1024)
|
|
expected = n_global * ctx * kv_per + n_swa * swa_cells * kv_per
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
|
|
|
|
def test_shared_floors_at_one_layer(self):
|
|
# Pathological: shared >= n_layers must not zero out the cache.
|
|
b = self._gqa_backend(_shared_kv_layers = 99)
|
|
ctx = 4096
|
|
kv_per = 8 * (128 + 128) * 2
|
|
assert b._estimate_kv_cache_bytes(ctx, "f16") == 1 * ctx * kv_per
|
|
|
|
def test_composes_with_n_parallel(self):
|
|
# Only the SWA portion of unshared layers scales by n_parallel;
|
|
# the global portion is constant.
|
|
b = self._gemma3n_backend()
|
|
ctx = 8192
|
|
swa = b._sliding_window
|
|
per_token = 4 * (256 + 256) * 2
|
|
unshared_pattern = b._sliding_window_pattern[:20] # 35 - 15 shared
|
|
sliding_in_unshared = sum(unshared_pattern)
|
|
global_in_unshared = len(unshared_pattern) - sliding_in_unshared
|
|
global_bytes = global_in_unshared * ctx * per_token
|
|
slots = 3
|
|
per_slot_ctx = max(1, ctx // slots)
|
|
swa_cells = min(ctx, 2 * swa, per_slot_ctx)
|
|
swa_bytes_per_slot = sliding_in_unshared * swa_cells * per_token
|
|
flagged = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False)
|
|
assert flagged == global_bytes + slots * swa_bytes_per_slot
|
|
|
|
def test_composes_with_ctx_checkpoints(self):
|
|
b = self._gemma3n_backend()
|
|
ctx = 8192
|
|
baseline = b._estimate_kv_cache_bytes(ctx, "f16")
|
|
with_cp = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4)
|
|
# Checkpoints count only over UNSHARED SWA layers (16 of them).
|
|
sliding_in_unshared = sum(b._sliding_window_pattern[:20])
|
|
per_cp_layer = 4 * 1024 * 4 * (256 + 256) * 2 # cps * swa * heads * (k+v) * bpe
|
|
assert with_cp == baseline + sliding_in_unshared * per_cp_layer
|
|
|
|
def test_unload_resets_shared_kv_layers(self):
|
|
b = LlamaCppBackend()
|
|
b._shared_kv_layers = 12
|
|
b.unload_model()
|
|
assert b._shared_kv_layers is None
|
|
|
|
|
|
# K. Lifecycle Tests
|
|
|
|
|
|
class TestLifecycle:
|
|
"""Init, unload, and reparse field management."""
|
|
|
|
def test_init_fields_none(self):
|
|
b = LlamaCppBackend()
|
|
for attr in [
|
|
"_kv_key_length",
|
|
"_kv_value_length",
|
|
"_sliding_window",
|
|
"_sliding_window_pattern",
|
|
"_full_attention_interval",
|
|
"_kv_lora_rank",
|
|
"_key_length_mla",
|
|
"_kv_key_length_swa",
|
|
"_kv_value_length_swa",
|
|
"_ssm_inner_size",
|
|
"_ssm_state_size",
|
|
"_shared_kv_layers",
|
|
]:
|
|
assert getattr(b, attr) is None
|
|
assert b._n_kv_heads_by_layer is None
|
|
|
|
def test_unload_resets_fields(self):
|
|
b = LlamaCppBackend()
|
|
b._n_layers = 32
|
|
b._kv_key_length = 128
|
|
b._kv_lora_rank = 512
|
|
b._sliding_window = 1024
|
|
b._sliding_window_pattern = [True, False]
|
|
b._n_kv_heads_by_layer = [8, 2]
|
|
b._kv_key_length_swa = 64
|
|
b._kv_value_length_swa = 64
|
|
b._ssm_inner_size = 4096
|
|
b._full_attention_interval = 4
|
|
b._shared_kv_layers = 8
|
|
b.unload_model()
|
|
for attr in [
|
|
"_kv_key_length",
|
|
"_kv_value_length",
|
|
"_sliding_window",
|
|
"_sliding_window_pattern",
|
|
"_full_attention_interval",
|
|
"_kv_lora_rank",
|
|
"_key_length_mla",
|
|
"_kv_key_length_swa",
|
|
"_kv_value_length_swa",
|
|
"_ssm_inner_size",
|
|
"_ssm_state_size",
|
|
"_shared_kv_layers",
|
|
]:
|
|
assert getattr(b, attr) is None
|
|
assert b._n_kv_heads_by_layer is None
|
|
|
|
def test_end_to_end_synthetic_mla(self):
|
|
"""Round-trip: write GGUF -> parse -> estimate."""
|
|
b = _backend_from_gguf(
|
|
"deepseek2",
|
|
{
|
|
"context_length": 163840,
|
|
"block_count": 61,
|
|
"attention.head_count_kv": 1,
|
|
"attention.head_count": 128,
|
|
"embedding_length": 7168,
|
|
"attention.key_length": 576,
|
|
"attention.value_length": 512,
|
|
"attention.kv_lora_rank": 512,
|
|
"attention.key_length_mla": 192,
|
|
},
|
|
)
|
|
assert b._can_estimate_kv()
|
|
result = b._estimate_kv_cache_bytes(163840, "f16")
|
|
expected = 61 * 163840 * 1 * 576 * 2
|
|
assert result == expected
|
|
|
|
def test_end_to_end_synthetic_hybrid(self):
|
|
b = _backend_from_gguf(
|
|
"qwen35",
|
|
{
|
|
"context_length": 262144,
|
|
"block_count": 64,
|
|
"attention.head_count_kv": 4,
|
|
"attention.head_count": 24,
|
|
"embedding_length": 5120,
|
|
"attention.key_length": 256,
|
|
"attention.value_length": 256,
|
|
"full_attention_interval": 4,
|
|
"ssm.inner_size": 6144,
|
|
"ssm.state_size": 128,
|
|
},
|
|
)
|
|
assert b._can_estimate_kv()
|
|
result = b._estimate_kv_cache_bytes(262144, "f16")
|
|
n_attn = 64 // 4
|
|
expected = n_attn * 262144 * 4 * (256 + 256) * 2
|
|
assert result == expected
|
|
|
|
def test_end_to_end_synthetic_swa(self):
|
|
b = _backend_from_gguf(
|
|
"gemma3",
|
|
{
|
|
"context_length": 131072,
|
|
"block_count": 62,
|
|
"attention.head_count_kv": 16,
|
|
"attention.head_count": 32,
|
|
"embedding_length": 5376,
|
|
"attention.key_length": 128,
|
|
"attention.value_length": 128,
|
|
"attention.sliding_window": 1024,
|
|
},
|
|
)
|
|
assert b._can_estimate_kv()
|
|
result = b._estimate_kv_cache_bytes(131072, "f16")
|
|
# gemma3 -> period 6 from bootstrap; SWA cache double-buffered to
|
|
# 2 * sliding_window cells.
|
|
period = 6
|
|
kv_per = 16 * 256 * 2
|
|
expected = 0
|
|
for i in range(62):
|
|
is_swa = (i + 1) % period != 0
|
|
layer_ctx = min(131072, 2 * 1024) if is_swa else 131072
|
|
expected += layer_ctx * kv_per
|
|
assert result == expected
|
|
|
|
def test_end_to_end_synthetic_shared_kv_round_trip(self):
|
|
# Mirrors gemma3n_text: 35 layers, 15 shared, sliding_window=1024.
|
|
b = _backend_from_gguf(
|
|
"gemma3n_text",
|
|
{
|
|
"context_length": 32768,
|
|
"block_count": 35,
|
|
"attention.head_count_kv": 4,
|
|
"attention.head_count": 8,
|
|
"embedding_length": 2048,
|
|
"attention.key_length": 256,
|
|
"attention.value_length": 256,
|
|
"attention.sliding_window": 1024,
|
|
"attention.shared_kv_layers": 15,
|
|
},
|
|
)
|
|
assert b._can_estimate_kv()
|
|
assert b._shared_kv_layers == 15
|
|
# Bootstrap for gemma3n_text -> period 5; resolver synthesises a
|
|
# 35-entry bool array. Only the first 20 (n_layers - shared)
|
|
# allocate KV.
|
|
result = b._estimate_kv_cache_bytes(8192, "f16")
|
|
assert result > 0
|
|
# Sanity: shared back to 0 -> strictly larger estimate (more
|
|
# layers allocate).
|
|
b._shared_kv_layers = 0
|
|
unshared = b._estimate_kv_cache_bytes(8192, "f16")
|
|
assert unshared > result
|
|
|
|
def test_end_to_end_synthetic_gqa(self):
|
|
b = _backend_from_gguf(
|
|
"qwen3",
|
|
{
|
|
"context_length": 40960,
|
|
"block_count": 28,
|
|
"attention.head_count_kv": 8,
|
|
"attention.head_count": 16,
|
|
"embedding_length": 1024,
|
|
"attention.key_length": 128,
|
|
"attention.value_length": 128,
|
|
},
|
|
)
|
|
assert b._can_estimate_kv()
|
|
result = b._estimate_kv_cache_bytes(40960, "f16")
|
|
expected = 28 * 40960 * 8 * 256 * 2
|
|
assert result == expected
|
|
|
|
def test_end_to_end_synthetic_legacy(self):
|
|
b = _backend_from_gguf(
|
|
"llama",
|
|
{
|
|
"context_length": 4096,
|
|
"block_count": 32,
|
|
"attention.head_count_kv": 8,
|
|
"attention.head_count": 32,
|
|
"embedding_length": 4096,
|
|
},
|
|
)
|
|
assert b._can_estimate_kv()
|
|
result = b._estimate_kv_cache_bytes(4096, "f16")
|
|
head_dim = 4096 // 32
|
|
expected = int(2 * 8 * head_dim * 32 * 4096 * 2)
|
|
assert result == expected
|