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."""