"""Assemble a full VRAM breakdown and a llama.cpp launch-command preview.""" from __future__ import annotations from dataclasses import dataclass, field from .quant import weight_bytes, QUANT_BPW from .kv import kv_cache_bytes, compute_scratch_bytes, cache_dtype_bytes from .yarn import yarn_effective_context, yarn_warnings from .gpu import gpu_split, GpuSplitResult, GpuSpec from .draft import DraftInputs, draft_bytes from .mmproj import Mmproj @dataclass class ModelArch: """Architecture parameters (fetched from GGUF or entered manually).""" name: str = "" architecture: str = "" n_layer: int = 0 n_embd: int = 0 n_head: int = 0 n_head_kv: int = 0 training_ctx: int = 0 params: int = 0 # total params including MoE experts rope_freq_base: float = 10000.0 n_expert: int = 0 # MoE n_expert_used: int = 0 n_mtp: int = 0 # MTP heads (e.g. DeepSeek-V3 = 1) # Hybrid-attention models (Qwen3.6 qwen3_5/qwen3_5_moe, Ornith, Gemma4): # only `full_attention` layers + MTP heads bear O(n_ctx) KV; the interleaved # linear/recurrent (Gated DeltaNet) layers use a constant-size state. When # n_full_attn_layers > 0 it overrides n_layer for KV layer counts; 0 ⇒ all # layers carry KV (the old, pure-attention behavior). n_full_attn_layers: int = 0 head_dim: int = 0 # explicit head dim (0 ⇒ derive n_embd // n_head) full_attention_interval: int = 0 # informational; GGUF-derived [L..F] interval def _full_attn_layers(arch: ModelArch) -> int: """Resolved count of layers that allocate KV cache (MTP accounted separately). 0 ⇒ all layers (pure-attention / unknown-architecture fallback). """ return arch.n_full_attn_layers or arch.n_layer @dataclass class Inputs: quant: str = "Q4_K_M" n_ctx: int = 8192 cache_dtype: str = "f16" flash_attn: bool = True compute_dtype: str = "f16" n_batch: int = 512 n_ubatch: int = 512 n_prompt: int = 0 # informational only — NOT read by estimate(). llama.cpp # pre-allocates KV for the full -c window regardless of # the actual prompt length. Used by the UI for prefill-pass # info + prompt-cache hints. # Server slots (--parallel). -c is the *total* KV budget across all slots, so # N slots × C ctx each needs -c = N*C. parallel_sizing toggles whether the # predictor sizes KV for that (True) or leaves KV at n_ctx (False, default). n_parallel: int = 1 parallel_sizing: bool = False # YaRN rope_freq_scale: float = 1.0 yarn_ext_factor: float = -1.0 yarn_attn_factor: float = 1.0 yarn_beta_fast: float = 32.0 yarn_beta_slow: float = 1.0 # GPU / multi-GPU gpu_vram_gb: list[float] = field(default_factory=lambda: [24.0]) split_mode: str = "layer" # layer | tensor | none main_gpu: int = 0 tensor_split: list[float] | None = None # explicit --tensor-split ratios unified_flags: list[bool] | None = None # per-device is_unified # MTP mtp_cache_dtype: str | None = None # MTP extra-layer KV dtype (else cache_dtype) # Speculative decoding draft: DraftInputs | None = None # Multimodal mmproj: Mmproj | None = None # margin safety_margin_pct: float = 5.0 _QUANTIZED_CACHE = {"q8_0", "q8_1", "q4_0", "q4_1", "q5_0", "q5_1"} @dataclass class Breakdown: weights_bytes: float = 0.0 kv_cache_bytes: float = 0.0 compute_scratch_bytes: float = 0.0 mtp_overhead_bytes: float = 0.0 draft_bytes_: float = 0.0 mmproj_bytes_: float = 0.0 gguf_overhead_bytes: float = 0.0 safety_margin_bytes: float = 0.0 total_bytes: float = 0.0 effective_context: int = 0 effective_kv_ctx: int = 0 warnings: list[str] = field(default_factory=list) gpu: GpuSplitResult | None = None def estimate(arch: ModelArch, inp: Inputs) -> Breakdown: weights = weight_bytes(arch.params, inp.quant) mtp_dt = inp.mtp_cache_dtype or inp.cache_dtype # Effective KV context. --parallel N makes -c the *total* KV budget across # all slots, so N slots × C ctx each requires -c = N*C. When parallel_sizing # is on, the predictor sizes KV for that; otherwise KV stays at n_ctx. if inp.parallel_sizing and inp.n_parallel > 1: kv_ctx = inp.n_ctx * inp.n_parallel else: kv_ctx = inp.n_ctx # KV cache (target; MTP extra layers use mtp_dt when set) kv = kv_cache_bytes( n_layer=arch.n_layer, n_embd=arch.n_embd, n_head=arch.n_head, n_head_kv=arch.n_head_kv, n_ctx=kv_ctx, cache_dtype=inp.cache_dtype, n_mtp=0, # MTP KV accounted separately for the dual-dtype case flash_attn=inp.flash_attn, n_full_attn_layers=arch.n_full_attn_layers, head_dim_override=arch.head_dim, ) # compute / activation scratch (physical tile = ubatch; logits buffer = batch) scratch = compute_scratch_bytes( n_layer=arch.n_layer, n_embd=arch.n_embd, n_head=arch.n_head, n_head_kv=arch.n_head_kv, n_batch=inp.n_batch, compute_dtype=inp.compute_dtype, cache_dtype=inp.cache_dtype, flash_attn=inp.flash_attn, n_mtp=arch.n_mtp, n_ubatch=inp.n_ubatch, n_full_attn_layers=arch.n_full_attn_layers, head_dim_override=arch.head_dim, ) # MTP overhead: the extra layers' weights + their KV (routed through the # GPU split below, and also displayed separately). mtp_overhead = 0.0 if arch.n_mtp > 0: if arch.n_layer > 0: per_layer_params = arch.params / arch.n_layer mtp_weights = per_layer_params * (QUANT_BPW[inp.quant] / 8.0) else: mtp_weights = 0.0 head_dim = arch.head_dim or (arch.n_embd // arch.n_head if arch.n_head > 0 else 0) mtp_kv = ( arch.n_mtp * kv_ctx * 2 * arch.n_head_kv * head_dim * cache_dtype_bytes(mtp_dt) ) mtp_overhead = mtp_weights + mtp_kv # Draft model (speculative decoding): weights + its own KV. draft_bd = draft_bytes( draft=inp.draft or DraftInputs(), target_n_ctx=kv_ctx, target_n_embd=arch.n_embd, target_n_head=arch.n_head, target_n_head_kv=arch.n_head_kv, ) if inp.draft is not None else None draft_total = draft_bd.total_bytes if draft_bd and draft_bd.enabled else 0.0 # Multimodal projector. mmproj_total = 0.0 if inp.mmproj is not None and inp.mmproj.enabled: mmproj_total = inp.mmproj.bytes_ # GGUF header / alignment overhead: small constant per file, rough estimate gguf_overhead = max(arch.n_layer * 4096, 1 << 20) # >= 1 MiB # The split handles weights + KV (target) + MTP KV + draft (on main GPU) + # mmproj (on main GPU, if offloaded). scratch rides on the main GPU. split_weights = weights + mtp_weights_only(arch, inp) split_kv = kv + (mtp_overhead - mtp_weights_only(arch, inp)) extra_main = draft_total # draft rides on the main GPU if inp.mmproj is not None and inp.mmproj.enabled and inp.mmproj.offload: extra_main += mmproj_total subtotal = ( weights + kv + scratch + mtp_overhead + draft_total + mmproj_total + gguf_overhead ) margin = subtotal * (inp.safety_margin_pct / 100.0) total = subtotal + margin eff_ctx = yarn_effective_context(arch.training_ctx, inp.rope_freq_scale) warns = yarn_warnings( training_ctx=arch.training_ctx, target_ctx=inp.n_ctx, rope_freq_scale=inp.rope_freq_scale, yarn_ext_factor=inp.yarn_ext_factor, yarn_attn_factor=inp.yarn_attn_factor, ) # GPU split specs = _gpu_specs(inp) gpu = gpu_split( gpu_specs=specs, weights_bytes=split_weights, kv_bytes=split_kv, scratch_bytes=scratch + extra_main, split_mode=inp.split_mode, main_gpu=inp.main_gpu, tensor_split=inp.tensor_split, cache_dtype_quantized=inp.cache_dtype in _QUANTIZED_CACHE, ) warns.extend(gpu.warnings) if inp.split_mode == "tensor" and not inp.flash_attn: warns.append("--split-mode tensor requires flash attention on.") if inp.split_mode == "tensor" and inp.cache_dtype in _QUANTIZED_CACHE: warns.append( "--split-mode tensor disallows quantized KV cache; use f16/bf16/f32." ) # batch/ubatch constraint: physical batch must be <= logical batch. if inp.n_ubatch > inp.n_batch and inp.n_batch > 0: warns.append( f"-ub (ubatch={inp.n_ubatch}) must be ≤ -b (batch={inp.n_batch}); " f"clamping ubatch to {inp.n_batch}." ) # --parallel accounting notes. if inp.n_parallel > 1 and not inp.parallel_sizing: warns.append( f"--parallel {inp.n_parallel}: -c ({inp.n_ctx}) is the total KV " f"budget across all slots → per-slot context ≈ " f"{inp.n_ctx // inp.n_parallel}. Enable “size KV for N slots” " f"to size KV for {inp.n_parallel} slots × {inp.n_ctx} each " f"(-c {kv_ctx})." ) elif inp.parallel_sizing and inp.n_parallel > 1: warns.append( f"KV sized for {inp.n_parallel} slots × {inp.n_ctx} = " f"{kv_ctx} tokens (-c {kv_ctx} --parallel {inp.n_parallel}). " f"Per-slot context = {inp.n_ctx}." ) # MTP — single neutral hint, not documented-as-incompatible warnings. In # current llama.cpp the model's MTP heads are exposed only via # `--spec-type draft-mtp` (a draft source); there is no standalone # `--mtp N` server flag, and no documented incompatibility with # `--parallel` or `--mmproj`. We surface only the one factual note so users # know MTP weights + extra-layer KV are accounted for in the breakdown. if arch.n_mtp > 0: warns.append( f"Model has {arch.n_mtp} MTP head(s) — extra weights and KV are " f"included in the breakdown; enable via --spec-type draft-mtp " f"(no standalone --mtp server flag exists in current llama.cpp)." ) # KV-quant honesty warnings. if inp.cache_dtype in {"q4_0", "q4_1"} and inp.n_ctx > 32768: warns.append( f"Quantized {inp.cache_dtype} KV dequant overhead may make attention " f"slower than f16 at this context length (measured up to ~92% slower " f"at 64K); the VRAM saving may cost throughput." ) if arch.architecture and arch.architecture.startswith("gemma") \ and inp.cache_dtype in _QUANTIZED_CACHE: warns.append( "Gemma + quantized KV cache has a known GPU-util regression " "(util 20–30%, CPU spikes); consider f16/bf16 KV for Gemma models." ) return Breakdown( weights_bytes=weights, kv_cache_bytes=kv, compute_scratch_bytes=scratch, mtp_overhead_bytes=mtp_overhead, draft_bytes_=draft_total, mmproj_bytes_=mmproj_total, gguf_overhead_bytes=gguf_overhead, safety_margin_bytes=margin, total_bytes=total, effective_context=eff_ctx, effective_kv_ctx=kv_ctx, warnings=warns, gpu=gpu, ) def mtp_weights_only(arch: ModelArch, inp: Inputs) -> float: """The weight portion of the MTP overhead (the KV portion is split_kv).""" if arch.n_mtp <= 0 or arch.n_layer <= 0: return 0.0 per_layer_params = arch.params / arch.n_layer return per_layer_params * (QUANT_BPW[inp.quant] / 8.0) def _gpu_specs(inp: Inputs) -> list[GpuSpec]: flags = inp.unified_flags or ([False] * len(inp.gpu_vram_gb)) return [ GpuSpec(vram_gb=v, is_unified=(i < len(flags) and flags[i])) for i, v in enumerate(inp.gpu_vram_gb) ] def _fit_target_mib(inp: Inputs) -> int: """Per-device --fit-target margin in MiB, from safety_margin_pct. --fit-target is a *per-device* MiB reserve. We derive one from the safety margin as a share of the smallest configured GPU's VRAM (the tightest device is what binds). Returns 0 to omit the flag (margin disabled). """ if not inp.gpu_vram_gb: return 0 if inp.safety_margin_pct <= 0: return 0 smallest_gb = min(inp.gpu_vram_gb) margin_bytes = smallest_gb * (1 << 30) * (inp.safety_margin_pct / 100.0) return int(round(margin_bytes / (1 << 20))) def format_bytes(n: float) -> str: """Human-readable byte size.""" n = float(n) if n < 0: return "-" + format_bytes(-n) units = [("GiB", 1 << 30), ("MiB", 1 << 20), ("KiB", 1 << 10)] for label, size in units: if n >= size: return f"{n / size:.2f} {label}" return f"{n:.0f} B" # AMD ROCm floating-point block quants that stock llama.cpp cannot run. They # require the pinned ciru-ai/ROCmFPX runner (see # https://github.com/ciru-ai/ROCmFPX). Listed so command_preview can flag it. ROCMFP_QUANTS = {"ROCmFP4", "ROCmFPX"} # NVIDIA NVFP4 (GGML_TYPE_NVFP4=40; llama.cpp PR #22196, merged Apr 2026). # Stock llama.cpp builds run it; Blackwell gets the FP4 tensor-core speedup, # older cards (Ada/Ampere/Hopper) get VRAM savings only. Listed so # command_preview can flag the experimental/build caveat. NVFP_QUANTS = {"nvfp4"} def command_preview(arch: ModelArch, inp: Inputs) -> str: """Generate a llama.cpp launch command from the current inputs.""" is_rocmfp = inp.quant in ROCMFP_QUANTS runner = "rocmfpx-llama-server" if is_rocmfp else "llama-server" parts = [runner] parts.append(f"-m model-{inp.quant}.gguf") # --parallel sizes -c across slots. With parallel_sizing the KV context is # n_ctx * n_parallel; emit that as -c and add --parallel N. kv_ctx = inp.n_ctx * inp.n_parallel if (inp.parallel_sizing and inp.n_parallel > 1) else inp.n_ctx parts.append(f"-c {kv_ctx}") if inp.n_parallel > 1: parts.append(f"--parallel {inp.n_parallel}") parts.append("-ngl 999") # full offload assumption parts.append(f"-b {inp.n_batch}") if inp.n_ubatch != inp.n_batch: parts.append(f"-ub {inp.n_ubatch}") if inp.cache_dtype != "f16": parts.append(f"--cache-type-k {inp.cache_dtype}") parts.append(f"--cache-type-v {inp.cache_dtype}") if inp.flash_attn: parts.append("--flash-attn") # Free throughput at long context: KV-shift reuse of cached prompt chunks. # On by default off (default value 0); we suggest 256 for large windows. if inp.n_ctx >= 65536: parts.append("--cache-reuse 256") # Complement path: we emit fully explicit flags, then --fit on with a # per-device margin target so runtime --fit is a no-op safety net that # can only *reduce* usage, never exceed our estimate. The margin is the # user's "how scared am I" knob (safety_margin_pct). We never claim a # guarantee; the margin is the honest version. parts.append("--fit on") margin_mib = _fit_target_mib(inp) if margin_mib: parts.append(f"--fit-target {margin_mib}") # Multi-GPU: emit --split-mode / --main-gpu / --tensor-split only when the # user has more than one GPU or chose a non-default mode. multi = len(inp.gpu_vram_gb) > 1 if multi or inp.split_mode != "layer": parts.append(f"--split-mode {inp.split_mode}") if inp.split_mode == "none": parts.append(f"--main-gpu {inp.main_gpu}") # With no layer splitting the user is pinning the run to one device; the # companion flag keeps host-side tensor ops on the host (no device # offload of operator execution). parts.append("--no-op-offload") if multi and inp.tensor_split: parts.append( "--tensor-split " + ",".join(str(s) for s in inp.tensor_split) ) # MTP draft KV dtype (only when it differs from the target KV dtype) if arch.n_mtp > 0 and inp.mtp_cache_dtype and inp.mtp_cache_dtype != inp.cache_dtype: # llama.cpp applies a single -ctk/-ctv; we note the MTP draft dtype. parts.append( f"# MTP draft KV dtype: {inp.mtp_cache_dtype} (target: {inp.cache_dtype})" ) # Speculative decoding. Two branches: # - weightless (n-gram family): no draft model, ~zero extra VRAM. Emit # only --spec-type (no -md). # - weighted (draft-*): a draft model with weights + KV. Emit -md plus # the full --spec-draft-* namespace, not the old -md-only stub. if inp.draft is not None and inp.draft.enabled: d = inp.draft parts.append(f"--spec-type {d.spec_type}") from .draft import is_weightless if not is_weightless(d.spec_type): parts.append(f"-md draft-{d.quant}.gguf") parts.append("--spec-draft-ngl 999") # full draft offload parts.append("--spec-draft-device all") if d.n_max: parts.append(f"--spec-draft-n-max {d.n_max}") if d.p_min: parts.append(f"--spec-draft-p-min {d.p_min}") if d.p_split: parts.append(f"--spec-draft-p-split {d.p_split}") if d.n_min: parts.append(f"--spec-draft-n-min {d.n_min}") else: # Weightless n-gram: no -md. n-min is meaningful for ngram-mod (PR # #19164 — default n-min=0 triggers acceptance resets). if d.n_min: parts.append(f"--spec-draft-n-min {d.n_min}") elif d.spec_type == "ngram-mod": parts.append( "# ngram-mod: consider --spec-draft-n-min 8 to avoid " "acceptance-rate resets (PR #19164)." ) # Multimodal projector if inp.mmproj is not None and inp.mmproj.enabled: parts.append(f"--mmproj {inp.mmproj.filename or 'mmproj.gguf'}") if not inp.mmproj.offload: parts.append("--no-mmproj-offload") # YaRN / rope yarn_args = [] if inp.rope_freq_scale != 1.0: yarn_args.append(f"--rope-freq-scale {inp.rope_freq_scale}") if arch.rope_freq_base != 10000.0: yarn_args.append(f"--rope-freq-base {arch.rope_freq_base}") if inp.yarn_ext_factor >= 0.0: yarn_args.append(f"--yarn-ext-factor {inp.yarn_ext_factor}") if inp.yarn_attn_factor != 1.0: yarn_args.append(f"--yarn-attn-factor {inp.yarn_attn_factor}") if inp.yarn_beta_fast != 32.0: yarn_args.append(f"--yarn-beta-fast {inp.yarn_beta_fast}") if inp.yarn_beta_slow != 1.0: yarn_args.append(f"--yarn-beta-slow {inp.yarn_beta_slow}") if inp.n_ctx > arch.training_ctx and arch.training_ctx > 0: yarn_args.append("--rope-scaling yarn") if yarn_args: parts.extend(yarn_args) if arch.n_mtp > 0: # There is no standalone `--mtp N` server flag in llama.cpp. The only # way to use MTP heads at inference is as the spec source # `--spec-type draft-mtp`. Emit a comment so users know how to wire # the VRAM they've budgeted for. parts.append( f"# MTP heads ({arch.n_mtp}) — enable via --spec-type draft-mtp " f"(no standalone --mtp server flag exists in llama.cpp)." ) if is_rocmfp: parts.append( f"# NOTE: {inp.quant} needs the ciru-ai/ROCmFPX runner " f"(stock llama.cpp cannot read these tensor types)." ) if inp.quant in NVFP_QUANTS: parts.append( "# NOTE: nvfp4 — Blackwell GPUs get the FP4 tensor-core speedup; " "older cards (Ada/Ampere/Hopper) get VRAM savings only (kernels " "run, slower). Experimental; verify your llama.cpp build has " "GGML_TYPE_NVFP4 support." ) # Assumed-on server defaults (not emitted as flags to keep the runnable # command clean): continuous batching, jinja chat templates, prompt # caching, flash-attention auto. Listed so users copying from older # tutorials aren't surprised by the differing behavior. parts.append( "# server defaults assumed on: cont-batching, jinja, cache-prompt, " "flash-attn=auto" ) return " \\\n ".join(parts)