diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index be0b1596ad..47e46405be 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -43,6 +43,7 @@ import httpx from core.inference.llama_server_args import ( _LAYER_OFFLOAD_FLAGS, _effective_tensor_parallel, + _flag_name, _tensor_parallel_matches_loaded, extra_args_disable_mmproj, parse_cache_override, @@ -1509,6 +1510,21 @@ def _kv_bytes_per_elem(cache_type: Optional[str]) -> float: }.get((cache_type or "f16").strip().lower(), 2.0) +def _pad_kv_cells(cells: int) -> int: + return ((cells + 255) // 256) * 256 + + +def _kv_cache_cell_layout(n_ctx: int, n_parallel: int, kv_unified: bool) -> tuple[int, int, int]: + """Return llama.cpp's slot count, stream count, and cells per stream.""" + slots = max(1, n_parallel) + padded_ctx = _pad_kv_cells(n_ctx) + streams = 1 if kv_unified else slots + if padded_ctx <= 0: + return slots, streams, 0 + cells_per_stream = padded_ctx if kv_unified else _pad_kv_cells(padded_ctx // slots) + return slots, streams, cells_per_stream + + 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. Unsloth emits --cache-type only for the @@ -1541,6 +1557,39 @@ def _extra_args_main_cache_type_for_budget(extra_args: Optional[Iterable[str]]) return max(candidates, key = _kv_bytes_per_elem) +def _effective_main_cache_types( + args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> tuple[str, str]: + """Effective main K/V cache types after environment and CLI precedence.""" + source_env = os.environ if env is None else env + env_k = (source_env.get("LLAMA_ARG_CACHE_TYPE_K") or "f16").strip().lower() + env_v = (source_env.get("LLAMA_ARG_CACHE_TYPE_V") or "f16").strip().lower() + arg_k, arg_v = parse_cache_override_per_axis(args) + return ( + (arg_k or env_k).strip().lower(), + (arg_v or env_v).strip().lower(), + ) + + +def _planned_main_cache_types( + cache_type_kv: Optional[str], + extra_args: Optional[Iterable[str]], + env: Optional[Mapping[str, str]] = None, +) -> tuple[str, str]: + """Main K/V types the loader's managed flags and user extras will produce.""" + args = list(extra_args or ()) + emitted_type = _extra_args_main_cache_type_for_budget(args) or cache_type_kv + if emitted_type: + args = [ + "--cache-type-k", + emitted_type, + "--cache-type-v", + emitted_type, + *args, + ] + return _effective_main_cache_types(args, env) + + def _auto_mode_drops_mtp( req_mode: Optional[str], size_b: Optional[float], @@ -1584,26 +1633,79 @@ def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool: # set keeps detection and stripping from drifting. _GPU_OFFLOAD_OVERRIDE_FLAGS = _LAYER_OFFLOAD_FLAGS _THREAD_OVERRIDE_FLAGS = frozenset({"-t", "--threads"}) - - -def _extra_arg_flag_name(token: str) -> Optional[str]: - if not token.startswith("-") or token in {"-", "--"}: - return None - if len(token) >= 2 and (token[1].isdigit() or token[1] == "."): - return None - return token.split("=", 1)[0] +# common_params defaults in the bundled llama.cpp runtime. +_DEFAULT_LLAMA_N_BATCH = 2048 +_DEFAULT_LLAMA_N_UBATCH = 512 +_LLAMA_ARG_TRUE_VALUES = frozenset({"on", "enabled", "true", "1"}) +_LLAMA_ARG_FALSE_VALUES = frozenset({"off", "disabled", "false", "0"}) +_LLAMA_ARG_AUTO_VALUES = frozenset({"auto", "-1"}) +_LLAMA_ARG_TRUE_OR_AUTO_VALUES = _LLAMA_ARG_TRUE_VALUES | _LLAMA_ARG_AUTO_VALUES +_LLAMA_ARG_TRUE_FALSE_AUTO_VALUES = _LLAMA_ARG_TRUE_OR_AUTO_VALUES | _LLAMA_ARG_FALSE_VALUES def _extra_args_set_any_flag(extra_args: Optional[Iterable[str]], flags: Collection[str]) -> bool: if not extra_args: return False for raw in extra_args: - flag = _extra_arg_flag_name(str(raw)) + flag = _flag_name(str(raw)) if flag in flags: return True return False +def _swa_full_from_args_or_env( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """Whether llama.cpp receives the enable-only full-size SWA option.""" + if _extra_args_set_any_flag(extra_args, {"--swa-full"}): + return True + value = (os.environ if env is None else env).get("LLAMA_ARG_SWA_FULL") + return value in _LLAMA_ARG_TRUE_VALUES + + +def _kv_unified_from_args( + extra_args: Optional[Iterable[str]], + default: bool = False, + env: Optional[Mapping[str, str]] = None, +) -> bool: + """Resolve llama.cpp's environment and last-wins unified KV flags.""" + enabled = False + value = (os.environ if env is None else env).get("LLAMA_ARG_KV_UNIFIED") + if value in _LLAMA_ARG_TRUE_VALUES: + enabled = True + elif value in _LLAMA_ARG_FALSE_VALUES: + enabled = False + if default: + # Studio's managed --kv-unified flag is appended after environment + # parsing and before user extras. + enabled = True + for raw in extra_args or (): + flag = _flag_name(str(raw)) + if flag in {"-kvu", "--kv-unified"}: + enabled = True + elif flag in {"-no-kvu", "--no-kv-unified"}: + enabled = False + return enabled + + +def _flash_attn_enabled_from_args(args: Optional[Iterable[str]], default: bool = True) -> bool: + """Resolve llama.cpp's last-wins flash-attention CLI setting.""" + enabled = default + values = [str(arg) for arg in args] if args else [] + for i, raw in enumerate(values): + if _flag_name(raw) not in {"-fa", "--flash-attn"}: + continue + _, eq, inline = raw.partition("=") + value = inline if eq else "on" + if not eq and i + 1 < len(values) and values[i + 1] in _LLAMA_ARG_TRUE_FALSE_AUTO_VALUES: + value = values[i + 1] + if value in _LLAMA_ARG_FALSE_VALUES: + enabled = False + elif value in _LLAMA_ARG_TRUE_OR_AUTO_VALUES: + enabled = True + return enabled + + def _effective_spec_type( extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None ) -> Optional[str]: @@ -1615,7 +1717,8 @@ def _effective_spec_type( cli_present = False cli_value: Optional[str] = None for i, raw in enumerate(args): - flag, eq, inline = raw.partition("=") + flag = _flag_name(raw) + _, eq, inline = raw.partition("=") if flag == "--spec-default": cli_present = True cli_value = "default" @@ -1659,7 +1762,8 @@ def _extra_args_spec_draft_n_max(extra_args: Optional[Iterable[str]]) -> Optiona args = [str(a) for a in extra_args] found: Optional[int] = None for i, raw in enumerate(args): - flag, eq, inline = raw.partition("=") + flag = _flag_name(raw) + _, 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 "") @@ -1689,7 +1793,8 @@ def _extra_args_mtp_draft_path( 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("=") + flag = _flag_name(raw) + _, eq, inline = raw.partition("=") if flag not in flags: continue value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") @@ -1713,7 +1818,8 @@ def _extra_args_draft_cache_types( k_type: Optional[str] = None v_type: Optional[str] = None for i, raw in enumerate(args): - flag, eq, inline = raw.partition("=") + flag = _flag_name(raw) + _, 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 "") @@ -1745,7 +1851,8 @@ def _extra_args_draft_offloaded_to_cpu( last_ngl: Optional[str] = None last_dev: Optional[str] = None for i, raw in enumerate(args): - flag, eq, inline = raw.partition("=") + flag = _flag_name(raw) + _, 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 @@ -1767,31 +1874,61 @@ def _extra_args_draft_offloaded_to_cpu( def _extra_args_n_ubatch( - extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None + extra_args: Optional[Iterable[str]], + env: Optional[Mapping[str, str]] = None, + n_ctx: Optional[int] = 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.""" + """Effective ubatch after llama.cpp normalizes it, or None at defaults.""" + values = { + "batch": _DEFAULT_LLAMA_N_BATCH, + "ubatch": _DEFAULT_LLAMA_N_UBATCH, + } + source_env = os.environ if env is None else env + overridden = False + for key, env_name in ( + ("batch", "LLAMA_ARG_BATCH"), + ("ubatch", "LLAMA_ARG_UBATCH"), + ): + raw = source_env.get(env_name) + if raw: + try: + values[key] = int(raw) + overridden = True + except (TypeError, ValueError): + pass + args = [str(a) for a in extra_args] if extra_args else [] - found: Optional[int] = None + flags = { + "-b": "batch", + "--batch-size": "batch", + "-ub": "ubatch", + "--ubatch-size": "ubatch", + } for i, raw in enumerate(args): - flag, eq, inline = raw.partition("=") - if flag not in ("--ubatch-size", "-ub"): + flag = _flag_name(raw) + _, eq, inline = raw.partition("=") + key = flags.get(flag) + if key is None: continue value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") try: - found = int(value) + values[key] = int(value) + overridden = True 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 + if not overridden: + return None + + # common_params stores signed values, then llama_context_params converts + # them to uint32_t. A zero ubatch means "use batch"; the context then caps + # ubatch at batch size. + batch = values["batch"] & 0xFFFFFFFF + raw_ubatch = values["ubatch"] + ubatch = batch if raw_ubatch == 0 else raw_ubatch & 0xFFFFFFFF + effective = min(batch, ubatch) + if n_ctx is not None and n_ctx > 0: + effective = min(effective, n_ctx) + return effective def _build_ngram_mod_flags( @@ -2150,6 +2287,14 @@ class LlamaCppBackend: # save can tell whether the model files were swapped on disk since load. self._slot_loaded_identity: Optional[tuple] = None self._prompt_cache_disabled: bool = False + self._swa_full: bool = False + self._kv_cache_unified: bool = False + self._n_ubatch: int = self._DEFAULT_N_UBATCH + self._flash_attn_enabled: bool = True + self._effective_cache_types: tuple[str, str] = ("f16", "f16") + # Total KV allocation context across all slots. _effective_context_length + # becomes the per-slot request limit after /props reconciliation. + self._kv_cache_context_total: Optional[int] = None # True once a probe has completed; cleared on transient failure. self._is_audio: bool = False self._audio_type: Optional[str] = None @@ -2202,6 +2347,11 @@ class LlamaCppBackend: """True when the loaded GGUF is a block-diffusion model (DiffusionGemma).""" return self._is_diffusion + @property + def swa_full(self) -> bool: + """Whether the active llama-server received full-size SWA mode.""" + return self._swa_full + @property def hf_variant(self) -> Optional[str]: return self._hf_variant @@ -4057,6 +4207,32 @@ class LlamaCppBackend: is non-None here.""" return self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] + def _max_kv_value_width( + self, + default_len: int, + swa_len: Optional[int] = None, + ) -> int: + """llama.cpp's hparams.n_embd_v_gqa_max() over every model layer.""" + n_layers = self._n_layers or 1 + n_kv = self._n_kv_heads or self._n_heads or 1 + if self._sliding_window_pattern is None: + max_len = max(default_len, swa_len or default_len) + return max( + self._kv_heads_for_layer(layer_idx, n_kv) * max_len for layer_idx in range(n_layers) + ) + return max( + self._kv_heads_for_layer(layer_idx, n_kv) + * ( + (swa_len or default_len) + if ( + layer_idx < len(self._sliding_window_pattern) + and self._sliding_window_pattern[layer_idx] + ) + else default_len + ) + for layer_idx in range(n_layers) + ) + def _estimate_kv_cache_bytes( self, n_ctx: int, @@ -4065,22 +4241,26 @@ class LlamaCppBackend: swa_full: bool = False, n_parallel: int = 1, kv_unified: bool = True, + n_ubatch: Optional[int] = None, ctx_checkpoints: int = 0, + flash_attn: bool = True, ) -> int: """Estimate KV cache VRAM for a given context length. 5-path architecture-aware estimation: 1. MLA -- compressed KV latent + RoPE, K-only (no separate V) 2. Hybrid -- only attention layers need KV (Mamba layers don't) - 3. SWA -- sliding-window layers cache min(ctx, window) tokens + 3. SWA -- sliding-window layers use compact or full cache cells 4. GQA -- standard full KV with explicit key/value dimensions 5. Legacy -- fallback using embed // n_heads Server-flag knobs (mirror llama-server's CLI): swa_full -- --swa-full: 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). + n_parallel -- --parallel slots: controls per-slot stream padding. + kv_unified -- --kv-unified: one shared stream vs one per slot. + n_ubatch -- --ubatch-size: SWA cache's processing headroom. ctx_checkpoints -- --ctx-checkpoints: N SWA snapshots per slot. + flash_attn -- False pads variable-width V tensors to the model max. Returns 0 if metadata is insufficient. """ @@ -4095,9 +4275,17 @@ 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 = _kv_bytes_per_elem(cache_type_kv) + bpe_k = _kv_bytes_per_elem(cache_type_kv) + # The automatic FA-off retry rewrites an invalid quantized V cache to + # f16. Pricing that viable retry here avoids under-reserving it. + bpe_v = bpe_k if flash_attn else max(bpe_k, _kv_bytes_per_elem("f16")) - slots = max(1, n_parallel) + slots, streams, cells_per_stream = _kv_cache_cell_layout(n_ctx, n_parallel, kv_unified) + total_cells = cells_per_stream * streams + ubatch = max( + 0, + int(self._DEFAULT_N_UBATCH if n_ubatch is None else n_ubatch), + ) # Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5) # One compressed KV latent per token/layer (shared across heads); V is @@ -4108,7 +4296,7 @@ class LlamaCppBackend: n_kv_mla = self._n_kv_heads or 1 rope_dim = self._key_length_mla or 64 key_len = self._kv_key_length or (self._kv_lora_rank + rope_dim) - return int(n_layers_kv * n_ctx * n_kv_mla * key_len * bpe) + return int(n_layers_kv * total_cells * n_kv_mla * key_len * bpe_k) key_len = self._kv_key_length val_len = self._kv_value_length @@ -4119,16 +4307,18 @@ class LlamaCppBackend: fai = self._full_attention_interval n_attn = -(-n_layers // fai) if fai > 0 else n_layers # ceiling division if key_len is not None and val_len is not None: - return int(n_attn * n_ctx * n_kv * (key_len + val_len) * bpe) + v_width = n_kv * val_len if flash_attn else self._max_kv_value_width(val_len) + return int(n_attn * total_cells * (n_kv * key_len * bpe_k + v_width * bpe_v)) head_dim = self._legacy_head_dim() - return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe) + return int(n_attn * total_cells * n_kv * 2 * head_dim * bpe_k) # 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. + # non-SWA cells total n_ctx across streams. Compact SWA adds one processing + # micro-batch to the window allowance and pads to 256 cells; unified mode + # holds all slots in one stream, while non-unified mode has one stream per + # slot. --swa-full expands SWA to each stream's full context. if ( self._sliding_window is not None and self._sliding_window > 0 @@ -4136,15 +4326,19 @@ class LlamaCppBackend: and val_len is not None ): swa = self._sliding_window - per_slot_ctx = max(1, n_ctx // slots) - # --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) + if swa_full: + swa_cells_total = total_cells + else: + swa_limit = swa * (slots if kv_unified else 1) + ubatch + swa_cells_per_stream = min(cells_per_stream, swa_limit) + swa_cells_per_stream = _pad_kv_cells(swa_cells_per_stream) + swa_cells_total = swa_cells_per_stream * streams key_len_swa = self._kv_key_length_swa or key_len val_len_swa = self._kv_value_length_swa or val_len + padded_v_width = None if flash_attn else self._max_kv_value_width(val_len, val_len_swa) if self._sliding_window_pattern is not None: - global_bytes = 0.0 # constant across slots - swa_bytes_per_slot = 0.0 # multiplied by slots + global_bytes = 0.0 + swa_bytes = 0.0 checkpoint_extra_per_slot = 0.0 # Only layers that allocate their own KV; trailing shared layers # reuse earlier caches. @@ -4154,41 +4348,48 @@ class LlamaCppBackend: layer_idx < len(self._sliding_window_pattern) and self._sliding_window_pattern[layer_idx] ) + layer_key_bytes = layer_n_kv * (key_len_swa if is_swa else key_len) * bpe_k + layer_value_bytes = ( + layer_n_kv * (val_len_swa if is_swa else val_len) + if padded_v_width is None + else padded_v_width + ) * bpe_v + layer_kv_bytes = layer_key_bytes + layer_value_bytes if is_swa: - swa_bytes_per_slot += ( - swa_cells_per_slot * layer_n_kv * (key_len_swa + val_len_swa) * bpe - ) + swa_bytes += swa_cells_total * layer_kv_bytes if ctx_checkpoints > 0 and not swa_full: - checkpoint_extra_per_slot += ( - ctx_checkpoints - * swa - * layer_n_kv - * (key_len_swa + val_len_swa) - * bpe - ) + checkpoint_extra_per_slot += ctx_checkpoints * swa * layer_kv_bytes else: - global_bytes += n_ctx * layer_n_kv * (key_len + val_len) * bpe - return int(global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot)) + global_bytes += total_cells * layer_kv_bytes + return int(global_bytes + swa_bytes + slots * checkpoint_extra_per_slot) n_global = max(1, n_layers_kv // 4) n_swa = n_layers_kv - n_global - kv_per_token = n_kv * (key_len + val_len) * bpe - kv_per_token_swa = n_kv * (key_len_swa + val_len_swa) * bpe - global_bytes = n_global * n_ctx * kv_per_token - swa_bytes_per_slot = n_swa * swa_cells_per_slot * kv_per_token_swa + global_v_width = n_kv * val_len if padded_v_width is None else padded_v_width + swa_v_width = n_kv * val_len_swa if padded_v_width is None else padded_v_width + kv_per_token = n_kv * key_len * bpe_k + global_v_width * bpe_v + kv_per_token_swa = n_kv * key_len_swa * bpe_k + swa_v_width * bpe_v + global_bytes = n_global * total_cells * kv_per_token + swa_bytes = n_swa * swa_cells_total * kv_per_token_swa checkpoint_extra_per_slot = ( ctx_checkpoints * n_swa * swa * kv_per_token_swa if ctx_checkpoints > 0 and not swa_full else 0.0 ) - return int(global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot)) + return int(global_bytes + swa_bytes + slots * checkpoint_extra_per_slot) # Path 4: Standard GQA with explicit key/value dimensions if key_len is not None and val_len is not None: - return int(n_layers_kv * n_ctx * n_kv * (key_len + val_len) * bpe) + padded_v_width = None if flash_attn else self._max_kv_value_width(val_len) + bytes_per_cell = 0.0 + for layer_idx in range(n_layers_kv): + layer_n_kv = self._kv_heads_for_layer(layer_idx, n_kv) + v_width = layer_n_kv * val_len if padded_v_width is None else padded_v_width + bytes_per_cell += layer_n_kv * key_len * bpe_k + v_width * bpe_v + return int(total_cells * bytes_per_cell) # Path 5: Legacy fallback (old GGUFs without explicit dimensions) head_dim = self._legacy_head_dim() - return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe) + return int(2 * n_kv * head_dim * n_layers_kv * total_cells * bpe_k) def _draft_backend_for(self, drafter_path: str) -> Optional["LlamaCppBackend"]: """Lightweight backend with a drafter GGUF's metadata, to size its own KV @@ -4236,6 +4437,10 @@ class LlamaCppBackend: draft_cache_type_k: Optional[str] = None, draft_cache_type_v: Optional[str] = None, n_parallel: int = 1, + swa_full: bool = False, + kv_unified: bool = True, + n_ubatch: Optional[int] = None, + flash_attn: bool = True, ) -> 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 @@ -4249,12 +4454,23 @@ class LlamaCppBackend: db = self._draft_backend_for(drafter_path) if db is None or not db._can_estimate_kv(): return None + # Gemma 4 assistant layers share the target context's final global + # and SWA KV tensors, so only the drafter weights add memory. + if getattr(db, "_architecture", None) == "gemma4-assistant": + return 0 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 + # The drafter uses the main model's slot and stream layout, so its + # compact SWA and per-stream padding must follow the same settings. + kv = db._estimate_kv_cache_bytes( + n_ctx, + heavier, + n_parallel = n_parallel, + swa_full = swa_full, + kv_unified = kv_unified, + n_ubatch = n_ubatch, + flash_attn = flash_attn, + ) + return kv if kv > 0 else None nextn = self._nextn_predict_layers or 0 n_kv = self._n_kv_heads or self._n_heads k_len = self._kv_key_length @@ -4268,7 +4484,14 @@ class LlamaCppBackend: 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) + _, streams, cells_per_stream = _kv_cache_cell_layout(n_ctx, n_parallel, kv_unified) + v_width = n_kv * v_len + if not flash_attn: + v_width = self._max_kv_value_width( + v_len, + self._kv_value_length_swa, + ) + return int(nextn * (n_kv * k_len * bpe_k + v_width * bpe_v) * cells_per_stream * streams) def _estimate_mtp_overhead_bytes( self, @@ -4281,6 +4504,10 @@ class LlamaCppBackend: draft_weights_bytes: int = 0, n_parallel: int = 1, mtp_keeps_target_ctx: bool = True, + swa_full: bool = False, + kv_unified: bool = True, + n_ubatch: Optional[int] = None, + flash_attn: bool = True, ) -> Optional[int]: """MTP draft reserve at ``n_ctx`` = draft KV (grows with ctx) + separate- drafter weights + (MTP + MLA only) a duplicated target KV context. The @@ -4296,6 +4523,10 @@ class LlamaCppBackend: draft_cache_type_k = draft_cache_type_k, draft_cache_type_v = draft_cache_type_v, n_parallel = n_parallel, + swa_full = swa_full, + kv_unified = kv_unified, + n_ubatch = n_ubatch, + flash_attn = flash_attn, ) weights = max(0, draft_weights_bytes) # MLA models (GLM-5.x, DeepSeek, Kimi-K2) under MTP keep a *second* full copy @@ -4311,7 +4542,15 @@ class LlamaCppBackend: # rather than duplicating the target, so they must not be charged for it. target_ctx_copy = 0 if mtp_keeps_target_ctx and self._kv_lora_rank is not None: - target_ctx_copy = self._estimate_kv_cache_bytes(n_ctx, "f16", n_parallel = n_parallel) + target_ctx_copy = self._estimate_kv_cache_bytes( + n_ctx, + "f16", + n_parallel = n_parallel, + swa_full = swa_full, + kv_unified = kv_unified, + n_ubatch = n_ubatch, + flash_attn = flash_attn, + ) if draft_kv is None: # KV unsized (exotic/remote drafter): still reserve known weights + any # MLA target copy so a large config can't launch over budget (the small @@ -4321,7 +4560,7 @@ class LlamaCppBackend: return total if total > 0 else None return draft_kv + weights + target_ctx_copy - _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Unsloth does not override it + _DEFAULT_N_UBATCH = _DEFAULT_LLAMA_N_UBATCH _COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate # Soft VRAM the modeled terms omit; charged to the fit budget on tight tiers (#6682). _CUDA_CONTEXT_RESERVE_BYTES = 320 * 1024 * 1024 # CUDA ctx + cuBLAS workspace (~330 MiB) @@ -4379,7 +4618,10 @@ class LlamaCppBackend: 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)) + ub = max( + 1, + int(self._DEFAULT_N_UBATCH if n_ubatch is None else 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 @@ -4411,7 +4653,10 @@ class LlamaCppBackend: n_embd = self._embedding_length or 0 if n_embd <= 0 or n_ctx <= 0: return 0 - ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + ub = max( + 1, + int(self._DEFAULT_N_UBATCH if n_ubatch is None else n_ubatch), + ) if getattr(self, "_architecture", None) == "deepseek4": # DSV4 indexer/CSA buffer (see constants): flat + linear, ub-scaled. Fires # for any KV type -- the indexer scratch is present even with an f16 cache. @@ -4459,6 +4704,9 @@ class LlamaCppBackend: per_device_overhead_bytes: int, min_gpus: int, n_ubatch: Optional[int] = None, + swa_full: bool = False, + kv_unified: bool = True, + flash_attn: bool = True, ) -> tuple[Optional[list[int]], bool, int]: """Largest serving-slot count in [1, n_parallel) whose fully-on-GPU footprint fits, so Unsloth keeps the model on GPU (-ngl -1) instead of --fit on, which offloads layers @@ -4477,7 +4725,15 @@ class LlamaCppBackend: total = ( base_footprint_bytes + cb - + self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = slots) + + self._estimate_kv_cache_bytes( + effective_ctx, + cache_type_kv, + n_parallel = slots, + swa_full = swa_full, + kv_unified = kv_unified, + n_ubatch = n_ubatch, + flash_attn = flash_attn, + ) ) gpu_indices, use_fit = self._select_gpus( total, @@ -4502,7 +4758,9 @@ class LlamaCppBackend: swa_full: bool = False, n_parallel: int = 1, kv_unified: bool = True, + n_ubatch: Optional[int] = None, ctx_checkpoints: int = 0, + flash_attn: bool = True, kv_on_gpu: bool = True, mtp_engaged: bool = False, mtp_overhead_fn: Optional[Callable[[int], int]] = None, @@ -4539,7 +4797,9 @@ class LlamaCppBackend: swa_full = swa_full, n_parallel = n_parallel, kv_unified = kv_unified, + n_ubatch = n_ubatch, ctx_checkpoints = ctx_checkpoints, + flash_attn = flash_attn, ) # byte-accurate mtp_overhead_fn supersedes the flat fraction (the fallback @@ -5202,6 +5462,12 @@ class LlamaCppBackend: self._is_audio = False # clear any prior TTS/audio model's routing flag self._model_identifier = model_identifier self._cache_type_kv = None + self._swa_full = False + self._kv_cache_unified = False + self._n_ubatch = self._DEFAULT_N_UBATCH + self._flash_attn_enabled = True + self._effective_cache_types = ("f16", "f16") + self._kv_cache_context_total = None self._gpu_offload_active = True # Diffusion doesn't use the llama.cpp GPU-memory knobs; reset them to # defaults (the picked device is still recorded below) so /load, /status @@ -5943,6 +6209,9 @@ class LlamaCppBackend: total_by_idx: Optional[dict[int, int]] = None, n_ubatch: Optional[int] = None, soft_overhead_bytes: int = 0, + swa_full: bool = False, + kv_unified: bool = True, + flash_attn: bool = True, ) -> tuple[int, int, list[int], Optional[list[int]]]: """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed. @@ -6030,6 +6299,17 @@ class LlamaCppBackend: def _mtp_at(ctx: int) -> int: return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + def _kv_at(ctx: int) -> int: + return self._estimate_kv_cache_bytes( + ctx, + cache_type_kv, + n_parallel = n_parallel, + swa_full = swa_full, + kv_unified = kv_unified, + n_ubatch = n_ubatch, + flash_attn = flash_attn, + ) + # Context-linear compute buffer, summed over the split. Tensor mode # replicates the compute graph on EVERY device (measured: the per-device # buffer grows a flat n_ubatch*2 bytes/token, ~1024 B/tok on Qwen3.5-9B at @@ -6055,31 +6335,21 @@ class LlamaCppBackend: # 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)+compute(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) - + _cc_ctx(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) - total_at = kv_at + _cc_ctx(ctx) # both ~linear through the origin - if total_at <= kv_budget_b: + def _consumer(c: int) -> int: + return _kv_at(c) + _mtp_at(c) + _cc_ctx(c) + + if _consumer(ctx) <= kv_budget_b: return ctx - return max(ctx_floor, int(ctx * kv_budget_b / total_at)) + 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 size unknown -> can't prove a safe cap; floor. return min(4096, ctx) if ctx > 0 else 4096 @@ -6091,11 +6361,7 @@ class LlamaCppBackend: effective_ctx = min(_fit_ctx(target_ctx), max_available_ctx) 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 - ) + kv_bytes = _kv_at(effective_ctx) if (self._can_estimate_kv() and effective_ctx > 0) else 0 # 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 @@ -6220,21 +6486,6 @@ class LlamaCppBackend: cls._is_signal_crash(returncode) or cls._is_abort_exit(returncode) ) - @staticmethod - def _canonical_long_flag(name: str) -> str: - """Return ``name`` with llama.cpp's long-option underscore normalization. - - llama.cpp runs ``std::replace(arg.begin(), arg.end(), '_', '-')`` on any - argv token that starts with ``--`` before looking it up, so a legal - pass-through spelling like ``--cache_type_v`` parses as - ``--cache-type-v``. Mirror that here so managed-flag matching sees the - same canonical name. Short flags (``-ctv``) never carry underscores and - keep their exact spelling; pass only the flag name (no attached value). - """ - if name.startswith("--"): - return name.replace("_", "-") - return name - @staticmethod def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]: """Return cmd with flash attention forced off, or None when its effective @@ -6247,23 +6498,25 @@ class LlamaCppBackend: def explicit(i): nxt = out[i + 1] if i + 1 < len(out) else None - return nxt if nxt in ("on", "auto", "off") else None + return nxt if nxt in _LLAMA_ARG_TRUE_FALSE_AUTO_VALUES else None effective = None for i, tok in enumerate(out): - if tok.startswith(("--flash-attn=", "-fa=")): + name = _flag_name(tok) + if name in ("--flash-attn", "-fa") and "=" in tok: effective = tok.partition("=")[2] - elif tok in ("--flash-attn", "-fa"): + elif name in ("--flash-attn", "-fa"): effective = explicit(i) or "on" - if effective not in ("on", "auto"): + if effective not in _LLAMA_ARG_TRUE_OR_AUTO_VALUES: return None for i, tok in enumerate(out): - if tok.startswith(("--flash-attn=", "-fa=")): + name = _flag_name(tok) + if name in ("--flash-attn", "-fa") and "=" in tok: flag, _, value = tok.partition("=") - if value in ("on", "auto"): + if value in _LLAMA_ARG_TRUE_OR_AUTO_VALUES: out[i] = f"{flag}=off" - elif tok in ("--flash-attn", "-fa"): - if explicit(i) in ("on", "auto"): + elif name in ("--flash-attn", "-fa"): + if explicit(i) in _LLAMA_ARG_TRUE_OR_AUTO_VALUES: out[i + 1] = "off" elif explicit(i) is None: # bare flag (reads as on) -> explicit off out[i] = f"{tok}=off" @@ -6295,7 +6548,7 @@ class LlamaCppBackend: # quantized V cache. Canonicalize the flag name the same way so the # reset recognizes the underscore aliases too; short flags (-ctv) # and the type value are left untouched. - name = LlamaCppBackend._canonical_long_flag(tok.partition("=")[0]) + name = _flag_name(tok) if name not in _v_cache_flags: continue if "=" in tok: @@ -6740,6 +6993,8 @@ class LlamaCppBackend: # same message remote validation already shows. raise LlamaServerNotFoundError(LLAMA_SERVER_NOT_FOUND_DETAIL) + server_caps = self.probe_server_capabilities(binary) + # Outside ``self._lock`` so /unload, /cancel, /status aren't # blocked. ``unload_model`` also records the kill, so the # frontend /unload+/load Apply path engages the wait here even @@ -6760,6 +7015,18 @@ class LlamaCppBackend: # state to publish. ctx_override = parse_ctx_override(extra_args) requested_ctx = resolve_requested_ctx(extra_args, n_ctx) + swa_full = _swa_full_from_args_or_env(extra_args) + _effective_ubatch = _extra_args_n_ubatch( + extra_args, + n_ctx = (requested_ctx if requested_ctx > 0 else self._context_length), + ) + planned_kv_unified = _kv_unified_from_args( + extra_args, + default = n_parallel > 1 and server_caps.get("supports_kv_unified", False), + ) + # A hard-crash recovery may relaunch this same plan with FA off. + # Size that larger cache up front so the recovery cannot OOM. + planned_flash_attn = False cache_override = parse_cache_override(extra_args) # Budget the heavier of asymmetric --cache-type-k/-v extras (they # win per axis at launch, appended last); resolve_cache_type_kv only @@ -7190,6 +7457,10 @@ class LlamaCppBackend: draft_cache_type_k = _mtp_draft_ck, draft_cache_type_v = _mtp_draft_cv, n_parallel = n_parallel, + swa_full = swa_full, + kv_unified = planned_kv_unified, + n_ubatch = _effective_ubatch, + flash_attn = planned_flash_attn, ) if ( self._estimate_mtp_overhead_bytes( @@ -7201,6 +7472,10 @@ class LlamaCppBackend: draft_weights_bytes = _mtp_draft_weights, n_parallel = n_parallel, mtp_keeps_target_ctx = _engaged_is_mtp, + swa_full = swa_full, + kv_unified = planned_kv_unified, + n_ubatch = _effective_ubatch, + flash_attn = planned_flash_attn, ) is not None ): @@ -7217,6 +7492,10 @@ class LlamaCppBackend: _w: int = _mtp_draft_weights, _np: int = n_parallel, _mtp: bool = _engaged_is_mtp, + _swa_full: bool = swa_full, + _kv_unified: bool = planned_kv_unified, + _n_ubatch: Optional[int] = _effective_ubatch, + _flash_attn: bool = planned_flash_attn, ) -> int: v = self._estimate_mtp_overhead_bytes( ctx, @@ -7227,15 +7506,26 @@ class LlamaCppBackend: draft_weights_bytes = _w, n_parallel = _np, mtp_keeps_target_ctx = _mtp, + swa_full = _swa_full, + kv_unified = _kv_unified, + n_ubatch = _n_ubatch, + flash_attn = _flash_attn, ) 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) + def _kv_bytes(ctx: int) -> int: + return self._estimate_kv_cache_bytes( + ctx, + cache_type_kv, + n_parallel = n_parallel, + swa_full = swa_full, + kv_unified = planned_kv_unified, + n_ubatch = _effective_ubatch, + flash_attn = planned_flash_attn, + ) def _cc_bytes(ctx: int, n_gpus: int = 1) -> int: # Context-linear compute-buffer growth (flash-attn KQ mask + @@ -7475,6 +7765,9 @@ class LlamaCppBackend: total_by_idx = total_by_idx, n_ubatch = _effective_ubatch, soft_overhead_bytes = _soft_overhead, + swa_full = swa_full, + kv_unified = planned_kv_unified, + flash_attn = planned_flash_attn, ) use_fit = False elif gpus and self._can_estimate_kv() and effective_ctx > 0: @@ -7507,16 +7800,18 @@ class LlamaCppBackend: pool_budget, _ms, cache_type_kv, + swa_full = swa_full, n_parallel = n_parallel, + kv_unified = planned_kv_unified, + n_ubatch = _effective_ubatch, + flash_attn = planned_flash_attn, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, compute_ctx_bytes_fn = _cc_sub, budget_frac = 1.0, total_mib = None, ) - kv = self._estimate_kv_cache_bytes( - capped, cache_type_kv, n_parallel = n_parallel - ) + kv = _kv_bytes(capped) footprint_mib = ( _ms + kv + _mtp_bytes(capped) + _cc_sub(capped) ) / (1024 * 1024) @@ -7536,9 +7831,7 @@ class LlamaCppBackend: # on and let llama-server flex -ngl (CPU offload). requested_total = ( model_size_fit - + self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv, n_parallel = n_parallel - ) + + _kv_bytes(effective_ctx) + _mtp_bytes(effective_ctx) + _cc_bytes(effective_ctx) ) @@ -7590,16 +7883,18 @@ class LlamaCppBackend: pool_budget, _ms, cache_type_kv, + swa_full = swa_full, n_parallel = n_parallel, + kv_unified = planned_kv_unified, + n_ubatch = _effective_ubatch, + flash_attn = planned_flash_attn, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, compute_ctx_bytes_fn = _cc_sub, budget_frac = 1.0, total_mib = None, ) - kv = self._estimate_kv_cache_bytes( - capped, cache_type_kv, n_parallel = n_parallel - ) + kv = _kv_bytes(capped) footprint_mib = ( _ms + kv + _mtp_bytes(capped) + _cc_sub(capped) ) / (1024 * 1024) @@ -7616,11 +7911,7 @@ class LlamaCppBackend: if effective_ctx > 0: for n_gpus in range(_auto_min_gpus, len(ranked) + 1): subset = ranked[:n_gpus] - kv = self._estimate_kv_cache_bytes( - effective_ctx, - cache_type_kv, - n_parallel = n_parallel, - ) + kv = _kv_bytes(effective_ctx) footprint_mib = ( _subset_model_size(n_gpus) + kv @@ -7677,7 +7968,11 @@ class LlamaCppBackend: _apple_fit_budget_mib, model_size_fit, cache_type_kv, + swa_full = swa_full, n_parallel = n_parallel, + kv_unified = planned_kv_unified, + n_ubatch = _effective_ubatch, + flash_attn = planned_flash_attn, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, compute_ctx_bytes_fn = _cc_bytes, @@ -7685,12 +7980,7 @@ class LlamaCppBackend: total_mib = None, ) _cap_footprint_mib = ( - model_size_fit - + self._estimate_kv_cache_bytes( - cap, cache_type_kv, n_parallel = n_parallel - ) - + _mtp_bytes(cap) - + _cc_bytes(cap) + model_size_fit + _kv_bytes(cap) + _mtp_bytes(cap) + _cc_bytes(cap) ) / (1024 * 1024) # Fit returns the request unchanged when it fits OR weights # exceed budget; only the latter over-commits, so floor to 4096. @@ -7737,6 +8027,9 @@ class LlamaCppBackend: _pipeline_overhead_bytes + _cc_bytes(effective_ctx), _layer_min_gpus, _effective_ubatch, + swa_full = swa_full, + kv_unified = planned_kv_unified, + flash_attn = planned_flash_attn, ) if not _uf_slots: logger.info( @@ -7761,9 +8054,7 @@ class LlamaCppBackend: _mtp_note = "" if effective_ctx < original_ctx: - kv_est = self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv, n_parallel = n_parallel - ) + kv_est = _kv_bytes(effective_ctx) logger.info( f"Context auto-reduced: {original_ctx} -> {effective_ctx} " f"(model: {model_size / (1024**3):.1f} GB, " @@ -7772,9 +8063,7 @@ class LlamaCppBackend: + ")" ) - kv_cache_bytes = self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv, n_parallel = n_parallel - ) + kv_cache_bytes = _kv_bytes(effective_ctx) mmproj_note = ( f"mmproj: {mmproj_size / (1024**3):.1f} GB, " if mmproj_size else "" ) @@ -7939,7 +8228,6 @@ class LlamaCppBackend: cmd.extend(["-ngl", "-1", "--fit", "off"]) fully_gpu_offloaded = True - server_caps = self.probe_server_capabilities(binary) # Expose Prometheus /metrics for the engine-stats logger, only # when the binary advertises it (older/custom binaries may not). if server_caps.get("supports_metrics"): @@ -8011,6 +8299,11 @@ class LlamaCppBackend: "iq4_nl", "f32", } + # Normalize like the budget does (_planned_main_cache_types): a + # case-sensitive match drops "Q8_0", emitting no flag, so llama.cpp + # runs f16 while the estimate priced q8_0. Emit the normalized + # spelling; kv_cache_type_from_str is case-sensitive. + cache_type_kv = cache_type_kv.strip().lower() if cache_type_kv else cache_type_kv if ( cache_type_kv and cache_type_kv in _valid_cache_types @@ -8213,6 +8506,8 @@ class LlamaCppBackend: cmd.extend(str(a) for a in extra_args) logger.info(f"Appending user extra args to llama-server: {list(extra_args)}") + kv_cache_unified = _kv_unified_from_args(cmd) + logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}") # Library paths so llama-server finds its shared libs and CUDA DLLs. @@ -8727,6 +9022,20 @@ class LlamaCppBackend: self._healthy = True self._commit_effective_parallel_slots(n_parallel) + self._swa_full = swa_full + self._kv_cache_unified = kv_cache_unified + self._n_ubatch = max( + 0, + int(self._DEFAULT_N_UBATCH if _effective_ubatch is None else _effective_ubatch), + ) + self._flash_attn_enabled = ( + _flash_attn_enabled_from_args(_last_spawn_cmd) and self._architecture != "grok" + ) + self._effective_cache_types = _effective_main_cache_types( + _last_spawn_cmd, + env, + ) + self._kv_cache_context_total = effective_ctx if effective_ctx > 0 else None # Server is up: adopt the real per-request context it allocated # -- the length --fit chose, or a --parallel slot split -- so the @@ -8734,6 +9043,11 @@ class LlamaCppBackend: # before the spawn above always failed; the seeded value was the # requested/native length.) self._reconcile_effective_ctx_with_server() + if self._kv_cache_context_total is not None: + self._n_ubatch = min( + self._n_ubatch, + self._kv_cache_context_total, + ) # Commit caller intent only after _healthy=True so a failed start # can't poison the next inheritance check. None keeps prior, [] @@ -9190,7 +9504,6 @@ class LlamaCppBackend: if _norm(self._cache_type_kv) != _norm(cache_type_kv): return False - # 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 @@ -9214,6 +9527,9 @@ class LlamaCppBackend: # layer/MoE/split knobs), so a standing manual preference in the # request must not force a needless reload -- only the GPU pick matters. if not self._is_diffusion: + requested_extra_args = extra_args if extra_args is not None else self._extra_args + if self._swa_full != _swa_full_from_args_or_env(requested_extra_args): + return False # A GPU-memory-mode flip (Unsloth / manual) must always reload. if self._gpu_memory_mode != gpu_memory_mode: return False @@ -9341,7 +9657,8 @@ class LlamaCppBackend: last_draft: Optional[str] = None args = [str(arg) for arg in cmd] for index, raw in enumerate(args): - flag, equals, inline = raw.partition("=") + flag = _flag_name(raw) + _, equals, inline = raw.partition("=") if flag not in main_flags and flag not in draft_flags: continue value = inline if equals else (args[index + 1] if index + 1 < len(args) else "") @@ -9414,6 +9731,12 @@ class LlamaCppBackend: self._slot_save_binary = None self._slot_loaded_identity = None self._prompt_cache_disabled = False + self._swa_full = False + self._kv_cache_unified = False + self._n_ubatch = self._DEFAULT_N_UBATCH + self._flash_attn_enabled = True + self._effective_cache_types = ("f16", "f16") + self._kv_cache_context_total = None self._chat_template = None self._chat_template_override = None self._supports_reasoning = False @@ -9957,8 +10280,12 @@ class LlamaCppBackend: tuple(sidecars), self._requested_n_ctx, self._effective_context_length, - getattr(self, "_cache_type_kv", None), + self._effective_cache_types, self.effective_parallel_slots, + self._swa_full, + self._kv_cache_unified, + self._n_ubatch, + self._flash_attn_enabled, ) def _gguf_file_identity(self, path) -> Optional[tuple]: @@ -9989,7 +10316,8 @@ class LlamaCppBackend: args = [str(a).strip() for a in (self._extra_args or ())] files: list[str] = [] for i, arg in enumerate(args): - flag, sep, inline = arg.partition("=") + flag = _flag_name(arg) + _, sep, inline = arg.partition("=") if flag not in self._SIDECAR_WEIGHT_FLAGS: continue operand = inline if sep else (args[i + 1] if i + 1 < len(args) else "") @@ -10029,7 +10357,7 @@ class LlamaCppBackend: if os.environ.get("LLAMA_ARG_NO_CACHE_PROMPT") is not None: return True env = (os.environ.get("LLAMA_ARG_CACHE_PROMPT") or "").strip().lower() - return env in {"off", "disabled", "false", "0"} + return env in _LLAMA_ARG_FALSE_VALUES def save_slots_for_resume( self, should_abort: Optional[Callable[[], bool]] = None @@ -10041,6 +10369,17 @@ class LlamaCppBackend: or self._prompt_cache_off() ): return None + # Same predicate as the estimator's SWA path: a window alone is not enough. + # phi3 GGUFs carry attention.sliding_window but no key/value length, and + # llama.cpp forces them back to a non-SWA cache, so their slots do restore. + if ( + (self._sliding_window or 0) > 0 + and self._kv_key_length is not None + and self._kv_value_length is not None + and not self._swa_full + ): + logger.debug("Skipping slot save: compact SWA cache cannot be reused after restart") + return None save_dir = Path(self._slot_save_dir) gguf_stat = self._gguf_file_identity(self._gguf_path) if gguf_stat is None: @@ -10057,9 +10396,16 @@ class LlamaCppBackend: return None try: estimate = self._estimate_kv_cache_bytes( - self._effective_context_length or self._context_length or 0, - self._cache_type_kv, + self._kv_cache_context_total + or self._effective_context_length + or self._context_length + or 0, + max(self._effective_cache_types, key = _kv_bytes_per_elem), n_parallel = self.effective_parallel_slots, + swa_full = self._swa_full, + kv_unified = self._kv_cache_unified, + n_ubatch = self._n_ubatch, + flash_attn = self._flash_attn_enabled, ) # Skip before writing anything when the estimate alone blows the cap, # rather than fully writing a slot and discarding it afterwards. @@ -10415,6 +10761,8 @@ class LlamaCppBackend: actual_n_ctx = self._query_server_n_ctx() if not actual_n_ctx or actual_n_ctx <= 0: return + slots = 1 if self._kv_cache_unified else self.effective_parallel_slots + self._kv_cache_context_total = actual_n_ctx * slots if self._effective_context_length and actual_n_ctx < self._effective_context_length: logger.warning( "llama-server allocated a smaller per-request context than " diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 6f1b931a7f..2ecd7e3e2e 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -80,9 +80,10 @@ _DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS) def _flag_name(token: str) -> Optional[str]: """Flag name for ``token``, or None if it isn't a flag. - Peels `--key=value` to `--key`, treats `-1`/`-0.5` as values (shorts - always start with a letter), and normalises attached `-np8` / `-np-1` / - `-np8x` to `-np`. Mirrors the CLI's `_expand_attached_np_short`. + Peels `--key=value` to `--key`, normalises long-option underscores like + llama.cpp, treats `-1`/`-0.5` as values (shorts always start with a letter), + and normalises attached `-np8` / `-np-1` / `-np8x` to `-np`. Mirrors the + CLI's `_expand_attached_np_short`. """ token = token.strip() if not token.startswith("-") or token in {"-", "--"}: @@ -90,6 +91,8 @@ def _flag_name(token: str) -> Optional[str]: if len(token) >= 2 and (token[1].isdigit() or token[1] == "."): return None name = token.split("=", 1)[0] + if name.startswith("--"): + name = name.replace("_", "-") if len(name) > 3 and name.startswith("-np"): suffix = name[3:] if suffix[0].isdigit() or ( diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index e66adb789e..acd60dd0b9 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -254,6 +254,8 @@ class ValidateModelRequest(BaseModel): # /load; defaults preserve old behavior for callers that omit them. max_seq_length: int = Field(0, ge = 0, le = 1048576) load_in_4bit: bool = Field(True) + cache_type_kv: Optional[str] = Field(None) + tensor_parallel: bool = Field(False) gpu_ids: Optional[List[int]] = Field(None) gpu_memory_mode: Literal["auto", "manual"] = Field( "auto", diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 97149f7a17..8b15779a50 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1004,8 +1004,13 @@ try: _DEFAULT_MAX_TOKENS_FLOOR, _DEFAULT_STREAM_STALL_TIMEOUT_S, _canonicalize_spec_mode, + _extra_args_n_ubatch, _extra_args_set_spec_type, _hf_offline_if_dns_dead, + _kv_bytes_per_elem, + _kv_unified_from_args, + _planned_main_cache_types, + _swa_full_from_args_or_env, detect_reasoning_flags, ) from core.inference.llama_server_args import ( @@ -1043,8 +1048,13 @@ except ImportError: _DEFAULT_MAX_TOKENS_FLOOR, _DEFAULT_STREAM_STALL_TIMEOUT_S, _canonicalize_spec_mode, + _extra_args_n_ubatch, _extra_args_set_spec_type, _hf_offline_if_dns_dead, + _kv_bytes_per_elem, + _kv_unified_from_args, + _planned_main_cache_types, + _swa_full_from_args_or_env, detect_reasoning_flags, ) from core.inference.llama_server_args import ( @@ -3320,6 +3330,10 @@ def _request_matches_loaded_settings( strip_offload = request.gpu_memory_mode == "manual", ) ) + if not llama_backend.is_diffusion and llama_backend.swa_full != _swa_full_from_args_or_env( + effective_extra + ): + return False if not _tensor_parallel_matches_loaded( effective_extra, request.tensor_parallel, llama_backend.tensor_parallel ): @@ -4435,10 +4449,12 @@ def _estimate_gguf_kv_gb( max_seq_length: int, llama_extra_args: Optional[list[str]] = None, n_parallel: int = 1, + cache_type_kv: Optional[str] = None, + tensor_parallel: bool = False, ) -> float: """KV-cache VRAM (GB) at the larger of max_seq_length and any `--ctx-size`/`-c` - override, over n_parallel slots, with the default f16 cache so the estimate is - never below what the server allocates. 0 if metadata is unreadable.""" + override, over n_parallel slots, using the effective cache settings and managed + launcher defaults. 0 if metadata is unreadable.""" try: from core.inference.llama_server_args import parse_ctx_override @@ -4453,7 +4469,43 @@ def _estimate_gguf_kv_gb( ctx = max(max_seq_length or 0, ctx_override) or (probe._context_length or 0) if ctx <= 0: return 0.0 - kv = probe._estimate_kv_cache_bytes(ctx, n_parallel = max(1, n_parallel or 1)) + slots = max(1, n_parallel or 1) + managed_kv_unified = bool( + slots > 1 + and LlamaCppBackend.probe_server_capabilities().get("supports_kv_unified", False) + ) + planned_cache_types = _planned_main_cache_types( + cache_type_kv, + llama_extra_args, + ) + if tensor_parallel and any( + cache_type not in LlamaCppBackend._TENSOR_PARALLEL_KV_TYPES + for cache_type in planned_cache_types + ): + # Tensor mode strips quantized axes, but a layer fallback restores + # the original settings. Size for the larger successful outcome. + tensor_cache_types = _planned_main_cache_types(None, None) + cache_type_for_budget = max( + (*planned_cache_types, *tensor_cache_types, "f16"), + key = _kv_bytes_per_elem, + ) + else: + cache_type_for_budget = max( + planned_cache_types, + key = _kv_bytes_per_elem, + ) + kv = probe._estimate_kv_cache_bytes( + ctx, + cache_type_for_budget, + n_parallel = slots, + swa_full = _swa_full_from_args_or_env(llama_extra_args), + kv_unified = _kv_unified_from_args( + llama_extra_args, + default = managed_kv_unified, + ), + n_ubatch = _extra_args_n_ubatch(llama_extra_args, n_ctx = ctx), + flash_attn = False, + ) return kv / (1024**3) except Exception as e: logger.warning(f"Could not size GGUF KV cache for training guard: {e}") @@ -4466,6 +4518,8 @@ def _estimate_gguf_required_gb( max_seq_length: int = 0, llama_extra_args: Optional[list[str]] = None, n_parallel: int = 1, + cache_type_kv: Optional[str] = None, + tensor_parallel: bool = False, ) -> Optional[float]: """Approximate GGUF VRAM (GB): quantized weights + companions, plus the KV cache for local files (unreadable pre-download for remote). None when nothing @@ -4481,7 +4535,12 @@ def _estimate_gguf_required_gb( total_bytes += Path(f).stat().st_size if total_bytes > 0: return total_bytes / (1024**3) + _estimate_gguf_kv_gb( - main, max_seq_length, llama_extra_args, n_parallel + main, + max_seq_length, + llama_extra_args, + n_parallel, + cache_type_kv, + tensor_parallel, ) repo = getattr(config, "gguf_hf_repo", None) @@ -4622,6 +4681,8 @@ def _guard_chat_load_against_training( requested_gpu_ids: Optional[List[int]], llama_extra_args: Optional[list[str]] = None, n_parallel: int = 1, + cache_type_kv: Optional[str] = None, + tensor_parallel: bool = False, gpu_memory_mode: Literal["auto", "manual"] = "auto", ) -> None: """Protect active training from automatically placed chat-model loads. @@ -4676,6 +4737,11 @@ def _guard_chat_load_against_training( max_seq_length = max_seq_length, llama_extra_args = llama_extra_args, n_parallel = n_parallel, + cache_type_kv = cache_type_kv, + tensor_parallel = ( + _effective_tensor_parallel(llama_extra_args, tensor_parallel) + and (is_vulkan or LlamaCppBackend._effective_gpu_count(requested_gpu_ids) >= 2) + ), ) if is_gguf else None @@ -5416,6 +5482,8 @@ async def _load_model_impl( requested_gpu_ids = effective_gpu_ids, llama_extra_args = extra_llama_args, n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1), + cache_type_kv = request.cache_type_kv, + tensor_parallel = bool(request.tensor_parallel), gpu_memory_mode = request.gpu_memory_mode, ) @@ -6092,6 +6160,8 @@ async def validate_model( if fastapi_request is not None else 1 ), + cache_type_kv = request.cache_type_kv, + tensor_parallel = request.tensor_parallel, gpu_memory_mode = request.gpu_memory_mode, ) diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index f1d973f004..6ec9c44e88 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -451,6 +451,9 @@ class TestChatLoadGuardRoute(unittest.TestCase): decision, gpu_memory_mode = "auto", requested_gpu_ids = None, + llama_extra_args = None, + cache_type_kv = None, + tensor_parallel = False, ): config = config or SimpleNamespace(is_gguf = False, is_lora = False, path = None) with _stub_guard_deps( @@ -463,6 +466,9 @@ class TestChatLoadGuardRoute(unittest.TestCase): load_in_4bit = True, max_seq_length = 0, requested_gpu_ids = requested_gpu_ids, + llama_extra_args = llama_extra_args, + cache_type_kv = cache_type_kv, + tensor_parallel = tensor_parallel, gpu_memory_mode = gpu_memory_mode, ) @@ -597,6 +603,32 @@ class TestChatLoadGuardRoute(unittest.TestCase): self.assertEqual(captured[0]["is_gguf"], True) self.assertEqual(captured[0]["required_override_gb"], 12.5) + def test_vulkan_gguf_estimate_keeps_tensor_cache_coercion(self): + config = SimpleNamespace(is_gguf = True) + estimate_kwargs = {} + with ( + patch.object( + self.route, + "_estimate_gguf_required_gb", + side_effect = lambda *args, **kwargs: estimate_kwargs.update(kwargs) or 12.5, + ), + patch.object( + self.route.LlamaCppBackend, + "_effective_gpu_count", + return_value = 0, + ), + patch.object(self.route.LlamaCppBackend, "_is_vulkan_backend", return_value = True), + ): + self._guard( + config = config, + training_active = True, + decision = (True, {}), + llama_extra_args = ["--split-mode", "tensor"], + cache_type_kv = "q4_0", + ) + self.assertEqual(estimate_kwargs["cache_type_kv"], "q4_0") + self.assertTrue(estimate_kwargs["tensor_parallel"]) + class TestEffectiveLoadIn4bit(unittest.TestCase): @classmethod @@ -745,7 +777,12 @@ class TestValidateRefusesDuringTraining(unittest.TestCase): # /load then 409s after the frontend has already unloaded. from models.inference import ValidateModelRequest - request = ValidateModelRequest(model_path = "unsloth/Qwen3-1.7B", max_seq_length = 4096) + request = ValidateModelRequest( + model_path = "unsloth/Qwen3-1.7B", + max_seq_length = 4096, + cache_type_kv = "f32", + tensor_parallel = True, + ) cfg = SimpleNamespace( identifier = "unsloth/Qwen3-1.7B", display_name = "Qwen3-1.7B", @@ -774,6 +811,8 @@ class TestValidateRefusesDuringTraining(unittest.TestCase): asyncio.run(self.route.validate_model(request, current_subject = "u")) self.assertEqual(captured.get("llama_extra_args"), ["-c", "32768"]) self.assertIn("n_parallel", captured) + self.assertEqual(captured.get("cache_type_kv"), "f32") + self.assertTrue(captured.get("tensor_parallel")) def test_metadata_probe_skips_training_guard(self): # A header-only probe (include_context_length) allocates no VRAM, so the @@ -985,6 +1024,8 @@ class TestEstimateGgufRequiredGb(unittest.TestCase): class _FakeBackend: _context_length = 2048 + _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + supports_kv_unified = True def _read_gguf_metadata(self, path): pass @@ -992,13 +1033,27 @@ class TestEstimateGgufRequiredGb(unittest.TestCase): def _can_estimate_kv(self): return True + @classmethod + def probe_server_capabilities(cls): + return {"supports_kv_unified": cls.supports_kv_unified} + def _estimate_kv_cache_bytes( self, ctx, + cache_type = None, n_parallel = 1, + swa_full = False, + kv_unified = False, + n_ubatch = None, + flash_attn = True, ): seen["ctx"] = ctx + seen["cache_type"] = cache_type seen["n_parallel"] = n_parallel + seen["swa_full"] = swa_full + seen["kv_unified"] = kv_unified + seen["n_ubatch"] = n_ubatch + seen["flash_attn"] = flash_attn return ctx * n_parallel * (1024**2) # 1 MiB per ctx unit per slot with patch.object(self.route, "LlamaCppBackend", _FakeBackend): @@ -1009,6 +1064,8 @@ class TestEstimateGgufRequiredGb(unittest.TestCase): ) self.assertEqual(seen["ctx"], 131072) self.assertEqual(seen["n_parallel"], 1) # default single slot + self.assertFalse(seen["swa_full"]) + self.assertFalse(seen["flash_attn"]) # override below max_seq_length -> larger (max_seq_length) wins self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "1024"]), 4.0) self.assertEqual(seen["ctx"], 4096) @@ -1020,6 +1077,50 @@ class TestEstimateGgufRequiredGb(unittest.TestCase): # --parallel slots scale the cache the same way the launcher does self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, None, 4), 16.0) self.assertEqual(seen["n_parallel"], 4) + self.assertTrue(seen["kv_unified"]) + # User extras are appended after Studio's managed default. + r._estimate_gguf_kv_gb("m", 4096, ["--no-kv-unified"], 4) + self.assertFalse(seen["kv_unified"]) + # An older binary without the flag keeps separate KV streams. + _FakeBackend.supports_kv_unified = False + r._estimate_gguf_kv_gb("m", 4096, None, 4) + self.assertFalse(seen["kv_unified"]) + r._estimate_gguf_kv_gb("m", 4096, None, 1, "f32") + self.assertEqual(seen["cache_type"], "f32") + r._estimate_gguf_kv_gb("m", 4096, ["--cache-type-v", "f32"]) + self.assertEqual(seen["cache_type"], "f32") + with patch.dict(self.route.os.environ, {"LLAMA_ARG_CACHE_TYPE_K": "f32"}): + r._estimate_gguf_kv_gb("m", 4096) + self.assertEqual(seen["cache_type"], "f32") + with patch.dict( + self.route.os.environ, + { + "LLAMA_ARG_CACHE_TYPE_K": "q4_0", + "LLAMA_ARG_CACHE_TYPE_V": "q4_0", + }, + ): + r._estimate_gguf_kv_gb("m", 4096) + self.assertEqual(seen["cache_type"], "q4_0") + r._estimate_gguf_kv_gb( + "m", + 4096, + ["--cache-type-k", "q4_0", "--cache-type-v", "q4_0"], + tensor_parallel = True, + ) + self.assertEqual(seen["cache_type"], "f16") + r._estimate_gguf_kv_gb( + "m", + 4096, + ["--cache-type-k", "f32", "--cache-type-v", "q4_0"], + tensor_parallel = True, + ) + self.assertEqual(seen["cache_type"], "f32") + # Full SWA mode follows the same pass-through args as the launcher. + r._estimate_gguf_kv_gb("m", 4096, ["--swa_full"]) + self.assertTrue(seen["swa_full"]) + r._estimate_gguf_kv_gb("m", 4096, ["--kv_unified", "--ubatch_size", "256"]) + self.assertTrue(seen["kv_unified"]) + self.assertEqual(seen["n_ubatch"], 256) # ── load_model integration: authoritative 409, and no unload before refusal ── diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py index 4259171da9..43365bd3ca 100644 --- a/studio/backend/tests/test_gpu_memory_mode.py +++ b/studio/backend/tests/test_gpu_memory_mode.py @@ -183,11 +183,12 @@ def test_already_in_target_state_reloads_on_mode_change(loaded, requested): assert _target_state(_loaded_backend(loaded), requested) is False -def test_already_in_target_state_ignores_mode_for_diffusion(): +def test_already_in_target_state_ignores_mode_for_diffusion(monkeypatch): # The diffusion runner is mode-agnostic (always "auto"), so a standing manual # preference must not force a needless reload. backend = _loaded_backend("auto") backend._is_diffusion = True + monkeypatch.setenv("LLAMA_ARG_SWA_FULL", "1") assert _target_state(backend, "manual") is True diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py index 27e9d0f57a..3cf86cf0ca 100644 --- a/studio/backend/tests/test_kv_cache_estimation.py +++ b/studio/backend/tests/test_kv_cache_estimation.py @@ -76,6 +76,39 @@ from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend # Helpers +def _runtime_kv_cells( + n_ctx: int, + *, + slots: int = 1, + unified: bool = True, +) -> int: + """Total KV cells allocated by llama.cpp across all streams.""" + slots = max(1, slots) + padded_ctx = ((n_ctx + 255) // 256) * 256 + streams = 1 if unified else slots + cells_per_stream = padded_ctx if unified else ((max(1, padded_ctx // slots) + 255) // 256) * 256 + return cells_per_stream * streams + + +def _runtime_swa_cells( + n_ctx: int, + sliding_window: int, + *, + slots: int = 1, + unified: bool = True, + n_ubatch: int = 512, +) -> tuple[int, int]: + """Return total non-SWA and compact-SWA cells allocated by llama.cpp.""" + slots = max(1, slots) + streams = 1 if unified else slots + base_cells = _runtime_kv_cells(n_ctx, slots = slots, unified = unified) + cells_per_stream = base_cells // streams + swa_limit = sliding_window * (slots if unified else 1) + n_ubatch + swa_cells_per_stream = min(cells_per_stream, swa_limit) + swa_cells_per_stream = ((swa_cells_per_stream + 255) // 256) * 256 + return base_cells, swa_cells_per_stream * streams + + def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes: """Build a minimal GGUF v3 blob with the given KV metadata. @@ -789,7 +822,7 @@ class TestMLAEstimation: b = self._mla_backend() result = b._estimate_kv_cache_bytes(1000, "f16") # n_layers * ctx * 1 * key_len(576) * 2 - expected = 61 * 1000 * 1 * 576 * 2 + expected = 61 * _runtime_kv_cells(1000) * 1 * 576 * 2 assert result == expected def test_mla_fallback_when_no_key_length(self): @@ -797,14 +830,14 @@ class TestMLAEstimation: b = self._mla_backend(_kv_key_length = None) # default _key_length_mla=192, so rope_dim=192 result = b._estimate_kv_cache_bytes(1000, "f16") - expected = 61 * 1000 * 1 * (512 + 192) * 2 # 704 + expected = 61 * _runtime_kv_cells(1000) * 1 * (512 + 192) * 2 # 704 assert result == expected def test_mla_fallback_no_key_length_mla(self): """No key_length and no key_length_mla: fall back to +64.""" b = self._mla_backend(_kv_key_length = None, _key_length_mla = None) result = b._estimate_kv_cache_bytes(1000, "f16") - expected = 61 * 1000 * 1 * (512 + 64) * 2 # 576 + expected = 61 * _runtime_kv_cells(1000) * 1 * (512 + 64) * 2 # 576 assert result == expected def test_mla_defaults_n_kv_to_1_when_heads_absent(self): @@ -812,7 +845,7 @@ class TestMLAEstimation: b = self._mla_backend(_n_kv_heads = None) # n_heads=128 still set result = b._estimate_kv_cache_bytes(1000, "f16") # Uses n_kv_mla=1, NOT n_heads=128 - expected = 61 * 1000 * 1 * 576 * 2 + expected = 61 * _runtime_kv_cells(1000) * 1 * 576 * 2 assert result == expected def test_mla_q4_quantization(self): @@ -821,7 +854,7 @@ class TestMLAEstimation: result_q4 = b._estimate_kv_cache_bytes(1000, "q4_0") assert result_q4 < result_f16 # q4_0 bpe = 0.5625, f16 bpe = 2.0 - assert result_q4 == int(61 * 1000 * 1 * 576 * 0.5625) + assert result_q4 == int(61 * _runtime_kv_cells(1000) * 1 * 576 * 0.5625) # D. Path 2: Hybrid Mamba Estimation @@ -910,9 +943,8 @@ class TestSlidingWindowEstimation: n_global = max(1, 62 // 4) # 15 n_swa = 62 - n_global # 47 kv_per = 16 * (128 + 128) * 2 - # SWA cache is double-buffered: 2 * sliding_window cells, capped at n_ctx. - swa_cells = min(131072, 2 * 1024) - expected = int(n_global * 131072 * kv_per + n_swa * swa_cells * kv_per) + base_cells, swa_cells = _runtime_swa_cells(131072, 1024) + expected = int(n_global * base_cells * kv_per + n_swa * swa_cells * kv_per) assert b._estimate_kv_cache_bytes(131072, "f16") == expected def test_gpt_oss(self): @@ -929,8 +961,8 @@ class TestSlidingWindowEstimation: n_global = max(1, 24 // 4) # 6 n_swa = 24 - n_global # 18 kv_per = 8 * (64 + 64) * 2 - swa_cells = min(131072, 2 * 128) - expected = int(n_global * 131072 * kv_per + n_swa * swa_cells * kv_per) + base_cells, swa_cells = _runtime_swa_cells(131072, 128) + expected = int(n_global * base_cells * kv_per + n_swa * swa_cells * kv_per) assert b._estimate_kv_cache_bytes(131072, "f16") == expected def test_gemma4_per_layer_swa_metadata(self): @@ -952,21 +984,67 @@ class TestSlidingWindowEstimation: sliding_layers = 25 def expected(ctx): - full = full_layers * ctx * 2 * (512 + 512) * 2 - sliding = sliding_layers * min(ctx, 2 * 1024) * 8 * (256 + 256) * 2 + base_cells, swa_cells = _runtime_swa_cells(ctx, 1024) + full = full_layers * base_cells * 2 * (512 + 512) * 2 + sliding = sliding_layers * swa_cells * 8 * (256 + 256) * 2 return int(full + sliding) for ctx in (4096, 46500, 262144): assert b._estimate_kv_cache_bytes(ctx, "f16") == expected(ctx) + def test_gemma4_flash_attn_off_pads_v_to_model_max(self): + b = self._swa_backend( + _n_layers = 35, + _n_kv_heads = 1, + _n_heads = 8, + _embedding_length = 1536, + _kv_key_length = 512, + _kv_value_length = 512, + _sliding_window = 512, + _sliding_window_pattern = [True, True, True, True, False] * 7, + _kv_key_length_swa = 256, + _kv_value_length_swa = 256, + _shared_kv_layers = 20, + ) + ctx = 5000 + slots = 3 + base_cells, swa_cells = _runtime_swa_cells(ctx, 512, slots = slots, unified = True) + max_v_width = 512 + expected = ( + 3 * base_cells * (512 + max_v_width) * 2 + 12 * swa_cells * (256 + max_v_width) * 2 + ) + actual = b._estimate_kv_cache_bytes( + ctx, + "f16", + n_parallel = slots, + flash_attn = False, + ) + assert actual == expected + assert actual == 66 * 1024**2 + assert actual > b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots) + + def test_flash_attn_off_prices_quantized_v_retry_as_f16(self): + b = self._swa_backend( + _n_layers = 2, + _n_kv_heads = None, + _n_kv_heads_by_layer = [8, 2], + _sliding_window_pattern = [True, False], + _kv_key_length_swa = 64, + _kv_value_length_swa = 64, + ) + off = b._estimate_kv_cache_bytes(4096, "q4_0", flash_attn = False) + on = b._estimate_kv_cache_bytes(4096, "q4_0") + assert off > on + def test_ctx_smaller_than_window(self): - """When ctx < 2 * sliding_window, SWA cache caps at ctx.""" + """When context is smaller than the compact allowance, SWA caps at context.""" b = self._swa_backend(_sliding_window = 8192) n_global = max(1, 62 // 4) # 15 n_swa = 62 - n_global # 47 kv_per = 16 * (128 + 128) * 2 ctx = 4096 - expected = int(n_global * ctx * kv_per + n_swa * min(ctx, 2 * 8192) * kv_per) + base_cells, swa_cells = _runtime_swa_cells(ctx, 8192) + expected = int(n_global * base_cells * kv_per + n_swa * swa_cells * kv_per) assert b._estimate_kv_cache_bytes(ctx, "f16") == expected def test_odd_layer_count(self): @@ -974,7 +1052,8 @@ class TestSlidingWindowEstimation: n_global = max(1, 63 // 4) # 15 n_swa = 63 - n_global # 48 kv_per = 16 * (128 + 128) * 2 - expected = int(n_global * 1000 * kv_per + n_swa * min(1000, 2 * 1024) * kv_per) + base_cells, swa_cells = _runtime_swa_cells(1000, 1024) + expected = int(n_global * base_cells * kv_per + n_swa * swa_cells * kv_per) assert b._estimate_kv_cache_bytes(1000, "f16") == expected @@ -1086,8 +1165,7 @@ class TestPathPriority: b._full_attention_interval = 4 b._sliding_window = 1024 # Would trigger SWA - # MLA: 61 * 1000 * 1 * 576 * 2 - expected_mla = int(61 * 1000 * 1 * 576 * 2) + expected_mla = int(61 * _runtime_kv_cells(1000) * 1 * 576 * 2) assert b._estimate_kv_cache_bytes(1000, "f16") == expected_mla def test_hybrid_over_swa(self): @@ -1104,7 +1182,7 @@ class TestPathPriority: b._sliding_window = 1024 # Would trigger SWA n_attn = 64 // 4 - expected_hybrid = int(n_attn * 1000 * 4 * (256 + 256) * 2) + expected_hybrid = int(n_attn * _runtime_kv_cells(1000) * 4 * (256 + 256) * 2) assert b._estimate_kv_cache_bytes(1000, "f16") == expected_hybrid def test_all_paths_produce_different_values(self): @@ -1192,7 +1270,7 @@ class TestQuantization: b._kv_key_length = 64 b._kv_value_length = 64 result = b._estimate_kv_cache_bytes(1000, cache_type) - expected = int(10 * 1000 * 1 * (64 + 64) * expected_bpe) + expected = int(10 * _runtime_kv_cells(1000) * 1 * (64 + 64) * expected_bpe) assert result == expected @@ -1221,7 +1299,7 @@ class TestEdgeCases: b._kv_key_length = 64 b._kv_value_length = 64 result = b._estimate_kv_cache_bytes(1, "f16") - assert result == int(10 * 1 * 1 * (64 + 64) * 2) + assert result == int(10 * _runtime_kv_cells(1) * 1 * (64 + 64) * 2) def test_very_large_context(self): """1M context should not overflow or crash.""" @@ -1242,7 +1320,7 @@ class TestEdgeCases: b._kv_key_length = 64 b._kv_value_length = 64 result = b._estimate_kv_cache_bytes(100, "f16") - expected = int(10 * 100 * 8 * (64 + 64) * 2) + expected = int(10 * _runtime_kv_cells(100) * 8 * (64 + 64) * 2) assert result == expected def test_both_heads_none_falls_to_one(self): @@ -1253,7 +1331,7 @@ class TestEdgeCases: b._kv_key_length = 64 b._kv_value_length = 64 result = b._estimate_kv_cache_bytes(100, "f16") - expected = int(10 * 100 * 1 * (64 + 64) * 2) + expected = int(10 * _runtime_kv_cells(100) * 1 * (64 + 64) * 2) assert result == expected @@ -1335,12 +1413,21 @@ class TestServerFlags: assert with_cp_full == no_cp_full assert with_cp > b._estimate_kv_cache_bytes(8192, "f16") + def test_compact_swa_includes_ubatch_headroom_and_padding(self): + b = self._swa_backend(_sliding_window = 128) + ctx = 8192 + result = b._estimate_kv_cache_bytes(ctx, "f16", n_ubatch = 512) + per_token = 4 * (256 + 256) * 2 + n_swa = sum(b._sliding_window_pattern) + n_global = b._n_layers - n_swa + expected = n_global * ctx * per_token + n_swa * 768 * per_token + assert result == expected + # ── --parallel + --kv-unified ────────────────────────────────── # Verified against llama-server: non-SWA caches partition n_ctx across - # slots (total memory constant); only SWA layers scale with --parallel. - # --kv-unified is a no-op for memory math (kept for API forward-compat). + # non-unified streams. Compact SWA sizing depends on the stream layout. - def test_gqa_kv_constant_across_parallel(self): + def test_gqa_kv_constant_for_aligned_stream_divisions(self): b = self._gqa_backend() baseline = b._estimate_kv_cache_bytes(4096, "f16") for slots in (1, 2, 4, 8): @@ -1359,7 +1446,7 @@ class TestServerFlags: == baseline ) - def test_swa_path_scales_only_swa_portion(self): + def test_swa_path_matches_aligned_stream_layout(self): b = self._swa_backend() ctx = 8192 baseline = b._estimate_kv_cache_bytes(ctx, "f16") @@ -1367,27 +1454,27 @@ class TestServerFlags: swa = b._sliding_window per_token_global = 4 * (256 + 256) * 2 # n_kv * (k+v) * f16 per_token_swa = 4 * (256 + 256) * 2 # k_swa/val_swa fall back - per_slot_swa_cells = min(ctx, 2 * swa) # not clamped at parallel=1 + base_cells, swa_cells = _runtime_swa_cells(ctx, swa) global_bytes = sum( - ctx * per_token_global for f in b._sliding_window_pattern[: b._n_layers] if not f + base_cells * per_token_global for f in b._sliding_window_pattern[: b._n_layers] if not f ) - swa_bytes_per_slot = sum( - per_slot_swa_cells * per_token_swa - for f in b._sliding_window_pattern[: b._n_layers] - if f + swa_bytes = sum( + swa_cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f ) # Sanity: parallel=1 reproduces baseline exactly - assert global_bytes + swa_bytes_per_slot == baseline - # Only the SWA portion scales by parallel + assert global_bytes + swa_bytes == baseline for slots in (1, 2, 3, 4): scaled = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False) - # SWA cells clamp to per_slot_ctx when ctx/slots < 2*swa - per_slot_ctx = max(1, ctx // slots) - cells = min(ctx, 2 * swa, per_slot_ctx) - swa_bps = sum( - cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f + base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = False) + expected_global = sum( + base_cells * per_token_global + for f in b._sliding_window_pattern[: b._n_layers] + if not f ) - assert scaled == global_bytes + slots * swa_bps + expected_swa = sum( + swa_cells * per_token_swa for f in b._sliding_window_pattern[: b._n_layers] if f + ) + assert scaled == expected_global + expected_swa def test_mla_kv_constant_across_parallel(self): b = LlamaCppBackend() @@ -1444,19 +1531,17 @@ class TestServerFlags: ctx = 8192 swa = b._sliding_window per_token = 4 * (256 + 256) * 2 - global_bytes = sum( - ctx * per_token for f in b._sliding_window_pattern[: b._n_layers] if not f - ) n_swa_layers = sum(1 for f in b._sliding_window_pattern[: b._n_layers] if f) slots = 3 - per_slot_ctx = max(1, ctx // slots) - swa_cells = min(ctx, 2 * swa, per_slot_ctx) - swa_bytes_per_slot = n_swa_layers * swa_cells * per_token + base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = False) + n_global_layers = b._n_layers - n_swa_layers + global_bytes = n_global_layers * base_cells * per_token + swa_bytes = n_swa_layers * swa_cells * per_token cp_extra_per_slot = n_swa_layers * 4 * swa * per_token # 4 checkpoints flagged = b._estimate_kv_cache_bytes( ctx, "f16", ctx_checkpoints = 4, n_parallel = slots, kv_unified = False ) - assert flagged == global_bytes + slots * (swa_bytes_per_slot + cp_extra_per_slot) + assert flagged == global_bytes + swa_bytes + slots * cp_extra_per_slot # ── --kv-offload (kv_on_gpu) ─────────────────────────────────── @@ -1535,22 +1620,40 @@ class TestServerFlags: assert fitted_default == ctx assert fitted_full < ctx + def test_tensor_planner_threads_swa_full_through_estimator(self): + b = self._swa_backend() + estimate = b._estimate_kv_cache_bytes + calls = [] + + def record(*args, **kwargs): + calls.append(kwargs) + return estimate(*args, **kwargs) + + b._estimate_kv_cache_bytes = record + b._plan_tensor_parallel( + [(0, 32768), (1, 32768)], + 1024**3, + 8192, + cache_type_kv = "f16", + swa_full = True, + flash_attn = False, + ) + assert calls + assert all(call["swa_full"] is True for call in calls) + assert all(call["flash_attn"] is False for call in calls) + # J2.5. --parallel N memory accounting (per-layer-type scaling rule) class TestParallelSWAScaling: - """Per-layer-type scaling rule vs the closed form measured from - llama-server. Empirical formula on Gemma-3 270m at ctx=8192: - total_kv = 24 + parallel * 15 (MiB). + """Per-layer-type scaling rule measured from llama-server. Rule (verified vs ``llama-server`` log on real GGUFs): - * non-SWA layers: total cells = n_ctx, partitioned across slots, - memory CONSTANT in n_parallel. - * SWA layers: per-slot cells = 2 * sliding_window (clamped at - n_ctx and at per_slot_ctx); memory LINEAR in n_parallel. - * --kv-unified is a no-op for memory math; both modes give the - same total in measured cases. + * non-SWA layers use the padded per-stream context. + * compact SWA adds ubatch headroom and pads to 256 cells. + * unified mode uses one stream with all slot windows. + * non-unified mode allocates one stream per slot. """ def _gqa_backend(self, **overrides): @@ -1586,7 +1689,7 @@ class TestParallelSWAScaling: setattr(b, k, v) return b - # ── non-SWA paths: constant ──────────────────────────────────── + # ── non-SWA paths: constant when stream divisions are aligned ── def test_pure_gqa_constant_across_parallel(self): b = self._gqa_backend() @@ -1633,25 +1736,53 @@ class TestParallelSWAScaling: for slots in (1, 2, 4, 8): assert b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots) == baseline - # ── SWA paths: scale only the SWA portion ────────────────────── + def test_non_swa_paths_follow_unaligned_stream_padding(self): + mla = LlamaCppBackend() + mla._n_layers = 60 + mla._n_kv_heads = 1 + mla._kv_lora_rank = 512 + mla._key_length_mla = 64 + mla._kv_key_length = 576 - def test_swa_pattern_scales_only_swa_portion(self): + hybrid = LlamaCppBackend() + hybrid._n_layers = 64 + hybrid._n_kv_heads = 16 + hybrid._n_heads = 32 + hybrid._embedding_length = 4096 + hybrid._kv_key_length = 128 + hybrid._kv_value_length = 128 + hybrid._ssm_inner_size = 4096 + hybrid._full_attention_interval = 4 + + legacy = LlamaCppBackend() + legacy._n_layers = 32 + legacy._n_kv_heads = 8 + legacy._n_heads = 8 + legacy._embedding_length = 4096 + + for backend in (self._gqa_backend(), mla, hybrid, legacy): + bytes_per_cell = backend._estimate_kv_cache_bytes(256, "f16") // 256 + unified = backend._estimate_kv_cache_bytes(5000, "f16", n_parallel = 3, kv_unified = True) + separate = backend._estimate_kv_cache_bytes(5000, "f16", n_parallel = 3, kv_unified = False) + assert unified == 5120 * bytes_per_cell + assert separate == 5376 * bytes_per_cell + + # ── SWA paths: aligned stream scaling ────────────────────────── + + def test_swa_pattern_matches_aligned_stream_layout(self): b = self._swa_backend() ctx = 8192 swa = b._sliding_window per_token = 1 * (256 + 256) * 2 # n_kv * (k+v) * f16 n_global = sum(1 for f in b._sliding_window_pattern if not f) n_swa = sum(1 for f in b._sliding_window_pattern if f) - global_bytes = n_global * ctx * per_token for slots in (1, 2, 4, 8): - per_slot_ctx = max(1, ctx // slots) - cells = min(ctx, 2 * swa, per_slot_ctx) - swa_bps = n_swa * cells * per_token for unified in (True, False): + base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = unified) got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = unified) - assert got == global_bytes + slots * swa_bps + assert got == (n_global * base_cells * per_token + n_swa * swa_cells * per_token) - def test_swa_fallback_scales_only_swa_portion(self): + def test_swa_fallback_matches_aligned_stream_layout(self): # No per-layer pattern -> 1/4-global heuristic. b = self._swa_backend(_sliding_window_pattern = None) ctx = 8192 @@ -1660,34 +1791,28 @@ class TestParallelSWAScaling: n_global = max(1, n_layers // 4) n_swa = n_layers - n_global per_token = 1 * (256 + 256) * 2 - global_bytes = n_global * ctx * per_token for slots in (1, 2, 4, 8): - per_slot_ctx = max(1, ctx // slots) - cells = min(ctx, 2 * swa, per_slot_ctx) - swa_bps = n_swa * cells * per_token - got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots) - assert got == global_bytes + slots * swa_bps + for unified in (True, False): + base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = unified) + got = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = unified) + assert got == (n_global * base_cells * per_token + n_swa * swa_cells * per_token) def test_swa_per_slot_clamped_when_ctx_lt_slots_x_2window(self): - # ctx=4096 / slots=8 -> per_slot_ctx=512, but 2*sliding=1024. - # SWA cells clamp at per_slot_ctx (512), not 2*sliding. + # ctx=4096 / slots=8 gives a 512-cell stream, which caps compact SWA. b = self._swa_backend() ctx = 4096 per_slot_ctx_at_8 = ctx // 8 - assert per_slot_ctx_at_8 < 2 * b._sliding_window - # Build expected with the clamped formula n_swa = sum(1 for f in b._sliding_window_pattern if f) n_global = sum(1 for f in b._sliding_window_pattern if not f) per_token = 1 * (256 + 256) * 2 - global_bytes = n_global * ctx * per_token - cells = min(ctx, 2 * b._sliding_window, per_slot_ctx_at_8) - assert cells == per_slot_ctx_at_8 - expected = global_bytes + 8 * (n_swa * cells * per_token) - assert b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = 8) == expected + base_cells, swa_cells = _runtime_swa_cells(ctx, b._sliding_window, slots = 8, unified = False) + assert swa_cells == 8 * per_slot_ctx_at_8 + expected = n_global * base_cells * per_token + n_swa * swa_cells * per_token + assert b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = 8, kv_unified = False) == expected - def test_swa_full_does_not_scale_under_parallel(self): - # swa_full forces every layer to n_ctx -> all-global GQA-style - # total, constant in parallel. + def test_swa_full_constant_for_aligned_stream_divisions(self): + # swa_full forces every layer to n_ctx. This aligned context remains + # constant across the tested stream divisions. b = self._swa_backend() ctx = 8192 baseline = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True) @@ -1696,25 +1821,32 @@ class TestParallelSWAScaling: b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True, n_parallel = slots) == baseline ) - # ── kv_unified: no-op for memory math ────────────────────────── + # ── kv_unified stream layout ──────────────────────────────────── - def test_kv_unified_is_no_op_for_memory_math(self): - # unified=True and unified=False must give the same total bytes - # for every backend type and parallel value. - backends = [ - ("gqa", self._gqa_backend()), - ("swa", self._swa_backend()), - ] - for label, b in backends: - for slots in (1, 2, 4, 8): - u = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = True) - nu = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = False) - assert u == nu, f"{label} parallel={slots} unified-mismatch" + def test_kv_unified_changes_only_compact_swa_for_aligned_context(self): + gqa = self._gqa_backend() + swa = self._swa_backend() + for slots in (1, 2, 4, 8): + gqa_unified = gqa._estimate_kv_cache_bytes( + 8192, "f16", n_parallel = slots, kv_unified = True + ) + gqa_separate = gqa._estimate_kv_cache_bytes( + 8192, "f16", n_parallel = slots, kv_unified = False + ) + assert gqa_unified == gqa_separate + + swa_unified = swa._estimate_kv_cache_bytes( + 8192, "f16", n_parallel = slots, kv_unified = True + ) + swa_separate = swa._estimate_kv_cache_bytes( + 8192, "f16", n_parallel = slots, kv_unified = False + ) + assert (swa_unified == swa_separate) is (slots == 1) # ── Empirical Gemma-3 270m formula ───────────────────────────── def test_matches_empirical_gemma3_270m_formula(self): - """Exact match against the formula measured from llama-server: + """Exact match against the non-unified formula measured from llama-server: total_kv = 24 + parallel * 15 (MiB) at ctx=8192. Geometry: 18 layers (3 global + 15 SWA), n_kv=1, head_dim=256, @@ -1736,12 +1868,16 @@ class TestParallelSWAScaling: # Confirm pattern shape assert sum(b._sliding_window_pattern) == n_swa for slots, expected_mib in [(1, 39), (2, 54), (4, 84)]: - got_bytes = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots) + got_bytes = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = False) got_mib = got_bytes / (1024 * 1024) assert ( got_mib == expected_mib ), f"slots={slots}: got {got_mib} MiB, expected {expected_mib} MiB" + for slots, expected_mib in [(1, 39), (2, 46.5), (4, 61.5)]: + got_bytes = b._estimate_kv_cache_bytes(8192, "f16", n_parallel = slots, kv_unified = True) + assert got_bytes / (1024 * 1024) == expected_mib + # J3. shared_kv_layers (Gemma 3n / Gemma 4) @@ -1844,8 +1980,8 @@ class TestSharedKVLayers: assert sliding_in_unshared == 16 assert full_in_unshared == 4 kv_per = 4 * (256 + 256) * 2 - swa_cells = min(ctx, 2 * 1024) - expected = full_in_unshared * ctx * kv_per + sliding_in_unshared * swa_cells * kv_per + base_cells, swa_cells = _runtime_swa_cells(ctx, 1024) + expected = full_in_unshared * base_cells * kv_per + sliding_in_unshared * swa_cells * kv_per assert b._estimate_kv_cache_bytes(ctx, "f16") == expected def test_shared_layers_reduces_estimate(self): @@ -1875,8 +2011,8 @@ class TestSharedKVLayers: n_global = max(1, n_layers_kv // 4) # 5 n_swa = n_layers_kv - n_global # 15 kv_per = 4 * (256 + 256) * 2 - swa_cells = min(ctx, 2 * 1024) - expected = n_global * ctx * kv_per + n_swa * swa_cells * kv_per + base_cells, swa_cells = _runtime_swa_cells(ctx, 1024) + expected = n_global * base_cells * kv_per + n_swa * swa_cells * kv_per assert b._estimate_kv_cache_bytes(ctx, "f16") == expected def test_shared_floors_at_one_layer(self): @@ -1896,13 +2032,12 @@ class TestSharedKVLayers: unshared_pattern = b._sliding_window_pattern[:20] # 35 - 15 shared sliding_in_unshared = sum(unshared_pattern) global_in_unshared = len(unshared_pattern) - sliding_in_unshared - global_bytes = global_in_unshared * ctx * per_token slots = 3 - per_slot_ctx = max(1, ctx // slots) - swa_cells = min(ctx, 2 * swa, per_slot_ctx) - swa_bytes_per_slot = sliding_in_unshared * swa_cells * per_token + base_cells, swa_cells = _runtime_swa_cells(ctx, swa, slots = slots, unified = False) + global_bytes = global_in_unshared * base_cells * per_token + swa_bytes = sliding_in_unshared * swa_cells * per_token flagged = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False) - assert flagged == global_bytes + slots * swa_bytes_per_slot + assert flagged == global_bytes + swa_bytes def test_composes_with_ctx_checkpoints(self): b = self._gemma3n_backend() @@ -2036,14 +2171,14 @@ class TestLifecycle: ) assert b._can_estimate_kv() result = b._estimate_kv_cache_bytes(131072, "f16") - # gemma3 -> period 6 from bootstrap; SWA cache double-buffered to - # 2 * sliding_window cells. + # gemma3 uses period 6 from the bootstrap resolver. period = 6 kv_per = 16 * 256 * 2 + base_cells, swa_cells = _runtime_swa_cells(131072, 1024) expected = 0 for i in range(62): is_swa = (i + 1) % period != 0 - layer_ctx = min(131072, 2 * 1024) if is_swa else 131072 + layer_ctx = swa_cells if is_swa else base_cells expected += layer_ctx * kv_per assert result == expected diff --git a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py index 45c8bcb032..f39baddcb4 100644 --- a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py +++ b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py @@ -221,6 +221,18 @@ class TestFlashAttnOff: assert _flash_off(["llama-server", "-fa", "auto"]) == ["llama-server", "-fa", "off"] assert _flash_off(["llama-server", "-fa=on"]) == ["llama-server", "-fa=off"] + @pytest.mark.parametrize("value", ["on", "enabled", "true", "1", "auto", "-1"]) + def test_flips_every_enabled_value(self, value): + assert _flash_off(["llama-server", "--flash-attn", value]) == [ + "llama-server", + "--flash-attn", + "off", + ] + + @pytest.mark.parametrize("value", ["off", "disabled", "false", "0"]) + def test_none_for_every_disabled_value(self, value): + assert _flash_off(["llama-server", "--flash-attn", value]) is None + def test_flips_every_occurrence_last_wins(self): # extra_args can re-enable FA after Unsloth's flag; llama.cpp is last-wins, # so one leftover 'on' would re-crash the retry. Every enable must flip. @@ -384,6 +396,10 @@ class TestFlashAttnOffQuantizedKvCache: out = _flash_off(["llama-server", "--flash-attn=on", "--cache_type_v=q8_0"]) assert out == ["llama-server", "--flash-attn=off", "--cache_type_v=f16"] + def test_underscore_alias_flash_attn_is_disabled(self): + out = _flash_off(["llama-server", "--flash_attn=on"]) + assert out == ["llama-server", "--flash_attn=off"] + def test_underscore_value_not_normalized_for_nonquantized(self): # Only the flag name is canonicalized; a non-quantized type value is # matched verbatim and left untouched (no spurious reset). diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 27c1b17a85..8754b86b18 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -63,7 +63,9 @@ from core.inference.llama_cpp import ( _extra_args_set_any_flag, _extra_args_set_spec_type, _is_mtp_model_name, + _kv_unified_from_args, _mla_mtp_auto_enabled, + _swa_full_from_args_or_env, ) @@ -147,6 +149,41 @@ def test_is_mtp_model_name_handles_none(): assert _is_mtp_model_name("", "") is False +@pytest.mark.parametrize("flag", ["--swa-full", "--swa_full"]) +def test_swa_full_detects_llama_cpp_long_flag_spellings(flag): + assert _swa_full_from_args_or_env([flag], {}) is True + + +@pytest.mark.parametrize("value", ["on", "enabled", "true", "1"]) +def test_swa_full_detects_llama_cpp_env_truth_values(value): + assert _swa_full_from_args_or_env([], {"LLAMA_ARG_SWA_FULL": value}) is True + + +@pytest.mark.parametrize("value", ["", "off", "yes", "TRUE", " true ", "0"]) +def test_swa_full_rejects_values_llama_cpp_treats_as_false(value): + assert _swa_full_from_args_or_env([], {"LLAMA_ARG_SWA_FULL": value}) is False + + +def test_swa_full_cli_wins_when_env_is_false(): + assert _swa_full_from_args_or_env(["--swa-full"], {"LLAMA_ARG_SWA_FULL": "0"}) is True + + +@pytest.mark.parametrize("flag", ["--kv-unified", "--kv_unified", "-kvu"]) +def test_kv_unified_detects_enable_aliases(flag): + assert _kv_unified_from_args([flag]) is True + + +@pytest.mark.parametrize("flag", ["--no-kv-unified", "--no_kv_unified", "-no-kvu"]) +def test_kv_unified_detects_disable_aliases(flag): + assert _kv_unified_from_args(["--kv-unified", flag]) is False + + +def test_kv_unified_uses_environment_before_cli(): + assert _kv_unified_from_args([], env = {"LLAMA_ARG_KV_UNIFIED": "true"}) is True + assert _kv_unified_from_args([], default = True, env = {"LLAMA_ARG_KV_UNIFIED": "false"}) is True + assert _kv_unified_from_args(["--kv-unified"], env = {"LLAMA_ARG_KV_UNIFIED": "false"}) is True + + def test_is_mtp_model_name_detects_marker_in_filename(tmp_path): gguf = tmp_path / "Qwen3.6-27B-MTP-Q4_K_M.gguf" gguf.write_bytes(b"") diff --git a/studio/backend/tests/test_llama_cpp_props_readback.py b/studio/backend/tests/test_llama_cpp_props_readback.py index fe1e67edad..1dc8bae8c2 100644 --- a/studio/backend/tests/test_llama_cpp_props_readback.py +++ b/studio/backend/tests/test_llama_cpp_props_readback.py @@ -104,6 +104,9 @@ def _make_backend(effective_ctx = 98304, port = 51234): inst._port = port inst._effective_context_length = effective_ctx inst._context_length = 262144 + inst._effective_parallel_slots = 1 + inst._kv_cache_unified = False + inst._kv_cache_context_total = None return inst @@ -173,6 +176,31 @@ def test_fit_shrunk_ctx_overwrites_advertised_value(monkeypatch): assert inst.context_length == 67584 +def test_props_keeps_total_cache_context_for_slot_preflight(monkeypatch): + inst = _make_backend(effective_ctx = 32768) + inst._effective_parallel_slots = 4 + _stub_props( + monkeypatch, + body = {"default_generation_settings": {"n_ctx": 8192}}, + ) + inst._reconcile_effective_ctx_with_server() + assert inst._effective_context_length == 8192 + assert inst._kv_cache_context_total == 32768 + + +def test_props_does_not_multiply_unified_cache_context(monkeypatch): + inst = _make_backend(effective_ctx = 32768) + inst._effective_parallel_slots = 4 + inst._kv_cache_unified = True + _stub_props( + monkeypatch, + body = {"default_generation_settings": {"n_ctx": 32768}}, + ) + inst._reconcile_effective_ctx_with_server() + assert inst._effective_context_length == 32768 + assert inst._kv_cache_context_total == 32768 + + def test_matching_ctx_is_left_alone(monkeypatch): inst = _make_backend(effective_ctx = 98304) _stub_props( diff --git a/studio/backend/tests/test_llama_cpp_slot_resume.py b/studio/backend/tests/test_llama_cpp_slot_resume.py index 8b20c952c4..fc1222b2da 100644 --- a/studio/backend/tests/test_llama_cpp_slot_resume.py +++ b/studio/backend/tests/test_llama_cpp_slot_resume.py @@ -221,6 +221,34 @@ def test_fingerprint_tracks_effective_context_length(tmp_path): assert backend._slot_launch_fingerprint() != before +def test_fingerprint_tracks_swa_full_mode(tmp_path): + backend = _resume_backend(tmp_path) + before = backend._slot_launch_fingerprint() + backend._swa_full = True + assert backend._slot_launch_fingerprint() != before + + +def test_fingerprint_tracks_unified_cache_mode(tmp_path): + backend = _resume_backend(tmp_path) + before = backend._slot_launch_fingerprint() + backend._kv_cache_unified = True + assert backend._slot_launch_fingerprint() != before + + +def test_fingerprint_tracks_flash_attention_mode(tmp_path): + backend = _resume_backend(tmp_path) + before = backend._slot_launch_fingerprint() + backend._flash_attn_enabled = False + assert backend._slot_launch_fingerprint() != before + + +def test_fingerprint_tracks_effective_cache_types(tmp_path): + backend = _resume_backend(tmp_path) + before = backend._slot_launch_fingerprint() + backend._effective_cache_types = ("f32", "f16") + assert backend._slot_launch_fingerprint() != before + + def test_gguf_file_identity_covers_split_shards(tmp_path): backend = _resume_backend(tmp_path) first = tmp_path / "m-00001-of-00002.gguf" @@ -444,6 +472,81 @@ def test_save_skipped_when_estimate_exceeds_cap(monkeypatch, tmp_path): assert backend.save_slots_for_resume() is None +def test_save_estimate_uses_total_context_and_active_cache_settings(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path, n_slots = 4) + backend._effective_context_length = 8192 + backend._kv_cache_context_total = 32768 + backend._sliding_window = 4096 + backend._swa_full = True + backend._flash_attn_enabled = False + backend._effective_cache_types = ("f32", "f16") + calls = [] + + def estimate(ctx, cache_type, **kwargs): + calls.append((ctx, cache_type, kwargs)) + return 0 + + backend._estimate_kv_cache_bytes = estimate + _fake_disk(monkeypatch) + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: _Resp(200, {"n_saved": 1, "n_written": 1}), + raising = False, + ) + + assert backend.save_slots_for_resume() is not None + assert calls == [ + ( + 32768, + "f32", + { + "n_parallel": 4, + "swa_full": True, + "kv_unified": False, + "n_ubatch": 512, + "flash_attn": False, + }, + ) + ] + + +def test_compact_swa_slot_save_is_skipped(monkeypatch, tmp_path): + backend = _resume_backend(tmp_path) + backend._sliding_window = 4096 + backend._kv_key_length = 256 + backend._kv_value_length = 256 + backend._swa_full = False + backend._estimate_kv_cache_bytes = lambda *a, **k: (_ for _ in ()).throw(AssertionError) + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: (_ for _ in ()).throw(AssertionError), + raising = False, + ) + assert backend.save_slots_for_resume() is None + + +def test_window_without_kv_dims_still_saves(monkeypatch, tmp_path): + # phi3 reports a window but no key/value length, and llama.cpp runs it + # non-SWA, so the compact-SWA skip must not catch it. + backend = _resume_backend(tmp_path) + backend._sliding_window = 262144 + backend._kv_key_length = None + backend._kv_value_length = None + backend._swa_full = False + posted = [] + monkeypatch.setattr( + llama_cpp.httpx, + "post", + lambda *a, **k: posted.append(a) + or SimpleNamespace(status_code = 200, json = lambda: {"filename": "slot.bin"}), + raising = False, + ) + backend.save_slots_for_resume() + assert posted + + def test_save_skipped_when_model_file_changed_since_load(monkeypatch, tmp_path): # The GGUF/sidecars were swapped on disk after the server loaded them, so the # live KV belongs to the old weights: refuse to persist it (no POST at all). diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index d3ead7d9f2..b2ec5034ac 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -112,6 +112,11 @@ def test_value_with_equals_form_passes_through(): assert validate_extra_args(["--top-k=20"]) == ["--top-k=20"] +def test_managed_long_flag_underscore_alias_is_rejected(): + with pytest.raises(ValueError, match = "slot-save-path"): + validate_extra_args(["--slot_save_path", "/tmp/slots"]) + + def test_non_flag_token_passes_through(): # Bare positionals are passed through; llama-server can reject them. assert validate_extra_args(["foo"]) == ["foo"] diff --git a/studio/backend/tests/test_mtp_vram_budget.py b/studio/backend/tests/test_mtp_vram_budget.py index 6c8b74fc54..77ca76325f 100644 --- a/studio/backend/tests/test_mtp_vram_budget.py +++ b/studio/backend/tests/test_mtp_vram_budget.py @@ -76,7 +76,9 @@ from core.inference.llama_cpp import ( # noqa: E402 _extra_args_spec_draft_n_max, _effective_tensor_parallel, _env_main_cache_type_for_budget, + _effective_main_cache_types, _extra_args_main_cache_type_for_budget, + _flash_attn_enabled_from_args, _kv_bytes_per_elem, _tensor_parallel_matches_loaded, ) @@ -132,6 +134,7 @@ class _StubDrafter: def __init__(self, kv_per_token): self._kv_per_token = kv_per_token + self._architecture = "gemma3" def _can_estimate_kv(self): return True @@ -177,6 +180,14 @@ class TestEmbeddedDraftKv: two = _make_backend(nextn = 2)._mtp_draft_kv_bytes(65536) assert two == pytest.approx(2 * one) + def test_unaligned_context_follows_runtime_stream_padding(self): + b = _make_backend() + bytes_per_cell = b._mtp_draft_kv_bytes(256) // 256 + unified = b._mtp_draft_kv_bytes(5000, n_parallel = 3, kv_unified = True) + separate = b._mtp_draft_kv_bytes(5000, n_parallel = 3, kv_unified = False) + assert unified == 5120 * bytes_per_cell + assert separate == 5376 * bytes_per_cell + 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 @@ -201,6 +212,15 @@ class TestEmbeddedDraftKv: 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_flash_attn_off_uses_model_wide_v_width(self): + b = _make_backend(n_layers = 2) + b._n_kv_heads_by_layer = [4, 1] + b._sliding_window_pattern = [False, True] + b._kv_value_length_swa = 2048 + ctx = 4096 + expected_per_cell = 4 * 256 * 2 + 1 * 2048 * 2 + assert b._mtp_draft_kv_bytes(ctx, flash_attn = False) == ctx * expected_per_cell + 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 @@ -232,6 +252,30 @@ class TestSeparateDrafter: c = b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf") assert c == pytest.approx(4 * a) + def test_gemma4_assistant_shares_target_kv(self, monkeypatch): + b = _make_backend(nextn = None) + stub = _StubDrafter(kv_per_token = 2000) + stub._architecture = "gemma4-assistant" + monkeypatch.setattr(b, "_draft_backend_for", lambda path: stub) + + assert ( + b._mtp_draft_kv_bytes( + 65536, + drafter_path = "/m/mtp-gemma4.gguf", + swa_full = True, + ) + == 0 + ) + assert ( + b._estimate_mtp_overhead_bytes( + 65536, + drafter_path = "/m/mtp-gemma4.gguf", + draft_weights_bytes = GIB, + swa_full = True, + ) + == GIB + ) + 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 @@ -398,6 +442,7 @@ class TestExtraArgsMtpDetection: (["--spec-type", "mtp"], True), (["--spec-type", "ngram-mod,draft-mtp"], True), (["--spec-type=draft-mtp"], True), + (["--spec_type=draft-mtp"], True), (["--spec-type", "ngram-mod"], False), (["--spec-default"], False), (["-c", "131072"], False), @@ -579,6 +624,7 @@ class TestExtraArgsMtpDetection: (["--spec-draft-ngl", "0"], True), (["-ngld", "0"], True), (["--spec-draft-ngl=0"], True), + (["--spec_draft_ngl=0"], True), (["--n-gpu-layers-draft", "0"], True), (["--spec-draft-ngl", "20"], False), (["--spec-draft-device", "none"], True), @@ -623,6 +669,7 @@ class TestExtraArgsMtpDetection: [ (["--spec-draft-n-max", "4"], 4), (["--spec-draft-n-max=6"], 6), + (["--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), @@ -644,6 +691,7 @@ class TestExtraArgsMtpDetection: (["--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=/m/draft.gguf"], "/m/draft.gguf"), (["--model-draft", "--spec-type"], None), (["-c", "4096"], None), (None, None), @@ -689,6 +737,7 @@ class TestExtraArgsMtpDetection: (["--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_draft=q8_0"], ("q8_0", None)), (["--cache-type-k", "q8_0"], (None, None)), # main type, not draft (["-c", "4096"], (None, None)), (None, (None, None)), @@ -717,8 +766,17 @@ class TestExtraArgsMtpDetection: "args,expected", [ (["--ubatch-size", "1024"], 1024), - (["-ub", "4096"], 4096), + (["-ub", "4096"], 2048), + (["--ubatch-size", "0"], 2048), + (["--batch-size", "256", "--ubatch-size", "0"], 256), + (["--batch-size", "-1"], 512), + (["--ubatch-size", "-1"], 2048), (["--ubatch-size=512"], 512), + (["--ubatch_size=512"], 512), + (["--batch-size", "256"], 256), + (["--batch_size=256"], 256), + (["-b", "256", "-ub", "1024"], 256), + (["-b", "4096"], 512), (["--ubatch", "2048"], None), # not a real llama-server flag; ignore it (["-c", "4096"], None), (None, None), @@ -727,12 +785,76 @@ class TestExtraArgsMtpDetection: def test_n_ubatch(self, args, expected): assert _extra_args_n_ubatch(args, env = {}) == expected + def test_n_ubatch_signed_values_cap_at_context(self): + assert ( + _extra_args_n_ubatch( + ["--batch-size", "-1", "--ubatch-size", "-1"], + env = {}, + n_ctx = 4096, + ) + == 4096 + ) + + @pytest.mark.parametrize( + "args,expected", + [ + (None, True), + (["--flash-attn", "off"], False), + (["--flash-attn", "disabled"], False), + (["--flash-attn", "false"], False), + (["--flash-attn", "0"], False), + (["--flash-attn=off"], False), + (["--flash-attn=disabled"], False), + (["--flash-attn=false"], False), + (["--flash-attn=0"], False), + (["--flash_attn", "off"], False), + (["-fa", "off", "--flash-attn", "auto"], True), + (["-fa", "off", "--flash-attn", "-1"], True), + (["-fa", "off", "--flash-attn", "enabled"], True), + (["-fa", "off", "--flash-attn=true"], True), + (["-fa", "off", "--flash-attn=1"], True), + (["--flash-attn", "off", "-fa"], True), + ], + ) + def test_flash_attn_last_value_wins(self, args, expected): + assert _flash_attn_enabled_from_args(args) is expected + + def test_effective_main_cache_types_follow_env_then_cli(self): + env = { + "LLAMA_ARG_CACHE_TYPE_K": "f32", + "LLAMA_ARG_CACHE_TYPE_V": "q4_0", + } + assert _effective_main_cache_types([], env) == ("f32", "q4_0") + assert _effective_main_cache_types(["--cache-type-v", "f16"], env) == ("f32", "f16") + 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 + # Environment values apply first, then each command-line option overrides + # its own axis before llama.cpp caps ubatch at batch size. + assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "4096"}) == 2048 + assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_BATCH": "256"}) == 256 + assert ( + _extra_args_n_ubatch( + [], + env = { + "LLAMA_ARG_BATCH": "1024", + "LLAMA_ARG_UBATCH": "4096", + }, + ) + == 1024 + ) assert ( _extra_args_n_ubatch(["-ub", "1024"], env = {"LLAMA_ARG_UBATCH": "4096"}) == 1024 ) # CLI wins + assert ( + _extra_args_n_ubatch( + ["-b", "1024"], + env = { + "LLAMA_ARG_BATCH": "256", + "LLAMA_ARG_UBATCH": "4096", + }, + ) + == 1024 + ) assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "notint"}) is None def test_env_main_cache_type_for_budget(self): diff --git a/studio/backend/tests/test_slot_offload_fit.py b/studio/backend/tests/test_slot_offload_fit.py index d354c7e113..6344905332 100644 --- a/studio/backend/tests/test_slot_offload_fit.py +++ b/studio/backend/tests/test_slot_offload_fit.py @@ -36,6 +36,7 @@ def _backend( vocab = 248320, embd = 5120, kv_fixed_mib = 0, + kv_calls = None, ): """Backend with the dims the compute buffer reads; KV mocked to a fixed size so the only slot-dependent term is the compute buffer (485 MiB/slot f32 output x 1.15).""" @@ -43,7 +44,17 @@ def _backend( b._vocab_size = vocab b._embedding_length = embd b._key_length_mla = None - b._estimate_kv_cache_bytes = lambda ctx, t = None, **k: kv_fixed_mib * MIB + + def estimate( + ctx, + t = None, + **kwargs, + ): + if kv_calls is not None: + kv_calls.append(kwargs) + return kv_fixed_mib * MIB + + b._estimate_kv_cache_bytes = estimate b._can_estimate_kv = lambda: True return b @@ -55,6 +66,7 @@ def _run( gpus, total_by_idx, overhead_mib = 0, + swa_full = False, ): return b._slots_that_fit_on_gpu( n_parallel, @@ -66,7 +78,8 @@ def _run( FRAC, int(overhead_mib * MIB), 1, - 512, + n_ubatch = 512, + swa_full = swa_full, ) @@ -113,3 +126,16 @@ class TestSlotsThatFitOnGpu: # base 19500 (= 22500 total at par-independent terms) the same par3 fit holds. gi, use_fit, slots = _run(_backend(kv_fixed_mib = 3000), 4, 19500, [(0, 24576)], {0: 24576}) assert use_fit is False and slots == 3 + + def test_swa_full_is_used_for_every_candidate(self): + calls = [] + _run( + _backend(kv_calls = calls), + 4, + 22500, + [(0, 24576)], + {0: 24576}, + swa_full = True, + ) + assert calls + assert all(call["swa_full"] is True for call in calls) diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 23c70f8499..88be5d8976 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -209,6 +209,13 @@ def test_already_in_target_state_reloads_on_tensor_parallel_change(loaded, reque assert _target_state(_loaded_backend(loaded), requested) is False +def test_already_in_target_state_reloads_when_swa_full_env_changes(monkeypatch): + backend = _loaded_backend(False) + backend._swa_full = False + monkeypatch.setenv("LLAMA_ARG_SWA_FULL", "1") + assert _target_state(backend, False) is False + + def test_already_in_target_state_reconciles_split_mode_extras(): # Tensor engaged via --split-mode in extras (boolean omitted/default False) # must match a server already running tensor mode -- no spurious reload. diff --git a/studio/backend/tests/test_tp_vision_regression.py b/studio/backend/tests/test_tp_vision_regression.py index 5dfc38f9af..1781bd70ae 100644 --- a/studio/backend/tests/test_tp_vision_regression.py +++ b/studio/backend/tests/test_tp_vision_regression.py @@ -663,6 +663,29 @@ def test_tensor_off_echo_preserves_multi_gpu_fallback(): ) +def test_route_dedupe_reloads_when_swa_full_env_changes(monkeypatch): + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + backend = _fallback_loaded_backend(layer_preserves_tensor_intent = False) + monkeypatch.setenv("LLAMA_ARG_SWA_FULL", "1") + + request = LoadRequest(model_path = "owner/repo") + assert inference_routes._request_matches_loaded_settings(request, backend) is False + + +def test_route_dedupe_ignores_swa_full_for_diffusion(monkeypatch): + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + backend = _fallback_loaded_backend(layer_preserves_tensor_intent = False) + backend._is_diffusion = True + monkeypatch.setenv("LLAMA_ARG_SWA_FULL", "1") + + request = LoadRequest(model_path = "owner/repo") + assert inference_routes._request_matches_loaded_settings(request, backend) is True + + def test_explicit_split_mode_layer_extras_reloads_after_multi_gpu_fallback(): """Tensor intent can be dropped via extras too: an explicit --split-mode layer matches the stored fallback extras but must still reload (reviewer.py P1, #6659).""" diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 4e35f7b319..08d17f2a65 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1507,6 +1507,8 @@ async function autoLoadSmallestModel(): Promise<{ // The safetensors fallback omits both fields and uses HF auto-placement. gpu_ids?: number[]; gpu_memory_mode?: "auto" | "manual"; + cache_type_kv?: string | null; + tensor_parallel?: boolean | null; }): Promise { const validation = await validateModel({ ...payload, @@ -1595,6 +1597,8 @@ async function autoLoadSmallestModel(): Promise<{ max_seq_length: fitMaxSeqLength, is_lora: false, gguf_variant: candidate.ggufVariant, + cache_type_kv: config.kvCacheDtype, + tensor_parallel: config.tensorParallel, // The same remembered-derived GPU pick the load below sends. ...(candidate.kind === "gguf" ? { diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index a40867beea..60b737fb68 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -185,6 +185,8 @@ export async function validateModel( // /load. Default placement is sized against the selected GPUs. max_seq_length: payload.max_seq_length, load_in_4bit: payload.load_in_4bit, + cache_type_kv: payload.cache_type_kv ?? null, + tensor_parallel: payload.tensor_parallel ?? false, gpu_ids: payload.gpu_ids, // Manual placement is an explicit override: Auto layers use llama.cpp // --fit, while a pinned layer count is owned by the user. Tell validate diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index d4057591b0..bc7227e70d 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -817,6 +817,8 @@ export function useChatModelRuntime() { load_in_4bit: true, is_lora: isLora, gguf_variant: ggufVariant ?? null, + cache_type_kv: loadKvCacheDtype, + tensor_parallel: loadTensorParallel, gpu_ids: validateGpuIds ?? undefined, ...(isGguf ? { gpu_memory_mode: loadGpuMemoryMode } : {}), }); diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index a070c9cb1f..44436b92df 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -1122,6 +1122,8 @@ export function SharedComposer({ gguf_variant: sel.ggufVariant ?? null, trust_remote_code: loadTrustRemoteCode, chat_template_override: effectiveChatTemplateOverride, + cache_type_kv: ownConfig.kvCacheDtype ?? null, + tensor_parallel: effectiveTensorParallel, // Scope the validate to the picked GPUs. GGUF-only, like the load // below: a non-GGUF target must not inherit a hidden GGUF GPU pick. ...(targetIsGguf