80f75c0df2
NVFP4 runs on sub-sm_90 cards via the Marlin W4A16 weight-only fallback
since vLLM v0.24 (floor = capability 7.5, marlin_utils_fp4). Live-confirmed
on 2x RTX 3090 sm_86 2026-07-11: the shipped dual/nvfp4/mtp.yml boots at
262K + fp8 KV + MTP n=3 (accept 97%+), 69.7/85.5 decode TPS, and scores
8-pack think-off 110/150 — a statistical tie with the fp8 production
tier's 109. c3 hiding these slugs as "incompatible-hw" was factually wrong.
Semantic split — required_sm stays the NATIVE floor; new registry field
`fallback_sm` = the weight-only-fallback floor:
- kv-calc: in the band [fallback_sm, required_sm) the fit is PRICED
normally and annotated with `hw_fallback` {required_sm, card_sm, note}
instead of returning incompatible-hw (both --fit and --fit-all).
- c3: FitVerdict carries hw_fallback; fits-* + fallback -> "⚑" glyph,
VISIBLE by default. incompatible-hw hiding stays for true incompatibles.
- gates.py/compat.py: the hard SM floor becomes fallback_sm when present
(else required_sm, unchanged).
- Registry: fallback_sm=7.5 on the 4 NVFP4 slugs + status_note honesty
pass ("REFUSES on Ampere" -> fallback reality + measured numbers);
same corrections in the 4 compose headers.
- test-kv-calc-fit section (f) rewritten to the new contract.
On Ampere these slugs still carry the honest downside in the note: no
speed edge (~20% slower than the AutoRound tier for the same model) —
their sub-sm_90 value is models where NVFP4 is the only quant.
Guards green: kv-calc-fit, registry-disk, status-drift, switch/launch
parity, profiles-compat, pullgate, pull, registry-json/emit, diagnose,
launch-compat, default resolvers, kv-calc --calibration; c3 248 fast
tests + headless incompatible-hw hide tests.
Claude-Session: https://claude.ai/code/session_01EfF565T9eSLaqGzidyJ1Pm
Co-authored-by: noonghunna <10742901+noonghunna@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1823 lines
82 KiB
Python
Executable File
1823 lines
82 KiB
Python
Executable File
#!/bin/sh
|
||
''':'
|
||
exec python3 "$0" "$@"
|
||
':'''
|
||
from __future__ import annotations
|
||
|
||
__doc__ = """kv-calc.py — predict per-card VRAM budget for vLLM composes.
|
||
|
||
Predicts (per card, after TP split):
|
||
- Model weights
|
||
- KV pool (attention layers only — recurrent / SSM states show up in activation)
|
||
- Activation peak (model-specific: Qwen GDN forward, Gemma SWA + dense MLP)
|
||
- Cudagraph + workspace overhead
|
||
- Drafter overhead (MTP / DFlash)
|
||
- Total vs available VRAM
|
||
- Verdict: PASS / TIGHT / FAIL
|
||
|
||
Four models modelled:
|
||
- Qwen 3.6 27B (DeltaNet hybrid: 16 full_attention + 48 GDN)
|
||
- Qwen 3.6 35B-A3B (MoE + DeltaNet hybrid: 10 attention + 30 GDN)
|
||
- Gemma 4 31B (SWA + dense MLP: 10 full_attention + 50 sliding_attention)
|
||
- Gemma 4 26B-A4B (MoE + SWA: 5 full_attention + 25 sliding_attention)
|
||
|
||
vLLM rate-limits KV pool to fit available budget; this predictor models that
|
||
capping behavior. When the requested KV pool exceeds what fits, the verdict
|
||
is TIGHT (vLLM will cap pool — effective concurrency reduced) not FAIL.
|
||
|
||
Anchored to:
|
||
- PerfMamba (arxiv 2511.22849) — Qwen GDN block-wise state materialization scaling
|
||
https://arxiv.org/html/2511.22849
|
||
- TurboQuant (arxiv 2504.19874, ICLR 2026) — TQ3 byte savings
|
||
https://arxiv.org/abs/2504.19874
|
||
- PagedAttention (arxiv 2309.06180) — KV pool layout
|
||
https://arxiv.org/abs/2309.06180
|
||
|
||
Calibrated against measured BENCHMARKS.md rows per model. Coefficients reflect
|
||
club-3090's empirical findings. See docs/KV_MATH.md for the derivation +
|
||
calibration trace.
|
||
|
||
Usage:
|
||
bash tools/kv-calc.py --compose dual-turbo --vram 24 # Qwen (default model)
|
||
bash tools/kv-calc.py --model gemma-4-31b --compose gemma-dual-int8 --vram 24
|
||
bash tools/kv-calc.py --model gemma-4-31b --solve-max-ctx --kv-format int8_per_token_head --tp 2 --vram 24
|
||
bash tools/kv-calc.py --calibration # all calibrated models, grouped per-model
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import logging
|
||
import os
|
||
import sys
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||
if str(REPO_ROOT) not in sys.path:
|
||
sys.path.insert(0, str(REPO_ROOT))
|
||
|
||
from scripts.lib.profiles.compat import load_profiles
|
||
from scripts.lib.profiles.compose_registry import COMPOSE_REGISTRY
|
||
|
||
|
||
# =============================================================================
|
||
# Model specs
|
||
# =============================================================================
|
||
|
||
def _load_profiles_silent():
|
||
logger = logging.getLogger("compat")
|
||
old_disabled, old_env = logger.disabled, os.environ.get("CLUB3090_LOG_LEVEL")
|
||
logger.disabled = True
|
||
os.environ["CLUB3090_LOG_LEVEL"] = "CRITICAL"
|
||
try:
|
||
return load_profiles()
|
||
finally:
|
||
logger.disabled = old_disabled
|
||
if old_env is None:
|
||
os.environ.pop("CLUB3090_LOG_LEVEL", None)
|
||
else:
|
||
os.environ["CLUB3090_LOG_LEVEL"] = old_env
|
||
|
||
|
||
PROFILES = _load_profiles_silent()
|
||
|
||
|
||
def _weight_size(model, variant):
|
||
value = model.weights[variant]["size_gb"]
|
||
if not isinstance(value, (int, float)):
|
||
raise ValueError(f"weight size for {model.id}/{variant} is not numeric: {value}")
|
||
return float(value)
|
||
|
||
|
||
def _load_model_specs_from_yaml(profiles):
|
||
qwen, gemma = profiles.models["qwen3.6-27b"], profiles.models["gemma-4-31b"]
|
||
qwen_moe, gemma_moe = profiles.models["qwen3.6-35b-a3b"], profiles.models["gemma-4-26b-a4b"]
|
||
gemma12 = profiles.models["gemma-4-12b"]
|
||
q_fields = ("hidden_size", "num_hidden_layers", "num_gdn_layers", "num_attn_layers", "num_attn_heads", "num_kv_heads", "head_dim_attn", "linear_num_v_heads", "linear_num_k_heads", "linear_v_head_dim", "linear_k_head_dim", "linear_conv_kernel_dim", "max_ctx_supported", "attention_k_eq_v")
|
||
g_fields = ("hidden_size", "intermediate_size", "num_hidden_layers", "num_full_attn_layers", "num_sliding_attn_layers", "num_attn_heads", "num_kv_heads", "head_dim_sliding", "global_head_dim", "sliding_window", "max_ctx_supported", "attention_k_eq_v")
|
||
gm_fields = (*g_fields, "num_global_kv_heads", "num_experts", "num_experts_per_tok", "moe_intermediate_size", "active_params_b", "mtp_num_hidden_layers")
|
||
qm_fields = (*q_fields, "num_experts", "num_experts_per_tok", "moe_intermediate_size", "shared_expert_intermediate_size", "active_params_b", "mtp_num_hidden_layers")
|
||
qspec = {"model_id": qwen.id, "model_family": qwen.family, **{k: getattr(qwen, k) for k in q_fields}, "valid_tp": list(qwen.valid_tp), "weights_total_gb": _weight_size(qwen, qwen.default_weight_variant), "weights_nvfp4_gb": _weight_size(qwen, "nvfp4"), "mamba_state_bytes": 4, "chunk_size": 256, "mtp_n_default": profiles.drafters["qwen-mtp-builtin"].n_default}
|
||
qmspec = {"model_id": qwen_moe.id, "model_family": qwen_moe.family, **{k: getattr(qwen_moe, k) for k in qm_fields}, "valid_tp": list(qwen_moe.valid_tp), "weights_total_gb": _weight_size(qwen_moe, qwen_moe.default_weight_variant), "weights_gptq_gb": _weight_size(qwen_moe, "gptq_int4"), "weights_nvfp4_gb": _weight_size(qwen_moe, "nvfp4"), "mamba_state_bytes": 4, "chunk_size": 256, "mtp_n_default": profiles.drafters["qwen-mtp-builtin"].n_default}
|
||
gspec = {"model_id": gemma.id, "model_family": gemma.family, **{k: getattr(gemma, k) for k in g_fields}, "valid_tp": list(gemma.valid_tp), "weights_int4_gb": _weight_size(gemma, "autoround-int4"), "weights_awq_gb": _weight_size(gemma, "awq"), "weights_bf16_gb": _weight_size(gemma, "bf16"), "drafter_mtp_gb": float(profiles.drafters["gemma-it-assistant"].vram_footprint_gb), "drafter_dflash_gb": float(profiles.drafters["gemma-dflash"].vram_footprint_gb), "mtp_n_default": profiles.drafters["gemma-it-assistant"].n_default}
|
||
gmspec = {"model_id": gemma_moe.id, "model_family": gemma_moe.family, **{k: getattr(gemma_moe, k) for k in gm_fields}, "valid_tp": list(gemma_moe.valid_tp), "weights_int4_gb": _weight_size(gemma_moe, "autoround-int4-mixed"), "weights_awq_gb": _weight_size(gemma_moe, "awq"), "drafter_mtp_gb": float(profiles.drafters["gemma-26b-it-assistant"].vram_footprint_gb), "mtp_n_default": profiles.drafters["gemma-26b-it-assistant"].n_default}
|
||
# Gemma-4-12B (gemma4_unified arch). Its TEXT backbone is gemma4-swa-dense
|
||
# shaped (same SWA KV math family as gemma-4-31b), so it rides the SAME
|
||
# "gemma4-swa-dense" prediction path — we set model_family to that internal
|
||
# KV-family tag (NOT its ModelProfile family "gemma4-unified", which is the
|
||
# vLLM arch name). It ships bf16 weights + an Intel AutoRound INT8 variant
|
||
# (~13 GB, W8A16) + an unsloth QAT W4A16 int4 variant (~9.6 GB,
|
||
# compressed-tensors) + a small assistant drafter. weights_int4_gb/weights_awq_gb
|
||
# carry the REAL int4 (qat-w4a16) footprint; weights_int8_gb the INT8 footprint
|
||
# (vllm/gemma-12b-single-int8-mtp path). NOTE: int4/awq FORMERLY pointed at the
|
||
# bf16 blob (23.9 GB) as a placeholder when no int4 variant existed — that
|
||
# over-priced the int4 SINGLE-card path by ~14 GB (false-FAIL'd
|
||
# vllm/gemma-12b-qat-w4a16-single at 27.9 GB vs the ~22.6 GB it actually boots
|
||
# at; the dual path hid it via /tp). Shared activation/overhead + the measured
|
||
# `measured_kv_growing_bpt_tp1` term are UNCHANGED — this is a weight-SIZE fix only.
|
||
# Activation/overhead constants stay the SHARED Gemma dense constants (NOT
|
||
# re-tuned). The ONE measured calibration is the growing KV per-token —
|
||
# `measured_kv_growing_bpt_tp1` (see kv_pool_per_card_bytes): gemma4_unified's
|
||
# global-layer KV measured 1.44x LOWER than the 31B-derived global-only
|
||
# formula predicts, so we ride the measurement for the 12B only.
|
||
g12_bf16 = _weight_size(gemma12, "bf16")
|
||
g12_int8 = _weight_size(gemma12, "autoround-int8")
|
||
g12_int4 = _weight_size(gemma12, "qat-w4a16")
|
||
g12spec = {"model_id": gemma12.id, "model_family": "gemma4-swa-dense", **{k: getattr(gemma12, k) for k in g_fields}, "valid_tp": list(gemma12.valid_tp), "weights_int4_gb": g12_int4, "weights_awq_gb": g12_int4, "weights_bf16_gb": g12_bf16, "weights_int8_gb": g12_int8, "drafter_mtp_gb": float(profiles.drafters["gemma-12b-it-assistant"].vram_footprint_gb), "measured_kv_growing_bpt_tp1": 45632, "mtp_n_default": profiles.drafters["gemma-12b-it-assistant"].n_default}
|
||
# Agents-A1 (InternScience 35B agentic MoE): geometry verified byte-identical
|
||
# to qwen3.6-35b-a3b (same Qwen3-Next MoE arch class) — rides the SAME
|
||
# qwen3-next-moe KV math. Own weights footprint (FP8-dynamic 36 GB, Marlin
|
||
# weight-only on Ampere). No MTP head in the checkpoint (safetensors header
|
||
# scan 2026-07-03: 0 mtp tensors) — mtp_n_default is inert (drafter=None).
|
||
a1 = profiles.models["agents-a1"]
|
||
a1spec = {"model_id": a1.id, "model_family": a1.family, **{k: getattr(a1, k) for k in qm_fields}, "valid_tp": list(a1.valid_tp), "weights_total_gb": _weight_size(a1, a1.default_weight_variant), "mamba_state_bytes": 4, "chunk_size": 256, "mtp_n_default": profiles.drafters["qwen-mtp-builtin"].n_default}
|
||
return {
|
||
"qwen3.6-27b": qspec,
|
||
"qwen3.6-35b-a3b": qmspec,
|
||
"agents-a1": a1spec,
|
||
"gemma-4-31b": gspec,
|
||
"gemma-4-26b-a4b": gmspec,
|
||
"gemma-4-12b": g12spec,
|
||
}
|
||
|
||
|
||
MODEL_SPECS = _load_model_specs_from_yaml(PROFILES)
|
||
QWEN36_27B = MODEL_SPECS["qwen3.6-27b"]
|
||
QWEN36_35B_A3B = MODEL_SPECS["qwen3.6-35b-a3b"]
|
||
GEMMA4_31B = MODEL_SPECS["gemma-4-31b"]
|
||
GEMMA4_26B_A4B = MODEL_SPECS["gemma-4-26b-a4b"]
|
||
GEMMA4_12B = MODEL_SPECS["gemma-4-12b"]
|
||
|
||
|
||
# =============================================================================
|
||
# KV format bytes-per-element
|
||
# =============================================================================
|
||
# (one element = one head dim of one head; K and V counted separately for
|
||
# models where attention_k_eq_v=False. For Gemma the K==V tying halves this
|
||
# at the formula level — see kv_pool_per_card_bytes.)
|
||
# Source: vLLM/HF docs + TurboQuant paper + PR #40391 (INT8 per-token-head).
|
||
|
||
KV_FORMAT_BYTES = {
|
||
"fp16": 2.0,
|
||
"bf16": 2.0,
|
||
"fp8_e5m2": 1.0,
|
||
"fp8_e4m3": 1.0,
|
||
"int8_per_token_head": 1.01, # 1.0 int8 + per-token-head fp16 scale (~1% amortized)
|
||
"q4_0": 0.5 + 0.0625, # 4-bit + per-group scale
|
||
"k8v4": 0.75, # avg of K=int8 V=int4
|
||
"turboquant_3bit_nc": 0.375 + 0.05, # 3 bits + small QJL overhead
|
||
"nvfp4": 0.5 + 0.0625, # 4-bit elements + fp8 block scale per 16
|
||
# (PROJECTED — Blackwell-only, no measured
|
||
# boot on this stack yet; #246 A/B arm 3)
|
||
}
|
||
|
||
INDEXER_FORMAT_BYTES = {
|
||
"fp16": 2.0,
|
||
"bf16": 2.0,
|
||
"fp8_e5m2": 1.0,
|
||
"fp8_e4m3": 1.0,
|
||
"int8": 1.0,
|
||
"fp4": 0.5,
|
||
"int4": 0.5,
|
||
}
|
||
|
||
|
||
# =============================================================================
|
||
# Per-model activation coefficients
|
||
# =============================================================================
|
||
|
||
# ---- Qwen GDN activation-peak per-layer per-token coefficient (bytes) ----
|
||
# Calibrated empirically against measured BENCHMARKS rows. PerfMamba's
|
||
# O(γ·D·N·L) scaling sets the form; this captures fla.ops.chunk
|
||
# implementation details + KV-format-dependent dequant overhead.
|
||
QWEN_GDN_ACTIVATION_COEF = {
|
||
"fp16": 135,
|
||
"bf16": 135,
|
||
"fp8_e5m2": 130,
|
||
"fp8_e4m3": 130,
|
||
"int8_per_token_head": 130,
|
||
"q4_0": 155,
|
||
"k8v4": 155,
|
||
"turboquant_3bit_nc": 165,
|
||
"nvfp4": 130, # PROJECTED — mirror fp8 until a Blackwell boot calibrates it (#246)
|
||
}
|
||
|
||
# ---- Qwen MoE activation + built-in MTP workspace ----
|
||
# Path-B low-anchor fit from the two v0.7.3 preview rows. The per-token GDN
|
||
# coefficient follows the dense-Qwen shape; the small constant captures MoE
|
||
# expert dispatch/router buffers. NOTE: vLLM shards both attention and MoE
|
||
# expert weights across TP ranks, so weights ARE divided by TP for
|
||
# qwen3-next-moe in _weights_per_card_gb() (fixed in #260). The ≈budget peak
|
||
# on the live rows is the elastic KV pool filling spare VRAM, not resident
|
||
# full-quant weights — the old no-/tp assumption over-predicted long-ctx
|
||
# configs (e.g. 262K) into a false FAIL.
|
||
QWEN_MOE_ACTIVATION_COEF = {
|
||
"fp16": 110,
|
||
"bf16": 110,
|
||
"fp8_e5m2": 105,
|
||
"fp8_e4m3": 105,
|
||
"int8_per_token_head": 105,
|
||
"q4_0": 130,
|
||
"k8v4": 130,
|
||
"turboquant_3bit_nc": 140,
|
||
"nvfp4": 105, # PROJECTED — mirror fp8 until a Blackwell boot calibrates it (#246)
|
||
}
|
||
QWEN_MOE_EXPERT_DISPATCH_GB = 0.20
|
||
QWEN_MOE_BUILTIN_MTP_WORKSPACE_GB = 0.10
|
||
|
||
# ---- Gemma activation peak (mostly constant in ctx) ----
|
||
# Unlike Qwen GDN, Gemma's activation peak comes from dense MLP forward +
|
||
# SWA windowed-attention prefill, both bounded by chunked-prefill chunk_size.
|
||
# Result: roughly CONSTANT in max_ctx. Calibrated as a per-TP base (GB) plus
|
||
# a small per-token term to capture residual ctx scaling.
|
||
GEMMA_ACTIVATION_CONST_GB = 1.5 # per card at TP=1 — calibrated, ~scales as 1/TP
|
||
GEMMA_ACTIVATION_PER_TOKEN_BYTES = 8 # tiny ctx scaling term to keep solver well-behaved
|
||
|
||
# ---- Gemma MoE activation peak ----
|
||
# Low-anchor fit from awq.yml and awq-mtp.yml. MoE dispatch is folded into the
|
||
# constant term; external assistant weights are modelled separately via
|
||
# drafter_gb.
|
||
GEMMA_MOE_ACTIVATION_CONST_GB = 1.8
|
||
GEMMA_MOE_ACTIVATION_PER_TOKEN_BYTES = 12
|
||
|
||
# ---- generic-dense activation peak (conservative lower-bound on fit) ----
|
||
# P1 / v0.8.0 Pull-Gate stratum-5/6: for an uncurated standard-dense
|
||
# (MHA/GQA) model we have no measured anchor, so we deliberately OVER-price
|
||
# the activation peak to guarantee generic-dense never predicts more fit
|
||
# than a curated exact branch would for an equivalent shape.
|
||
#
|
||
# The form is a per-token coefficient × hidden_size × max_ctx (per card,
|
||
# divided by TP). The coefficient is chosen to be at least as conservative
|
||
# as the most conservative curated dense activation term in this file. The
|
||
# most conservative curated dense per-token term is Qwen GDN (turboquant)
|
||
# at 165 bytes/layer/token; expressed per hidden-unit that is
|
||
# 165 * num_gdn_layers / hidden_size ≈ 165 * 64 / 2048 ≈ 5.16 bytes per
|
||
# hidden-unit-token for the dense Qwen 27B shape. We round UP to 6.0 to stay
|
||
# strictly on the conservative (lower-bound-on-fit) side for arbitrary
|
||
# dense shapes, and add a fixed per-card floor at least as large as Gemma's
|
||
# constant dense activation term (1.5 GB) so small-ctx configs are never
|
||
# under-priced relative to the curated Gemma dense branch.
|
||
GENERIC_DENSE_ACTIVATION_COEF_PER_HIDDEN = 6.0 # bytes per (hidden_size × token); ≥ most conservative curated dense
|
||
GENERIC_DENSE_ACTIVATION_FLOOR_GB = 1.5 # ≥ Gemma dense constant activation term
|
||
|
||
|
||
# =============================================================================
|
||
# Compose presets (per-model)
|
||
# =============================================================================
|
||
COMPOSE_ALIAS_TEXT = {
|
||
"qwen3.6-27b": "minimal=vllm/minimal dual=vllm/dual nvfp4-single=vllm/qwen-27b-single-nvfp4 nvfp4-dual=vllm/qwen-27b-dual-nvfp4",
|
||
"qwen3.6-35b-a3b": "qwen-a3b-preview-single=vllm/qwen-a3b-preview-single qwen-35b-a3b-dual=vllm/qwen-35b-a3b-dual nvfp4-single=vllm/qwen-35b-a3b-single-nvfp4 nvfp4-dual=vllm/qwen-35b-a3b-dual-nvfp4",
|
||
"agents-a1": "agents-a1-dual=vllm/agents-a1-dual",
|
||
"gemma-4-31b": "gemma-dual=vllm/gemma-bf16-mtp gemma-dual-int8=vllm/gemma-int8-mtp gemma-single=vllm/gemma-mtp-tp1",
|
||
# gemma-4-12b legacy alias namespace is keyed by model id, so reusing the
|
||
# bare `gemma-dual` string here is harmless — compat + the CLI always pass
|
||
# an explicit --model, and the reverse map is keyed by (unique) registry
|
||
# slug. `gemma-dual` → the MTP dual; `gemma-no-mtp` → the no-drafter dual.
|
||
"gemma-4-12b": "gemma-dual=vllm/gemma-12b-dual-bf16-mtp gemma-single-int8-mtp=vllm/gemma-12b-single-int8-mtp",
|
||
"gemma-4-26b-a4b": "gemma-26ba4b-single=vllm/gemma-26ba4b-single gemma-26ba4b-dual=vllm/gemma-26ba4b-dual",
|
||
}
|
||
COMPOSE_ALIASES = {model: tuple(part.split("=", 1) for part in text.split()) for model, text in COMPOSE_ALIAS_TEXT.items()}
|
||
|
||
REGISTRY_TO_LEGACY_COMPOSE = {
|
||
registry: legacy
|
||
for aliases in COMPOSE_ALIASES.values()
|
||
for legacy, registry in aliases
|
||
}
|
||
|
||
COMPOSE_COMPAT_OVERRIDES = {
|
||
("qwen3.6-27b", "minimal"): {"max_num_seqs": 4, "mem_util": 0.90},
|
||
("qwen3.6-27b", "dual"): {"mem_util": 0.95},
|
||
("gemma-4-31b", "gemma-single"): {"kv_format": "fp8_e5m2"},
|
||
}
|
||
|
||
|
||
def _compose_cfg_from_registry(profiles, model_id, legacy_name, registry_name):
|
||
entry = COMPOSE_REGISTRY[registry_name]
|
||
drafter = profiles.drafters[entry["drafter"]] if entry.get("drafter") else None
|
||
cfg = {k: entry[k] for k in ("max_ctx", "max_num_seqs", "tp", "kv_format", "mem_util")}
|
||
cfg["mtp"] = drafter is not None and drafter.spec_method in ("mtp", "mtp_assistant")
|
||
if drafter is not None and drafter.spec_method == "dflash":
|
||
cfg.update({"mtp": False, "dflash_draft_gb": float(drafter.vram_footprint_gb)})
|
||
if drafter is not None:
|
||
cfg["mtp_n"] = int(drafter.n_default)
|
||
# gemma-4-12b rides the shared Gemma dense path (model_family
|
||
# gemma4-swa-dense in MODEL_SPECS); treat its drafter + weights the same way.
|
||
if model_id in ("gemma-4-31b", "gemma-4-12b") and drafter is not None:
|
||
cfg["drafter_gb"] = float(drafter.vram_footprint_gb)
|
||
if model_id in ("gemma-4-31b", "gemma-4-12b"):
|
||
cfg["weights_variant"] = {"awq": "awq", "bf16": "bf16", "autoround-int8": "int8"}.get(entry["weights_variant"], "int4")
|
||
if model_id == "gemma-4-26b-a4b":
|
||
cfg["weights_variant"] = "awq" if entry["weights_variant"] == "awq" else "int4"
|
||
if drafter is not None:
|
||
cfg["drafter_gb"] = float(drafter.vram_footprint_gb)
|
||
if model_id == "qwen3.6-35b-a3b":
|
||
cfg["weights_variant"] = {"gptq_int4": "gptq", "nvfp4": "nvfp4"}.get(entry["weights_variant"], "default")
|
||
if model_id == "qwen3.6-27b":
|
||
cfg["weights_variant"] = "nvfp4" if entry["weights_variant"] == "nvfp4" else "default"
|
||
cfg.update(COMPOSE_COMPAT_OVERRIDES.get((model_id, legacy_name), {}))
|
||
return cfg
|
||
|
||
|
||
COMPOSES = {
|
||
model_id: {legacy: _compose_cfg_from_registry(PROFILES, model_id, legacy, registry) for legacy, registry in aliases}
|
||
for model_id, aliases in COMPOSE_ALIASES.items()
|
||
}
|
||
|
||
|
||
# =============================================================================
|
||
# Calibration: measured BENCHMARKS rows (peak per-card VRAM during bench)
|
||
# =============================================================================
|
||
|
||
CALIBRATION = {
|
||
model_id: [
|
||
(REGISTRY_TO_LEGACY_COMPOSE[row["compose"]], row["vram_gb"], row["measured_peak_gb"], row.get("ctx_override"), row.get("source", ""))
|
||
for row in cal.rows
|
||
]
|
||
for model_id, cal in PROFILES.calibration.items()
|
||
}
|
||
|
||
|
||
# =============================================================================
|
||
# Prediction
|
||
# =============================================================================
|
||
|
||
@dataclass
|
||
class Prediction:
|
||
model: str
|
||
weights_gb: float
|
||
kv_pool_requested_gb: float
|
||
kv_pool_actual_gb: float # capped at available budget (vLLM behavior)
|
||
kv_pool_sliding_fixed_gb: float # Gemma sliding-window fixed term (0 for Qwen)
|
||
activation_gb: float
|
||
cudagraph_overhead_gb: float
|
||
drafter_gb: float
|
||
total_gb: float
|
||
vram_gb: float
|
||
budget_gb: float
|
||
pct_of_vram: float
|
||
verdict: str
|
||
notes: list[str]
|
||
|
||
|
||
@dataclass
|
||
class CacheBreakdown:
|
||
layout: str
|
||
sequences: int
|
||
gpus: int
|
||
tp: int
|
||
attention_kv_growing_gb: float
|
||
attention_kv_sliding_fixed_gb: float
|
||
recurrent_state_gb: float
|
||
compressed_kv_gb: float
|
||
indexer_cache_gb: float
|
||
draft_kv_gb: float
|
||
total_cache_gb: float
|
||
notes: list[str]
|
||
|
||
|
||
def _weights_per_card_gb(spec, tp, weights_variant="default"):
|
||
"""Return per-card weights footprint in GB after TP split."""
|
||
if spec["model_family"] == "qwen3-next-hybrid":
|
||
# nvfp4 (nvidia modelopt, 21.9 GB) vs the default AutoRound INT4 —
|
||
# the community Hopper/Blackwell slugs (vllm/qwen-27b-*-nvfp4).
|
||
if weights_variant == "nvfp4":
|
||
return spec["weights_nvfp4_gb"] / tp
|
||
return spec["weights_total_gb"] / tp
|
||
elif spec["model_family"] == "qwen3-next-moe":
|
||
# vLLM shards attention + MoE expert weights across TP ranks, so
|
||
# per-card weights are /tp (#260). The old no-/tp assumption was
|
||
# mis-inferred from the live ≈budget peak — but that peak is the
|
||
# elastic KV pool filling spare VRAM (vLLM allocates all available
|
||
# KV memory as blocks), not resident full-quant weights. Not dividing
|
||
# over-predicted the fixed footprint and turned long-ctx configs
|
||
# (e.g. 262K dual, which fits) into a false FAIL.
|
||
if weights_variant == "gptq":
|
||
return spec["weights_gptq_gb"] / tp
|
||
if weights_variant == "nvfp4":
|
||
return spec["weights_nvfp4_gb"] / tp
|
||
return spec["weights_total_gb"] / tp
|
||
elif spec["model_family"] == "gemma4-swa-dense":
|
||
if weights_variant == "awq":
|
||
return spec["weights_awq_gb"] / tp
|
||
elif weights_variant == "bf16":
|
||
return spec["weights_bf16_gb"] / tp
|
||
elif weights_variant == "int8":
|
||
return spec["weights_int8_gb"] / tp
|
||
else: # int4 default
|
||
return spec["weights_int4_gb"] / tp
|
||
elif spec["model_family"] == "gemma4-swa-moe":
|
||
# NOT /tp'd (unlike qwen3-next-moe, fixed in #260) — deliberately left
|
||
# pending its own re-calibration. The v0.7.3 AWQ rows peak at 23.45–
|
||
# 23.50 GB/card (TP=2, ~0.95 GB above the 22.08 budget), so the current
|
||
# no-/tp footprint produces a *protective* TIGHT verdict there. Dividing
|
||
# by TP without a measured fixed-footprint anchor would relax that to
|
||
# PASS for a config that genuinely runs at 23.5/24. Re-do with a proper
|
||
# boot-log anchor (KV-pool tokens → fixed = peak − pool) before /tp'ing.
|
||
if weights_variant == "int4":
|
||
return spec["weights_int4_gb"]
|
||
return spec["weights_awq_gb"]
|
||
elif spec["model_family"] == "generic-dense":
|
||
# Weight authority is the caller-supplied summed selected-blob size
|
||
# (HF siblings, populated by P2). generic-dense NEVER recomputes
|
||
# weights from params — it trusts the provided blob-size GB. Dense
|
||
# weights are sharded evenly across TP ranks.
|
||
return spec["weights_total_gb"] / tp
|
||
raise ValueError(f"Unknown model_family: {spec['model_family']}")
|
||
|
||
|
||
def kv_pool_per_card_bytes(spec, kv_format, max_ctx, max_num_seqs, tp, mtp_n=0):
|
||
"""Per-card KV pool bytes (growing portion only).
|
||
|
||
Returns a tuple (growing_per_card_bytes, sliding_fixed_per_card_bytes).
|
||
Sliding term is zero for models without sliding-window layers.
|
||
|
||
For Qwen 3.6 (DeltaNet hybrid):
|
||
Only the 16 full_attention layers grow KV. GDN layers have a fixed-size
|
||
recurrent state (not seq-len-dependent), so they show up in activation.
|
||
K and V stored independently → ×2 factor.
|
||
|
||
For Gemma 4 (SWA + dense MLP):
|
||
Only the full_attention layers grow KV (at global_head_dim); the count is
|
||
read from the spec, not hardcoded — 10/50 full/sliding on the 31B, 8/40 on
|
||
the 12B. The sliding_attention layers hold a FIXED window (constant in ctx,
|
||
separate small term). K==V tying IS exploited by vLLM's allocator → ×1
|
||
factor (calibrated against BENCHMARKS data; see docs/KV_MATH.md).
|
||
gemma4_unified (12B) overrides the growing per-token with a MEASURED
|
||
constant (`measured_kv_growing_bpt_tp1`) — see the branch below.
|
||
"""
|
||
bpe = KV_FORMAT_BYTES[kv_format]
|
||
|
||
if spec["model_family"] == "qwen3-next-hybrid":
|
||
# K and V stored independently
|
||
per_token = (
|
||
spec["num_attn_layers"]
|
||
* spec["num_kv_heads"]
|
||
* spec["head_dim_attn"]
|
||
* 2 # K + V
|
||
* bpe
|
||
)
|
||
effective_ctx = max_ctx + mtp_n * 32
|
||
growing = (per_token / tp) * effective_ctx * max_num_seqs
|
||
return growing, 0.0
|
||
|
||
elif spec["model_family"] == "qwen3-next-moe":
|
||
# K and V stored independently. GDN recurrent state is fixed-size and
|
||
# per-stream, not context-linear.
|
||
per_token = (
|
||
spec["num_attn_layers"]
|
||
* spec["num_kv_heads"]
|
||
* spec["head_dim_attn"]
|
||
* 2
|
||
* bpe
|
||
)
|
||
effective_ctx = max_ctx + mtp_n * 32
|
||
growing = (per_token / tp) * effective_ctx * max_num_seqs
|
||
recurrent_per_stream = (
|
||
spec["num_gdn_layers"]
|
||
* (
|
||
spec["linear_num_v_heads"] * spec["linear_v_head_dim"]
|
||
+ spec["linear_num_k_heads"] * spec["linear_k_head_dim"]
|
||
+ spec["linear_conv_kernel_dim"] * spec["hidden_size"]
|
||
)
|
||
* 2 # recurrent state kept in bf16 on this stack
|
||
)
|
||
recurrent_fixed = recurrent_per_stream * max_num_seqs
|
||
return growing, recurrent_fixed
|
||
|
||
elif spec["model_family"] == "gemma4-swa-dense":
|
||
# K==V tied → ×1 storage
|
||
measured_bpt_tp1 = spec.get("measured_kv_growing_bpt_tp1")
|
||
if measured_bpt_tp1 is not None:
|
||
# gemma4_unified (gemma-4-12b) MEASURED calibration. The 31B-derived
|
||
# global-only formula (below) predicts 65,536 B/tok TP1 for the 12B
|
||
# (8 full x 8 kv x 512 global_head_dim x bpe2), but the LIVE
|
||
# gemma4_unified pool measured 22,816 B/tok/card = 45,632 TP1
|
||
# (8.16 GiB available KV / 384,019 tokens @ 131K / TP2 / mem-util
|
||
# 0.90, MTP, 2026-06-04) — 1.44x LOWER. Cause: gemma4_unified's
|
||
# unified-KV global layers + vLLM's hybrid-SWA pool accounting differ
|
||
# from the 31B's gemma4-swa-dense. bf16 baseline; scale by bpe for
|
||
# quantized KV. Single-anchor measured calibration; the precise arch
|
||
# decomposition is a TODO (needs vLLM kv_cache_utils source or >1
|
||
# anchor). Set ONLY on the 12B spec, so the 31B keeps the formula.
|
||
per_token_growing = measured_bpt_tp1 * (bpe / 2.0)
|
||
else:
|
||
per_token_growing = (
|
||
spec["num_full_attn_layers"]
|
||
* spec["num_kv_heads"]
|
||
* spec["global_head_dim"]
|
||
* 1 # K==V tied; vLLM stores once
|
||
* bpe
|
||
)
|
||
# No MTP draft-token bump on Gemma — drafter is a separate model
|
||
growing = (per_token_growing / tp) * max_ctx * max_num_seqs
|
||
|
||
# Sliding-window fixed term — 50 layers × window × head_dim × 1 × bpe
|
||
sliding_fixed_total = (
|
||
spec["num_sliding_attn_layers"]
|
||
* spec["num_kv_heads"]
|
||
* spec["head_dim_sliding"]
|
||
* 1 # K==V tied here too
|
||
* bpe
|
||
* spec["sliding_window"]
|
||
)
|
||
sliding_per_card = sliding_fixed_total / tp
|
||
return growing, sliding_per_card
|
||
|
||
elif spec["model_family"] == "gemma4-swa-moe":
|
||
# K==V tied. Global layers use their own KV-head count; sliding
|
||
# layers keep the windowed KV head count.
|
||
per_token_growing = (
|
||
spec["num_full_attn_layers"]
|
||
* spec["num_global_kv_heads"]
|
||
* spec["global_head_dim"]
|
||
* 1
|
||
* bpe
|
||
)
|
||
growing = (per_token_growing / tp) * max_ctx * max_num_seqs
|
||
sliding_fixed_total = (
|
||
spec["num_sliding_attn_layers"]
|
||
* spec["num_kv_heads"]
|
||
* spec["head_dim_sliding"]
|
||
* 1
|
||
* bpe
|
||
* spec["sliding_window"]
|
||
)
|
||
sliding_per_card = sliding_fixed_total / tp
|
||
return growing, sliding_per_card
|
||
|
||
elif spec["model_family"] == "generic-dense":
|
||
# Standard MHA/GQA: every hidden layer grows KV. K and V stored
|
||
# independently (no K==V tying assumption for an arbitrary uncurated
|
||
# model — the conservative choice → ×2). Same shape as the curated
|
||
# Qwen dense branch (num_layers × num_kv_heads × head_dim × 2 × bpe),
|
||
# generalized to all layers since a pure-dense model has no GDN/SWA
|
||
# split. No sliding-window fixed term for generic-dense.
|
||
per_token = (
|
||
spec["num_hidden_layers"]
|
||
* spec["num_kv_heads"]
|
||
* spec["head_dim_attn"]
|
||
* 2 # K + V stored independently (conservative; no K==V tying)
|
||
* bpe
|
||
)
|
||
growing = (per_token / tp) * max_ctx * max_num_seqs
|
||
return growing, 0.0
|
||
|
||
raise ValueError(f"Unknown model_family: {spec['model_family']}")
|
||
|
||
|
||
|
||
def _auto_kv_layout(spec) -> str:
|
||
family = spec["model_family"]
|
||
if family in {"qwen3-next-hybrid", "qwen3-next-moe"}:
|
||
return "hybrid_mamba"
|
||
if family in {"gemma4-swa-dense", "gemma4-swa-moe"}:
|
||
return "sliding_window"
|
||
return "dense_gqa"
|
||
|
||
|
||
def _recurrent_state_per_card_bytes(spec, max_num_seqs, tp):
|
||
"""Non-attention recurrent/SSM state, reported separately from KV.
|
||
|
||
Qwen3-Next GDN/DeltaNet state is not attention KV. It is replicated per
|
||
running sequence and should be visible in planning output because it
|
||
competes for the same per-card VRAM budget as the attention KV pool.
|
||
"""
|
||
if spec["model_family"] not in {"qwen3-next-hybrid", "qwen3-next-moe"}:
|
||
return 0.0
|
||
required = (
|
||
"num_gdn_layers", "linear_num_v_heads", "linear_v_head_dim",
|
||
"linear_num_k_heads", "linear_k_head_dim", "linear_conv_kernel_dim",
|
||
"hidden_size",
|
||
)
|
||
if any(spec.get(k) is None for k in required):
|
||
return 0.0
|
||
bytes_per_elem = spec.get("mamba_state_bytes", 2)
|
||
recurrent_per_stream = (
|
||
spec["num_gdn_layers"]
|
||
* (
|
||
spec["linear_num_v_heads"] * spec["linear_v_head_dim"]
|
||
+ spec["linear_num_k_heads"] * spec["linear_k_head_dim"]
|
||
+ spec["linear_conv_kernel_dim"] * spec["hidden_size"]
|
||
)
|
||
* bytes_per_elem
|
||
)
|
||
# Runtime state is per rank, not sharded like dense weights.
|
||
return recurrent_per_stream * max_num_seqs
|
||
|
||
|
||
def _indexer_cache_per_card_bytes(
|
||
*,
|
||
max_ctx,
|
||
max_num_seqs,
|
||
tp,
|
||
indexer_ratio_layers=0,
|
||
indexer_compress_ratio=4,
|
||
indexer_head_dim=None,
|
||
indexer_format="fp4",
|
||
):
|
||
if not indexer_ratio_layers:
|
||
return 0.0
|
||
if indexer_head_dim is None or indexer_head_dim <= 0:
|
||
raise ValueError("--indexer-ratio-layers requires --indexer-head-dim")
|
||
if indexer_compress_ratio <= 0:
|
||
raise ValueError("--indexer-compress-ratio must be positive")
|
||
bpe = INDEXER_FORMAT_BYTES[indexer_format]
|
||
indexed_tokens = max_ctx // indexer_compress_ratio
|
||
return (
|
||
indexer_ratio_layers
|
||
* indexed_tokens
|
||
* indexer_head_dim
|
||
* bpe
|
||
* max_num_seqs
|
||
/ tp
|
||
)
|
||
|
||
|
||
def _compressed_kv_per_card_bytes(
|
||
*,
|
||
max_ctx,
|
||
max_num_seqs,
|
||
tp,
|
||
compressed_layers=0,
|
||
compression_ratio=128,
|
||
compressed_head_dim=None,
|
||
kv_format="fp8_e5m2",
|
||
):
|
||
if not compressed_layers:
|
||
return 0.0
|
||
if compressed_head_dim is None or compressed_head_dim <= 0:
|
||
raise ValueError("--compressed-layers requires --compressed-head-dim")
|
||
if compression_ratio <= 0:
|
||
raise ValueError("--compression-ratio must be positive")
|
||
bpe = KV_FORMAT_BYTES[kv_format]
|
||
compressed_tokens = max_ctx // compression_ratio
|
||
return (
|
||
compressed_layers
|
||
* compressed_tokens
|
||
* compressed_head_dim
|
||
* bpe
|
||
* max_num_seqs
|
||
/ tp
|
||
)
|
||
|
||
|
||
def architecture_cache_breakdown(
|
||
spec,
|
||
kv_format,
|
||
max_ctx,
|
||
max_num_seqs,
|
||
tp,
|
||
*,
|
||
gpus=1,
|
||
layout="auto",
|
||
mtp=False,
|
||
include_draft_kv=False,
|
||
draft_kv_gb=0.0,
|
||
compressed_layers=0,
|
||
compression_ratio=128,
|
||
compressed_head_dim=None,
|
||
indexer_ratio_layers=0,
|
||
indexer_compress_ratio=4,
|
||
indexer_head_dim=None,
|
||
indexer_format="fp4",
|
||
) -> CacheBreakdown:
|
||
"""Raw per-card cache/state breakdown for home-rig planning.
|
||
|
||
This is intentionally reporting-only. predict() remains the calibrated
|
||
fit authority used by existing scripts; this function exposes the extra
|
||
architecture buckets that generic KV calculators show explicitly.
|
||
"""
|
||
_validate_tp_for_spec(spec, tp)
|
||
if gpus < 1 or gpus > 8:
|
||
raise ValueError("--gpus must be between 1 and 8 for single-node home/workstation planning")
|
||
if max_num_seqs < 1:
|
||
raise ValueError("sequences/max_num_seqs must be positive")
|
||
|
||
mtp_n = int(spec.get("mtp_n_default", 3)) if mtp else 0
|
||
growing_b, fixed_b = kv_pool_per_card_bytes(
|
||
spec, kv_format, max_ctx, max_num_seqs, tp, mtp_n=mtp_n,
|
||
)
|
||
|
||
family = spec["model_family"]
|
||
sliding_fixed_b = fixed_b if family in {"gemma4-swa-dense", "gemma4-swa-moe"} else 0.0
|
||
recurrent_b = fixed_b if family == "qwen3-next-moe" else _recurrent_state_per_card_bytes(spec, max_num_seqs, tp)
|
||
|
||
compressed_b = _compressed_kv_per_card_bytes(
|
||
max_ctx=max_ctx,
|
||
max_num_seqs=max_num_seqs,
|
||
tp=tp,
|
||
compressed_layers=compressed_layers,
|
||
compression_ratio=compression_ratio,
|
||
compressed_head_dim=compressed_head_dim,
|
||
kv_format=kv_format,
|
||
)
|
||
indexer_b = _indexer_cache_per_card_bytes(
|
||
max_ctx=max_ctx,
|
||
max_num_seqs=max_num_seqs,
|
||
tp=tp,
|
||
indexer_ratio_layers=indexer_ratio_layers,
|
||
indexer_compress_ratio=indexer_compress_ratio,
|
||
indexer_head_dim=indexer_head_dim,
|
||
indexer_format=indexer_format,
|
||
)
|
||
|
||
draft_kv_b = draft_kv_gb * 1e9 if include_draft_kv else 0.0
|
||
layout = _auto_kv_layout(spec) if layout == "auto" else layout
|
||
total_b = growing_b + sliding_fixed_b + recurrent_b + compressed_b + indexer_b + draft_kv_b
|
||
|
||
notes = []
|
||
if include_draft_kv and draft_kv_gb <= 0:
|
||
notes.append("draft KV requested, but no --draft-kv-gb override was provided; built-in MTP token bump is already included when --mtp is set")
|
||
if compressed_layers or indexer_ratio_layers:
|
||
notes.append("compressed/indexer buckets are architecture-level estimates; calibrated fit verdict still comes from the main prediction path")
|
||
if gpus != tp:
|
||
notes.append(f"planning {gpus} GPU(s) with TP={tp}; per-card memory is sharded by TP, not by idle estate cards")
|
||
|
||
return CacheBreakdown(
|
||
layout=layout,
|
||
sequences=max_num_seqs,
|
||
gpus=gpus,
|
||
tp=tp,
|
||
attention_kv_growing_gb=growing_b / 1e9,
|
||
attention_kv_sliding_fixed_gb=sliding_fixed_b / 1e9,
|
||
recurrent_state_gb=recurrent_b / 1e9,
|
||
compressed_kv_gb=compressed_b / 1e9,
|
||
indexer_cache_gb=indexer_b / 1e9,
|
||
draft_kv_gb=draft_kv_b / 1e9,
|
||
total_cache_gb=total_b / 1e9,
|
||
notes=notes,
|
||
)
|
||
|
||
|
||
def activation_peak_per_card_bytes(spec, kv_format, max_ctx, tp):
|
||
"""Per-card peak activation during prefill forward.
|
||
|
||
For Qwen 3.6 (DeltaNet GDN): linear in seq_len, KV-format-dependent
|
||
coefficient (PerfMamba O(γ·D·N·L) form, fla.ops.chunk implementation
|
||
details calibrated empirically).
|
||
|
||
For Gemma 4 (dense MLP + SWA): mostly CONSTANT in seq_len because chunked
|
||
prefill bounds the MLP intermediate. Small per-token residual to keep
|
||
the solver smooth.
|
||
"""
|
||
if spec["model_family"] == "qwen3-next-hybrid":
|
||
coef = QWEN_GDN_ACTIVATION_COEF[kv_format]
|
||
return (coef * spec["num_gdn_layers"] * max_ctx) / tp
|
||
|
||
elif spec["model_family"] == "qwen3-next-moe":
|
||
coef = QWEN_MOE_ACTIVATION_COEF[kv_format]
|
||
return (coef * spec["num_gdn_layers"] * max_ctx) / tp + QWEN_MOE_EXPERT_DISPATCH_GB * 1e9
|
||
|
||
elif spec["model_family"] == "gemma4-swa-dense":
|
||
const_bytes = GEMMA_ACTIVATION_CONST_GB * 1e9
|
||
per_token = GEMMA_ACTIVATION_PER_TOKEN_BYTES * max_ctx
|
||
return (const_bytes + per_token) / tp
|
||
|
||
elif spec["model_family"] == "gemma4-swa-moe":
|
||
const_bytes = GEMMA_MOE_ACTIVATION_CONST_GB * 1e9
|
||
per_token = GEMMA_MOE_ACTIVATION_PER_TOKEN_BYTES * max_ctx
|
||
return (const_bytes + per_token) / tp
|
||
|
||
elif spec["model_family"] == "generic-dense":
|
||
# Conservative hidden_size-based coefficient × max_ctx, plus a fixed
|
||
# per-card floor ≥ the most conservative curated dense constant
|
||
# (Gemma 1.5 GB). The coefficient (6.0 B / hidden-unit / token) is
|
||
# ≥ the most conservative curated dense per-token term re-expressed
|
||
# per hidden-unit (Qwen GDN turboquant ≈ 5.16). This guarantees a
|
||
# lower-bound on fit vs. an equivalent curated dense shape.
|
||
per_token = (
|
||
GENERIC_DENSE_ACTIVATION_COEF_PER_HIDDEN
|
||
* spec["hidden_size"]
|
||
* max_ctx
|
||
)
|
||
floor_bytes = GENERIC_DENSE_ACTIVATION_FLOOR_GB * 1e9
|
||
return (floor_bytes + per_token) / tp
|
||
|
||
raise ValueError(f"Unknown model_family: {spec['model_family']}")
|
||
|
||
|
||
def cudagraph_overhead_gb(mem_util, tp):
|
||
"""vLLM cudagraph capture + flashinfer workspace overhead per card.
|
||
Roughly linear with mem_util (higher mem-util → more graphs captured).
|
||
TP increases per-card overhead slightly due to NCCL workspaces.
|
||
"""
|
||
base = 0.5 + 1.0 * mem_util
|
||
tp_bump = 0.0 if tp == 1 else 0.3 * (tp - 1)
|
||
return base + tp_bump
|
||
|
||
|
||
def _validate_tp_for_spec(spec, tp):
|
||
valid_tp = spec.get("valid_tp")
|
||
if valid_tp and tp not in valid_tp:
|
||
raise ValueError(
|
||
f"TP={tp} invalid for {spec['model_id']} "
|
||
f"(num_kv_heads={spec['num_kv_heads']} cannot be divided across TP cleanly). "
|
||
f"Valid TP values: {valid_tp}"
|
||
)
|
||
|
||
|
||
def predict(
|
||
spec=QWEN36_27B,
|
||
kv_format="fp8_e5m2",
|
||
max_ctx=180000,
|
||
max_num_seqs=1,
|
||
tp=1,
|
||
mem_util=0.95,
|
||
vram_gb=24,
|
||
dflash_draft_gb=0.0,
|
||
drafter_gb=0.0,
|
||
mtp=False,
|
||
weights_variant="default",
|
||
) -> Prediction:
|
||
"""Predict per-card VRAM usage.
|
||
|
||
vLLM caps KV pool to (budget - fixed_components), so the prediction
|
||
reflects what actually gets allocated. When requested > available,
|
||
verdict is TIGHT with a note about effective concurrency reduction.
|
||
|
||
Args:
|
||
drafter_gb: total drafter weight (MTP / DFlash) — split by TP.
|
||
dflash_draft_gb: legacy alias — folded into drafter_gb if set.
|
||
"""
|
||
_validate_tp_for_spec(spec, tp)
|
||
|
||
weights_gb = _weights_per_card_gb(spec, tp, weights_variant)
|
||
|
||
mtp_n = int(spec.get("mtp_n_default", 3)) if mtp else 0
|
||
growing_b, sliding_b = kv_pool_per_card_bytes(
|
||
spec, kv_format, max_ctx, max_num_seqs, tp,
|
||
mtp_n=mtp_n,
|
||
)
|
||
kv_pool_requested_gb = growing_b / 1e9
|
||
kv_pool_sliding_fixed_gb = sliding_b / 1e9
|
||
|
||
activation_gb = activation_peak_per_card_bytes(spec, kv_format, max_ctx, tp) / 1e9
|
||
overhead_gb = cudagraph_overhead_gb(mem_util, tp)
|
||
|
||
# Drafter: prefer drafter_gb; fall back to legacy dflash_draft_gb.
|
||
drafter_total = drafter_gb if drafter_gb > 0 else dflash_draft_gb
|
||
if mtp and spec["model_family"] == "qwen3-next-moe":
|
||
drafter_total += QWEN_MOE_BUILTIN_MTP_WORKSPACE_GB
|
||
drafter_per_card = drafter_total / tp if tp > 1 else drafter_total
|
||
|
||
fixed_gb = weights_gb + activation_gb + overhead_gb + drafter_per_card + kv_pool_sliding_fixed_gb
|
||
budget_gb = mem_util * vram_gb
|
||
available_for_kv = max(0.0, budget_gb - fixed_gb)
|
||
|
||
# vLLM caps the KV pool to fit available budget (PagedAttention allocator).
|
||
kv_pool_actual_gb = min(kv_pool_requested_gb, available_for_kv)
|
||
|
||
total_gb = fixed_gb + kv_pool_actual_gb
|
||
pct = 100 * total_gb / budget_gb if budget_gb > 0 else 999.0
|
||
|
||
notes = []
|
||
|
||
# Verdict logic:
|
||
# - FAIL: fixed components alone exceed budget (no room even for minimum KV).
|
||
# - TIGHT: requested KV pool exceeds available — vLLM will cap, effective
|
||
# concurrency reduced (BOOT OK, but `--max-num-seqs` may not be
|
||
# honored at full max_ctx).
|
||
# - PASS: requested KV fits with room to spare.
|
||
# vLLM's boot pre-check is token-capacity based, NOT an absolute byte floor:
|
||
# the (capped) growing KV pool must hold at least ONE max_model_len sequence.
|
||
# Growing KV scales linearly with max_num_seqs, so one sequence's growing KV
|
||
# is kv_pool_requested_gb / max_num_seqs. KV-light families (gemma4 SWA, Qwen3
|
||
# -Next MoE) legitimately boot with a sub-GB growing pool — the old flat 1.0 GB
|
||
# floor false-FAILed them (e.g. gemma-26ba4b-single: 0.78 GB pool holds 17,490
|
||
# tok ≥ 16,384 max_ctx, boots+serves live, #326). So lower the floor to the
|
||
# per-sequence requirement — which also generalizes the old qwen3-next-moe 0.05
|
||
# special-case. CAP it at the legacy 1.0 GB though: for dense/long-KV configs a
|
||
# single sequence needs many GB, and a hard per-seq threshold would false-FAIL
|
||
# measured-working configs sitting inside the estimator's ±1.5 GB band (e.g.
|
||
# gemma-dual-int8 @262K: 10.71 GB avail vs 10.84 GB/seq, 1.2% short but boots).
|
||
# Result: relax the floor for KV-light models, leave dense on its prior ≥1 GB
|
||
# behavior. Still FAILs when even one (small) sequence won't fit / fixed alone
|
||
# exceeds budget (available_for_kv → 0).
|
||
per_sequence_kv_gb = kv_pool_requested_gb / max(max_num_seqs, 1)
|
||
MIN_KV_GB = max(0.01, min(1.0, per_sequence_kv_gb))
|
||
if available_for_kv < MIN_KV_GB:
|
||
verdict = "FAIL"
|
||
notes.append(
|
||
f"fixed components ({fixed_gb:.1f} GB) leave only {available_for_kv:.2f} GB for KV — "
|
||
f"below the {MIN_KV_GB:.2f} GB that one max_ctx={max_ctx:,} sequence needs; vLLM "
|
||
f"pre-check will refuse to boot — lower max_ctx, drop a drafter, or raise mem_util"
|
||
)
|
||
elif kv_pool_requested_gb > available_for_kv * 1.05:
|
||
verdict = "TIGHT"
|
||
notes.append(
|
||
f"requested KV pool ({kv_pool_requested_gb:.1f} GB) > available ({available_for_kv:.1f} GB) — "
|
||
f"vLLM will cap to {available_for_kv:.1f} GB; effective concurrency may be lower than "
|
||
f"--max-num-seqs={max_num_seqs} at full max_ctx={max_ctx:,}"
|
||
)
|
||
else:
|
||
verdict = "PASS"
|
||
|
||
# Model-specific advisory notes (preserved from v1)
|
||
if kv_format == "turboquant_3bit_nc" and vram_gb < 24 and spec["model_family"] == "qwen3-next-hybrid":
|
||
notes.append("⚠ TQ3 KV on <24 GB cards: consider --kv-format fp8_e5m2 (see docs/HARDWARE.md, #47)")
|
||
if max_ctx > 50000 and tp == 1 and spec["model_family"] == "qwen3-next-hybrid" and kv_format != "fp16":
|
||
notes.append("⚠ single-card vLLM at >50K single-prompt: Cliff 2 territory (DeltaNet GDN forward); see docs/CLIFFS.md")
|
||
if spec["model_family"] == "gemma4-swa-dense" and kv_format == "fp8_e4m3":
|
||
notes.append("⚠ fp8_e4m3 on Ampere (sm_86): Triton `fp8e4nv` kernel unsupported; use int8_per_token_head instead (PR #40391 via #42102)")
|
||
if spec["model_family"] == "gemma4-swa-dense" and tp == 1 and vram_gb < 32:
|
||
notes.append("⚠ Gemma 4 31B TP=1 needs ≥32 GB VRAM; 24 GB Ampere boot-OOMs (model weights + drafter + min KV)")
|
||
if spec["model_family"] in {"qwen3-next-moe", "gemma4-swa-moe"}:
|
||
notes.append("MoE projection uses low-anchor calibration; add max_ctx/max_num_seqs A/B rows before treating this as production-grade.")
|
||
if (
|
||
spec["model_family"] == "qwen3-next-moe"
|
||
and kv_pool_requested_gb > 0
|
||
and available_for_kv > kv_pool_requested_gb * 1.5
|
||
):
|
||
notes.append(
|
||
f"DeltaNet-hybrid KV is cheap — one max_ctx={max_ctx:,} seq needs only "
|
||
f"{kv_pool_requested_gb:.2f} GB, but vLLM fills the ~{available_for_kv:.1f} GB available pool "
|
||
f"to budget (≈{available_for_kv / kv_pool_requested_gb:.1f}× concurrency headroom). The sub-budget "
|
||
f"'predicted total' is the single-seq floor, not spare VRAM you can repurpose."
|
||
)
|
||
if tp > 4:
|
||
notes.append("TP > 4 predictions are extrapolated; report deltas via scripts/report.sh --bench")
|
||
|
||
return Prediction(
|
||
model=spec["model_id"],
|
||
weights_gb=weights_gb,
|
||
kv_pool_requested_gb=kv_pool_requested_gb,
|
||
kv_pool_actual_gb=kv_pool_actual_gb,
|
||
kv_pool_sliding_fixed_gb=kv_pool_sliding_fixed_gb,
|
||
activation_gb=activation_gb,
|
||
cudagraph_overhead_gb=overhead_gb,
|
||
drafter_gb=drafter_per_card,
|
||
total_gb=total_gb,
|
||
vram_gb=vram_gb,
|
||
budget_gb=budget_gb,
|
||
pct_of_vram=pct,
|
||
verdict=verdict,
|
||
notes=notes,
|
||
)
|
||
|
||
|
||
# =============================================================================
|
||
# v0.8.0 Pull-Gate — generic-dense eligibility predicate + verdict adapter
|
||
# =============================================================================
|
||
#
|
||
# Import path for P2/P3/P4: `tools/kv-calc.py` is not a dotted-importable
|
||
# module name (hyphen). Load it via importlib, e.g.:
|
||
#
|
||
# import importlib.util, pathlib, sys
|
||
# _p = pathlib.Path(REPO_ROOT) / "tools" / "kv-calc.py"
|
||
# spec = importlib.util.spec_from_file_location("kv_calc", _p)
|
||
# kv_calc = importlib.util.module_from_spec(spec)
|
||
# sys.modules["kv_calc"] = kv_calc # required: @dataclass resolves
|
||
# # cls.__module__ via sys.modules
|
||
# spec.loader.exec_module(kv_calc)
|
||
# kv_calc.is_generic_dense_eligible(config) # stratum-5 predicate
|
||
# kv_calc.raw_verdict(...) # [B]→[C1] adapter
|
||
#
|
||
# (This mirrors how the existing test harness loads the module.)
|
||
|
||
# Markers that make a config NOT a standard dense (MHA/GQA) transformer.
|
||
_SSM_RWKV_CONFIG_KEYS = (
|
||
"ssm_cfg", # mamba / mamba2
|
||
"mamba_d_state",
|
||
"mamba_expand",
|
||
"mamba_n_heads",
|
||
"rwkv", # rwkv-family marker
|
||
"time_mix_extra_dim",
|
||
"linear_attn_config", # qwen3-next / deltanet hybrid
|
||
"linear_num_value_heads",
|
||
"linear_num_key_heads",
|
||
"full_attention_interval", # hybrid attention layout
|
||
)
|
||
_SSM_RWKV_MODEL_TYPE_SUBSTRINGS = (
|
||
"mamba", "rwkv", "ssm", "qwen3_next", "qwen3next",
|
||
"jamba", "zamba", "recurrentgemma", "griffin", "hgrn",
|
||
)
|
||
|
||
|
||
def is_generic_dense_eligible(config: dict) -> bool:
|
||
"""Stratum-5 pre-`[B]` eligibility: True iff `config` is a standard
|
||
dense MHA/GQA transformer the generic-dense pricing branch can price.
|
||
|
||
Eligible iff ALL of:
|
||
- has hidden_size, num_hidden_layers, num_attention_heads,
|
||
num_key_value_heads
|
||
- head_dim present OR derivable as hidden_size / num_attention_heads
|
||
with no remainder
|
||
AND NONE of (ineligible — return False, never emit a number):
|
||
- SSM / Mamba / RWKV / linear-attention (DeltaNet hybrid) markers
|
||
- MoE (`num_local_experts` or `num_experts`)
|
||
- sliding-window-only attention (`sliding_window` set without any
|
||
full-attention layers)
|
||
"""
|
||
if not isinstance(config, dict):
|
||
return False
|
||
|
||
required = ("hidden_size", "num_hidden_layers",
|
||
"num_attention_heads", "num_key_value_heads")
|
||
for key in required:
|
||
v = config.get(key)
|
||
if not isinstance(v, int) or isinstance(v, bool) or v <= 0:
|
||
return False
|
||
|
||
hidden_size = config["hidden_size"]
|
||
num_attn_heads = config["num_attention_heads"]
|
||
|
||
# head_dim explicit OR cleanly derivable.
|
||
head_dim = config.get("head_dim")
|
||
if not (isinstance(head_dim, int) and not isinstance(head_dim, bool) and head_dim > 0):
|
||
if num_attn_heads <= 0 or hidden_size % num_attn_heads != 0:
|
||
return False
|
||
|
||
# --- exclusions ---
|
||
model_type = str(config.get("model_type", "")).lower()
|
||
architectures = " ".join(
|
||
str(a) for a in config.get("architectures", []) if a is not None
|
||
).lower()
|
||
for sub in _SSM_RWKV_MODEL_TYPE_SUBSTRINGS:
|
||
if sub in model_type or sub in architectures:
|
||
return False
|
||
for key in _SSM_RWKV_CONFIG_KEYS:
|
||
if config.get(key) is not None:
|
||
return False
|
||
|
||
# MoE.
|
||
if config.get("num_local_experts") is not None:
|
||
return False
|
||
if config.get("num_experts") is not None:
|
||
return False
|
||
|
||
# Sliding-window-only attention: a sliding_window is set and there is no
|
||
# evidence of any full-attention layers. (A model that mixes full +
|
||
# sliding — e.g. has a full-attention interval — is already excluded
|
||
# above via the hybrid-layout markers; here we reject pure-SWA.)
|
||
sliding_window = config.get("sliding_window")
|
||
if sliding_window is not None and sliding_window:
|
||
# Heuristics that indicate at least some full-attention layers.
|
||
has_full_attn = bool(
|
||
config.get("sliding_window_pattern") # e.g. Gemma "every Nth is global"
|
||
or config.get("max_window_layers") # Qwen2-style: first N layers full
|
||
or config.get("global_attn_layers")
|
||
or config.get("layer_types") # explicit per-layer attn type list
|
||
)
|
||
if not has_full_attn:
|
||
return False
|
||
|
||
return True
|
||
|
||
|
||
# Map predict()'s coarse verdict → the design's [B]→[C1] vocabulary.
|
||
_RAW_VERDICT_MAP = {
|
||
"FAIL": "wont-fit",
|
||
"TIGHT": "fits-constrained",
|
||
"PASS": "fits-clean",
|
||
}
|
||
|
||
|
||
def raw_verdict(
|
||
spec,
|
||
kv_format="fp8_e5m2",
|
||
max_ctx=180000,
|
||
max_num_seqs=1,
|
||
tp=1,
|
||
mem_util=0.95,
|
||
vram_gb=24,
|
||
dflash_draft_gb=0.0,
|
||
drafter_gb=0.0,
|
||
mtp=False,
|
||
weights_variant="default",
|
||
) -> dict:
|
||
"""Thin `[B]`→`[C1]` adapter over predict().
|
||
|
||
Returns a JSON-serializable dict exposing predict()'s result in the
|
||
design's vocabulary:
|
||
FAIL → "wont-fit"
|
||
TIGHT → "fits-constrained"
|
||
PASS → "fits-clean"
|
||
|
||
The structured dict (verdict + key GB breakdown + notes) is what the
|
||
P3/P4 `[C1]` §4.1 total function consumes alongside the confidence tier.
|
||
Confidence assignment is NOT done here (that is the deriver / stratum-5
|
||
layer in later STEPs); this adapter only translates the raw fit verdict.
|
||
"""
|
||
p = predict(
|
||
spec=spec,
|
||
kv_format=kv_format,
|
||
max_ctx=max_ctx,
|
||
max_num_seqs=max_num_seqs,
|
||
tp=tp,
|
||
mem_util=mem_util,
|
||
vram_gb=vram_gb,
|
||
dflash_draft_gb=dflash_draft_gb,
|
||
drafter_gb=drafter_gb,
|
||
mtp=mtp,
|
||
weights_variant=weights_variant,
|
||
)
|
||
return {
|
||
"raw_verdict": _RAW_VERDICT_MAP[p.verdict],
|
||
"predict_verdict": p.verdict,
|
||
"model": p.model,
|
||
"breakdown_gb": {
|
||
"weights": round(p.weights_gb, 4),
|
||
"kv_pool_requested": round(p.kv_pool_requested_gb, 4),
|
||
"kv_pool_actual": round(p.kv_pool_actual_gb, 4),
|
||
"kv_pool_sliding_fixed": round(p.kv_pool_sliding_fixed_gb, 4),
|
||
"activation": round(p.activation_gb, 4),
|
||
"cudagraph_overhead": round(p.cudagraph_overhead_gb, 4),
|
||
"drafter": round(p.drafter_gb, 4),
|
||
"total": round(p.total_gb, 4),
|
||
},
|
||
"budget_gb": round(p.budget_gb, 4),
|
||
"vram_gb": p.vram_gb,
|
||
"pct_of_vram": round(p.pct_of_vram, 2),
|
||
"notes": list(p.notes),
|
||
}
|
||
|
||
|
||
|
||
def fmt_cache_breakdown(b: CacheBreakdown) -> str:
|
||
lines = []
|
||
lines.append("Cache architecture breakdown")
|
||
lines.append("----------------------------")
|
||
lines.append(f" Layout: {b.layout}")
|
||
lines.append(f" Scope: {b.gpus} GPU(s), TP={b.tp}, sequences={b.sequences}")
|
||
lines.append(f" Attention KV — growing: {b.attention_kv_growing_gb:>6.2f} GB / card")
|
||
if b.attention_kv_sliding_fixed_gb > 0.001:
|
||
lines.append(f" Attention KV — sliding: {b.attention_kv_sliding_fixed_gb:>6.2f} GB / card")
|
||
if b.recurrent_state_gb > 0.001:
|
||
lines.append(f" Recurrent / SSM state: {b.recurrent_state_gb:>6.2f} GB / card")
|
||
if b.compressed_kv_gb > 0.001:
|
||
lines.append(f" Compressed KV estimate: {b.compressed_kv_gb:>6.2f} GB / card")
|
||
if b.indexer_cache_gb > 0.001:
|
||
lines.append(f" Indexer cache estimate: {b.indexer_cache_gb:>6.2f} GB / card")
|
||
if b.draft_kv_gb > 0.001:
|
||
lines.append(f" Draft KV override: {b.draft_kv_gb:>6.2f} GB / card")
|
||
lines.append(f" Cache/state subtotal: {b.total_cache_gb:>6.2f} GB / card")
|
||
for note in b.notes:
|
||
lines.append(f" Note: {note}")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def fmt_prediction(p: Prediction, header: str = "") -> str:
|
||
lines = []
|
||
if header:
|
||
lines.append(header)
|
||
lines.append("-" * len(header))
|
||
lines.append(f" Model: {p.model}")
|
||
lines.append(f" Weights: {p.weights_gb:>6.2f} GB / card")
|
||
if p.kv_pool_sliding_fixed_gb > 0.01:
|
||
lines.append(f" KV pool — sliding fixed: {p.kv_pool_sliding_fixed_gb:>6.2f} GB / card (constant, doesn't grow with ctx)")
|
||
if abs(p.kv_pool_requested_gb - p.kv_pool_actual_gb) > 0.05:
|
||
lines.append(f" KV pool — growing (req): {p.kv_pool_requested_gb:>6.2f} GB / card (requested)")
|
||
lines.append(f" KV pool — growing (cap): {p.kv_pool_actual_gb:>6.2f} GB / card (vLLM-capped to fit)")
|
||
else:
|
||
lines.append(f" KV pool — growing: {p.kv_pool_actual_gb:>6.2f} GB / card")
|
||
lines.append(f" Activation peak: {p.activation_gb:>6.2f} GB / card")
|
||
lines.append(f" Cudagraph + workspace: {p.cudagraph_overhead_gb:>6.2f} GB / card")
|
||
if p.drafter_gb > 0:
|
||
lines.append(f" Drafter (MTP / DFlash): {p.drafter_gb:>6.2f} GB / card")
|
||
lines.append(f" ─────────────────────────────────────")
|
||
lines.append(f" Predicted total: {p.total_gb:>6.2f} GB / card ({p.pct_of_vram:.0f}% of {p.budget_gb:.1f} GB engine budget)")
|
||
lines.append(f" Verdict: {p.verdict}")
|
||
for note in p.notes:
|
||
lines.append(f" Note: {note}")
|
||
return "\n".join(lines)
|
||
|
||
|
||
# =============================================================================
|
||
# Calibration runner
|
||
# =============================================================================
|
||
|
||
def _resolve_compose_for_predict(model_key, compose_id, vram, ctx_override=None):
|
||
"""Resolve a compose preset to predict() kwargs, applying optional ctx override."""
|
||
spec = MODEL_SPECS[model_key]
|
||
cfg = COMPOSES[model_key][compose_id]
|
||
max_ctx = ctx_override if ctx_override is not None else cfg["max_ctx"]
|
||
|
||
kwargs = dict(
|
||
spec=spec,
|
||
kv_format=cfg["kv_format"],
|
||
max_ctx=max_ctx,
|
||
max_num_seqs=cfg["max_num_seqs"],
|
||
tp=cfg["tp"],
|
||
mem_util=cfg["mem_util"],
|
||
vram_gb=vram,
|
||
mtp=cfg.get("mtp", False),
|
||
weights_variant=cfg.get("weights_variant", "default"),
|
||
drafter_gb=cfg.get("drafter_gb", 0.0),
|
||
dflash_draft_gb=cfg.get("dflash_draft_gb", 0.0),
|
||
)
|
||
return kwargs
|
||
|
||
|
||
def _calibration_block(model_key: str) -> tuple[int, int]:
|
||
"""Print calibration table for one model. Returns (correct, total)."""
|
||
rows = CALIBRATION.get(model_key, [])
|
||
if not rows:
|
||
return 0, 0
|
||
|
||
spec = MODEL_SPECS[model_key]
|
||
print(f"== {spec['model_id']} ==")
|
||
print(f" {'compose':<26s} {'predicted':>10s} {'budget':>9s} {'measured':>10s} {'verdict':>8s}")
|
||
print(f" {'─'*25:<26s} {'─'*9:>10s} {'─'*8:>9s} {'─'*9:>10s} {'─'*7:>8s}")
|
||
|
||
correct = 0
|
||
for row in rows:
|
||
compose, vram, measured, ctx_override, _src = row
|
||
kwargs = _resolve_compose_for_predict(model_key, compose, vram, ctx_override)
|
||
p = predict(**kwargs)
|
||
# Verdict is "correct" if (PASS/TIGHT and measured fits) or (FAIL and would OOM).
|
||
# We don't have negative (FAIL) data points in BENCHMARKS — every row booted —
|
||
# so verdict_correct simplifies to: PASS/TIGHT and measured < vram.
|
||
if p.verdict in ("PASS", "TIGHT") and measured < vram:
|
||
mark = "✓"
|
||
correct += 1
|
||
elif p.verdict == "FAIL" and measured >= vram:
|
||
mark = "✓"
|
||
correct += 1
|
||
else:
|
||
mark = "⨯"
|
||
compose_disp = compose if ctx_override is None else f"{compose}@{ctx_override//1024}K"
|
||
print(f" {compose_disp:<26s} {p.total_gb:>8.2f} GB {p.budget_gb:>7.2f} GB {measured:>8.2f} GB {p.verdict:>7s} {mark}")
|
||
|
||
print()
|
||
print(f" Verdict accuracy: {correct}/{len(rows)} ({100*correct/len(rows):.0f}%)")
|
||
print()
|
||
return correct, len(rows)
|
||
|
||
|
||
def run_calibration():
|
||
print("=" * 88)
|
||
print("Calibration — predicted per-card VRAM vs measured BENCHMARKS rows")
|
||
print("=" * 88)
|
||
print()
|
||
print(" Predicted = weights + activation + overhead + drafter + (KV capped at available).")
|
||
print(" Budget = mem_util × VRAM. Measured = nvidia-smi peak during bench (target ≈ budget).")
|
||
print(" Verdict ✓ iff PASS/TIGHT and measured < VRAM (boot OK).")
|
||
print()
|
||
|
||
total_c, total_n = 0, 0
|
||
for model_key in MODEL_SPECS:
|
||
c, n = _calibration_block(model_key)
|
||
total_c += c
|
||
total_n += n
|
||
|
||
if total_n > 0:
|
||
print(f"Overall: {total_c}/{total_n} ({100*total_c/total_n:.0f}%)")
|
||
print()
|
||
print("Notes:")
|
||
print(" - This is a directional estimator (±1.5 GB error band on the breakdown).")
|
||
print(" - vLLM's `gpu_worker.py` boot log is the authoritative source.")
|
||
print(" - If predicted PASS but measured > budget, file an issue with `scripts/report.sh --bench`.")
|
||
|
||
|
||
# =============================================================================
|
||
# Max-ctx solver
|
||
# =============================================================================
|
||
|
||
def solve_max_ctx(spec, kv_format, max_num_seqs, tp, mem_util, vram_gb,
|
||
drafter_gb=0.0, dflash_draft_gb=0.0, mtp=False, weights_variant="default"):
|
||
"""Binary search for the largest max_ctx that keeps the verdict at PASS or TIGHT."""
|
||
lo, hi = 1024, spec.get("max_ctx_supported", 262144)
|
||
best = 0
|
||
while lo <= hi:
|
||
mid = (lo + hi) // 2
|
||
mid = (mid // 1024) * 1024 # round to nearest 1024 for cleaner numbers
|
||
if mid == 0:
|
||
break
|
||
p = predict(
|
||
spec=spec, kv_format=kv_format, max_ctx=mid, max_num_seqs=max_num_seqs,
|
||
tp=tp, mem_util=mem_util, vram_gb=vram_gb,
|
||
drafter_gb=drafter_gb, dflash_draft_gb=dflash_draft_gb,
|
||
mtp=mtp, weights_variant=weights_variant,
|
||
)
|
||
if p.verdict in ("PASS", "TIGHT"):
|
||
best = mid
|
||
lo = mid + 1024
|
||
else:
|
||
hi = mid - 1024
|
||
return best
|
||
|
||
|
||
# =============================================================================
|
||
# --fit — structured verdict for a registry slug or bare model on a given card
|
||
# =============================================================================
|
||
#
|
||
# Strictly-additive wrapper. REUSES the SAME pricing the pull.sh gate relies
|
||
# on (predict() / solve_max_ctx() / _RAW_VERDICT_MAP via raw_verdict()), the
|
||
# SAME registry→predict()-kwargs conversion the calibration path uses
|
||
# (_compose_cfg_from_registry), and the SAME default-slug resolver the
|
||
# launchers use (compose_registry.curated_default_target). No new math.
|
||
|
||
# Known per-card VRAM (GB) for the `--card` flag. A bare numeric string is
|
||
# also accepted (e.g. `--card 48`) so this never blocks an uncatalogued card.
|
||
CARD_VRAM_GB = {
|
||
"rtx3090": 24.0,
|
||
"3090": 24.0,
|
||
"rtx4090": 24.0,
|
||
"4090": 24.0,
|
||
"rtx3090ti": 24.0,
|
||
"a6000": 48.0,
|
||
"rtxa6000": 48.0,
|
||
"a100-40": 40.0,
|
||
"a100-80": 80.0,
|
||
"a100": 80.0,
|
||
"h100": 80.0,
|
||
"l40s": 48.0,
|
||
"l40": 48.0,
|
||
"rtx5090": 32.0,
|
||
"5090": 32.0,
|
||
"rtx3080": 10.0,
|
||
"rtx3060": 12.0,
|
||
# Canonical hyphenated hardware-profile ids — the exact form switch.sh
|
||
# --explain (explain_detect_card) and the rest of the stack emit. Values
|
||
# mirror scripts/lib/profiles/hardware/*.yml.
|
||
"rtx-3090": 24.0,
|
||
"rtx-3090-ti": 24.0,
|
||
"rtx-4090": 24.0,
|
||
"rtx-5090": 32.0,
|
||
"rtx-a5000": 24.0,
|
||
"a5000": 24.0,
|
||
"a100-40gb": 40.0,
|
||
"h100-80gb": 80.0,
|
||
"rtx-3060-12gb": 12.0,
|
||
"rtx-6000-pro-blackwell": 96.0,
|
||
}
|
||
|
||
# Documented estimator error band (the file states "±1.5 GB error band" in
|
||
# run_calibration() + the solver epilogue). Surfaced verbatim so consumers
|
||
# can compare vram_est to budget with the same tolerance the gate uses.
|
||
FIT_BAND_GB = 1.5
|
||
|
||
# Per-card SM (compute capability) for the required_sm gate — derived from the
|
||
# hardware profiles (the SAME source compat.py's C3 floor uses), keyed like
|
||
# CARD_VRAM_GB (canonical hyphenated id + dehyphenated alias). A bare-number
|
||
# --card describes VRAM only, so it has no SM → the gate is skipped
|
||
# (permissive by design: an uncatalogued card is not assumed incompatible).
|
||
CARD_SM = {}
|
||
for _hid, _h in PROFILES.hardware.items():
|
||
CARD_SM[_hid] = float(_h.sm)
|
||
CARD_SM[_hid.replace("-", "")] = float(_h.sm)
|
||
|
||
|
||
def _resolve_card_sm(card: Optional[str]) -> Optional[float]:
|
||
if card is None:
|
||
return None
|
||
raw = str(card).strip()
|
||
try:
|
||
float(raw)
|
||
return None # bare VRAM number — carries no arch info
|
||
except ValueError:
|
||
pass
|
||
key = raw.lower().replace(" ", "").replace("_", "").replace("/", "-")
|
||
return CARD_SM.get(key) or CARD_SM.get(key.replace("-", ""))
|
||
|
||
|
||
def _resolve_card_vram_gb(card: Optional[str]) -> Optional[float]:
|
||
"""Map a `--card` value to per-card VRAM in GB. Accepts a known card key
|
||
(case/sep-insensitive) or a bare positive number. Returns None when the
|
||
value is neither — the caller emits an `unknown` verdict, never crashes."""
|
||
if card is None:
|
||
return None
|
||
raw = str(card).strip()
|
||
try:
|
||
v = float(raw)
|
||
return v if v > 0 else None
|
||
except ValueError:
|
||
pass
|
||
key = raw.lower().replace(" ", "").replace("_", "").replace("/", "-")
|
||
# Try the normalized key, then a hyphen-stripped variant, so a hyphenated
|
||
# profile-id ('rtx-3090') also matches the de-hyphenated keys ('rtx3090').
|
||
return CARD_VRAM_GB.get(key) or CARD_VRAM_GB.get(key.replace("-", ""))
|
||
|
||
|
||
def _resolve_fit_slug(target: str) -> tuple[Optional[str], Optional[str], Optional[str]]:
|
||
"""Resolve a `--fit` argument to a priceable vLLM registry slug.
|
||
|
||
Accepts either a registry slug (e.g. `vllm/dual`) or a bare model id
|
||
(e.g. `qwen3.6-27b` → its curated default, preferring dual then single).
|
||
|
||
Returns (slug, model_id, error). On success error is None; on failure
|
||
slug/model_id are None and error explains why (so the caller emits an
|
||
`unknown` verdict rather than raising)."""
|
||
from scripts.lib.profiles.compose_registry import (
|
||
COMPOSE_REGISTRY,
|
||
curated_default_target,
|
||
model_of_slug,
|
||
model_set,
|
||
)
|
||
|
||
if target in COMPOSE_REGISTRY:
|
||
slug = target
|
||
model_id = model_of_slug(slug)
|
||
elif target in model_set():
|
||
# Bare model → its curated default slug. Prefer a priceable (vLLM)
|
||
# topology: dual first, then single.
|
||
slug = None
|
||
for topo in ("dual", "single"):
|
||
cand = curated_default_target(target, topo)
|
||
if cand and COMPOSE_REGISTRY.get(cand, {}).get("kvcalc_key") not in (None, "SKIP"):
|
||
slug = cand
|
||
break
|
||
if slug is None:
|
||
# Fall back to ANY priceable vLLM slug for this model.
|
||
for s, e in COMPOSE_REGISTRY.items():
|
||
if e.get("model") == target and s.startswith("vllm/") and e.get("kvcalc_key") not in (None, "SKIP"):
|
||
slug = s
|
||
break
|
||
if slug is None:
|
||
return None, target, f"no vLLM (kv-calc priceable) compose for model {target!r}"
|
||
model_id = target
|
||
else:
|
||
return None, None, f"unknown slug/model {target!r} (not a registry slug or catalogued model)"
|
||
|
||
entry = COMPOSE_REGISTRY[slug]
|
||
if entry.get("kvcalc_key") in (None, "SKIP"):
|
||
return None, model_id, f"slug {slug!r} is not vLLM (kv-calc prices vLLM only; engine sets kvcalc_key=SKIP)"
|
||
if model_id not in MODEL_SPECS:
|
||
return None, model_id, f"model {model_id!r} has no calibrated kv-calc spec (not in MODEL_SPECS)"
|
||
return slug, model_id, None
|
||
|
||
|
||
def fit_verdict(target: str, card: Optional[str], vram_default: float) -> dict:
|
||
"""Structured fit verdict for `--fit <slug|model> --card <gpu>`.
|
||
|
||
Wraps the EXISTING predict()/solve_max_ctx() pricing — does NOT
|
||
reimplement the math. Returns a dict shaped:
|
||
{"verdict": "fits-clean"|"fits-constrained"|"wont-fit",
|
||
"vram_est_gb": float, "band_gb": float, "max_ctx": int}
|
||
or {"verdict": "unknown", "error": "..."} when an input can't be resolved.
|
||
"""
|
||
vram = _resolve_card_vram_gb(card)
|
||
if card is not None and vram is None:
|
||
return {"verdict": "unknown", "error": f"unrecognized --card {card!r} (pass a known card name or a GB number)"}
|
||
if vram is None:
|
||
vram = float(vram_default)
|
||
|
||
try:
|
||
slug, model_id, err = _resolve_fit_slug(target)
|
||
except Exception as exc: # registry import / lookup robustness
|
||
return {"verdict": "unknown", "error": f"could not resolve {target!r}: {exc}"}
|
||
if err is not None:
|
||
return {"verdict": "unknown", "error": err}
|
||
|
||
# required_sm gate (compat C3's arch floor, surfaced in the FIT verdict so
|
||
# the cockpit can hide/warn hardware-incompatible slugs — e.g. NVFP4 on
|
||
# Ampere: the VRAM would "fit" but the NATIVE kernels don't exist below
|
||
# sm 9.0). Since v0.24 some formats have a weight-only FALLBACK kernel
|
||
# below required_sm (registry `fallback_sm` — NVFP4 -> Marlin W4A16,
|
||
# floor sm 7.5; live-confirmed on 2x3090 2026-07-11). In the band
|
||
# [fallback_sm, required_sm) we PRICE the fit normally and attach an
|
||
# `hw_fallback` annotation instead of hiding — the cockpit badges it.
|
||
_entry_reg = COMPOSE_REGISTRY[slug]
|
||
_req_sm = _entry_reg.get("required_sm")
|
||
_fb_sm = _entry_reg.get("fallback_sm")
|
||
_card_sm = _resolve_card_sm(card)
|
||
_hw_fallback = None
|
||
if _req_sm is not None and _card_sm is not None and _card_sm < float(_req_sm):
|
||
if _fb_sm is not None and _card_sm >= float(_fb_sm):
|
||
_hw_fallback = {
|
||
"required_sm": float(_req_sm),
|
||
"card_sm": _card_sm,
|
||
"note": (
|
||
f"no native kernels below sm {float(_req_sm):g} — runs via "
|
||
f"weight-only fallback (Marlin); works, no speed edge vs "
|
||
f"native 4-bit quants on this card"
|
||
),
|
||
}
|
||
else:
|
||
return {
|
||
"verdict": "incompatible-hw",
|
||
"required_sm": float(_req_sm),
|
||
"card_sm": _card_sm,
|
||
"error": f"requires sm >= {float(_req_sm):g} (this card is sm_{_card_sm:g})",
|
||
}
|
||
|
||
try:
|
||
spec = MODEL_SPECS[model_id]
|
||
cfg = _compose_cfg_from_registry(PROFILES, model_id, "__fit__", slug)
|
||
predict_kwargs = dict(
|
||
spec=spec,
|
||
kv_format=cfg["kv_format"],
|
||
max_ctx=cfg["max_ctx"],
|
||
max_num_seqs=cfg["max_num_seqs"],
|
||
tp=cfg["tp"],
|
||
mem_util=cfg["mem_util"],
|
||
vram_gb=vram,
|
||
mtp=cfg.get("mtp", False),
|
||
weights_variant=cfg.get("weights_variant", "default"),
|
||
drafter_gb=cfg.get("drafter_gb", 0.0),
|
||
dflash_draft_gb=cfg.get("dflash_draft_gb", 0.0),
|
||
)
|
||
pred = predict(**predict_kwargs)
|
||
solver_kwargs = {k: v for k, v in predict_kwargs.items() if k != "max_ctx"}
|
||
max_ctx_fit = solve_max_ctx(**solver_kwargs)
|
||
except Exception as exc: # any pricing-path failure → honest unknown
|
||
return {"verdict": "unknown", "error": f"pricing {slug!r} failed: {exc}"}
|
||
|
||
res = {
|
||
"verdict": _RAW_VERDICT_MAP[pred.verdict],
|
||
"vram_est_gb": round(pred.total_gb, 4),
|
||
"band_gb": FIT_BAND_GB,
|
||
"max_ctx": int(max_ctx_fit),
|
||
}
|
||
if _hw_fallback is not None:
|
||
res["hw_fallback"] = _hw_fallback
|
||
return res
|
||
|
||
|
||
def fit_all_verdicts(card: Optional[str], vram_default: float) -> dict:
|
||
"""Batch fit verdicts for EVERY registry slug on one card, in a SINGLE
|
||
process — one registry/profile load instead of N subprocesses.
|
||
|
||
Mirrors `--fit` per slug (same verdict shape); non-vLLM slugs
|
||
(kvcalc_key SKIP) get ``{"verdict": "skip"}``. Shaped:
|
||
{"card": <str|None>, "card_vram_gb": <float>,
|
||
"variants": {<slug>: <fit_verdict dict>, ...}}
|
||
The cockpit catalog consumes this instead of fanning out N `--fit` calls.
|
||
"""
|
||
from scripts.lib.profiles.compose_registry import COMPOSE_REGISTRY
|
||
|
||
vram = _resolve_card_vram_gb(card)
|
||
if card is not None and vram is None:
|
||
return {
|
||
"card": card,
|
||
"error": f"unrecognized --card {card!r} (pass a known card name or a GB number)",
|
||
"variants": {},
|
||
}
|
||
if vram is None:
|
||
vram = float(vram_default)
|
||
|
||
variants: dict = {}
|
||
_card_sm = _resolve_card_sm(card)
|
||
for slug, entry in COMPOSE_REGISTRY.items():
|
||
# required_sm gate first — applies to SKIP (non-vLLM) slugs too, and
|
||
# short-circuits the pricing for slugs whose kernels can't exist here.
|
||
# Slugs with a `fallback_sm` at/below this card fall THROUGH to normal
|
||
# pricing (fit_verdict re-checks and attaches the hw_fallback note).
|
||
_req_sm = entry.get("required_sm")
|
||
_fb_sm = entry.get("fallback_sm")
|
||
if (
|
||
_req_sm is not None
|
||
and _card_sm is not None
|
||
and _card_sm < float(_req_sm)
|
||
and not (_fb_sm is not None and _card_sm >= float(_fb_sm))
|
||
):
|
||
variants[slug] = {
|
||
"verdict": "incompatible-hw",
|
||
"required_sm": float(_req_sm),
|
||
"card_sm": _card_sm,
|
||
"error": f"requires sm >= {float(_req_sm):g} (this card is sm_{_card_sm:g})",
|
||
}
|
||
continue
|
||
if entry.get("kvcalc_key") in (None, "SKIP"):
|
||
variants[slug] = {"verdict": "skip"}
|
||
continue
|
||
variants[slug] = fit_verdict(slug, card, vram_default)
|
||
return {"card": card, "card_vram_gb": round(float(vram), 4), "variants": variants}
|
||
|
||
|
||
# =============================================================================
|
||
# CLI
|
||
# =============================================================================
|
||
|
||
def _all_compose_choices() -> list[str]:
|
||
"""Flat list of compose names across all models for argparse choices."""
|
||
out = []
|
||
for model_key in COMPOSES:
|
||
out.extend(COMPOSES[model_key].keys())
|
||
return sorted(set(out))
|
||
|
||
|
||
def _resolve_compose_model(compose_name: str, explicit_model: Optional[str]) -> str:
|
||
"""Infer model from compose name if --model not given.
|
||
|
||
Composes are namespaced by prefix; Qwen uses bare names, Gemma uses gemma-*.
|
||
"""
|
||
if explicit_model:
|
||
return explicit_model
|
||
for model_key, composes in COMPOSES.items():
|
||
if compose_name in composes:
|
||
return model_key
|
||
return "qwen3.6-27b" # back-compat default
|
||
|
||
|
||
def main():
|
||
p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
|
||
p.add_argument("--model", choices=sorted(MODEL_SPECS.keys()),
|
||
help="Which model to predict for. Default: qwen3.6-27b (back-compat) or inferred from --compose.")
|
||
p.add_argument("--compose", choices=_all_compose_choices(),
|
||
help="Use a shipped compose's defaults. Override individual flags below.")
|
||
p.add_argument("--kv-format", choices=sorted(KV_FORMAT_BYTES.keys()),
|
||
help="KV cache format. Default: from --compose, or fp8_e5m2.")
|
||
p.add_argument("--max-ctx", type=int, help="max_model_len. Default: from --compose, or 180000.")
|
||
p.add_argument("--max-num-seqs", type=int, help="max_num_seqs. Default: from --compose, or 1.")
|
||
p.add_argument("--sequences", type=int, help="Alias for --max-num-seqs, for KV planning calculators.")
|
||
p.add_argument("--tp", type=int, choices=[1, 2, 4, 8, 16], help="tensor_parallel_size. Default: from --compose, or 1.")
|
||
p.add_argument("--gpus", type=int, default=None, help="Physical GPUs in the local single-node rig (1-8). Informational; TP controls sharding. Default: TP.")
|
||
p.add_argument("--mem-util", type=float, help="gpu_memory_utilization. Default: from --compose, or 0.95.")
|
||
p.add_argument("--vram", type=float, default=24, help="VRAM per card in GB. Default 24.")
|
||
p.add_argument("--mtp", action="store_true", default=None, help="MTP enabled (Qwen: n=3 built-in; Gemma: external drafter).")
|
||
p.add_argument("--no-mtp", dest="mtp", action="store_false")
|
||
p.add_argument("--drafter-gb", type=float, default=None,
|
||
help="Drafter model size in GB (MTP / DFlash). 0 if not using a drafter.")
|
||
p.add_argument("--dflash-draft-gb", type=float, default=None,
|
||
help="(deprecated alias for --drafter-gb)")
|
||
p.add_argument("--weights-variant", choices=["default", "int4", "awq", "bf16", "int8", "nvfp4"], default=None,
|
||
help="Gemma 4 only: which weight quant variant. Default: from --compose, or int4.")
|
||
p.add_argument("--calibration", action="store_true", help="Print predicted vs measured for all calibrated models.")
|
||
p.add_argument("--fit", metavar="SLUG|MODEL",
|
||
help="Structured fit verdict for a registry compose slug (e.g. vllm/dual) or a bare "
|
||
"catalogued model (resolved to its curated vLLM default). Use with --card and --json.")
|
||
p.add_argument("--fit-all", action="store_true",
|
||
help="Batch --fit for EVERY registry slug on --card, as one JSON object "
|
||
"{card, card_vram_gb, variants:{slug: verdict}}. One process, no per-slug fan-out "
|
||
"(the cockpit catalog's fit column). Always emits JSON.")
|
||
p.add_argument("--card", metavar="GPU",
|
||
help="Per-card GPU for --fit: a known card name (e.g. rtx3090, a6000) or a GB number. "
|
||
"Default: --vram.")
|
||
p.add_argument("--solve-max-ctx", action="store_true", help="Binary-search for the largest max_ctx that fits.")
|
||
p.add_argument("--json", action="store_true", help="Output prediction as JSON.")
|
||
p.add_argument("--kv-breakdown", action="store_true",
|
||
help="Also print/report raw architecture cache buckets (attention KV, SSM state, optional indexer/draft estimates).")
|
||
p.add_argument("--kv-layout", choices=["auto", "dense_gqa", "hybrid_mamba", "sliding_window", "compressed_sparse"],
|
||
default="auto", help="Reporting label for --kv-breakdown. Default: infer from model family.")
|
||
p.add_argument("--include-draft-kv", action="store_true",
|
||
help="Include --draft-kv-gb in --kv-breakdown. Built-in MTP token bump still comes from --mtp.")
|
||
p.add_argument("--draft-kv-gb", type=float, default=0.0,
|
||
help="Per-card draft KV estimate to include with --include-draft-kv. Default 0.")
|
||
p.add_argument("--compressed-layers", type=int, default=0,
|
||
help="Optional compressed/sparse KV layer count for --kv-breakdown reporting. Default 0.")
|
||
p.add_argument("--compression-ratio", type=int, default=128,
|
||
help="Token compression ratio for --compressed-layers. Default 128.")
|
||
p.add_argument("--compressed-head-dim", type=int,
|
||
help="Effective compressed KV head dimension for --compressed-layers.")
|
||
p.add_argument("--indexer-ratio-layers", type=int, default=0,
|
||
help="Optional indexer-cache layer count for --kv-breakdown reporting. Default 0.")
|
||
p.add_argument("--indexer-compress-ratio", type=int, default=4,
|
||
help="Token compression ratio for indexer cache. Default 4.")
|
||
p.add_argument("--indexer-head-dim", type=int,
|
||
help="Effective indexer head dimension for --indexer-ratio-layers.")
|
||
p.add_argument("--indexer-format", choices=sorted(INDEXER_FORMAT_BYTES.keys()), default="fp4",
|
||
help="Indexer precision for --kv-breakdown. Default fp4.")
|
||
args = p.parse_args()
|
||
|
||
if args.sequences is not None and args.max_num_seqs is not None and args.sequences != args.max_num_seqs:
|
||
print("ERROR: --sequences and --max-num-seqs disagree; pass only one value.", file=sys.stderr)
|
||
return 2
|
||
if args.sequences is not None:
|
||
args.max_num_seqs = args.sequences
|
||
|
||
if args.calibration:
|
||
run_calibration()
|
||
return 0
|
||
|
||
if args.fit is not None:
|
||
result = fit_verdict(args.fit, args.card, args.vram)
|
||
if args.json:
|
||
print(json.dumps(result, indent=2))
|
||
else:
|
||
print(f"Fit verdict — {args.fit}"
|
||
+ (f" on {args.card}" if args.card is not None else f" on {args.vram} GB/card"))
|
||
if result["verdict"] == "unknown":
|
||
print(f" Verdict: unknown ({result['error']})")
|
||
else:
|
||
print(f" Verdict: {result['verdict']}")
|
||
print(f" VRAM est: {result['vram_est_gb']:.2f} GB / card (±{result['band_gb']:.1f} GB band)")
|
||
print(f" Max ctx fit: {result['max_ctx']:,} tokens")
|
||
return 0 if result["verdict"] != "unknown" else 2
|
||
|
||
if args.fit_all:
|
||
result = fit_all_verdicts(args.card, args.vram)
|
||
print(json.dumps(result, indent=2))
|
||
return 0 if "error" not in result else 2
|
||
|
||
# Resolve model: explicit --model > inferred from --compose > qwen3.6-27b
|
||
model_key = _resolve_compose_model(args.compose, args.model) if args.compose else (args.model or "qwen3.6-27b")
|
||
spec = MODEL_SPECS[model_key]
|
||
|
||
# Resolve compose-derived defaults
|
||
if args.compose:
|
||
# Compose must belong to the resolved model
|
||
if args.compose not in COMPOSES[model_key]:
|
||
print(f"ERROR: --compose {args.compose} is not in --model {model_key}'s compose list.", file=sys.stderr)
|
||
print(f" Available for {model_key}: {', '.join(sorted(COMPOSES[model_key].keys()))}", file=sys.stderr)
|
||
return 2
|
||
cfg = COMPOSES[model_key][args.compose]
|
||
kv_format = args.kv_format or cfg["kv_format"]
|
||
max_ctx = args.max_ctx or cfg["max_ctx"]
|
||
max_num_seqs = args.max_num_seqs or cfg["max_num_seqs"]
|
||
tp = args.tp or cfg["tp"]
|
||
mem_util = args.mem_util if args.mem_util is not None else cfg["mem_util"]
|
||
mtp = args.mtp if args.mtp is not None else cfg.get("mtp", False)
|
||
drafter_gb = args.drafter_gb if args.drafter_gb is not None else cfg.get("drafter_gb", 0.0)
|
||
dflash_gb = args.dflash_draft_gb if args.dflash_draft_gb is not None else cfg.get("dflash_draft_gb", 0.0)
|
||
weights_variant = args.weights_variant or cfg.get("weights_variant", "default")
|
||
header = f"Predicted budget — {model_key} / {args.compose} on {args.vram} GB VRAM (kv={kv_format}, ctx={max_ctx:,}, seqs={max_num_seqs}, TP={tp}, mem={mem_util})"
|
||
else:
|
||
kv_format = args.kv_format or "fp8_e5m2"
|
||
max_ctx = args.max_ctx or 180000
|
||
max_num_seqs = args.max_num_seqs or 1
|
||
tp = args.tp or 1
|
||
mem_util = args.mem_util if args.mem_util is not None else 0.95
|
||
mtp = bool(args.mtp) if args.mtp is not None else False
|
||
drafter_gb = args.drafter_gb or 0.0
|
||
dflash_gb = args.dflash_draft_gb or 0.0
|
||
weights_variant = args.weights_variant or "default"
|
||
header = f"Predicted budget — {model_key} custom config on {args.vram} GB VRAM (kv={kv_format}, ctx={max_ctx:,}, seqs={max_num_seqs}, TP={tp}, mem={mem_util})"
|
||
|
||
try:
|
||
_validate_tp_for_spec(spec, tp)
|
||
except ValueError as exc:
|
||
print(f"ERROR: {exc}", file=sys.stderr)
|
||
return 2
|
||
|
||
gpus = args.gpus if args.gpus is not None else tp
|
||
if gpus < 1 or gpus > 8:
|
||
print("ERROR: --gpus must be between 1 and 8 for single-node home/workstation planning.", file=sys.stderr)
|
||
return 2
|
||
|
||
if args.solve_max_ctx:
|
||
best = solve_max_ctx(
|
||
spec, kv_format=kv_format, max_num_seqs=max_num_seqs,
|
||
tp=tp, mem_util=mem_util, vram_gb=args.vram,
|
||
drafter_gb=drafter_gb, dflash_draft_gb=dflash_gb, mtp=mtp,
|
||
weights_variant=weights_variant,
|
||
)
|
||
if best > 0:
|
||
pred_at_best = predict(
|
||
spec=spec, kv_format=kv_format, max_ctx=best, max_num_seqs=max_num_seqs,
|
||
tp=tp, mem_util=mem_util, vram_gb=args.vram,
|
||
drafter_gb=drafter_gb, dflash_draft_gb=dflash_gb, mtp=mtp,
|
||
weights_variant=weights_variant,
|
||
)
|
||
if args.json:
|
||
out = pred_at_best.__dict__.copy()
|
||
out["solved_max_ctx"] = best
|
||
print(json.dumps(out, indent=2))
|
||
else:
|
||
print(f"Max-ctx solver — {model_key} / {kv_format}, seqs={max_num_seqs}, TP={tp}, mem_util={mem_util}, VRAM={args.vram} GB")
|
||
print(f" Largest max_ctx that fits: {best:,} tokens")
|
||
print(f" At that ctx: predicted = {pred_at_best.total_gb:.2f} GB / card ({pred_at_best.pct_of_vram:.0f}% of budget)")
|
||
print(f" Verdict at that ctx: {pred_at_best.verdict}")
|
||
for note in pred_at_best.notes:
|
||
print(f" Note: {note}")
|
||
print()
|
||
print("Note: this is a directional estimate (±1.5 GB error band). The vLLM engine")
|
||
print("pre-check (gpu_worker.py boot log) is authoritative.")
|
||
else:
|
||
print(f"No max_ctx fits at this config on {args.vram} GB. Reduce TP, swap KV format, or get bigger cards.")
|
||
return 0
|
||
|
||
pred = predict(
|
||
spec=spec, kv_format=kv_format, max_ctx=max_ctx, max_num_seqs=max_num_seqs,
|
||
tp=tp, mem_util=mem_util, vram_gb=args.vram,
|
||
drafter_gb=drafter_gb, dflash_draft_gb=dflash_gb, mtp=mtp,
|
||
weights_variant=weights_variant,
|
||
)
|
||
|
||
breakdown = None
|
||
if args.kv_breakdown:
|
||
try:
|
||
breakdown = architecture_cache_breakdown(
|
||
spec=spec,
|
||
kv_format=kv_format,
|
||
max_ctx=max_ctx,
|
||
max_num_seqs=max_num_seqs,
|
||
tp=tp,
|
||
gpus=gpus,
|
||
layout=args.kv_layout,
|
||
mtp=mtp,
|
||
include_draft_kv=args.include_draft_kv,
|
||
draft_kv_gb=args.draft_kv_gb,
|
||
compressed_layers=args.compressed_layers,
|
||
compression_ratio=args.compression_ratio,
|
||
compressed_head_dim=args.compressed_head_dim,
|
||
indexer_ratio_layers=args.indexer_ratio_layers,
|
||
indexer_compress_ratio=args.indexer_compress_ratio,
|
||
indexer_head_dim=args.indexer_head_dim,
|
||
indexer_format=args.indexer_format,
|
||
)
|
||
except ValueError as exc:
|
||
print(f"ERROR: {exc}", file=sys.stderr)
|
||
return 2
|
||
|
||
if args.json:
|
||
out = pred.__dict__.copy()
|
||
if breakdown is not None:
|
||
out["cache_breakdown"] = breakdown.__dict__
|
||
print(json.dumps(out, indent=2))
|
||
else:
|
||
print(fmt_prediction(pred, header=header))
|
||
if breakdown is not None:
|
||
print()
|
||
print(fmt_cache_breakdown(breakdown))
|
||
print()
|
||
print("Run `tools/kv-calc.py --calibration` to see predicted-vs-measured for all anchors.")
|
||
print("Run `tools/kv-calc.py --solve-max-ctx ...` to find the largest max_ctx that fits.")
|
||
print("See docs/KV_MATH.md for math + per-model architecture details.")
|
||
return 0 if pred.verdict in ("PASS", "TIGHT") else 1
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|