From 0e69614878096fee53ee0fbe724011327e4bb601 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 17 Jun 2026 03:10:22 -0700 Subject: [PATCH] studio: deterministic VRAM auto-fit for GGUF (MTP reserve, compute buffer, total-based budget) (#6312) * 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> --- studio/backend/core/inference/llama_cpp.py | 1261 ++++++++++++++--- .../core/inference/llama_server_args.py | 92 +- .../backend/core/inference/tensor_fallback.py | 22 +- studio/backend/routes/inference.py | 7 +- studio/backend/tests/test_compute_buffer.py | 152 ++ .../backend/tests/test_kv_cache_estimation.py | 8 +- .../tests/test_llama_cpp_context_fit.py | 44 +- .../test_llama_cpp_max_context_threshold.py | 4 +- .../tests/test_llama_cpp_mtp_detection.py | 14 +- .../backend/tests/test_llama_server_args.py | 67 + studio/backend/tests/test_mtp_vram_budget.py | 1002 +++++++++++++ studio/backend/tests/test_tensor_parallel.py | 203 ++- .../tests/test_windows_gpu_detection_mock.py | 16 + 13 files changed, 2656 insertions(+), 236 deletions(-) create mode 100644 studio/backend/tests/test_compute_buffer.py create mode 100644 studio/backend/tests/test_mtp_vram_budget.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 7b6c0b0668..9221323bd9 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -23,18 +23,19 @@ import sys import threading import time from pathlib import Path -from typing import Callable, Collection, Generator, Iterable, List, Optional, Union +from typing import Callable, Collection, Generator, Iterable, List, Mapping, Optional, Union import httpx from core.inference.llama_server_args import ( + _effective_tensor_parallel, + _tensor_parallel_matches_loaded, extra_args_disable_mmproj, parse_cache_override, + parse_cache_override_per_axis, parse_ctx_override, parse_split_mode_override, - resolve_cache_type_kv, resolve_requested_ctx, - resolve_tensor_parallel, strip_shadowing_flags, strip_split_mode_only, ) @@ -724,18 +725,63 @@ def _gguf_files_for_variant(files: Iterable[str], variant: str) -> list[str]: # _build_speculative_flags); auto mode drops MTP under it. _MTP_MIN_SIZE_B = 3.0 -# Context-fit VRAM budget: tighter than _GPU_PIN_VRAM_FRACTION (0.95) on -# purpose -- over-promising context OOMs at runtime (#5106). -_CTX_FIT_VRAM_FRACTION = 0.90 +# Cap total GPU occupancy at this fraction of the card. The fit reserves an +# absolute (1 - frac) * total per GPU when total VRAM is known, else a fraction +# of free (see _fit_context_to_vram), plus a byte-accurate MTP draft reserve. +_CTX_FIT_VRAM_FRACTION = 0.95 -# Extra VRAM fraction reserved when MTP will engage: the draft model's -# weights, KV cache, and compute buffers live outside the main model's -# estimate. Applied to BOTH the ctx-fit budget and the GPU pin thresholds -- -# tightening only the fit lets a load whose weights land between the two -# fractions pin without any room for the drafter. +# Flat MTP reserve, used only when GGUF dims are too sparse for the byte-accurate +# reserve (_estimate_mtp_overhead_bytes). Applied to both the fit budget and pin. _MTP_VRAM_RESERVE_FRAC = 0.05 +def _kv_bytes_per_elem(cache_type: Optional[str]) -> float: + """Bytes per KV-cache element for a llama.cpp cache type (f16 default).""" + return { + "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, + }.get((cache_type or "f16").strip().lower(), 2.0) + + +def _env_main_cache_type_for_budget(env: Optional[Mapping[str, str]] = None) -> Optional[str]: + """Heavier of the inherited LLAMA_ARG_CACHE_TYPE_K/_V env types when it + exceeds the f16 default, else None. Studio emits --cache-type only for the + param/extras path, so a heavier env (f32) would otherwise reach the child + unbudgeted; quantized env types stay over-reserved by f16 (-> None).""" + e = os.environ if env is None else env + f16_bpe = _kv_bytes_per_elem("f16") + heaviest: Optional[str] = None + heaviest_bpe = f16_bpe + for var in ("LLAMA_ARG_CACHE_TYPE_K", "LLAMA_ARG_CACHE_TYPE_V"): + raw = (e.get(var) or "").strip().lower() + if not raw: + continue + bpe = _kv_bytes_per_elem(raw) + if bpe > heaviest_bpe: + heaviest, heaviest_bpe = raw, bpe + return heaviest + + +def _extra_args_main_cache_type_for_budget(extra_args: Optional[Iterable[str]]) -> Optional[str]: + """Heavier (max bytes/elem) of the explicit --cache-type-k/-v extras, or None. + + Extras are appended last and win per axis, so an asymmetric K=f32,V=f16 must be + budgeted by its heavier axis. resolve_cache_type_kv returns only the last-wins + single type, which under-reserves the heavier axis when the lighter one is last.""" + k, v = parse_cache_override_per_axis(extra_args) + candidates = [c for c in (k, v) if c] + if not candidates: + return None + return max(candidates, key = _kv_bytes_per_elem) + + def _auto_mode_drops_mtp( req_mode: Optional[str], size_b: Optional[float], @@ -779,6 +825,196 @@ def _extra_args_set_any_flag(extra_args: Optional[Iterable[str]], flags: Collect return False +def _effective_spec_type( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> Optional[str]: + """The --spec-type llama-server will use: the last CLI --spec-type (or + --spec-default, which resolves non-MTP), else LLAMA_ARG_SPEC_TYPE. A CLI flag + overrides the env (matching llama.cpp), so a stale MTP env can't make the + budget reserve a drafter the launch won't load. None if neither sets it.""" + args = [str(a) for a in extra_args] if extra_args else [] + cli_present = False + cli_value: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag == "--spec-default": + cli_present = True + cli_value = "default" + continue + if flag != "--spec-type": + continue + cli_present = True + cli_value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if cli_present: + return cli_value + return (os.environ if env is None else env).get("LLAMA_ARG_SPEC_TYPE") + + +def _extra_args_requests_mtp( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True if the effective --spec-type selects MTP (mtp/draft-mtp), so the + budget must reserve for it.""" + value = _effective_spec_type(extra_args, env) + if not value: + return False + return any(p.strip().lower() in ("mtp", "draft-mtp") for p in value.split(",")) + + +def _extra_args_requests_separate_draft( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True if the effective --spec-type selects a non-MTP model draft mode + (draft-simple/draft-eagle3), which loads a separate draft model the budget + must reserve (draft-mtp -> _extra_args_requests_mtp; ngram-* load no model).""" + value = _effective_spec_type(extra_args, env) + if not value: + return False + return any(p.strip().lower() in ("draft-simple", "draft-eagle3") for p in value.split(",")) + + +def _extra_args_spec_draft_n_max(extra_args: Optional[Iterable[str]]) -> Optional[int]: + """Draft depth from extras (``--spec-draft-n-max`` or legacy ``--draft-max``), else None.""" + if not extra_args: + return None + args = [str(a) for a in extra_args] + found: Optional[int] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in ("--spec-draft-n-max", "--draft-max"): + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + try: + found = int(value) + except (TypeError, ValueError): + continue + return found + + +def _extra_args_mtp_draft_path( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> Optional[str]: + """Separate drafter path from extras (local --model-draft/-md or HF + --spec-draft-hf/-hfd/...), else the LLAMA_ARG_SPEC_DRAFT_MODEL/_HF_REPO env, + else None. An HF repo isn't a local file, so the budget can't size it (falls + back to the flat reserve), but recognizing it avoids sizing the wrong one.""" + flags = { + "--model-draft", + "--spec-draft-model", + "-md", + "--spec-draft-hf", + "-hfd", + "-hfrd", + "--hf-repo-draft", + } + args = [str(a) for a in extra_args] if extra_args else [] + found: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in flags: + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if value and not value.startswith("-"): + found = value + if found is not None: + return found + e = os.environ if env is None else env + return e.get("LLAMA_ARG_SPEC_DRAFT_MODEL") or e.get("LLAMA_ARG_SPEC_DRAFT_HF_REPO") or None + + +def _extra_args_draft_cache_types( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> tuple[Optional[str], Optional[str]]: + """Draft KV cache types (k_type, v_type), each from extras else the + LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K/_V env, else None (f16). K and V are + independent: a one-sided override must not apply to both.""" + args = [str(a) for a in extra_args] if extra_args else [] + k_flags = {"--cache-type-k-draft", "--spec-draft-type-k", "-ctkd"} + v_flags = {"--cache-type-v-draft", "--spec-draft-type-v", "-ctvd"} + k_type: Optional[str] = None + v_type: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in k_flags and flag not in v_flags: + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if not value or value.startswith("-"): + continue + if flag in k_flags: + k_type = value + else: + v_type = value + e = os.environ if env is None else env + if k_type is None: + k_type = e.get("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K") or None + if v_type is None: + v_type = e.get("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V") or None + return k_type, v_type + + +def _extra_args_draft_offloaded_to_cpu( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True if the SEPARATE draft model is on CPU (so the budget must not charge + its weights+KV): --spec-draft-ngl 0, or --spec-draft-device naming only + cpu/none, else the LLAMA_ARG_N_GPU_LAYERS_DRAFT env the child honors (the + device flag has no env). An embedded MTP head follows the main -ngl, so these + draft-only flags don't move it. Last-wins, so only each flag's final value counts.""" + ngl_flags = {"--spec-draft-ngl", "-ngld", "--gpu-layers-draft", "--n-gpu-layers-draft"} + dev_flags = {"--spec-draft-device", "-devd", "--device-draft"} + args = [str(a) for a in extra_args] if extra_args else [] + last_ngl: Optional[str] = None + last_dev: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if flag in ngl_flags: + last_ngl = value + elif flag in dev_flags: + last_dev = value + if last_ngl is None: + last_ngl = (os.environ if env is None else env).get("LLAMA_ARG_N_GPU_LAYERS_DRAFT") + if last_ngl is not None: + try: + if int(last_ngl) == 0: + return True + except (TypeError, ValueError): + pass + if last_dev is not None: + devs = [d.strip().lower() for d in last_dev.split(",") if d.strip()] + if devs and all(d in ("cpu", "none") for d in devs): + return True + return False + + +def _extra_args_n_ubatch( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> Optional[int]: + """Physical micro-batch from extras (--ubatch-size/-ub) else the LLAMA_ARG_UBATCH + env, else None. It sizes the compute-graph buffer, so an override must reach + the VRAM reserve.""" + args = [str(a) for a in extra_args] if extra_args else [] + found: Optional[int] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in ("--ubatch-size", "-ub"): + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + try: + found = int(value) + except (TypeError, ValueError): + continue + if found is not None: + return found + raw = (os.environ if env is None else env).get("LLAMA_ARG_UBATCH") + if raw: + try: + return int(raw) + except (TypeError, ValueError): + pass + return None + + def _build_ngram_mod_flags( caps: Optional[dict], n_match: int = 24, @@ -954,6 +1190,9 @@ class LlamaCppBackend: self._n_kv_heads_by_layer: Optional[list[int]] = None self._n_heads: Optional[int] = None self._embedding_length: Optional[int] = None + # For the compute-graph buffer estimate; vocab from the tokens array len. + self._feed_forward_length: Optional[int] = None + self._vocab_size: Optional[int] = None # Architecture-aware KV fields for 5-path estimation self._kv_key_length: Optional[int] = None self._kv_value_length: Optional[int] = None @@ -1729,7 +1968,14 @@ class LlamaCppBackend: @staticmethod def _get_gpu_free_memory() -> list[tuple[int, int]]: - """Query free memory per GPU. + """Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by + index; empty if no supported GPU is reachable. Thin wrapper over + ``_get_gpu_memory`` for callers that only need free VRAM.""" + return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()] + + @staticmethod + def _get_gpu_memory() -> list[tuple[int, int, int]]: + """Query free AND total memory per GPU. Order: 1. ``nvidia-smi`` (NVIDIA CUDA hosts) -- respects @@ -1740,15 +1986,15 @@ class LlamaCppBackend: probe returned [] on AMD) and NVIDIA hosts missing ``nvidia-smi`` from PATH. - Returns list of (gpu_index, free_mib) sorted by index; empty if no - supported GPU is reachable. + Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no + supported GPU is reachable. ``total`` lets the fit reserve absolute headroom. """ # ── NVIDIA via nvidia-smi ──────────────────────────────────── try: result = subprocess.run( [ "nvidia-smi", - "--query-gpu=index,memory.free", + "--query-gpu=index,memory.free,memory.total", "--format=csv,noheader,nounits", ], capture_output = True, @@ -1768,15 +2014,30 @@ class LlamaCppBackend: allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip()) except ValueError: pass - gpus: list[tuple[int, int]] = [] + gpus: list[tuple[int, int, int]] = [] for line in result.stdout.strip().splitlines(): - parts = line.split(",") - if len(parts) == 2: - idx = int(parts[0].strip()) - free_mib = int(parts[1].strip()) - if allowed is not None and idx not in allowed: - continue - gpus.append((idx, free_mib)) + parts = [p.strip() for p in line.split(",")] + if len(parts) < 2: + continue + # Index and free required; skip a bad line rather than abandon + # the probe to the torch fallback. + try: + idx = int(parts[0]) + free_mib = int(parts[1]) + except ValueError: + continue + # Total parsed separately: a two-column line or a non-integer + # total ("N/A" on MIG/vGPU) keeps the GPU at total 0 (fit uses + # the free*frac fallback) instead of dropping it. + total_mib = 0 + if len(parts) >= 3 and parts[2]: + try: + total_mib = int(parts[2]) + except ValueError: + total_mib = 0 + if allowed is not None and idx not in allowed: + continue + gpus.append((idx, free_mib, total_mib)) # Match the docstring's sort-by-id guarantee (driver order isn't). gpus.sort(key = lambda g: g[0]) if gpus: @@ -1821,13 +2082,13 @@ class LlamaCppBackend: physical_ids = None gpus = [] for ordinal in range(torch.cuda.device_count()): - free_bytes, _total_bytes = torch.cuda.mem_get_info(ordinal) + free_bytes, total_bytes = torch.cuda.mem_get_info(ordinal) idx = ( physical_ids[ordinal] if physical_ids is not None and ordinal < len(physical_ids) else ordinal ) - gpus.append((idx, free_bytes // (1024 * 1024))) + gpus.append((idx, free_bytes // (1024 * 1024), total_bytes // (1024 * 1024))) # Match the nvidia-smi path's docstring guarantee of sorted-by-id. return sorted(gpus, key = lambda g: g[0]) except Exception as e: @@ -1905,19 +2166,17 @@ class LlamaCppBackend: # buffers; 0.90 dropped 91-94% fits to CPU offload (#5106). _GPU_PIN_VRAM_FRACTION = 0.95 - # Per-GPU compute-graph buffer to reserve in tensor mode (MiB). This is the - # logits buffer (n_batch x vocab) + activation scratch that llama.cpp sizes - # via graph_reserve -- it is roughly EQUAL on every device (not proportional - # to the tensor split) and independent of context. Measured ~2.3 GB - # (gemma-3-27B) to ~3.8 GB (gemma-4-31B) on a 256k-vocab model; we reserve a - # conservative headroom above that. It is (a) subtracted from each GPU's free - # VRAM before computing --tensor-split, so the roomier GPU absorbs more - # weight and the smallest GPU keeps room for KV, and (b) reserved per device - # when capping context. The auto-fallback to layer split covers any - # underestimate. NOTE: scales with the model's vocab / batch size; tune if a - # large-vocab model OOMs at load. + # Fallback per-device tensor-mode compute buffer (MiB), used only when GGUF + # dims are unavailable so _estimate_compute_buffer_bytes (the primary, derived + # path) returns 0. _TENSOR_PARALLEL_BUFFER_RESERVE_MIB = 5120 + # Fixed per-device overhead on every GPU of a LAYER split (CUDA context + + # scratch), beyond the conserved slot-scaling buffer. ~0.9 GB/device measured + # (Qwen3.6-27B, b9625), independent of --parallel; reserved per extra GPU so a + # tight layer split can't advertise a context that OOMs at load. + _PIPELINE_PER_DEVICE_OVERHEAD_MIB = 1024 + # KV cache types llama.cpp accepts in tensor mode. A quantized KV cache # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) @@ -2085,6 +2344,8 @@ class LlamaCppBackend: model_size_bytes: int, gpus: list[tuple[int, int]], usable_fraction: Optional[float] = None, + total_by_idx: Optional[dict[int, int]] = None, + per_device_overhead_bytes: int = 0, ) -> tuple[Optional[list[int]], bool]: """Pick GPU(s) for a model from estimated VRAM and free memory. @@ -2092,6 +2353,11 @@ class LlamaCppBackend: ``usable_fraction`` (default ``_GPU_PIN_VRAM_FRACTION``) provides headroom for compute buffers, CUDA context, and other runtime overhead; callers lower it when MTP reserves VRAM for a draft model. + ``total_by_idx`` (index -> total MiB) makes the headroom an ABSOLUTE + ``(1 - fraction) * total`` per GPU instead of a fraction of free. + ``per_device_overhead_bytes`` is the fixed layer-split cost per GPU beyond + the first; a k-GPU pin must hold ``model + (k-1) * overhead`` or it can OOM + a device after -ngl -1 (no --fit fallback). Single-GPU adds none. Returns (gpu_indices, use_fit): - ([1], False) fits on 1 GPU at the headroom threshold @@ -2105,20 +2371,31 @@ class LlamaCppBackend: if usable_fraction is None: usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION - # Sort GPUs by free memory descending - ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + # Per-GPU usable budget: free - (1-frac)*total when total is known, else + # the legacy free*frac (also covers a total-0 two-column probe). + def _usable(idx: int, free_mib: int) -> float: + t = total_by_idx.get(idx, 0) if total_by_idx else 0 + if t > 0: + return max(0.0, free_mib - (1.0 - usable_fraction) * t) + return free_mib * usable_fraction + + # Rank by usable budget (free - reserve), not raw free: a more-used large + # card can have less usable room than a less-used small one. + ranked = sorted(gpus, key = lambda g: _usable(g[0], g[1]), reverse = True) # Try 1 GPU at the usable-VRAM threshold. - if ranked[0][1] * usable_fraction >= model_size_mib: + if _usable(ranked[0][0], ranked[0][1]) >= model_size_mib: return [ranked[0][0]], False - # Try N GPUs (accumulate free memory from most-free) - cumulative = 0 + # Try N GPUs (accumulate usable memory from most-free). Each GPU past the + # first adds a fixed per-device overhead the pool must hold. + overhead_mib = per_device_overhead_bytes / (1024 * 1024) + cumulative = 0.0 selected = [] for idx, free_mib in ranked: selected.append(idx) - cumulative += free_mib * usable_fraction - if cumulative >= model_size_mib: + cumulative += _usable(idx, free_mib) + if cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib: return sorted(selected), False # Too large even for all GPUs; let --fit handle it @@ -2173,14 +2450,10 @@ class LlamaCppBackend: 5. Legacy -- fallback using embed // n_heads Server-flag knobs (mirror llama-server's CLI): - swa_full -- ``--swa-full``: force SWA layers to cache full - ``n_ctx`` (collapses path 3 to path 4 for them). - n_parallel -- ``--parallel`` slots: non-SWA layers stay constant - (cells split across slots), SWA layers scale linearly. - kv_unified -- ``--kv-unified`` (default on): no-op for memory math; - kept for API forward-compat. - ctx_checkpoints -- ``--ctx-checkpoints`` (PR #15293): N SWA snapshots - per slot, one sliding-window of state per SWA layer. + swa_full -- --swa-full: SWA layers cache full n_ctx (path 3->4). + n_parallel -- --parallel slots: non-SWA constant, SWA scale linearly. + kv_unified -- --kv-unified: memory no-op (API forward-compat). + ctx_checkpoints -- --ctx-checkpoints: N SWA snapshots per slot. Returns 0 if metadata is insufficient. """ @@ -2195,17 +2468,7 @@ class LlamaCppBackend: n_kv = self._n_kv_heads or self._n_heads or 1 # type: ignore[assignment] # Bytes per element depends on KV cache quantization - 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, - }.get(cache_type_kv or "f16", 2.0) + bpe = _kv_bytes_per_elem(cache_type_kv) slots = max(1, n_parallel) @@ -2233,15 +2496,12 @@ class LlamaCppBackend: head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe) - # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). - # Pattern filled by the resolver at parse time; if absent, falls through - # to the legacy 1/4-global heuristic below. Per-layer-type --parallel N - # accounting (verified against llama-server): - # * non-SWA layers: total cells = n_ctx split across slots -> CONSTANT. - # * SWA layers: per-slot cells = 2*sliding_window (capped at n_ctx - # and per_slot_ctx) -> grows LINEARLY in slots. - # --swa-full forces full n_ctx for SWA layers; --ctx-checkpoints N adds - # N snapshots per SWA layer per slot. + # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). Pattern + # from the resolver; if absent, falls through to the legacy 1/4-global + # heuristic. --parallel N accounting (verified against llama-server): + # non-SWA cells = n_ctx split across slots (CONSTANT); SWA per-slot cells + # = 2*sliding_window (capped at n_ctx/per_slot_ctx) -> LINEAR in slots. + # --swa-full forces full n_ctx for SWA; --ctx-checkpoints N adds snapshots. if ( self._sliding_window is not None and self._sliding_window > 0 @@ -2250,8 +2510,7 @@ class LlamaCppBackend: ): swa = self._sliding_window per_slot_ctx = max(1, n_ctx // slots) - # --swa-full caches full context like non-SWA (per-slot cells = - # per_slot_ctx, collapsing to constant n_ctx total); otherwise SWA + # --swa-full caches full per_slot_ctx (constant n_ctx total); else SWA # caches 2*sliding_window per slot, clamped at per-slot ctx. swa_cells_per_slot = per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx) key_len_swa = self._kv_key_length_swa or key_len @@ -2304,6 +2563,149 @@ class LlamaCppBackend: head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe) + def _draft_backend_for(self, drafter_path: str) -> Optional["LlamaCppBackend"]: + """Lightweight backend with a drafter GGUF's metadata, to size its own KV + via _estimate_kv_cache_bytes. Cached per path; None if unreadable.""" + cache = getattr(self, "_draft_backend_cache", None) + if cache is not None and cache[0] == drafter_path: + return cache[1] + db: Optional[LlamaCppBackend] = None + try: + db = LlamaCppBackend.__new__(LlamaCppBackend) + for attr in ( + "_context_length", + "_n_layers", + "_n_kv_heads", + "_n_heads", + "_embedding_length", + "_kv_key_length", + "_kv_value_length", + "_kv_lora_rank", + "_sliding_window", + "_sliding_window_pattern", + "_ssm_inner_size", + "_full_attention_interval", + "_key_length_mla", + "_n_kv_heads_by_layer", + "_kv_key_length_swa", + "_kv_value_length_swa", + "_shared_kv_layers", + "_nextn_predict_layers", + ): + setattr(db, attr, None) + db._model_identifier = "mtp-draft" + db._read_gguf_metadata(drafter_path) + except Exception as e: # unreadable drafter -> caller falls back + logger.debug(f"Could not read drafter GGUF for MTP budget: {e}") + db = None + self._draft_backend_cache = (drafter_path, db) + return db + + def _mtp_draft_kv_bytes( + self, + n_ctx: int, + *, + drafter_path: Optional[str] = None, + draft_cache_type_k: Optional[str] = None, + draft_cache_type_v: Optional[str] = None, + n_parallel: int = 1, + ) -> Optional[int]: + """Draft KV cache bytes at n_ctx, sized from GGUF dims (K and V types are + independent). Separate drafter (Gemma): its own KV via _estimate_kv_cache_bytes + at the heavier type. Embedded head (Qwen): nextn_predict_layers attention + layers from the main dims. None when dims are missing (flat fallback).""" + if n_ctx <= 0: + return None + bpe_k = _kv_bytes_per_elem(draft_cache_type_k) + bpe_v = _kv_bytes_per_elem(draft_cache_type_v) + if drafter_path: + db = self._draft_backend_for(drafter_path) + if db is None or not db._can_estimate_kv(): + return None + heavier = draft_cache_type_k if bpe_k >= bpe_v else draft_cache_type_v + # The drafter is served under the same --parallel slot count as the + # main model, so price its KV per slot too: a sliding-window drafter + # (Gemma) grows KV with slots and would otherwise be under-reserved. + kv = db._estimate_kv_cache_bytes(n_ctx, heavier, n_parallel = n_parallel) + return kv or None + nextn = self._nextn_predict_layers or 0 + n_kv = self._n_kv_heads or self._n_heads + k_len = self._kv_key_length + v_len = self._kv_value_length + if not (nextn and n_kv and k_len and v_len): + return None + # The embedded MTP head is one draft layer, so a quantized draft KV can't + # amortize its overhead and fits *less* context than f16 (llama.cpp#24102). + # Floor it at f16: a quantized override is priced as f16, f32 keeps its 4 + # bytes. The separate-drafter branch is multi-layer, so it keeps its type. + f16_bpe = _kv_bytes_per_elem("f16") + bpe_k = max(bpe_k, f16_bpe) + bpe_v = max(bpe_v, f16_bpe) + return int(nextn * n_kv * (k_len * bpe_k + v_len * bpe_v) * n_ctx) + + def _estimate_mtp_overhead_bytes( + self, + n_ctx: int, + *, + spec_draft_n_max: int = 0, + draft_cache_type_k: Optional[str] = None, + draft_cache_type_v: Optional[str] = None, + drafter_path: Optional[str] = None, + draft_weights_bytes: int = 0, + n_parallel: int = 1, + ) -> Optional[int]: + """MTP draft reserve at ``n_ctx`` = draft KV (grows with ctx) + separate- + drafter weights. The verify buffer rides in the ctx-fit headroom (no tuned + constant). None when the draft KV can't be sized (caller keeps the flat + fallback). ``draft_weights_bytes`` is the drafter file size (0 for embedded).""" + draft_kv = self._mtp_draft_kv_bytes( + n_ctx, + drafter_path = drafter_path, + draft_cache_type_k = draft_cache_type_k, + draft_cache_type_v = draft_cache_type_v, + n_parallel = n_parallel, + ) + weights = max(0, draft_weights_bytes) + if draft_kv is None: + # KV unsized (exotic/remote drafter): still reserve known weights so a + # large drafter can't launch over budget (the small unsized KV rides in + # the cushion). Nothing known -> None, so the caller keeps the flat + # fallback. + return weights if weights > 0 else None + return draft_kv + weights + + _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Studio does not override it + _COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate + + def _estimate_compute_buffer_bytes( + self, + *, + n_ubatch: Optional[int] = None, + n_parallel: int = 1, + per_device_tensor: bool = False, + ) -> int: + """Per-device compute-graph buffer (bytes) from GGUF dims: a vocab-width + output buffer + activation scratch. Context-independent; scales with + ``--parallel`` (serving slots). Tensor mode materializes it on every device. + A slight upper bound over measured allocations; 0 when dims are missing.""" + n_vocab = self._vocab_size or 0 + n_embd = self._embedding_length or 0 + if n_vocab <= 0 or n_embd <= 0: + return 0 + ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + par = max(1, int(n_parallel)) + out_buffer = n_vocab * ub * 4 # f32 output/logits buffer + act_scratch = 4 * n_embd * ub * 4 # a few resident hidden-width buffers + if per_device_tensor: + # Output + comm/staging materialized on every device, every slot. + compute = 2 * act_scratch + out_buffer * par + else: + # Each extra concurrent slot adds one output buffer (chat decode sizes + # ~one logit row per slot; would under-count embeddings/--logits-all, + # not run here). Matches measured {1:36,2:492,4:1388,8:3220} MiB. + compute = act_scratch + out_buffer * max(0, par - 1) + return int(compute * self._COMPUTE_BUFFER_SAFETY) + def _fit_context_to_vram( self, requested_ctx: int, @@ -2318,13 +2720,15 @@ class LlamaCppBackend: ctx_checkpoints: int = 0, kv_on_gpu: bool = True, mtp_engaged: bool = False, + mtp_overhead_fn: Optional[Callable[[int], int]] = None, budget_frac: Optional[float] = None, + total_mib: Optional[int] = None, ) -> int: """Return the largest context length that fits in GPU VRAM. - Uses 90% of available VRAM as the ctx-fit budget -- tighter than - ``_GPU_PIN_VRAM_FRACTION`` on purpose (over-promising context OOMs at - runtime). If the weights alone don't fit, returns ``requested_ctx``. + Budget caps occupancy at ``_CTX_FIT_VRAM_FRACTION`` of the card: an + absolute ``free - (1 - frac) * total`` when ``total_mib`` is given, else + ``free * frac``. Weights alone over budget returns ``requested_ctx``. ``kv_on_gpu`` mirrors ``--kv-offload`` (default on); when False the KV cache lives in CPU RAM and the requested context is honored verbatim. @@ -2352,17 +2756,25 @@ class LlamaCppBackend: ctx_checkpoints = ctx_checkpoints, ) - # MTP engaged: carve the drafter's reserve out of the fit budget. Callers - # can override outright (tensor-parallel mode passes a fatter margin), so - # only compute a default when none was supplied. + # byte-accurate mtp_overhead_fn supersedes the flat fraction (the fallback + # when dims can't size the draft KV); callers may override budget_frac. if budget_frac is None: - budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if mtp_engaged else 0.0) - budget_bytes = available_mib * 1024 * 1024 * budget_frac + flat_mtp = mtp_engaged and mtp_overhead_fn is None + budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if flat_mtp else 0.0) + # Absolute reserve off total when known, else fraction-of-free; clamp >=0. + if total_mib is not None and total_mib > 0: + budget_mib = max(0.0, available_mib - (1.0 - budget_frac) * total_mib) + else: + budget_mib = available_mib * budget_frac + budget_bytes = budget_mib * 1024 * 1024 model_footprint = model_size_bytes + def _mtp_at(ctx: int) -> int: + return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + # Already fits? kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs) - if model_footprint + kv <= budget_bytes: + if model_footprint + kv + _mtp_at(requested_ctx) <= budget_bytes: return requested_ctx # Weights alone exceed budget -- reducing ctx can't help; --fit handles it. @@ -2375,7 +2787,7 @@ class LlamaCppBackend: ) return requested_ctx - # Binary search for max context that fits + # Binary search for max context that fits (KV + MTP draft reserve at that ctx) remaining = budget_bytes - model_footprint effective_min = min(min_ctx, requested_ctx) lo, hi = effective_min, requested_ctx @@ -2383,7 +2795,7 @@ class LlamaCppBackend: while lo <= hi: mid = (lo + hi) // 2 kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs) - if kv <= remaining: + if kv + _mtp_at(mid) <= remaining: best = mid lo = mid + 1 else: @@ -2562,6 +2974,8 @@ class LlamaCppBackend: self._n_kv_heads_by_layer = None self._n_heads = None self._embedding_length = None + self._feed_forward_length = None + self._vocab_size = None self._kv_key_length = None self._kv_value_length = None self._sliding_window = None @@ -2583,6 +2997,8 @@ class LlamaCppBackend: WANTED = { "general.architecture", "tokenizer.chat_template", + # Vocab size = tokens array length (no vocab_size key in many GGUFs). + "tokenizer.ggml.tokens", # Block-diffusion marker (DiffusionGemma); routes to the diffusion runner. "diffusion.canvas_length", # Source-repo hints for the SWA resolver's HF fallback. @@ -2646,6 +3062,7 @@ class LlamaCppBackend: f"{arch}.attention.head_count_kv": "n_kv_heads", f"{arch}.attention.head_count": "n_heads", f"{arch}.embedding_length": "embedding_length", + f"{arch}.feed_forward_length": "feed_forward_length", f"{arch}.attention.key_length": "kv_key_length", f"{arch}.attention.value_length": "kv_value_length", f"{arch}.attention.sliding_window": "sliding_window", @@ -2679,6 +3096,9 @@ class LlamaCppBackend: elif vtype == 9: # ARRAY atype = struct.unpack(" tuple[int, int, list[int], Optional[list[int]]]: """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed. @@ -3470,22 +3894,41 @@ class LlamaCppBackend: Policy (assumes >= 2 GPUs; the caller drops the toggle below that): - Cap context to the KV that fits the pooled VRAM after the weights and - one per-device compute-graph buffer (``_TENSOR_PARALLEL_BUFFER_RESERVE_MIB``). + one per-device compute-graph buffer (``_estimate_compute_buffer_bytes``, + deterministic from dims; flat fallback when dims are unavailable). llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only cap, honored even for an explicit ``-c``. It is more accurate than the 0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused. - ``tensor_split`` is None (llama.cpp's even default, safe for every arch incl. Gemma 3n which GGML_ASSERTs on a weighted split) when an even - share fits the smallest GPU; otherwise it is weighted by - ``(free - buffer)`` so the roomier GPU absorbs more weight and the - smallest GPU keeps room for KV. + share fits the smallest GPU; otherwise it is weighted by usable budget + so the roomier GPU absorbs more weight and the smallest keeps room for KV. + ``total_by_idx`` enables the total-based occupancy cap; ``n_ubatch`` sizes + the compute buffer. """ - # Drop GPUs that can't hold the per-device compute-graph buffer; they'd - # OOM in tensor mode. load_model already filters before calling, so this - # is defense-in-depth that also keeps the pure function self-contained - # (and unit-testable without a GPU). - reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB - usable_gpus = [g for g in gpus if g[1] >= reserve_mib] + + # Per-GPU usable budget: free - (1-frac)*total, else (unknown total, e.g. a + # two-column probe) the legacy free*frac. Mirrors _select_gpus and + # _gpu_usable so the 5% cushion is kept on every path, not dropped here. + def _usable(idx: int, free_mib: int) -> float: + t = total_by_idx.get(idx, 0) if total_by_idx else 0 + if t > 0: + return max(0.0, free_mib - (1.0 - _CTX_FIT_VRAM_FRACTION) * t) + return max(0.0, free_mib * _CTX_FIT_VRAM_FRACTION) + + # Drop GPUs whose usable budget can't hold the per-device compute-graph + # buffer; they'd OOM in tensor mode. Admitting on raw free would let a + # partly-used big card in with no budget left. Defense-in-depth (load_model + # gates too). Derived per-device reserve; flat fallback. + _reserve_bytes = self._estimate_compute_buffer_bytes( + n_ubatch = n_ubatch, n_parallel = n_parallel, per_device_tensor = True + ) + reserve_mib = ( + _reserve_bytes // (1024 * 1024) + if _reserve_bytes > 0 + else self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + ) + usable_gpus = [g for g in gpus if _usable(g[0], g[1]) >= reserve_mib] gpu_indices = sorted(idx for idx, _ in usable_gpus) if len(gpu_indices) < 2: # Tensor parallelism is meaningless on <2 GPUs (the caller drops the @@ -3497,21 +3940,50 @@ class LlamaCppBackend: None, ) free_by_idx = {idx: free for idx, free in usable_gpus} - pool_mib = sum(free_by_idx.values()) - kv_budget_b = (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - if mtp_engaged: - # MTP keeps a draft model + its own KV cache on GPU. - kv_budget_b -= 2 * 1024**3 + usable_by_idx = {idx: _usable(idx, free_by_idx[idx]) for idx in gpu_indices} + pool_mib = sum(usable_by_idx.values()) + # MTP reserve: byte-accurate per-ctx inside _fit_ctx (mtp_overhead_fn) plus + # a flat cushion that the byte fn can't size -- 2 GiB when dims are wholly + # unavailable (no fn), or mtp_flat_reserve_bytes when the fn is weights-only + # because the draft KV couldn't be sized (_mtp_kv_unsized). Without this the + # binary search spends the unsized-KV cushion on main context and OOMs. + flat_mtp_bytes = max(0, mtp_flat_reserve_bytes) + if mtp_engaged and mtp_overhead_fn is None: + flat_mtp_bytes = max(flat_mtp_bytes, 2 * 1024**3) + kv_budget_b = ( + (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - flat_mtp_bytes + ) + + def _mtp_at(ctx: int) -> int: + return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 def _fit_ctx(ctx: int) -> int: - # Largest context whose KV fits the pooled budget. Floors small, but - # never raises an explicit ctx above what was asked. + # Largest context whose KV (+ MTP draft reserve) fits the pooled + # budget. Floors small, but never raises an explicit ctx above asked. if self._can_estimate_kv() and ctx > 0: ctx_floor = min(2048, ctx) if kv_budget_b <= 0: # Weights + buffers exceed the pool -> floor; the load then # falls back to layer split. return ctx_floor + if mtp_overhead_fn is not None: + # kv(ctx)+mtp(ctx) is not single-linear, so binary search. + def _consumer(c: int) -> int: + return self._estimate_kv_cache_bytes( + c, cache_type_kv, n_parallel = n_parallel + ) + _mtp_at(c) + + if _consumer(ctx) <= kv_budget_b: + return ctx + lo, hi, best = ctx_floor, ctx, ctx_floor + while lo <= hi: + mid = (lo + hi) // 2 + if _consumer(mid) <= kv_budget_b: + best = mid + lo = mid + 1 + else: + hi = mid - 1 + return best kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel) if kv_at <= kv_budget_b: return ctx @@ -3526,16 +3998,19 @@ class LlamaCppBackend: max_available_ctx = _fit_ctx(max_ctx_target) effective_ctx = min(_fit_ctx(target_ctx), max_available_ctx) - min_free_mib = min(free_by_idx.values()) + min_usable_mib = min(usable_by_idx.values()) kv_bytes = ( self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = n_parallel) if (self._can_estimate_kv() and effective_ctx > 0) else 0 ) - even_share_mib = (model_size + kv_bytes) / len(gpu_indices) / (1024 * 1024) + # The MTP reserve also has to fit the even split (mirror the pooled budget): + # byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above. + mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes + even_share_mib = (model_size + kv_bytes + mtp_bytes) / len(gpu_indices) / (1024 * 1024) tensor_split: Optional[list[int]] = None - if even_share_mib > (min_free_mib - reserve_mib): - adj = [max(0, int(free_by_idx[i] - reserve_mib)) for i in gpu_indices] + if even_share_mib > (min_usable_mib - reserve_mib): + adj = [max(0, int(usable_by_idx[i] - reserve_mib)) for i in gpu_indices] if sum(adj) > 0: tensor_split = adj return effective_ctx, max_available_ctx, gpu_indices, tensor_split @@ -3862,27 +4337,58 @@ class LlamaCppBackend: ctx_override = parse_ctx_override(extra_args) requested_ctx = resolve_requested_ctx(extra_args, n_ctx) cache_override = parse_cache_override(extra_args) - cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv) - # A user --split-mode in extras last-wins-overrides the - # toggle, so reconcile it back into tensor_parallel state. + # Budget the heavier of asymmetric --cache-type-k/-v extras (they + # win per axis at launch, appended last); resolve_cache_type_kv only + # returns the last-wins type, which under-reserves the heavier axis. + # The user's extras still set the real (possibly asymmetric) child + # cache, so this only affects the reserve, not the emitted command. + _extras_cache = _extra_args_main_cache_type_for_budget(extra_args) + cache_type_kv = _extras_cache if _extras_cache is not None else cache_type_kv + _cache_type_from_env = False + if cache_type_kv is None: + # Param/extras set nothing, so the child inherits + # LLAMA_ARG_CACHE_TYPE_K/_V. Adopt a heavier env type (f32) for + # the reserve only; the launch does NOT re-emit it (that would + # rewrite an asymmetric K=f32,V=f16 env into symmetric flags), + # so _cache_type_from_env keeps it out of the emitted flags. + cache_type_kv = _env_main_cache_type_for_budget() + _cache_type_from_env = cache_type_kv is not None + # A user --split-mode in extras last-wins-overrides the toggle, and + # an inherited tensor LLAMA_ARG_SPLIT_MODE flips it on (the child + # would run tensor unbudgeted otherwise). The duplicate-load matchers + # use the same helper so a healthy env-driven tensor server matches. split_mode_override = parse_split_mode_override(extra_args) - tensor_parallel = resolve_tensor_parallel(extra_args, tensor_parallel) + tensor_parallel = _effective_tensor_parallel(extra_args, tensor_parallel) # Tensor mode aborts on a quantized KV cache, so drop it for the # tensor attempt (and strip any inherited/explicit --cache-type - # that would re-impose it when appended last). The layer-split - # fallback re-runs with tensor_parallel False and keeps the type. - if ( - tensor_parallel - and cache_type_kv - and cache_type_kv.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES - ): + # that would re-impose it when appended last). Layer split does + # support it, so remember the dropped type and the original extras + # to restore (verbatim, incl. an asymmetric K/V) if we later fall + # back to layer split below. + _tensor_dropped_cache_type_kv: Optional[str] = None + _tensor_dropped_extra_args: Optional[list] = None + # Tensor mode rejects any quantized axis. cache_type_kv is the + # heavier-by-bytes budget type, which can mask a quantized axis (an + # f16 budget hides a paired q4_0), so also test each explicit + # --cache-type-k/-v extra, not just the budget type. + _ck_extra, _cv_extra = parse_cache_override_per_axis(extra_args) + _cache_non_tensor_safe = any( + c and c.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES + for c in (cache_type_kv, _ck_extra, _cv_extra) + ) + if tensor_parallel and _cache_non_tensor_safe: logger.info( "Tensor parallelism requires a non-quantized KV cache; " "ignoring cache type %s for the tensor attempt.", cache_type_kv, ) + _tensor_dropped_cache_type_kv = cache_type_kv cache_type_kv = None if extra_args: + # Keep the originals so a layer downgrade restores the real + # (possibly asymmetric) --cache-type-k/-v the layer path + # supports, not just the scalar heavier type. + _tensor_dropped_extra_args = list(extra_args) extra_args = strip_shadowing_flags( extra_args, strip_context = False, @@ -3891,10 +4397,24 @@ class LlamaCppBackend: strip_template = False, strip_split_mode = False, ) + # The launch keeps an inherited tensor-safe env cache type (the + # env cleanup only pops quantized ones), so re-adopt a heavier + # env type (f32) for the budget here too -- mirrors the initial + # adoption, which was skipped because the param/extras set the + # (now-dropped) quantized type. Else the child allocates f32 KV + # against an f16 budget. + _env_tensor_cache = _env_main_cache_type_for_budget() + if _env_tensor_cache is not None: + cache_type_kv = _env_tensor_cache + _cache_type_from_env = True if ctx_override is not None and ctx_override > 0: logger.info(f"User --ctx-size {ctx_override} honored; skipping auto-reduce") if cache_override is not None: - logger.info(f"User --cache-type-k/-v {cache_override} honored for KV estimate") + _ck, _cv = parse_cache_override_per_axis(extra_args) + logger.info( + f"User --cache-type-k/-v (k={_ck}, v={_cv}) honored; " + "KV estimate budgets the heavier axis" + ) if split_mode_override is not None: logger.info( f"User --split-mode {split_mode_override} honored; " @@ -3926,7 +4446,27 @@ class LlamaCppBackend: self._mmproj_vram_bytes(launch_mmproj_path) if effective_is_vision else 0 ) model_size = gguf_size + mmproj_size - gpus = self._get_gpu_free_memory() + # 2-tuple gpus for existing logic + a total map for the absolute + # per-GPU headroom (correct when the GPU is already partly used). + _gpu_mem = self._get_gpu_memory() + gpus = [(idx, free) for idx, free, _t in _gpu_mem] + total_by_idx = {idx: total for idx, _f, total in _gpu_mem} + + def _gpu_usable(g, frac = _CTX_FIT_VRAM_FRACTION): + # Per-GPU usable budget for ranking: free - (1-frac)*total. + # Callers pass the ACTIVE fraction so the ranking matches the + # budget the fit then tests (else mixed totals mis-order). + idx, free = g + t = total_by_idx.get(idx, 0) + if t > 0: + return free - (1.0 - frac) * t + return free * frac + + def _pool_budget_mib(subset, frac): + # Sum each GPU's own usable budget. Pooling free and total + # separately would let an unknown-total GPU (MIG/vGPU/N/A) + # add full free with no cushion among known-total GPUs. + return sum(max(0.0, _gpu_usable(g, frac)) for g in subset) # Resolve effective context: 0 means let llama-server use # the model's native length. Only expand to a known native @@ -3942,12 +4482,10 @@ class LlamaCppBackend: # GPU/VRAM-fit logic below may shrink it on limited HW. max_available_ctx = self._context_length or effective_ctx - # Will MTP engage on this load? If so, auto-fit reserves - # extra VRAM for the draft model. Mirrors - # _build_speculative_flags' resolver: forced mtp / mtp+ngram - # always engage; auto only on an MTP model >= 3B; ngram / - # ngram-simple / off never engage MTP. A separate drafter - # (Gemma) counts as an MTP model just like a baked-in head. + # Will MTP engage? If so, auto-fit reserves draft-model VRAM. + # Mirrors _build_speculative_flags: forced mtp/mtp+ngram always + # engage; auto only on an MTP model >= 3B; ngram/off never. A + # separate drafter (Gemma) counts as an MTP model. _mtp_canonical = _canonicalize_spec_mode(speculative_type) _mtp_effective = _mtp_canonical or "auto" _mtp_size_for_fit = _extract_model_size_b(model_identifier) @@ -3958,49 +4496,229 @@ class LlamaCppBackend: and _mtp_size_for_fit < _MTP_MIN_SIZE_B and not bool(mtp_draft_path) ) + # LLAMA_ARG_SPEC_TYPE only reaches the child when neither extras + # nor Studio emit a spec flag (mode "off", no user --spec-type), + # since _build_speculative_flags emits one for every other mode. + # Consult the env for the reserve only then, else a stale MTP env + # would over-reserve. + _spec_env: Mapping[str, str] = ( + os.environ + if (not _extra_args_set_spec_type(extra_args) and _mtp_canonical == "off") + else {} + ) + # Extras can run MTP even when Studio suppresses its own emission. + _user_mtp_via_extras = _extra_args_requests_mtp(extra_args, env = _spec_env) + # A non-MTP model-based draft mode (draft-simple/draft-eagle3) in + # extras also loads a separate draft model that needs reserving; + # engage only when extras actually name a drafter for it. + _user_draft_via_extras = _extra_args_requests_separate_draft( + extra_args, env = _spec_env + ) and bool(_extra_args_mtp_draft_path(extra_args)) + # Mirror _build_speculative_flags: reserve only for MTP the launch + # resolver will actually emit (needs a head/drafter and a binary + # that supports --spec-type mtp). + _mtp_model_for_fit = bool( + self._nextn_predict_layers + or _is_mtp_model_name(model_identifier, model_path) + or bool(mtp_draft_path) + ) + _mtp_binary_ok = True + if not _user_mtp_via_extras: + try: + _mtp_binary_ok = bool( + (self.probe_server_capabilities(binary) or {}).get("mtp_token") + ) + except Exception: + _mtp_binary_ok = False _mtp_will_engage = bool( - not _extra_args_set_spec_type(extra_args) - and ( - _mtp_effective in ("mtp", "mtp+ngram") - or ( - _mtp_effective == "auto" - and ( - bool(self._nextn_predict_layers) - or _is_mtp_model_name(model_identifier, model_path) - or bool(mtp_draft_path) - ) - and not _mtp_sub_3b_for_fit + _user_mtp_via_extras + or _user_draft_via_extras + or ( + not _extra_args_set_spec_type(extra_args) + and _mtp_binary_ok + and _mtp_model_for_fit + and ( + _mtp_effective in ("mtp", "mtp+ngram") + or (_mtp_effective == "auto" and not _mtp_sub_3b_for_fit) ) ) ) - # Auto-cap context to fit GPU VRAM and select GPUs. Two - # policies by whether the user set n_ctx: - # Explicit n_ctx: honor it. Try the full context with - # _select_gpus (as many GPUs as needed); cap only if it - # doesn't fit on any combination. - # Auto n_ctx=0 (native): prefer fewer GPUs with reduced - # context, since multi-GPU is slower. + # Effective draft depth: extras win (last-wins at launch), else + # the field, else the platform default (2 GPU / 3 CPU). + _extra_n_max = _extra_args_spec_draft_n_max(extra_args) + _mtp_eff_n_max = _extra_n_max if _extra_n_max is not None else spec_draft_n_max + if _mtp_eff_n_max is None: + _mtp_eff_n_max = 2 if gpus else 3 + # Separate-drafter weights live on GPU (an embedded head is + # already in model_size). Size the drafter the launch loads, by + # precedence: extras --model-draft (last-wins), else Studio's + # emitted mtp_draft_path, else the env drafter. Sizing the wrong + # one would under-reserve and OOM. + _cli_draft_for_budget = _extra_args_mtp_draft_path(extra_args, env = {}) + _studio_draft_for_budget = ( + mtp_draft_path + if ( + _mtp_will_engage + and mtp_draft_path + and not _extra_args_set_spec_type(extra_args) + ) + else None + ) + _env_draft_for_budget = _extra_args_mtp_draft_path([], env = os.environ) + _mtp_draft_for_budget = ( + _cli_draft_for_budget or _studio_draft_for_budget or _env_draft_for_budget + ) + # Drafter offloaded to CPU keeps its weights+KV off the GPU, so + # drop it from the budget (an embedded head stays in the model). + # Consult the env too: the child honors LLAMA_ARG_N_GPU_LAYERS_DRAFT. + _draft_on_cpu = _extra_args_draft_offloaded_to_cpu(extra_args, env = os.environ) + if _draft_on_cpu: + _mtp_draft_for_budget = None + _mtp_draft_weights = 0 + if _mtp_draft_for_budget: + try: + _mtp_draft_weights = self._get_gguf_size_bytes(_mtp_draft_for_budget) + except Exception: + _mtp_draft_weights = 0 + # Draft K/V types (f16 by default; independent extras overrides). + _mtp_draft_ck, _mtp_draft_cv = _extra_args_draft_cache_types(extra_args) + + # Byte-accurate reserve when dims allow, else None -> flat fallback. + mtp_overhead_fn: Optional[Callable[[int], int]] = None + # True when the byte reserve is the drafter weights ONLY because + # its KV couldn't be sized; the flat fraction must then stay on + # as the cushion for that unsized draft KV (it is not covered by + # the weights-only mtp_overhead_fn). + _mtp_kv_unsized = False + if _mtp_will_engage: + _probe_ctx = self._context_length or ( + effective_ctx if effective_ctx > 0 else 4096 + ) + _draft_kv_probe = self._mtp_draft_kv_bytes( + _probe_ctx, + drafter_path = _mtp_draft_for_budget, + draft_cache_type_k = _mtp_draft_ck, + draft_cache_type_v = _mtp_draft_cv, + n_parallel = n_parallel, + ) + if ( + self._estimate_mtp_overhead_bytes( + _probe_ctx, + spec_draft_n_max = _mtp_eff_n_max, + draft_cache_type_k = _mtp_draft_ck, + draft_cache_type_v = _mtp_draft_cv, + drafter_path = _mtp_draft_for_budget, + draft_weights_bytes = _mtp_draft_weights, + n_parallel = n_parallel, + ) + is not None + ): + # Reserve is weights-only when the draft KV is unsizable. + _mtp_kv_unsized = _draft_kv_probe is None + + # Closure binding this load's draft params; ctx varies. + def mtp_overhead_fn( + ctx: int, + _n: int = _mtp_eff_n_max, + _ck: Optional[str] = _mtp_draft_ck, + _cv: Optional[str] = _mtp_draft_cv, + _dp: Optional[str] = _mtp_draft_for_budget, + _w: int = _mtp_draft_weights, + _np: int = n_parallel, + ) -> int: + v = self._estimate_mtp_overhead_bytes( + ctx, + spec_draft_n_max = _n, + draft_cache_type_k = _ck, + draft_cache_type_v = _cv, + drafter_path = _dp, + draft_weights_bytes = _w, + n_parallel = _np, + ) + return v if v is not None else 0 + + def _mtp_bytes(ctx: int) -> int: + return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + + # Effective micro-batch (a user --ubatch override scales the + # compute buffer); None -> the 512 default in the estimate. + _effective_ubatch = _extra_args_n_ubatch(extra_args) + + # Layer-split compute buffer (one lump; tensor mode reserves it + # per device in _plan_tensor_parallel). Context-independent, so + # fold it into the model footprint for the branches below. Falls + # back to the flat reserve when dims are missing (returns 0), a + # safe upper bound since the tensor buffer >= the layer one. + _compute_buffer_pipeline = self._estimate_compute_buffer_bytes( + n_ubatch = _effective_ubatch, + n_parallel = n_parallel, + per_device_tensor = False, + ) + if _compute_buffer_pipeline <= 0: + _compute_buffer_pipeline = ( + self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024 + ) + model_size_fit = model_size + _compute_buffer_pipeline + + # Layer split adds a fixed per-device overhead on every GPU. The + # folded buffer covers one device; reserve the extra devices' + # share so a k-GPU split can't pin a context that OOMs a device + # (k=1 adds nothing). + _pipeline_overhead_bytes = self._PIPELINE_PER_DEVICE_OVERHEAD_MIB * 1024 * 1024 + + def _subset_model_size(n_gpus: int) -> int: + return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes + + # Auto-cap context to fit VRAM and select GPUs. Explicit n_ctx: + # honor it, cap only if it fits no combination. Auto (native): + # prefer fewer GPUs with reduced context (multi-GPU is slower). gpu_indices, use_fit = None, True # Per-GPU weight proportions for tensor mode (None = even). tp_tensor_split: Optional[list[int]] = None explicit_ctx = requested_ctx > 0 - # MTP draft model lives outside the main estimates; carve - # its reserve out of every fit budget and pin threshold so - # a load can't pin into the drafter's headroom. - _mtp_reserve = _MTP_VRAM_RESERVE_FRAC if _mtp_will_engage else 0.0 - _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _mtp_reserve + # Flat MTP reserve fraction: used only as the fallback when the + # byte-accurate mtp_overhead_fn can't size the draft KV (dims + # unavailable, or _mtp_kv_unsized = weights-only). A separate + # drafter on CPU uses no GPU (no reserve); an embedded head is on + # GPU regardless of draft-offload flags (keep its reserve). + _flat_mtp_engages = _mtp_will_engage and ( + mtp_overhead_fn is None or _mtp_kv_unsized + ) + _draft_cpu_no_embedded = _draft_on_cpu and not self._nextn_predict_layers + # MTP reserves GPU VRAM unless its only drafter is a separate + # CPU-offloaded one (an embedded head stays on GPU). The tensor + # path reserves like the layer path; gate both on this. + _mtp_reserves_gpu = _mtp_will_engage and not _draft_cpu_no_embedded + _flat_mtp_reserve = ( + _MTP_VRAM_RESERVE_FRAC + if (_flat_mtp_engages and not _draft_cpu_no_embedded) + else 0.0 + ) + _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve - # Tensor mode allocates a compute-graph buffer on every - # participating GPU, so a GPU with less free VRAM than that - # reserve can't host it and would OOM at load. Drop those - # from the tensor-parallel set up front (gpu_indices below - # becomes the CUDA_VISIBLE_DEVICES mask, so they're excluded - # from llama-server entirely, not just given zero weight). + # Tensor mode replicates a compute buffer on every GPU, so drop + # GPUs below that reserve from the set up front (gpu_indices + # becomes the CUDA_VISIBLE_DEVICES mask, fully excluding them). tp_gpus = gpus if tensor_parallel: - reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB - tp_gpus = [g for g in gpus if g[1] >= reserve_mib] + # Deterministic per-device compute buffer (replicated on + # every device in tensor mode); flat fallback when dims + # are unavailable. _plan_tensor_parallel uses the same. + _tp_reserve_bytes = self._estimate_compute_buffer_bytes( + n_ubatch = _effective_ubatch, + n_parallel = n_parallel, + per_device_tensor = True, + ) + reserve_mib = ( + _tp_reserve_bytes // (1024 * 1024) + if _tp_reserve_bytes > 0 + else self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + ) + # Admit by usable budget (free - (1-frac)*total), not raw + # free: a partly-used big card can clear the reserve on raw + # free yet have no budget left. + tp_gpus = [g for g in gpus if _gpu_usable(g) >= reserve_mib] if tensor_parallel and len(tp_gpus) < 2: # Tensor parallelism needs >= 2 usable GPUs. On a single @@ -4016,21 +4734,81 @@ class LlamaCppBackend: len(gpus), ) tensor_parallel = False - # A user --split-mode tensor in extras is appended after - # Studio's flags, so it would still reach llama-server and - # fail here; strip it so the downgrade actually applies. - extra_args = strip_split_mode_only(extra_args) + # Layer split supports a quantized KV the tensor attempt + # dropped; restore it and re-emit it (clear the env flag the + # tensor re-adoption may have set, so the restored type wins + # over a stale inherited env on the layer launch). + if _tensor_dropped_cache_type_kv is not None: + cache_type_kv = _tensor_dropped_cache_type_kv + _cache_type_from_env = False + # Restore the original extras (with the real, possibly + # asymmetric, --cache-type-k/-v the tensor attempt stripped), + # then drop the user --split-mode tensor so the downgrade + # actually applies (extras are appended last). + extra_args = strip_split_mode_only( + _tensor_dropped_extra_args + if _tensor_dropped_extra_args is not None + else extra_args + ) if tensor_parallel and tp_gpus: - # Tensor-parallel allocation: use all usable GPUs, weight - # the split by (free - buffer), and cap context to the - # pooled VRAM after weights + per-device compute-graph - # buffers. See _plan_tensor_parallel for the policy. + # Pooled usable budget (after each device's compute buffer) + # must hold the non-shrinkable footprint: weights + the MTP + # reserve. The planner can shrink ctx/KV, not these. + _tp_weight_budget_mib = ( + sum(_gpu_usable(g) for g in tp_gpus) - len(tp_gpus) * reserve_mib + ) + _tp_flat_mtp = 2 * 1024**3 # flat reserve when dims unavailable + if not _mtp_reserves_gpu: + # No MTP, or its only drafter is CPU-offloaded (no GPU). + _tp_mtp_floor = 0 + elif mtp_overhead_fn is not None and not _mtp_kv_unsized: + _tp_mtp_floor = _mtp_bytes( + min(2048, effective_ctx) if effective_ctx > 0 else 2048 + ) + else: + # Dims unavailable / weights-only: tensor mode has no + # --fit valve, so keep the flat reserve as the unsized-KV + # cushion, never below the known byte reserve. + _tp_mtp_floor = max( + _tp_flat_mtp, + _mtp_bytes(min(2048, effective_ctx) if effective_ctx > 0 else 2048), + ) + _tp_required_mib = (model_size + _tp_mtp_floor) / (1024 * 1024) + if _tp_weight_budget_mib <= _tp_required_mib: + logger.info( + "Tensor parallelism requested but the pooled VRAM " + "budget cannot hold the weights, MTP reserve, and " + "per-device compute buffers; falling back to layer split." + ) + tensor_parallel = False + # Restore the dropped quantized KV (layer split supports + # it); clear the env flag so the restored type is emitted. + if _tensor_dropped_cache_type_kv is not None: + cache_type_kv = _tensor_dropped_cache_type_kv + _cache_type_from_env = False + # Restore the original (possibly asymmetric) cache extras + # too, dropping only the user --split-mode tensor. + extra_args = strip_split_mode_only( + _tensor_dropped_extra_args + if _tensor_dropped_extra_args is not None + else extra_args + ) + + if tensor_parallel and tp_gpus: + # Tensor-parallel allocation; see _plan_tensor_parallel. target_ctx = ( effective_ctx if explicit_ctx else (self._context_length or effective_ctx) ) + # When the draft KV couldn't be sized (weights-only reserve), + # the planner's mtp_overhead_fn is non-None but covers only + # weights, so pass the flat cushion for the unsized KV (else + # the binary search spends it on context). + _tp_unsized_mtp_reserve = ( + 2 * 1024**3 if (_mtp_reserves_gpu and _mtp_kv_unsized) else 0 + ) ( effective_ctx, max_available_ctx, @@ -4042,10 +4820,14 @@ class LlamaCppBackend: target_ctx, cache_type_kv = cache_type_kv, n_parallel = n_parallel, - mtp_engaged = _mtp_will_engage, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + mtp_flat_reserve_bytes = _tp_unsized_mtp_reserve, # Report the UI ceiling from native ctx, not the # explicit small request. max_target_ctx = self._context_length or target_ctx, + total_by_idx = total_by_idx, + n_ubatch = _effective_ubatch, ) use_fit = False elif gpus and self._can_estimate_kv() and effective_ctx > 0: @@ -4054,24 +4836,38 @@ class LlamaCppBackend: # bounds), independent of the currently requested context. native_ctx_for_cap = self._context_length or effective_ctx if native_ctx_for_cap > 0: - ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) + ranked_for_cap = sorted( + gpus, + key = lambda g: _gpu_usable( + g, _CTX_FIT_VRAM_FRACTION - _flat_mtp_reserve + ), + reverse = True, + ) best_cap = 0 + _cap_fraction = _CTX_FIT_VRAM_FRACTION - _flat_mtp_reserve for n_gpus in range(1, len(ranked_for_cap) + 1): subset = ranked_for_cap[:n_gpus] - pool_mib = sum(free for _, free in subset) + # Per-GPU-consistent pool budget (fixes mixed + # known/unknown totals); pass it as an absolute + # budget so the fit and the check below agree. + pool_budget = _pool_budget_mib(subset, _cap_fraction) + _ms = _subset_model_size(n_gpus) capped = self._fit_context_to_vram( native_ctx_for_cap, - pool_mib, - model_size, + pool_budget, + _ms, cache_type_kv, n_parallel = n_parallel, - mtp_engaged = _mtp_will_engage, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + budget_frac = 1.0, + total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * (_CTX_FIT_VRAM_FRACTION - _mtp_reserve): + footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + if footprint_mib <= pool_budget: best_cap = max(best_cap, capped) if best_cap > 0: max_available_ctx = best_cap @@ -4085,34 +4881,49 @@ class LlamaCppBackend: # Honor the requested context verbatim. If it fits, # pin GPUs and skip --fit; else ship -c --fit # on and let llama-server flex -ngl (CPU offload). - requested_total = model_size + self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv, n_parallel = n_parallel + requested_total = ( + model_size_fit + + self._estimate_kv_cache_bytes( + effective_ctx, cache_type_kv, n_parallel = n_parallel + ) + + _mtp_bytes(effective_ctx) ) gpu_indices, use_fit = self._select_gpus( - requested_total, gpus, usable_fraction = _pin_fraction + requested_total, + gpus, + usable_fraction = _pin_fraction, + total_by_idx = total_by_idx, + per_device_overhead_bytes = _pipeline_overhead_bytes, ) # No silent shrink: effective_ctx stays == requested_ctx. else: # Auto context: prefer fewer GPUs, cap to fit. Same - # headroom threshold as _select_gpus (#5106). - ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + # headroom threshold as _select_gpus (#5106). Rank by the + # active pin fraction so the order matches the fit budget. pin_fraction = _pin_fraction + ranked = sorted( + gpus, key = lambda g: _gpu_usable(g, pin_fraction), reverse = True + ) for n_gpus in range(1, len(ranked) + 1): subset = ranked[:n_gpus] - pool_mib = sum(free for _, free in subset) + pool_budget = _pool_budget_mib(subset, pin_fraction) + _ms = _subset_model_size(n_gpus) capped = self._fit_context_to_vram( effective_ctx, - pool_mib, - model_size, + pool_budget, + _ms, cache_type_kv, n_parallel = n_parallel, - mtp_engaged = _mtp_will_engage, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + budget_frac = 1.0, + total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * pin_fraction: + footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + if footprint_mib <= pool_budget: effective_ctx = capped gpu_indices = sorted(idx for idx, _ in subset) use_fit = False @@ -4125,14 +4936,17 @@ class LlamaCppBackend: if effective_ctx > 0: for n_gpus in range(1, len(ranked) + 1): subset = ranked[:n_gpus] - pool_mib = sum(free for _, free in subset) kv = self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv, n_parallel = n_parallel, ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * pin_fraction: + footprint_mib = ( + _subset_model_size(n_gpus) + + kv + + _mtp_bytes(effective_ctx) + ) / (1024 * 1024) + if footprint_mib <= _pool_budget_mib(subset, pin_fraction): gpu_indices = sorted(idx for idx, _ in subset) use_fit = False break @@ -4144,14 +4958,36 @@ class LlamaCppBackend: "Falling back to file-size-only GPU selection", model_size_gb = round(model_size / (1024**3), 2), ) + # Add the byte-accurate MTP reserve here too when it is + # available; otherwise _pin_fraction carries the flat + # fallback (the two are mutually exclusive by design). + _fs_total = model_size_fit + _mtp_bytes( + self._context_length or effective_ctx or 4096 + ) gpu_indices, use_fit = self._select_gpus( - model_size, gpus, usable_fraction = _pin_fraction + _fs_total, + gpus, + usable_fraction = _pin_fraction, + total_by_idx = total_by_idx, + per_device_overhead_bytes = _pipeline_overhead_bytes, ) if use_fit and not explicit_ctx: # Weights don't fit on any subset; default UI to 4096 # so the slider isn't on an unusable native ctx. effective_ctx = min(4096, effective_ctx) if effective_ctx > 0 else 4096 + # MTP reserve at the final context, for the logs below. + _mtp_reserve_bytes = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0 + if _mtp_will_engage: + _mtp_note = ( + f"MTP reserve: {_mtp_reserve_bytes / (1024**3):.2f} GB " + f"(draft KV @ {effective_ctx} + verify n_max={_mtp_eff_n_max}" + + (", flat-frac fallback" if mtp_overhead_fn is None else "") + + "), " + ) + else: + _mtp_note = "" + if effective_ctx < original_ctx: kv_est = self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv, n_parallel = n_parallel @@ -4159,7 +4995,9 @@ class LlamaCppBackend: logger.info( f"Context auto-reduced: {original_ctx} -> {effective_ctx} " f"(model: {model_size / (1024**3):.1f} GB, " - f"est. KV cache: {kv_est / (1024**3):.1f} GB)" + f"est. KV cache: {kv_est / (1024**3):.1f} GB, " + f"{_mtp_note}".rstrip(", ") + + ")" ) kv_cache_bytes = self._estimate_kv_cache_bytes( @@ -4172,6 +5010,7 @@ class LlamaCppBackend: f"GGUF size: {gguf_size / (1024**3):.1f} GB, " f"{mmproj_note}" f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, " + f"{_mtp_note}" f"context: {effective_ctx}, " f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}" ) @@ -4266,7 +5105,11 @@ class LlamaCppBackend: "iq4_nl", "f32", } - if cache_type_kv and cache_type_kv in _valid_cache_types: + if ( + cache_type_kv + and cache_type_kv in _valid_cache_types + and not _cache_type_from_env + ): cmd.extend( [ "--cache-type-k", @@ -4278,6 +5121,8 @@ class LlamaCppBackend: self._cache_type_kv = cache_type_kv logger.info(f"KV cache type: {cache_type_kv}") else: + # An env-only type is left inherited (untouched) so an + # asymmetric K/V env reaches the child as set. self._cache_type_kv = None # Tensor parallelism: split the model across GPUs by tensor @@ -4435,6 +5280,31 @@ class LlamaCppBackend: if "--threads" not in cmd: env.pop("LLAMA_ARG_THREADS", None) + # Reconcile the inherited LLAMA_ARG_* env with Studio's final + # decision: stripping CLI extras on a tensor->layer downgrade + # can't remove env vars, so the child could run a mode/KV Studio + # didn't budget. + if not tensor_parallel: + # Layer split: clear a non-layer inherited split mode (and any + # paired tensor-split) so the child can't override the layer plan. + _inherited_sm = (env.get("LLAMA_ARG_SPLIT_MODE") or "").strip().lower() + if _inherited_sm and _inherited_sm != "layer": + env.pop("LLAMA_ARG_SPLIT_MODE", None) + env.pop("LLAMA_ARG_TENSOR_SPLIT", None) + else: + # Studio owns the tensor split: it emits --tensor-split when it + # picks an uneven one (CLI wins) and nothing when an even split + # is safe. Clear any inherited LLAMA_ARG_TENSOR_SPLIT so the even + # case can't be overridden by a stale env (the layer branch above + # clears it too). + env.pop("LLAMA_ARG_TENSOR_SPLIT", None) + # Tensor split aborts on a quantized KV; clear an inherited + # quantized cache type so the child uses the tensor-safe default. + for _ct_var in ("LLAMA_ARG_CACHE_TYPE_K", "LLAMA_ARG_CACHE_TYPE_V"): + _ct_raw = (env.get(_ct_var) or "").strip().lower() + if _ct_raw and _ct_raw not in self._TENSOR_PARALLEL_KV_TYPES: + env.pop(_ct_var, None) + # Windows + full offload: PASSIVE OMP + 2 threads stop # spin-wait burning CPU. CPU/partial offload keeps default # OMP parallelism. #5692. @@ -4885,6 +5755,11 @@ class LlamaCppBackend: "run `unsloth studio update`. Loading without " "speculative decoding." ) + # Override an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI wins + # over env) so the child matches the binary-capability gate and + # the no-MTP budget, like the sibling no-head/non-MTP fallbacks. + flags.append("--spec-default") + self._speculative_type = "default" self._spec_fallback_reason = "binary_no_mtp" return False draft_n_max = _resolved_draft_n_max() @@ -5071,10 +5946,12 @@ class LlamaCppBackend: if _norm(self._cache_type_kv) != _norm(cache_type_kv): return False - # Reconcile a user --split-mode in extras (load_model does the same), so - # an extras-driven tensor load isn't seen as a mismatch that needlessly - # kills/reloads a healthy server. - if self._tensor_parallel != resolve_tensor_parallel(extra_args, tensor_parallel): + # Reconcile a user --split-mode in extras AND an inherited tensor + # LLAMA_ARG_SPLIT_MODE env, but only against a server that actually + # launched tensor: if load_model downgraded to layer split it scrubbed + # the child env, so the env must not force an endless reload of a healthy + # server. An identical request would downgrade the same way. + if not _tensor_parallel_matches_loaded(extra_args, tensor_parallel, self._tensor_parallel): return False # Compare on the canonical requested mode. With --spec-type in diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 69a86fa3ba..b42be5ee0d 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -13,7 +13,8 @@ Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md from __future__ import annotations -from typing import Iterable, Optional +import os +from typing import Iterable, Mapping, Optional # Each group = every alias (short + long) of one hard-denied flag. # Extend the matching group when llama.cpp adds a new alias. @@ -124,7 +125,9 @@ def is_managed_flag(flag: str) -> bool: # from inherited extras so they can't last-wins-override an Apply that # re-sets the same field. _CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"}) -_CACHE_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"}) +_CACHE_TYPE_K_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k"}) +_CACHE_TYPE_V_FLAGS: frozenset[str] = frozenset({"-ctv", "--cache-type-v"}) +_CACHE_FLAGS: frozenset[str] = _CACHE_TYPE_K_FLAGS | _CACHE_TYPE_V_FLAGS _SPEC_FLAGS: frozenset[str] = frozenset( { "--spec-default", @@ -133,13 +136,22 @@ _SPEC_FLAGS: frozenset[str] = frozenset( "--spec-ngram-size", "--draft-min", "--draft-max", - # MTP path (llama.cpp #22673). --model-draft and aliases are - # Studio-managed since the separate-drafter support (Gemma 4): an - # inherited copy must not last-wins-override the auto-detected - # drafter. Explicit extras for the current load are never stripped. + # MTP path (llama.cpp #22673). The drafter selectors (local --model-draft + # and HF --spec-draft-hf aliases) are Studio-managed since the separate- + # drafter support (Gemma 4): an inherited copy must not last-wins-override + # the auto-detected drafter. Explicit extras for the current load are never + # stripped. The per-drafter tuning knobs (--spec-draft-type-*, -ngld, + # --spec-draft-device) are deliberately NOT stripped: 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. "--model-draft", "-md", "--spec-draft-model", + "--spec-draft-hf", + "-hfd", + "-hfrd", + "--hf-repo-draft", "--spec-draft-n-max", "--spec-draft-n-min", "--spec-draft-p-min", @@ -274,6 +286,20 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: return _last_flag_value(args, _CACHE_FLAGS) +def parse_cache_override_per_axis( + args: Optional[Iterable[str]], +) -> tuple[Optional[str], Optional[str]]: + """Last-wins --cache-type-k / --cache-type-v values kept apart, as (k, v). + + parse_cache_override collapses both axes to one last-wins value; this keeps + them separate so an asymmetric K/V can be budgeted by its heavier axis. + """ + return ( + _last_flag_value(args, _CACHE_TYPE_K_FLAGS), + _last_flag_value(args, _CACHE_TYPE_V_FLAGS), + ) + + def resolve_cache_type_kv( args: Optional[Iterable[str]], fallback_cache_type_kv: Optional[str] ) -> Optional[str]: @@ -309,6 +335,60 @@ def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_paral return override.strip().lower() == "tensor" +def _env_split_mode_is_tensor(env: Optional[Mapping[str, str]] = None) -> bool: + """True when the inherited LLAMA_ARG_SPLIT_MODE env selects tensor. Studio + emits --split-mode only on its tensor branch, so a tensor env on the layer + path would run the child tensor-parallel unbudgeted; this flips the budget + to tensor. Only tensor is heavier, so other modes are ignored.""" + raw = (os.environ if env is None else env).get("LLAMA_ARG_SPLIT_MODE") + return bool(raw) and raw.strip().lower() == "tensor" + + +def _effective_tensor_parallel( + extra_args: Optional[Iterable[str]], + tensor_parallel: bool, + env: Optional[Mapping[str, str]] = None, +) -> bool: + """Tensor-parallel decision including the inherited LLAMA_ARG_SPLIT_MODE env. + + resolve_tensor_parallel (extras + toggle), flipped on when extras set no split + mode but the child inherits a tensor split env. Shared by load_model (which + budgets and launches it) and the tensor-fallback wrapper (so an env-only + tensor crash still retries layer split).""" + resolved = resolve_tensor_parallel(extra_args, tensor_parallel) + if ( + not resolved + and parse_split_mode_override(extra_args) is None + and _env_split_mode_is_tensor(env) + ): + return True + return resolved + + +def _tensor_parallel_matches_loaded( + extra_args: Optional[Iterable[str]], + requested_tensor_parallel: bool, + loaded_tensor_parallel: bool, + env: Optional[Mapping[str, str]] = None, +) -> bool: + """Whether a duplicate load request matches a loaded server's tensor state. + + Env-only tensor mode is a launch hint load_model may downgrade to layer split + (capacity/buffer), scrubbing the child env. So only let an inherited tensor env + raise a match against a server that *actually* launched tensor; on a downgraded + (layer) server the env is ignored, and an identical request would downgrade the + same way -- avoiding an endless reload of a healthy server.""" + requested = resolve_tensor_parallel(extra_args, requested_tensor_parallel) + if ( + loaded_tensor_parallel + and not requested + and parse_split_mode_override(extra_args) is None + and _env_split_mode_is_tensor(env) + ): + requested = True + return requested == loaded_tensor_parallel + + _MMPROJ_DISABLE_FLAGS: frozenset[str] = frozenset({"--no-mmproj", "--no-mmproj-auto"}) _MMPROJ_ENABLE_FLAGS: frozenset[str] = frozenset({"--mmproj-auto"}) diff --git a/studio/backend/core/inference/tensor_fallback.py b/studio/backend/core/inference/tensor_fallback.py index 73687165b8..3ceb1a268a 100644 --- a/studio/backend/core/inference/tensor_fallback.py +++ b/studio/backend/core/inference/tensor_fallback.py @@ -13,7 +13,7 @@ import logging from typing import Awaitable, Callable, Optional from core.inference.llama_server_args import ( - resolve_tensor_parallel, + _effective_tensor_parallel, strip_split_mode_only, ) @@ -34,18 +34,20 @@ async def load_with_tensor_fallback( True on success; it *raises* on a hard crash (llama-server aborts on some archs / older builds), which is treated the same as a False return. - Tensor mode can be requested by the toggle or by a ``--split-mode tensor`` - in ``extra_args`` (an allowed shadow flag), so the retry is keyed on whether - tensor mode is actually engaged, and it strips ``--split-mode`` from the - extras so the layer retry can't relaunch the same failing tensor load. A - non-tensor load keeps its original contract and propagates exceptions. + Tensor mode can be requested by the toggle, by a ``--split-mode tensor`` in + ``extra_args`` (an allowed shadow flag), or by an inherited + ``LLAMA_ARG_SPLIT_MODE=tensor`` env (load_model engages it the same way), so + the retry is keyed on whether tensor mode is actually engaged, and it forces + ``--split-mode layer`` on the retry so neither leftover extras nor the + inherited tensor env can relaunch the same failing tensor load. A non-tensor + load keeps its original contract and propagates exceptions. ``cancelled()`` distinguishes a real tensor-start failure from a user cancellation: ``attempt_load`` also returns False when the load was cancelled, so without this the helper would restart a load the user just cancelled. """ - tensor_requested = resolve_tensor_parallel(extra_args, requested_tensor) + tensor_requested = _effective_tensor_parallel(extra_args, requested_tensor) try: success = await attempt_load(requested_tensor, extra_args) except Exception as exc: @@ -67,4 +69,8 @@ async def load_with_tensor_fallback( "(this model may not support tensor parallelism)", label, ) - return await attempt_load(False, strip_split_mode_only(extra_args)) + # Force --split-mode layer (CLI wins over env) so neither leftover extras nor + # an inherited LLAMA_ARG_SPLIT_MODE=tensor can re-engage tensor and re-crash + # the retry; load_model and the child both honor the explicit layer override. + layer_extras = strip_split_mode_only(extra_args) or [] + return await attempt_load(False, [*layer_extras, "--split-mode", "layer"]) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index da266bf768..2193d562e3 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -612,6 +612,7 @@ try: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + _tensor_parallel_matches_loaded, resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, @@ -646,6 +647,7 @@ except ImportError: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + _tensor_parallel_matches_loaded, resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, @@ -1813,9 +1815,8 @@ def _request_matches_loaded_settings( strip_split_mode = _should_strip_split_mode(request, backend_extra), ) ) - if ( - resolve_tensor_parallel(effective_extra, request.tensor_parallel) - != llama_backend.tensor_parallel + if not _tensor_parallel_matches_loaded( + effective_extra, request.tensor_parallel, llama_backend.tensor_parallel ): return False # Spec decoding works on vision models too (MTP is mmproj-compatible, diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py new file mode 100644 index 0000000000..42c400383e --- /dev/null +++ b/studio/backend/tests/test_compute_buffer.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for ``_estimate_compute_buffer_bytes``: it scales with ``--parallel``, +tensor exceeds pipeline, and it is a safe upper bound on the allocations measured +on real hardware (Qwen3.6-27B-MTP: parallel 1/2/4/8 -> 36/492/1388/3220 MiB single +GPU, ~600 MiB/device tensor). No GPU, subprocess, or GGUF I/O.""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +_structlog_stub = _types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +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 in tests collected after +# this one (the stub leaks via sys.modules for the whole session). +try: + import httpx as _httpx_real # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "RequestError", + ): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) + _httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) + _httpx_stub.Response = type("Response", (), {}) + _httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__exit__": lambda s, *a: None, + }, + ) + sys.modules["httpx"] = _httpx_stub + +from core.inference.llama_cpp import LlamaCppBackend + +MIB = 1024 * 1024 + + +def _backend(vocab = 248320, embd = 5120): + """Backend with just the dims the compute-buffer estimate reads.""" + b = LlamaCppBackend.__new__(LlamaCppBackend) + b._vocab_size = vocab + b._embedding_length = embd + return b + + +# Measured ground truth (MiB) the estimate must upper-bound. +_PIPELINE_MEASURED = {1: 36, 2: 492, 4: 1388, 8: 3220} +_TENSOR_MEASURED_PER_DEVICE = 600 + + +class TestSafeUpperBound: + """The estimate must be >= every measured allocation (never under-reserve).""" + + @pytest.mark.parametrize("parallel,measured", sorted(_PIPELINE_MEASURED.items())) + def test_pipeline_upper_bounds_measured(self, parallel, measured): + est = _backend()._estimate_compute_buffer_bytes(n_parallel = parallel) / MIB + assert est >= measured, f"under-reserved at parallel={parallel}: {est:.0f} < {measured}" + + @pytest.mark.parametrize("parallel,measured", sorted(_PIPELINE_MEASURED.items())) + def test_pipeline_not_wildly_over(self, parallel, measured): + # Stay within ~2x of measured so we don't waste context (the point of + # replacing the flat reserve). parallel=1 is tiny in absolute terms. + est = _backend()._estimate_compute_buffer_bytes(n_parallel = parallel) / MIB + assert est <= max(measured * 2.0, 128) + + def test_tensor_upper_bounds_measured(self): + est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) / MIB + assert est >= _TENSOR_MEASURED_PER_DEVICE + + def test_tensor_far_below_old_flat_reserve(self): + # The whole point: deterministic estimate << flat 5120 for this model. + est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) / MIB + assert est < LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + + +class TestScaling: + def test_grows_with_serving_slots(self): + b = _backend() + vals = [b._estimate_compute_buffer_bytes(n_parallel = p) for p in (1, 2, 4, 8)] + assert vals == sorted(vals) and vals[0] < vals[-1] + + def test_parallel_1_is_small(self): + # Single-token decode: a few tens of MiB, not gigabytes. + est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1) / MIB + assert est < 128 + + def test_tensor_exceeds_pipeline_at_same_parallel(self): + b = _backend() + pipe = b._estimate_compute_buffer_bytes(n_parallel = 1) + tens = b._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) + assert tens > pipe + + def test_scales_with_vocab(self): + small = _backend(vocab = 32000)._estimate_compute_buffer_bytes(n_parallel = 4) + big = _backend(vocab = 256000)._estimate_compute_buffer_bytes(n_parallel = 4) + assert big > small + + def test_scales_with_ubatch(self): + b = _backend() + lo = b._estimate_compute_buffer_bytes(n_parallel = 4, n_ubatch = 256) + hi = b._estimate_compute_buffer_bytes(n_parallel = 4, n_ubatch = 1024) + assert hi > lo + + +class TestFallback: + def test_zero_when_vocab_missing(self): + assert _backend(vocab = None)._estimate_compute_buffer_bytes(n_parallel = 4) == 0 + + def test_zero_when_embd_missing(self): + assert _backend(embd = None)._estimate_compute_buffer_bytes(n_parallel = 4) == 0 + + def test_zero_lets_tensor_plan_use_flat_fallback(self): + # When dims are missing, _plan_tensor_parallel must fall back to the flat + # reserve (defense-in-depth) rather than reserving 0 and OOMing. + b = _backend(vocab = None, embd = None) + b._n_layers = None # can't estimate KV -> floors ctx, still returns a plan + ec, mac, gi, ts = b._plan_tensor_parallel([(0, 48000), (1, 48000)], 8 * 1024**3, 8192) + assert gi == [0, 1] # both GPUs usable under the flat fallback + + +class TestParallel1Default: + """At Studio's default --parallel 1 the buffer is negligible in pipeline.""" + + def test_default_n_parallel(self): + est = _backend()._estimate_compute_buffer_bytes() / MIB + assert est < 128 diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py index cd834b345b..27e9d0f57a 100644 --- a/studio/backend/tests/test_kv_cache_estimation.py +++ b/studio/backend/tests/test_kv_cache_estimation.py @@ -71,7 +71,7 @@ except ImportError: ) sys.modules["httpx"] = _httpx_stub -from core.inference.llama_cpp import LlamaCppBackend +from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend # Helpers @@ -1484,8 +1484,8 @@ class TestServerFlags: assert fitted < 32_768 def test_fit_mtp_engaged_returns_smaller_or_equal_context(self): - # MTP budget is 0.85 of available, non-MTP is 0.90; on a tight - # budget MTP must yield <= non-MTP. + # 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, @@ -1518,7 +1518,7 @@ class TestServerFlags: 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) / 0.90 + 1 + 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), diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py index abacc6e953..9866fc4ae1 100644 --- a/studio/backend/tests/test_llama_cpp_context_fit.py +++ b/studio/backend/tests/test_llama_cpp_context_fit.py @@ -67,7 +67,11 @@ _httpx_stub.Client = type( ) sys.modules.setdefault("httpx", _httpx_stub) -from core.inference.llama_cpp import LlamaCppBackend, classify_gpu_offload_lines +from core.inference.llama_cpp import ( + _CTX_FIT_VRAM_FRACTION, + LlamaCppBackend, + classify_gpu_offload_lines, +) from core.inference.llama_server_args import parse_ctx_override, resolve_requested_ctx @@ -171,7 +175,7 @@ def _drive( ) kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv) total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * 0.90: + if total_mib <= pool_mib * _CTX_FIT_VRAM_FRACTION: best_cap = max(best_cap, capped) if best_cap > 0: max_available_ctx = best_cap @@ -664,3 +668,39 @@ class TestClassifyGpuOffload: def test_module_level_no_signal_returns_none(self): assert classify_gpu_offload_lines(["INFO starting server"]) is None + + +def test_select_gpus_ranks_by_usable_not_raw_free(): + # 80 GB card (30 GB free -> 25.9 GB usable) vs 32 GB card (29 GB free -> 27.4 + # GB usable). A 27 GB model fits the 32 GB card alone; raw-free ranking would + # try the 80 GB card first and split across both. Usable ranking picks [1]. + gpus = [(0, 30000), (1, 29000)] + totals = {0: 81920, 1: 32607} + model = int(27000 * 1024 * 1024) + idxs, use_fit = LlamaCppBackend._select_gpus(model, gpus, total_by_idx = totals) + assert idxs == [1] and use_fit is False + + +def test_select_gpus_reserves_per_device_overhead(): + # Two 16 GB cards, ~15181 MiB usable each at 0.95 -> 30362 MiB pooled. A 30000 + # MiB model fits the pool with no per-device overhead, but a layer split also + # pays ~1 GiB/extra-GPU; that pushes the 2-GPU need to 31024 MiB > pool, so a + # pin would OOM -> must fall back to --fit. Single-GPU fits add no overhead + # (Finding F1, the explicit/file-size multi-GPU pin gap). + gpus = [(0, 16000), (1, 16000)] + totals = {0: 16384, 1: 16384} + gib = 1024 * 1024 * 1024 + model = int(30000 * 1024 * 1024) + idxs, use_fit = LlamaCppBackend._select_gpus(model, gpus, total_by_idx = totals) + assert idxs == [0, 1] and use_fit is False # fits 2 GPUs without overhead + idxs2, use_fit2 = LlamaCppBackend._select_gpus( + model, gpus, total_by_idx = totals, per_device_overhead_bytes = gib + ) + assert idxs2 is None and use_fit2 is True # overhead tips it past the pool + # A single-GPU fit is unchanged by the overhead (k=1 adds nothing). + small = int(15000 * 1024 * 1024) + a, _ = LlamaCppBackend._select_gpus(small, gpus, total_by_idx = totals) + b, _ = LlamaCppBackend._select_gpus( + small, gpus, total_by_idx = totals, per_device_overhead_bytes = gib + ) + assert a == [0] and b == [0] diff --git a/studio/backend/tests/test_llama_cpp_max_context_threshold.py b/studio/backend/tests/test_llama_cpp_max_context_threshold.py index 310aaf6c0f..545b9c4e7a 100644 --- a/studio/backend/tests/test_llama_cpp_max_context_threshold.py +++ b/studio/backend/tests/test_llama_cpp_max_context_threshold.py @@ -74,7 +74,7 @@ _httpx_stub.Client = type( ) sys.modules.setdefault("httpx", _httpx_stub) -from core.inference.llama_cpp import LlamaCppBackend +from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend # Helpers @@ -140,7 +140,7 @@ def _compute_max_available_ctx( ) kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv) total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * 0.90: + if total_mib <= pool_mib * _CTX_FIT_VRAM_FRACTION: best_cap = max(best_cap, capped) if best_cap > 0: max_available_ctx = best_cap diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 14c3848d33..063a9ce2cd 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -1268,9 +1268,10 @@ def test_build_speculative_flags_user_draft_n_max_override(monkeypatch): assert backend.spec_draft_n_max == 5 -def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch): - # Outdated llama-server with no MTP support: forced MTP must degrade - # to spec-off (warned) rather than emit a bad --spec-type. +def test_build_speculative_flags_mtp_token_missing_emits_spec_default(monkeypatch): + # Outdated llama-server with no MTP support: forced MTP must degrade (warned) + # and emit --spec-default so an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI + # wins over env) can't make the child attempt MTP the gate budgeted off. backend = _resolver_backend(monkeypatch, mtp_token = None) flags = backend._build_speculative_flags( speculative_type = "mtp", @@ -1282,10 +1283,11 @@ def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch): binary = "/fake/llama-server", ) assert "--spec-type" not in flags - # _speculative_type stays None (resolved emission was none); the user's - # choice is still reflected in _requested_spec_mode. + assert "--spec-default" in flags + # Degraded to non-speculative; the user's choice is still reflected. + assert backend.speculative_type == "default" assert backend.requested_spec_mode == "mtp" - assert backend.speculative_type is None + assert backend.spec_fallback_reason == "binary_no_mtp" def test_forced_mtp_on_non_mtp_model_defaults_back(monkeypatch): diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index f3a3ea1ec4..deeb228026 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -24,6 +24,7 @@ _lsa = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(_lsa) is_managed_flag = _lsa.is_managed_flag parse_cache_override = _lsa.parse_cache_override +parse_cache_override_per_axis = _lsa.parse_cache_override_per_axis parse_ctx_override = _lsa.parse_ctx_override parse_split_mode_override = _lsa.parse_split_mode_override resolve_cache_type_kv = _lsa.resolve_cache_type_kv @@ -467,6 +468,25 @@ def test_parse_cache_override_rejects_malformed_values(args): parse_cache_override(args) +@pytest.mark.parametrize( + "args, expected", + [ + (["--cache-type-k", "f32", "--cache-type-v", "f16"], ("f32", "f16")), + (["-ctk", "q8_0", "-ctv", "q4_0"], ("q8_0", "q4_0")), + (["--cache-type-k=f32"], ("f32", None)), + (["--cache-type-v", "f16"], (None, "f16")), + (["-c", "4096"], (None, None)), + (None, (None, None)), + # Last-wins is kept per axis. + (["-ctk", "f16", "-ctk", "f32"], ("f32", None)), + ], +) +def test_parse_cache_override_per_axis(args, expected): + # Unlike parse_cache_override (collapses both axes to one last-wins value), + # this keeps K and V apart so an asymmetric cache can be budgeted per axis. + assert parse_cache_override_per_axis(args) == expected + + def test_resolve_cache_type_kv_uses_override_when_present(): assert resolve_cache_type_kv(["--cache-type-k", "q8_0"], "f16") == "q8_0" @@ -649,6 +669,53 @@ def test_strip_shadowing_flags_drops_model_draft_with_spec(): assert out == ["--top-k", "20"] +@pytest.mark.parametrize( + "selector", + [ + ["--spec-draft-hf", "org/repo"], + ["-hfd", "org/repo"], + ["-hfrd", "org/repo"], + ["--hf-repo-draft", "org/repo"], + ["--spec-draft-hf=org/repo"], + ], +) +def test_strip_shadowing_flags_drops_hf_drafter_selectors_with_spec(selector): + # HF drafter selectors must reset on inherit like local --model-draft, or a + # stale inherited HF drafter last-wins over Studio's re-derived spec choice. + out = strip_shadowing_flags( + selector + ["--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = True, + strip_template = False, + ) + assert out == ["--top-k", "20"] + + +def test_strip_shadowing_flags_keeps_draft_tuning_with_spec(): + # Per-drafter tuning knobs are deliberately preserved: the VRAM budget reads + # them via the same parsers the child honors (so they stay consistent on + # inherit), and stripping --spec-draft-ngl would move a CPU drafter to GPU. + keep = [ + "--spec-draft-type-k", + "q4_0", + "--spec-draft-type-v", + "q4_0", + "--spec-draft-ngl", + "0", + "--spec-draft-device", + "cpu", + ] + out = strip_shadowing_flags( + list(keep), + strip_context = False, + strip_cache = False, + strip_spec = True, + strip_template = False, + ) + assert out == keep + + def test_strip_shadowing_flags_keeps_split_mode_when_not_requested(): # No tensor_parallel field supplied on the Apply -> an inherited # --split-mode survives (mirrors the chat-template keep behavior). diff --git a/studio/backend/tests/test_mtp_vram_budget.py b/studio/backend/tests/test_mtp_vram_budget.py new file mode 100644 index 0000000000..f61ffa3c21 --- /dev/null +++ b/studio/backend/tests/test_mtp_vram_budget.py @@ -0,0 +1,1002 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the deterministic MTP VRAM reserve used by load-time auto-fit. + +reserve(ctx) = draft_KV(ctx, draft_cache_type) + separate_drafter_weights, sized +from GGUF dims (embedded head from the main model's dims; separate drafter from +its own KV). Anchors checked against real llama-server measurements. Pure: no +GPU, network, subprocess, or GGUF I/O.""" + +from __future__ import annotations + +import inspect +import os +import sys +import types as _types +from pathlib import Path + +import pytest + +# --------------------------------------------------------------------------- +# Stub heavy/unavailable deps before importing the module under test. +# --------------------------------------------------------------------------- + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +sys.modules.setdefault("structlog", _types.ModuleType("structlog")) + +# 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 in tests collected after +# this one (the stub leaks via sys.modules for the whole session). +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,), {})) + _httpx_stub.Timeout = type("Timeout", (), {"__init__": lambda self, *a, **kw: None}) + _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 ( # noqa: E402 + _CTX_FIT_VRAM_FRACTION, + LlamaCppBackend, + _extra_args_draft_cache_types, + _extra_args_draft_offloaded_to_cpu, + _extra_args_mtp_draft_path, + _extra_args_n_ubatch, + _extra_args_requests_mtp, + _extra_args_requests_separate_draft, + _extra_args_spec_draft_n_max, + _effective_tensor_parallel, + _env_main_cache_type_for_budget, + _extra_args_main_cache_type_for_budget, + _kv_bytes_per_elem, + _tensor_parallel_matches_loaded, +) +from core.inference.llama_server_args import _env_split_mode_is_tensor # noqa: E402 + +MIB = 1024 * 1024 +GIB = 1024**3 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_backend( + *, + nextn = 1, + n_kv_heads = 4, + n_heads = 24, + kv_key_length = 256, + kv_value_length = 256, + embedding_length = 5120, + n_layers = 65, + native_ctx = 262144, +): + """Qwen3.6-27B-MTP-class backend (embedded head) with the MTP-math dims.""" + b = LlamaCppBackend.__new__(LlamaCppBackend) + b._nextn_predict_layers = nextn + b._n_kv_heads = n_kv_heads + b._n_heads = n_heads + b._kv_key_length = kv_key_length + b._kv_value_length = kv_value_length + b._embedding_length = embedding_length + b._n_layers = n_layers + b._context_length = native_ctx + # Hybrid attention/Mamba (qwen35 path) + remaining KV-estimator fields. + b._shared_kv_layers = 0 + b._kv_lora_rank = None + b._sliding_window = None + b._sliding_window_pattern = None + b._ssm_inner_size = 6144 + b._full_attention_interval = 4 + b._key_length_mla = None + b._n_kv_heads_by_layer = None + b._kv_key_length_swa = None + b._kv_value_length_swa = None + b._draft_backend_cache = None + return b + + +class _StubDrafter: + """Stand-in for a separate drafter backend (no GGUF I/O).""" + + def __init__(self, kv_per_token): + self._kv_per_token = kv_per_token + + def _can_estimate_kv(self): + return True + + def _estimate_kv_cache_bytes( + self, + n_ctx, + cache_type = None, + n_parallel = 1, + **_k, + ): + bpe = _kv_bytes_per_elem(cache_type) + # n_parallel scales like a sliding-window drafter's per-slot KV. + return 0 if n_ctx <= 0 else int(n_ctx * self._kv_per_token * bpe / 2.0 * n_parallel) + + +# --------------------------------------------------------------------------- +# Embedded draft KV: deterministic from nextn dims, scales with ctx + draft type +# --------------------------------------------------------------------------- + + +class TestEmbeddedDraftKv: + def test_scales_linearly_with_context(self): + b = _make_backend() + kv_8k = b._mtp_draft_kv_bytes(8192) + kv_16k = b._mtp_draft_kv_bytes(16384) + kv_64k = b._mtp_draft_kv_bytes(65536) + assert kv_8k and kv_16k and kv_64k + assert kv_16k == pytest.approx(2 * kv_8k) + assert kv_64k == pytest.approx(8 * kv_8k) + + def test_value_matches_dim_formula_f16(self): + # nextn(1) * n_kv(4) * (256+256) * 2(f16) * ctx -- no magic safety factor. + b = _make_backend() + ctx = 131072 + expected = int(1 * 4 * 512 * 2.0 * ctx) + assert b._mtp_draft_kv_bytes(ctx) == expected + # And that is 512 MiB, matching the measured 27B draft-KV slope (~4 MiB/1k). + assert b._mtp_draft_kv_bytes(ctx) / MIB == pytest.approx(512, abs = 1) + + def test_scales_with_nextn_predict_layers(self): + one = _make_backend(nextn = 1)._mtp_draft_kv_bytes(65536) + two = _make_backend(nextn = 2)._mtp_draft_kv_bytes(65536) + assert two == pytest.approx(2 * one) + + def test_embedded_draft_kv_floored_at_f16(self): + # The embedded MTP head is one layer, so llama.cpp's quantized-KV + # overhead is not amortized: a quantized draft KV fits LESS context than + # f16, not more (ggml-org/llama.cpp#24102). The embedded reserve floors a + # quantized draft type at f16 (never under-reserved); f32 still costs more. + b = _make_backend() + f16 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "f16", draft_cache_type_v = "f16") + q8 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "q8_0", draft_cache_type_v = "q8_0") + q4 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "q4_0", draft_cache_type_v = "q4_0") + f32 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "f32", draft_cache_type_v = "f32") + assert q8 == f16 and q4 == f16 # quantized draft KV priced as f16, not less + assert f32 == pytest.approx(f16 * 2.0) # f32 genuinely larger, not floored + + def test_draft_kv_split_axes_no_under_reserve(self): + # A quantized draft type on either or both axes never reserves below the + # all-f16 value for the single-layer embedded head (the f16 floor; #24102). + b = _make_backend() + both_q4 = b._mtp_draft_kv_bytes( + 131072, draft_cache_type_k = "q4_0", draft_cache_type_v = "q4_0" + ) + k_only = b._mtp_draft_kv_bytes(131072, draft_cache_type_k = "q4_0") # V defaults f16 + both_f16 = b._mtp_draft_kv_bytes(131072, draft_cache_type_k = "f16", draft_cache_type_v = "f16") + assert both_q4 == k_only == both_f16 # floored at f16, never under-reserved + + def test_none_when_dims_missing(self): + assert _make_backend(nextn = 0)._mtp_draft_kv_bytes(65536) is None + assert _make_backend(kv_key_length = None)._mtp_draft_kv_bytes(65536) is None + assert _make_backend()._mtp_draft_kv_bytes(0) is None + + +# --------------------------------------------------------------------------- +# Separate drafter (Gemma): sized from the drafter GGUF's own dims + weights +# --------------------------------------------------------------------------- + + +class TestSeparateDrafter: + def test_uses_drafter_kv_and_weights(self, monkeypatch): + b = _make_backend(nextn = None) # main has no embedded head + stub = _StubDrafter(kv_per_token = 2000) + monkeypatch.setattr(b, "_draft_backend_for", lambda path: stub) + ctx = 65536 + kv = b._mtp_draft_kv_bytes(ctx, drafter_path = "/m/draft.gguf") + assert kv == stub._estimate_kv_cache_bytes(ctx) + total = b._estimate_mtp_overhead_bytes( + ctx, drafter_path = "/m/draft.gguf", draft_weights_bytes = GIB + ) + assert total == kv + GIB + + def test_drafter_kv_scales_with_context(self, monkeypatch): + b = _make_backend(nextn = None) + monkeypatch.setattr(b, "_draft_backend_for", lambda path: _StubDrafter(2000)) + a = b._mtp_draft_kv_bytes(16384, drafter_path = "/m/d.gguf") + c = b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf") + assert c == pytest.approx(4 * a) + + def test_drafter_kv_scales_with_parallel_slots(self, monkeypatch): + # The drafter is served under the same --parallel slots as the main model, + # so a sliding-window drafter's KV grows per slot; the reserve must thread + # n_parallel or it under-reserves (Finding G1). + b = _make_backend(nextn = None) + monkeypatch.setattr(b, "_draft_backend_for", lambda path: _StubDrafter(2000)) + one = b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf", n_parallel = 1) + four = b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf", n_parallel = 4) + assert four == pytest.approx(4 * one) + # And it threads through the overhead estimate too. + ov1 = b._estimate_mtp_overhead_bytes( + 65536, drafter_path = "/m/d.gguf", draft_weights_bytes = GIB, n_parallel = 1 + ) + ov4 = b._estimate_mtp_overhead_bytes( + 65536, drafter_path = "/m/d.gguf", draft_weights_bytes = GIB, n_parallel = 4 + ) + assert (ov4 - GIB) == pytest.approx(4 * (ov1 - GIB)) # KV scales, weights flat + + def test_none_when_drafter_unreadable(self, monkeypatch): + b = _make_backend(nextn = None) + monkeypatch.setattr(b, "_draft_backend_for", lambda path: None) + assert b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf") is None + assert b._estimate_mtp_overhead_bytes(65536, drafter_path = "/m/d.gguf") is None + + def test_keeps_weights_when_drafter_kv_unsizable(self, monkeypatch): + # KV can't be sized (exotic/remote drafter), but the local weights are + # known: reserve the weights so a drafter larger than the flat fallback + # cushion can't slip through and OOM (Finding C). Nothing known -> None. + b = _make_backend(nextn = None) + monkeypatch.setattr(b, "_draft_backend_for", lambda path: None) + assert b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf") is None + assert ( + b._estimate_mtp_overhead_bytes( + 65536, drafter_path = "/m/d.gguf", draft_weights_bytes = 3 * GIB + ) + == 3 * GIB + ) + assert b._estimate_mtp_overhead_bytes(65536, drafter_path = "/m/d.gguf") is None + + +# --------------------------------------------------------------------------- +# Total overhead = draft KV (+ separate drafter weights); no verify constant +# --------------------------------------------------------------------------- + + +class TestOverheadTotal: + def test_equals_draft_kv_for_embedded(self): + b = _make_backend() + for ctx in (16384, 65536, 131072): + assert b._estimate_mtp_overhead_bytes(ctx) == b._mtp_draft_kv_bytes(ctx) + + def test_does_not_depend_on_n_max(self): + # The verify buffer (the only n_max-dependent term) rides in headroom now. + b = _make_backend() + assert b._estimate_mtp_overhead_bytes( + 65536, spec_draft_n_max = 2 + ) == b._estimate_mtp_overhead_bytes(65536, spec_draft_n_max = 6) + + def test_none_when_draft_kv_unsizable(self): + assert _make_backend(nextn = 0)._estimate_mtp_overhead_bytes(65536) is None + + def test_includes_separate_drafter_weights(self): + b = _make_backend() + base = b._estimate_mtp_overhead_bytes(65536) + with_w = b._estimate_mtp_overhead_bytes(65536, draft_weights_bytes = GIB) + assert with_w - base == GIB + + @pytest.mark.parametrize( + "ctx,measured_draft_kv_mib", + # Measured 27B MTP delta minus the (headroom-covered) ~500 MiB verify + # buffer leaves the draft KV; the deterministic estimate must match it. + [(16384, 64), (65536, 256), (131072, 512)], + ) + def test_draft_kv_matches_measured(self, ctx, measured_draft_kv_mib): + b = _make_backend() + pred = b._estimate_mtp_overhead_bytes(ctx) / MIB + assert pred == pytest.approx(measured_draft_kv_mib, abs = 2) + + +# --------------------------------------------------------------------------- +# _fit_context_to_vram: MTP reserve lowers the chosen context +# --------------------------------------------------------------------------- + + +class TestFitContextWithMtp: + def _fit_backend(self, kv_per_token = 325_000): + b = _make_backend() + b._can_estimate_kv = lambda: True + b._estimate_kv_cache_bytes = lambda n, _t = None, **_k: (0 if n <= 0 else n * kv_per_token) + return b + + def test_overhead_fn_lowers_context(self): + b = self._fit_backend() + avail_mib = 24_000 + model = 8 * GIB + without = b._fit_context_to_vram(131072, avail_mib, model) + with_mtp = b._fit_context_to_vram( + 131072, + avail_mib, + model, + mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes(c) or 0, + ) + assert 0 < with_mtp < without + + def test_quantized_embedded_draft_kv_does_not_inflate_context(self): + # For the single-layer embedded head, quantizing the draft KV does NOT + # buy more context (it fits less in practice; ggml-org/llama.cpp#24102), + # so the f16-floored reserve advertises the same context as f16 -- never + # a larger one off a smaller (unsafe) reserve. + b = self._fit_backend() + avail_mib, model = 24_000, 8 * GIB + f16 = b._fit_context_to_vram( + 131072, + avail_mib, + model, + mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes( + c, draft_cache_type_k = "f16", draft_cache_type_v = "f16" + ) + or 0, + ) + q4 = b._fit_context_to_vram( + 131072, + avail_mib, + model, + mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes( + c, draft_cache_type_k = "q4_0", draft_cache_type_v = "q4_0" + ) + or 0, + ) + assert 0 < q4 == f16 + + def test_no_mtp_unchanged(self): + b = self._fit_backend() + avail_mib, model = 24_000, 8 * GIB + a = b._fit_context_to_vram(131072, avail_mib, model) + bb = b._fit_context_to_vram( + 131072, avail_mib, model, mtp_engaged = False, mtp_overhead_fn = None + ) + assert a == bb + + def test_chosen_context_actually_fits_budget(self): + b = self._fit_backend() + avail_mib, model = 24_000, 8 * GIB + fn = lambda c: b._estimate_mtp_overhead_bytes(c) or 0 # noqa: E731 + ctx = b._fit_context_to_vram(131072, avail_mib, model, mtp_overhead_fn = fn) + budget = avail_mib * MIB * _CTX_FIT_VRAM_FRACTION + assert model + b._estimate_kv_cache_bytes(ctx) + fn(ctx) <= budget + + +# --------------------------------------------------------------------------- +# extra_args parsing: detect user-enabled MTP + draft depth + draft KV type +# --------------------------------------------------------------------------- + + +class TestExtraArgsMtpDetection: + @pytest.mark.parametrize( + "args,expected", + [ + (["--spec-type", "draft-mtp"], True), + (["--spec-type", "mtp"], True), + (["--spec-type", "ngram-mod,draft-mtp"], True), + (["--spec-type=draft-mtp"], True), + (["--spec-type", "ngram-mod"], False), + (["--spec-default"], False), + (["-c", "131072"], False), + (None, False), + ([], False), + ], + ) + def test_requests_mtp(self, args, expected): + assert _extra_args_requests_mtp(args, env = {}) is expected + + def test_requests_mtp_env(self): + # The child honors LLAMA_ARG_SPEC_TYPE; env-requested MTP must reserve too. + assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"}) is True + assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "ngram-mod,mtp"}) is True + assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"}) is False + assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "none"}) is False + + def test_requests_mtp_effective_spec_type(self): + # llama.cpp uses the LAST CLI --spec-type and ignores the env when any CLI + # --spec-type is present. The reserve must track that effective value, not + # any earlier/MTP-ish one, or it over-reserves a drafter the launch won't + # load (Finding B). + env_mtp = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"} + # Later CLI value overrides an earlier MTP one (last-wins). + assert ( + _extra_args_requests_mtp( + ["--spec-type", "draft-mtp", "--spec-type", "ngram-mod"], env = {} + ) + is False + ) + # A non-MTP CLI flag overrides a stale MTP env. + assert _extra_args_requests_mtp(["--spec-type", "ngram-mod"], env = env_mtp) is False + assert _extra_args_requests_mtp(["--spec-type", "none"], env = env_mtp) is False + # A later MTP CLI value still engages. + assert ( + _extra_args_requests_mtp( + ["--spec-type", "ngram-mod", "--spec-type", "draft-mtp"], env = {} + ) + is True + ) + # Same precedence for separate (draft-simple/eagle3) detection. + assert ( + _extra_args_requests_separate_draft( + ["--spec-type", "draft-simple", "--spec-type", "ngram-mod"], env = {} + ) + is False + ) + assert ( + _extra_args_requests_separate_draft( + ["--spec-type", "ngram-mod"], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"} + ) + is False + ) + + @pytest.mark.parametrize( + "args,expected", + [ + (["--spec-type", "draft-simple"], True), + (["--spec-type", "draft-eagle3"], True), + (["--spec-type=draft-eagle3"], True), + (["--spec-type", "draft-mtp"], False), # MTP path handles this one + (["--spec-type", "ngram-mod"], False), # loads no draft model + (["-c", "4096"], False), + (None, False), + ], + ) + def test_requests_separate_draft(self, args, expected): + assert _extra_args_requests_separate_draft(args, env = {}) is expected + + def test_requests_separate_draft_env(self): + assert ( + _extra_args_requests_separate_draft([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"}) + is True + ) + assert ( + _extra_args_requests_separate_draft([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"}) + is False + ) + + def test_load_model_reserves_for_non_mtp_draft_modes(self): + # load_model engages the draft reserve for a non-MTP model-based draft mode + # only when extras also name a drafter (else nothing is loaded to reserve). + # Strip all whitespace so the check survives any line-wrapping the + # formatter applies to the call (pre-commit black wraps long lines). + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_user_draft_via_extras" in compact + # called with extra_args (an env kwarg may follow); prefix match stays + # robust to that and to any formatter line-wrapping. + assert "_extra_args_requests_separate_draft(extra_args" in compact + assert "or_user_draft_via_extras" in compact # OR'd into the reserve gate + # The drafter check must NOT force extras-only (env={}); the default + # env=None lets it see an env LLAMA_ARG_SPEC_DRAFT_MODEL, so an env-only + # drafter still engages the reserve (codex review 4507014299). + assert "bool(_extra_args_mtp_draft_path(extra_args))" in compact + + def test_env_only_drafter_engages_separate_draft_reserve(self, monkeypatch): + # An env-provided drafter (no --model-draft in extras) must still engage + # the draft reserve, or auto-fit spends the drafter's VRAM and OOMs. Mirror + # load_model's _user_draft_via_extras gate (codex review 4507014299). + monkeypatch.setenv("LLAMA_ARG_SPEC_DRAFT_MODEL", "/large.gguf") + monkeypatch.delenv("LLAMA_ARG_SPEC_DRAFT_HF_REPO", raising = False) + ea = ["--spec-type", "draft-simple"] # _spec_env is {} (extras set spec-type) + assert _extra_args_requests_separate_draft(ea, env = {}) is True + assert _extra_args_mtp_draft_path(ea) == "/large.gguf" # env=None -> os.environ + # -> _user_draft_via_extras True; _env_draft_for_budget sizes the drafter. + assert _extra_args_mtp_draft_path([], env = dict(os.environ)) == "/large.gguf" + + def test_load_model_gates_env_spec_type_on_off_mode(self): + # LLAMA_ARG_SPEC_TYPE only reaches the child when Studio emits no spec + # flag (UI mode "off", no user --spec-type); otherwise the emitted + # --spec-type/--spec-default overrides the env, so the reserve must not + # consult it or a stale MTP env over-reserves (Finding F3). Whitespace- + # stripped so the check survives formatter line-wrapping. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert '_mtp_canonical=="off"' in compact # the env-reaches-child gate + assert "_extra_args_requests_mtp(extra_args,env=_spec_env)" in compact + + def test_spec_default_overrides_env_mtp(self): + # --spec-default is a CLI spec flag (resolves to the model default, + # non-MTP) that overrides a stale LLAMA_ARG_SPEC_TYPE env, so the reserve + # must not treat it as MTP (reviewer.py R4). + env_mtp = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"} + assert _extra_args_requests_mtp(["--spec-default"], env = env_mtp) is False + assert ( + _extra_args_requests_separate_draft( + ["--spec-default"], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"} + ) + is False + ) + # A later --spec-type still wins over an earlier --spec-default. + assert ( + _extra_args_requests_mtp(["--spec-default", "--spec-type", "draft-mtp"], env = {}) is True + ) + + def test_load_model_drafter_budget_precedence(self): + # The budget sizes the drafter the launch actually loads: CLI extras win, + # then Studio's emitted mtp_draft_path (overrides LLAMA_ARG_SPEC_DRAFT_MODEL), + # then the env drafter -- not the env before Studio's (reviewer.py R3). + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" in compact + assert "_env_draft_for_budget=_extra_args_mtp_draft_path([],env=os.environ)" in compact + assert "_cli_draft_for_budgetor_studio_draft_for_budgetor_env_draft_for_budget" in compact + + def test_load_model_drops_cpu_offloaded_drafter_from_budget(self): + # A SEPARATE drafter offloaded to CPU (--spec-draft-ngl 0 / + # --spec-draft-device none) consumes no GPU, so it must be dropped from the + # budget and get no flat reserve (Finding F2). But an embedded head is on + # GPU regardless of those draft-only flags, so the flat reserve is only + # suppressed when there is no embedded head (Finding G5). + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + # env-aware: also honors the inherited LLAMA_ARG_N_GPU_LAYERS_DRAFT. + assert ( + "_draft_on_cpu=_extra_args_draft_offloaded_to_cpu(extra_args,env=os.environ)" in compact + ) + assert "if_draft_on_cpu:_mtp_draft_for_budget=None" in compact + # flat reserve suppressed only for a CPU drafter with no embedded head + assert "_draft_cpu_no_embedded=_draft_on_cpuandnotself._nextn_predict_layers" in compact + assert "not_draft_cpu_no_embedded" in compact + + def test_load_model_keeps_flat_reserve_for_unsized_draft_kv(self): + # When only the drafter weights could be sized (KV unsizable), the flat + # fraction stays on as the cushion for the still-unsized draft KV, on top + # of the byte-accurate weights reserve (Finding G3). + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_mtp_kv_unsized" in compact + assert "mtp_overhead_fnisNoneor_mtp_kv_unsized" in compact + + def test_load_model_ranks_subsets_by_active_pin_fraction(self): + # Auto/cap subset ranking uses the active budget fraction (lowered by the + # flat MTP reserve), not a hard-coded 0.95, so the ranking order matches + # the fit budget that is then tested (Finding G4). + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_gpu_usable(g,pin_fraction)" in compact + assert "_gpu_usable(g,_CTX_FIT_VRAM_FRACTION-_flat_mtp_reserve)" in compact + + @pytest.mark.parametrize( + "args,expected", + [ + (["--spec-draft-ngl", "0"], True), + (["-ngld", "0"], True), + (["--spec-draft-ngl=0"], True), + (["--n-gpu-layers-draft", "0"], True), + (["--spec-draft-ngl", "20"], False), + (["--spec-draft-device", "none"], True), + (["--spec-draft-device", "CPU"], True), + (["-devd", "cpu,none"], True), + (["--spec-draft-device", "CUDA0"], False), + (["--spec-draft-device", "CUDA0,CPU"], False), # any GPU -> on GPU + (["-c", "4096"], False), + (None, False), + # last-wins: only the final value of each flag counts (Finding G2) + (["--spec-draft-ngl", "0", "--spec-draft-ngl", "-1"], False), # last = GPU + (["--spec-draft-ngl", "-1", "--spec-draft-ngl", "0"], True), # last = CPU + (["--spec-draft-device", "CUDA0", "--spec-draft-device", "none"], True), + (["--spec-draft-device", "none", "--spec-draft-device", "CUDA0"], False), + ], + ) + def test_draft_offloaded_to_cpu(self, args, expected): + assert _extra_args_draft_offloaded_to_cpu(args, env = {}) is expected + + def test_draft_offloaded_to_cpu_env(self): + # The child honors LLAMA_ARG_N_GPU_LAYERS_DRAFT; an env-only CPU offload + # must drop the drafter from the budget too (review run3 #3). CLI wins. + assert ( + _extra_args_draft_offloaded_to_cpu([], env = {"LLAMA_ARG_N_GPU_LAYERS_DRAFT": "0"}) + is True + ) + assert ( + _extra_args_draft_offloaded_to_cpu([], env = {"LLAMA_ARG_N_GPU_LAYERS_DRAFT": "-1"}) + is False + ) + # CLI --spec-draft-ngl wins over the env (last-wins is CLI-only). + assert ( + _extra_args_draft_offloaded_to_cpu( + ["--spec-draft-ngl", "-1"], env = {"LLAMA_ARG_N_GPU_LAYERS_DRAFT": "0"} + ) + is False + ) + assert _extra_args_draft_offloaded_to_cpu([], env = {}) is False + + @pytest.mark.parametrize( + "args,expected", + [ + (["--spec-draft-n-max", "4"], 4), + (["--spec-draft-n-max=6"], 6), + (["--spec-type", "draft-mtp", "--spec-draft-n-max", "3"], 3), + (["--spec-draft-n-max", "2", "--spec-draft-n-max", "5"], 5), # last wins + (["--spec-draft-n-max", "notanint"], None), + (["-c", "4096"], None), + (None, None), + (["--draft-max", "6"], 6), + (["--draft-max=4"], 4), + (["--spec-type", "draft-mtp", "--draft-max", "6"], 6), + (["--spec-draft-n-max", "2", "--draft-max", "5"], 5), + ], + ) + def test_spec_draft_n_max(self, args, expected): + assert _extra_args_spec_draft_n_max(args) == expected + + @pytest.mark.parametrize( + "args,expected", + [ + (["--model-draft", "/m/draft.gguf"], "/m/draft.gguf"), + (["--spec-draft-model", "/m/draft.gguf"], "/m/draft.gguf"), + (["-md", "/m/draft.gguf"], "/m/draft.gguf"), + (["--model-draft=/m/draft.gguf"], "/m/draft.gguf"), + (["--model-draft", "--spec-type"], None), + (["-c", "4096"], None), + (None, None), + ], + ) + def test_mtp_draft_path(self, args, expected): + # env={} isolates pure-CLI behavior from a polluted test environment. + assert _extra_args_mtp_draft_path(args, env = {}) == expected + + @pytest.mark.parametrize( + "args,expected", + [ + # HF draft repo flags are real llama-server flags; the budget must see them. + (["--spec-draft-hf", "big/repo:Q8_0"], "big/repo:Q8_0"), + (["-hfd", "big/repo"], "big/repo"), + (["-hfrd", "big/repo"], "big/repo"), + (["--hf-repo-draft=big/repo"], "big/repo"), + ], + ) + def test_mtp_draft_path_hf_flags(self, args, expected): + assert _extra_args_mtp_draft_path(args, env = {}) == expected + + def test_mtp_draft_path_env_fallback(self): + # The child honors LLAMA_ARG_SPEC_DRAFT_MODEL / _HF_REPO; CLI wins over env. + assert ( + _extra_args_mtp_draft_path([], env = {"LLAMA_ARG_SPEC_DRAFT_MODEL": "/m/e.gguf"}) + == "/m/e.gguf" + ) + assert _extra_args_mtp_draft_path([], env = {"LLAMA_ARG_SPEC_DRAFT_HF_REPO": "x/y"}) == "x/y" + assert ( + _extra_args_mtp_draft_path( + ["-md", "/m/cli.gguf"], env = {"LLAMA_ARG_SPEC_DRAFT_HF_REPO": "x/y"} + ) + == "/m/cli.gguf" + ) + + @pytest.mark.parametrize( + "args,expected", + [ + (["--cache-type-k-draft", "q8_0"], ("q8_0", None)), + (["--spec-draft-type-k", "q4_0"], ("q4_0", None)), + (["-ctkd", "q8_0"], ("q8_0", None)), + (["--cache-type-v-draft", "q4_0"], (None, "q4_0")), # K stays f16, V only + (["--cache-type-k-draft", "q4_0", "--cache-type-v-draft", "q8_0"], ("q4_0", "q8_0")), + (["--cache-type-k-draft=q8_0"], ("q8_0", None)), + (["--cache-type-k", "q8_0"], (None, None)), # main type, not draft + (["-c", "4096"], (None, None)), + (None, (None, None)), + ], + ) + def test_draft_cache_types(self, args, expected): + assert _extra_args_draft_cache_types(args, env = {}) == expected + + def test_draft_cache_types_env_fallback(self): + # The child honors LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K/_V per axis; CLI wins. + assert _extra_args_draft_cache_types( + [], env = {"LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K": "q8_0"} + ) == ("q8_0", None) + assert _extra_args_draft_cache_types( + [], + env = { + "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K": "q8_0", + "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V": "q4_0", + }, + ) == ("q8_0", "q4_0") + assert _extra_args_draft_cache_types( + ["-ctkd", "q4_0"], env = {"LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K": "q8_0"} + ) == ("q4_0", None) + + @pytest.mark.parametrize( + "args,expected", + [ + (["--ubatch-size", "1024"], 1024), + (["-ub", "4096"], 4096), + (["--ubatch-size=512"], 512), + (["--ubatch", "2048"], None), # not a real llama-server flag; ignore it + (["-c", "4096"], None), + (None, None), + ], + ) + def test_n_ubatch(self, args, expected): + assert _extra_args_n_ubatch(args, env = {}) == expected + + def test_n_ubatch_env_fallback(self): + # The child honors LLAMA_ARG_UBATCH; it must reach the compute-buffer reserve. + assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "4096"}) == 4096 + assert ( + _extra_args_n_ubatch(["-ub", "1024"], env = {"LLAMA_ARG_UBATCH": "4096"}) == 1024 + ) # CLI wins + assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "notint"}) is None + + def test_env_main_cache_type_for_budget(self): + # The child inherits LLAMA_ARG_CACHE_TYPE_K/_V, but Studio emits no + # --cache-type when neither param nor extras set it -> a heavier env + # main KV (f32) must be adopted so the reserve matches the child. + assert _env_main_cache_type_for_budget(env = {}) is None + # f32 exceeds the f16 default -> adopt it (lower-cased so the launch + # re-emits it via _valid_cache_types). + assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "f32"}) == "f32" + assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_V": "F32"}) == "f32" + # Heavier of K/V (single knob; over-reserves the lighter axis). + assert ( + _env_main_cache_type_for_budget( + env = {"LLAMA_ARG_CACHE_TYPE_K": "f32", "LLAMA_ARG_CACHE_TYPE_V": "f16"} + ) + == "f32" + ) + # Quantized env types are <= f16 -> already over-reserved by the default. + assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "q4_0"}) is None + assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_V": "q8_0"}) is None + assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "f16"}) is None + # Unknown env type self-neutralizes (treated as f16 by _kv_bytes_per_elem). + assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "wat"}) is None + + def test_load_model_adopts_env_main_cache_type(self): + # Source-level: load_model budgets the heavier of asymmetric --cache-type + # extras, then (only when neither param nor extras set it) adopts the env + # main KV type, so the reserve covers a child that inherits a heavier + # LLAMA_ARG_CACHE_TYPE_*. Whitespace-stripped to survive formatter wraps. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_extra_args_main_cache_type_for_budget(extra_args)" in compact + assert "ifcache_type_kvisNone:" in compact + assert "cache_type_kv=_env_main_cache_type_for_budget()" in compact + + def test_env_split_mode_is_tensor(self): + # The child inherits LLAMA_ARG_SPLIT_MODE, but Studio emits --split-mode + # only on its tensor branch -> a tensor env must flip the budget so the + # heavier per-device compute buffer is reserved (not layer overhead). + assert _env_split_mode_is_tensor(env = {}) is False + assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "tensor"}) is True + assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "Tensor"}) is True + # Other modes are not a runtime-heavier surprise -> not acted on. + assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "layer"}) is False + assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "row"}) is False + assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "none"}) is False + + def test_effective_tensor_parallel_env_flip(self): + # Shared by load_model and both duplicate-load matchers, so they agree. + tensor_env = {"LLAMA_ARG_SPLIT_MODE": "tensor"} + # No extras, toggle off, tensor env -> flips on. + assert _effective_tensor_parallel(None, False, env = tensor_env) is True + # Extras override (any --split-mode) beats the env, even if non-tensor. + assert _effective_tensor_parallel(["--split-mode", "layer"], False, env = tensor_env) is False + # Explicit extras/toggle tensor stays on regardless of env. + assert _effective_tensor_parallel(["--split-mode", "tensor"], False, env = {}) is True + assert _effective_tensor_parallel(None, True, env = {}) is True + # One-directional: a non-tensor env never downgrades, and no env -> no flip. + assert _effective_tensor_parallel(None, False, env = {}) is False + assert ( + _effective_tensor_parallel(None, False, env = {"LLAMA_ARG_SPLIT_MODE": "layer"}) is False + ) + + def test_tensor_parallel_matches_loaded_env_downgrade(self): + # Env-only tensor matches a server that actually launched tensor, but a + # server load_model downgraded to layer (env scrubbed) must still match + # an identical request -- not reload forever (#6312). + tensor_env = {"LLAMA_ARG_SPLIT_MODE": "tensor"} + # Launched tensor: env-only request matches. + assert _tensor_parallel_matches_loaded(None, False, True, env = tensor_env) is True + # Downgraded to layer: same env-only request still matches (no reload loop). + assert _tensor_parallel_matches_loaded(None, False, False, env = tensor_env) is True + # No env: a plain request matches a layer server and mismatches a tensor one. + assert _tensor_parallel_matches_loaded(None, False, False, env = {}) is True + assert _tensor_parallel_matches_loaded(None, False, True, env = {}) is False + # Explicit tensor request stays strict: must have a tensor server. + assert _tensor_parallel_matches_loaded(None, True, False, env = {}) is False + assert _tensor_parallel_matches_loaded(None, True, True, env = {}) is True + # An explicit non-tensor --split-mode beats the env (no flip). + assert ( + _tensor_parallel_matches_loaded(["--split-mode", "layer"], False, True, env = tensor_env) + is False + ) + + def test_route_matcher_uses_tensor_parallel_matches_loaded(self): + # Fix: the route duplicate-load matcher must use the downgrade-aware + # helper, or an env-driven tensor server (or its layer downgrade) is + # needlessly reloaded (#6312). Read from disk (importing routes.inference + # drags in heavy deps). + routes_src = ( + Path(__file__).resolve().parent.parent / "routes" / "inference.py" + ).read_text() + start = routes_src.index("def _request_matches_loaded_settings") + end = routes_src.index("\ndef ", start + 1) + body = "".join(routes_src[start:end].split()) + assert ( + "_tensor_parallel_matches_loaded(effective_extra," + "request.tensor_parallel,llama_backend.tensor_parallel)" in body + ) + + def test_extra_args_main_cache_type_heavier_axis(self): + # Asymmetric --cache-type-k/-v must budget the heavier axis (extras win + # per axis at launch), not the last-wins single type that under-reserves. + H = _extra_args_main_cache_type_for_budget + assert H(["--cache-type-k", "f32", "--cache-type-v", "f16"]) == "f32" + assert H(["--cache-type-v", "f16", "--cache-type-k", "f32"]) == "f32" # order-free + assert H(["--cache-type-k=f32", "--cache-type-v=f16"]) == "f32" # = form + assert H(["-ctk", "q4_0", "-ctv", "q8_0"]) == "q8_0" # heavier quant + assert H(["--cache-type-k", "q8_0"]) == "q8_0" # single axis honored as-is + assert H(["-c", "4096"]) is None # no cache flags + assert H(None) is None + + def test_load_model_budgets_heavier_asymmetric_cache_axis(self): + # load_model must reserve from the heavier of asymmetric cache extras, or + # an f32 K against an f16 budget over-advertises context and can OOM. + load = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_extra_args_main_cache_type_for_budget(extra_args)" in load + + def test_load_model_tensor_drops_any_quantized_cache_axis(self): + # The heavier-by-bytes budget type can mask a quantized axis (an f16 + # budget hides a paired q4_0), so the tensor-safety drop must test each + # --cache-type-k/-v extra, not just cache_type_kv -- else the quantized + # axis survives into tensor mode and crashes the load (#6312). + load = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_ck_extra,_cv_extra=parse_cache_override_per_axis(extra_args)" in load + assert "forcin(cache_type_kv,_ck_extra,_cv_extra)" in load + assert "iftensor_paralleland_cache_non_tensor_safe:" in load + + def test_load_model_layer_downgrade_restores_original_cache_extras(self): + # Tensor mode strips asymmetric --cache-type-k/-v (it rejects quantized), + # but layer split supports them, so a downgrade must restore the ORIGINAL + # extras, not just the scalar heavier type (else q4_0/f16 silently becomes + # f16/f16 on the layer fallback) (#6312). + load = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_tensor_dropped_extra_args=list(extra_args)" in load + # Both tensor->layer downgrade points restore the saved originals. + assert load.count("strip_split_mode_only(_tensor_dropped_extra_argsif") == 2 + + def test_load_model_tensor_skips_reserve_for_cpu_drafter(self): + # A separate CPU-offloaded drafter (no embedded head) uses no GPU, so the + # tensor reserve must be suppressed like the layer path -- else tensor mode + # subtracts a phantom flat MTP reserve and under-advertises context (#6312). + load = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_mtp_will_engageandnot_draft_cpu_no_embedded" in load + assert "ifnot_mtp_reserves_gpu:" in load + assert "mtp_engaged=_mtp_reserves_gpu" in load + + def test_load_model_adopts_env_tensor_split_mode(self): + # load_model delegates the tensor decision to _effective_tensor_parallel, + # which flips to tensor only one-directionally: extras set no split mode, + # none is overridden, and the env selects tensor (an existing tensor plan + # is never downgraded). Whitespace-stripped to survive formatter wrapping. + load = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "tensor_parallel=_effective_tensor_parallel(extra_args,tensor_parallel)" in load + helper = "".join(inspect.getsource(_effective_tensor_parallel).split()) + assert "notresolved" in helper + assert "parse_split_mode_override(extra_args)isNone" in helper + assert "_env_split_mode_is_tensor(env)" in helper + + def test_load_model_does_not_emit_env_only_cache_type(self): + # Cluster C: an env-only (budget) cache type must not be re-emitted as + # --cache-type flags (that would rewrite an asymmetric K/V env). Emission + # is guarded by `not _cache_type_from_env`, set when the value came from + # _env_main_cache_type_for_budget(). Whitespace-stripped for formatter. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "cache_type_kv=_env_main_cache_type_for_budget()" in compact + assert "_cache_type_from_env=cache_type_kvisnotNone" in compact + assert "andnot_cache_type_from_env" in compact + + def test_load_model_clears_inherited_split_mode_on_layer(self): + # Cluster A: when the final decision is layer split, an inherited + # non-layer LLAMA_ARG_SPLIT_MODE (and paired LLAMA_ARG_TENSOR_SPLIT) must + # be popped from the child env so the child cannot run tensor/row/none + # against Studio's layer budget. Whitespace-stripped for formatter. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert 'env.get("LLAMA_ARG_SPLIT_MODE")' in compact + assert '_inherited_sm!="layer"' in compact + assert 'env.pop("LLAMA_ARG_SPLIT_MODE",None)' in compact + assert 'env.pop("LLAMA_ARG_TENSOR_SPLIT",None)' in compact + + def test_load_model_clears_quantized_kv_env_for_tensor(self): + # Cluster B: tensor mode aborts on quantized KV. An inherited quantized + # LLAMA_ARG_CACHE_TYPE_K/_V must be popped from the child env so it cannot + # crash the tensor child (and matches the tensor-safe budget). + # Whitespace-stripped for formatter. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert '("LLAMA_ARG_CACHE_TYPE_K","LLAMA_ARG_CACHE_TYPE_V")' in compact + assert "_ct_rawnotinself._TENSOR_PARALLEL_KV_TYPES" in compact + assert "env.pop(_ct_var,None)" in compact + + def test_load_model_clears_tensor_split_env_in_tensor_mode(self): + # review run3 #2: Studio owns the tensor split. When it emits no + # --tensor-split (even split), a stale inherited LLAMA_ARG_TENSOR_SPLIT must + # be cleared in the TENSOR branch too (not just the layer downgrade), or the + # child runs a split Studio didn't budget. The else (tensor) branch pops it. + src = inspect.getsource(LlamaCppBackend.load_model) + compact = "".join(src.split()) + # appears in both the layer branch and the tensor branch. + assert compact.count('env.pop("LLAMA_ARG_TENSOR_SPLIT",None)') >= 2 + + def test_load_model_layer_compute_buffer_fallback(self): + # review run3 #4: when GGUF dims are missing the compute-buffer estimate is + # 0; the layer path must still reserve the flat fallback (tensor buffer >= + # layer buffer), not fold 0, or it under-reserves at high --parallel. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "if_compute_buffer_pipeline<=0:" in compact + # The RHS expression only (no `_compute_buffer_pipeline=` prefix) so the + # match survives the formatter wrapping it in parens. It's unique within + # load_model (the other use of this constant is in MiB, no *1024*1024). + assert "self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB*1024*1024" in compact + + def test_load_model_passes_unsized_mtp_reserve_to_tensor_planner(self): + # review run3 #1/#5: a weights-only (KV-unsized) MTP reserve must flow into + # _plan_tensor_parallel as a flat cushion, else its binary search spends the + # unsized draft KV on context and OOMs. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_tp_unsized_mtp_reserve=" in compact + # Gated on _mtp_reserves_gpu so a CPU-offloaded drafter reserves nothing. + assert "(_mtp_reserves_gpuand_mtp_kv_unsized)" in compact + assert "mtp_flat_reserve_bytes=_tp_unsized_mtp_reserve" in compact + + def test_pool_budget_sums_per_gpu_usable(self): + # Finding #1: the multi-GPU pooled budget must sum each GPU's own usable + # budget (so an unknown-total GPU gets the free*frac cushion) rather than + # pooling free and total separately. The fit calls pass the precomputed + # budget as an absolute (budget_frac=1.0, total_mib=None) so fit and check + # agree. Whitespace-stripped for formatter. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "def_pool_budget_mib(subset,frac):" in compact + assert "sum(max(0.0,_gpu_usable(g,frac))forginsubset)" in compact + # No revert to the pooled free/total form. + assert "def_pool_total(" not in compact + assert "budget_frac=1.0" in compact + + +# --------------------------------------------------------------------------- +# Regression: the reported Qwen3.6-27B MTP / 24 GB scenario +# --------------------------------------------------------------------------- + + +def test_qwen36_class_regression_picks_lower_ctx_with_mtp(): + """A 24 GB card that auto-picks a high context without MTP must pick a + strictly lower one once the MTP draft reserve is accounted for.""" + b = _make_backend() + b._can_estimate_kv = lambda: True + b._estimate_kv_cache_bytes = lambda n, _t = None, **_k: (0 if n <= 0 else int(n * 66_000)) + avail_mib = 24_000 + model = int(17.9 * GIB) # UD-Q4_K_XL weights + no_mtp = b._fit_context_to_vram(262144, avail_mib, model) + with_mtp = b._fit_context_to_vram( + 262144, + avail_mib, + model, + mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes(c) or 0, + ) + assert 0 < with_mtp < no_mtp + + +def test_mtp_draft_budget_prefers_user_extras_drafter(): + # A user --model-draft in extras is appended last and wins at launch, so the + # VRAM budget must size it first; then Studio's emitted mtp_draft_path (which + # overrides LLAMA_ARG_SPEC_DRAFT_MODEL), then the env drafter (load_model is too + # entangled to drive end-to-end; assert the precedence at the source level). + # Whitespace-stripped so the check survives any formatter line-wrapping. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + # CLI extras sized first (env={} so the env doesn't pre-empt Studio's drafter). + assert "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" in compact + # Order: CLI extras, then Studio's mtp_draft_path, then the env drafter. + assert "_cli_draft_for_budgetor_studio_draft_for_budgetor_env_draft_for_budget" in compact + # The env must not be consulted before Studio's resolved drafter. + assert "_extra_args_mtp_draft_path(extra_args)ormtp_draft_path" not in compact diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 30bfb91a08..8ec629bd2c 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -62,8 +62,11 @@ _httpx_stub.Client = type( sys.modules.setdefault("httpx", _httpx_stub) from core.inference import llama_cpp as llama_cpp_module -from core.inference.llama_cpp import LlamaCppBackend -from core.inference.llama_server_args import resolve_tensor_parallel +from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend +from core.inference.llama_server_args import ( + _effective_tensor_parallel, + resolve_tensor_parallel, +) from core.inference.tensor_fallback import load_with_tensor_fallback from models.inference import ( InferenceStatusResponse, @@ -330,16 +333,21 @@ def _plan( def _kv_budget_b(model_gb, gpus = _ASYM): + # No totals here, so usable is the legacy free*frac (keeps the 5% cushion). reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB - return (sum(f for _, f in gpus) - len(gpus) * reserve) * 1024 * 1024 - int(model_gb * _GB) + usable = sum(f * _CTX_FIT_VRAM_FRACTION for _, f in gpus) + return (usable - len(gpus) * reserve) * 1024 * 1024 - int(model_gb * _GB) def test_tp_plan_weighted_split_on_asymmetric_big_model(): b, (ec, mac, gi, ts) = _plan(50) reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB assert gi == [0, 1] - # split weighted by (free - buffer), not raw free - assert ts == [48000 - reserve, 24000 - reserve] + # split weighted by (usable - buffer); with no totals usable is free*frac + assert ts == [ + int(48000 * _CTX_FIT_VRAM_FRACTION - reserve), + int(24000 * _CTX_FIT_VRAM_FRACTION - reserve), + ] assert ec < 131072 # capped below native @@ -442,11 +450,9 @@ def test_tp_plan_drops_gpu_below_buffer_reserve(): # ── route auto-fallback survives a *raised* tensor-load crash ───────── -# A tensor-incompatible model makes load_model RAISE (Gemma 3n aborts) rather -# than return False. The /load fallback helper must catch that and retry with -# layer split -- stripping any --split-mode from the extras so the retry can't -# relaunch tensor -- while a non-tensor load propagates its exception. These -# exercise the real helper with a fake loader (no GPU, no llama-server). +# A tensor-incompatible model makes load_model RAISE (not return False); the +# /load fallback must catch it and retry with layer split (stripping --split-mode +# so the retry can't relaunch tensor), while a non-tensor load propagates. class _RecordingLoader: @@ -555,15 +561,46 @@ def test_tensor_fallback_skips_layer_retry_when_cancelled(): ) def test_tensor_fallback_strips_split_mode_from_extras_on_retry(extras): # Tensor engaged via extras (boolean False); the retry must drop every - # --split-mode form (long/short, space/=) but keep the user's other flags, - # else resolve_tensor_parallel re-enables tensor and relaunches the crash. + # --split-mode form (long/short, space/=) and force layer, keeping the user's + # other flags, else tensor is re-enabled and relaunches the crash. loader = _RecordingLoader() ok = asyncio.run( load_with_tensor_fallback(loader, requested_tensor = False, extra_args = extras, label = "m") ) assert ok is True assert len(loader.calls) == 2 - assert loader.calls[1][1] == ["-c", "4096"] # split-mode stripped, -c kept + # User --split-mode replaced by an explicit layer override; -c kept. + assert loader.calls[1][1] == ["-c", "4096", "--split-mode", "layer"] + + +def test_tensor_fallback_env_tensor_retry_forces_layer(monkeypatch): + # Env-only tensor (toggle off, no --split-mode extra): load_model engages + # tensor via LLAMA_ARG_SPLIT_MODE and a tensor-incompatible model crashes. The + # wrapper must (1) recognise the env tensor request and retry, and (2) force + # --split-mode layer so the retry doesn't re-engage tensor via the still-set + # env and crash again (#6312). + monkeypatch.setenv("LLAMA_ARG_SPLIT_MODE", "tensor") + calls: list = [] + + async def _crash_when_effectively_tensor(tensor_parallel, extra_args): + calls.append(list(extra_args) if extra_args else extra_args) + # Mirror real load_model: env-aware tensor engagement crashes. + if _effective_tensor_parallel(extra_args, tensor_parallel): + raise RuntimeError("llama-server failed to start (tensor)") + return True + + ok = asyncio.run( + load_with_tensor_fallback( + _crash_when_effectively_tensor, + requested_tensor = False, + extra_args = None, + label = "m", + ) + ) + assert ok is True + assert len(calls) == 2 + # The forced layer override neutralises the inherited tensor env on retry. + assert calls[1] == ["--split-mode", "layer"] def test_tensor_fallback_propagates_non_tensor_crash(): @@ -576,3 +613,143 @@ def test_tensor_fallback_propagates_non_tensor_crash(): _always_raise, requested_tensor = False, extra_args = None, label = "m" ) ) + + +# ── _plan_tensor_parallel: total-based headroom + ubatch (review fixes) ── + + +def test_tensor_caps_context_to_total_vram_budget(): + # Partly-used 80 GB cards: 20 GB free each. With total_by_idx the planner must + # cap occupancy at 0.95*total (not spend the cushion the layer-split paths keep). + b = _kv_seeded_backend() + gpus = [(0, 20000), (1, 20000)] + totals = {0: 81920, 1: 81920} + model = int(18 * _GB) + with_total, *_ = b._plan_tensor_parallel(gpus, model, 131072, total_by_idx = totals) + without, *_ = b._plan_tensor_parallel(gpus, model, 131072) + assert with_total < without # total cap tightens the chosen context + + MIB = 1024 * 1024 + reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB # flat (no vocab dims) + pool_usable = sum(f - (1.0 - _CTX_FIT_VRAM_FRACTION) * totals[i] for i, f in gpus) + foot_total = (model + b._estimate_kv_cache_bytes(with_total, None)) / MIB + len(gpus) * reserve + foot_free = (model + b._estimate_kv_cache_bytes(without, None)) / MIB + len(gpus) * reserve + assert foot_total <= pool_usable + 2 # fix: fits the total-based budget + assert foot_free > pool_usable # old behavior over-spent the cushion + + +def test_tensor_unknown_total_keeps_fraction_cushion(): + # A two-column nvidia-smi probe yields total 0. The planner must fall back to + # free*frac (keep the 5% cushion), like _select_gpus/_gpu_usable, not raw free, + # or it over-advertises context exactly where the PR is hardening the budget. + b = _kv_seeded_backend() + gpus = [(0, 20000), (1, 20000)] + MIB = 1024 * 1024 + reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + model = int(18 * _GB) + ec_zero, *_ = b._plan_tensor_parallel(gpus, model, 131072, total_by_idx = {0: 0, 1: 0}) + ec_none, *_ = b._plan_tensor_parallel(gpus, model, 131072) + assert ec_zero == ec_none # total 0 == total absent: both use free*frac + pool_free = sum(f for _, f in gpus) + foot = (model + b._estimate_kv_cache_bytes(ec_zero, None)) / MIB + len(gpus) * reserve + assert foot <= pool_free * _CTX_FIT_VRAM_FRACTION + 2 # within free*frac, not raw free + + +def test_tensor_reserve_scales_with_ubatch(): + # A user --ubatch override must enlarge the per-device reserve -> less ctx room. + b = _kv_seeded_backend() + b._vocab_size = 152064 # enable the deterministic compute-buffer estimate + gpus = [(0, 16000), (1, 16000)] + model = int(18 * _GB) + small_ub, *_ = b._plan_tensor_parallel(gpus, model, 131072, n_ubatch = 512) + big_ub, *_ = b._plan_tensor_parallel(gpus, model, 131072, n_ubatch = 4096) + assert big_ub < small_ub + + +def test_plan_tensor_carries_unsized_mtp_flat_reserve(): + # review run3 #1/#5: with a weights-only (KV-unsized) MTP reserve, the planner + # gets a non-None mtp_overhead_fn but must still subtract the flat unsized-KV + # cushion, or its binary search spends it on context. Passing the reserve must + # pick a strictly smaller context than passing 0. + b = _kv_seeded_backend() + gpus = [(0, 14000), (1, 14000)] # tight pool so the context is actually capped + model = int(8 * _GB) + weights_only = lambda c: 3 * _GB # noqa: E731 -- constant drafter weights, no KV term + ctx_no_flat, *_ = b._plan_tensor_parallel( + gpus, + model, + 131072, + mtp_engaged = True, + mtp_overhead_fn = weights_only, + mtp_flat_reserve_bytes = 0, + ) + ctx_flat, *_ = b._plan_tensor_parallel( + gpus, + model, + 131072, + mtp_engaged = True, + mtp_overhead_fn = weights_only, + mtp_flat_reserve_bytes = 2 * _GB, + ) + assert 0 < ctx_flat < ctx_no_flat + + +def test_tensor_admission_drops_gpu_below_usable_budget(): + # A partly-used big card can clear the buffer reserve on raw free yet have no + # usable budget left (free - 0.05*total). Admit by usable budget: GPU 0 here is + # 6000 free on an 80 GB card -> usable 1904 < flat reserve 5120, so it's dropped + # (leaving <2 -> no split). Without total_by_idx, raw free 6000 >= 5120 admits it. + b = _kv_seeded_backend() + gpus = [(0, 6000), (1, 40000)] + totals = {0: 81920, 1: 81920} + _ec, _mac, gi, ts = b._plan_tensor_parallel(gpus, int(8 * _GB), 8192, total_by_idx = totals) + assert gi == [1] and ts is None # GPU 0 excluded on usable budget + _ec2, _mac2, gi_raw, _ts2 = b._plan_tensor_parallel(gpus, int(8 * _GB), 8192) + assert gi_raw == [0, 1] # raw free would have admitted both + + +def test_load_model_tensor_admission_and_capacity_gate_use_usable_budget(): + # load_model is too entangled (subprocess + GPU probe) to drive end-to-end, so + # assert at the source level that the tensor prefilter admits on the usable + # budget (_gpu_usable), not raw free, and downgrades to layer split when the + # pooled budget can't hold weights + per-device compute buffers. + src = inspect.getsource(LlamaCppBackend.load_model) + assert "_gpu_usable(g) >= reserve_mib" in src # admit by usable budget + assert "g[1] >= reserve_mib" not in src # not raw free + assert "_tp_weight_budget_mib" in src # pooled-weight capacity gate + assert "falling back to layer split" in src # downgrade on overcommit + # The gate's required footprint must include the non-shrinkable MTP reserve, + # not weights alone, or a separate-drafter MTP load can still overcommit. + assert "_tp_mtp_floor" in src + assert "model_size + _tp_mtp_floor" in src + + +def test_load_model_tensor_floor_keeps_flat_reserve_for_weights_only(): + # Tensor mode has no --fit valve, so a weights-only drafter (KV unsized) must + # keep the flat reserve as the draft-KV cushion, not just the byte weights + # (Finding H1, the tensor analog of the layer-split _mtp_kv_unsized handling). + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + # byte-only floor used only when KV is sizable (not the weights-only case) + assert "mtp_overhead_fnisnotNoneandnot_mtp_kv_unsized" in compact + # weights-only / dims-unavailable: flat reserve, never below the byte floor + assert "_tp_mtp_floor=max(" in compact + + +def test_load_model_reserves_pipeline_per_device_overhead(): + # Layer split must reserve the fixed per-device overhead per EXTRA device so a + # tight multi-GPU split can't pin a context that OOMs a device (Finding A); k=1 + # adds nothing. + assert LlamaCppBackend._PIPELINE_PER_DEVICE_OVERHEAD_MIB > 0 + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "def_subset_model_size(n_gpus:int)->int:" in compact + assert "max(0,n_gpus-1)*_pipeline_overhead_bytes" in compact + assert "_subset_model_size(n_gpus)" in compact # used in the layer-split fit + + +def test_load_model_restores_quantized_kv_on_tensor_downgrade(): + # A quantized KV dropped for the tensor attempt must be restored if tensor + # downgrades to layer split (Finding D); captured once, restored at both the + # GPU-count and capacity-gate downgrades. + compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split()) + assert "_tensor_dropped_cache_type_kv=cache_type_kv" in compact # captured pre-null + assert compact.count("cache_type_kv=_tensor_dropped_cache_type_kv") >= 2 # restored diff --git a/studio/backend/tests/test_windows_gpu_detection_mock.py b/studio/backend/tests/test_windows_gpu_detection_mock.py index 88a1a28d14..a4ec3f3fb5 100644 --- a/studio/backend/tests/test_windows_gpu_detection_mock.py +++ b/studio/backend/tests/test_windows_gpu_detection_mock.py @@ -211,6 +211,22 @@ class TestWindowsGpuDetectionAfter5106Fix: gpus = LlamaCppBackend._get_gpu_free_memory() assert gpus == [(1, 24576)], gpus + def test_get_gpu_memory_parses_three_and_two_column(self, monkeypatch): + """Total is parsed when present; a legacy two-column line or a non-integer + total ("N/A") yields total 0 (back-compat) rather than dropping the GPU, + which would silently spill to CPU.""" + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + with _mock_nvidia_smi_run("0, 22805, 24576\n"): + assert LlamaCppBackend._get_gpu_memory() == [(0, 22805, 24576)] + with _mock_nvidia_smi_run("0, 22805\n"): + assert LlamaCppBackend._get_gpu_memory() == [(0, 22805, 0)] + # A non-integer total must keep the GPU (total 0), not drop it. + with _mock_nvidia_smi_run("0, 22805, N/A\n"): + assert LlamaCppBackend._get_gpu_memory() == [(0, 22805, 0)] + # A bad free still skips that line (free is required). + with _mock_nvidia_smi_run("0, N/A, 24576\n1, 22805, 24576\n"): + assert LlamaCppBackend._get_gpu_memory() == [(1, 22805, 24576)] + def test_windows_install_dir_has_all_three_cudart_dlls(self, tmp_path): """All three bundle DLLs must land in install_dir/build/bin/ Release; any missing one breaks ggml-cuda.dll's PE import chain."""