diff --git a/docs/cli.md b/docs/cli.md index fb5e613f7b..d30b6e73cd 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -7,6 +7,7 @@ ft [args] | Command | Purpose | |---|---| | `ft serve` | Start the API server (OpenAI `/v1/*`, Anthropic `/v1/messages`, Responses) | +| `ft info` | Forecast a model's GPU memory for a serve configuration, without loading it | | `ft shell` | Chat with a server in the terminal | | `ft ctl` | Query and manage a running server over HTTP | | `ft launch` | Configure and launch a coding agent against a server | @@ -69,6 +70,7 @@ ft serve --model ... --gpu GPU-9e8d7c6b # the same card by UUID (a unique prefi | Flag | Default | Meaning | |---|---|---| | `--memory-ratio` | 0.9 | Fraction of free VRAM the engine may use (weights + MoE cache + KV) | +| `--skip-preflight` | off | Load the weights even when the pre-load memory forecast says the configuration cannot fit; see [ft info](#ft-info) | | `--num-pages` / `--num-tokens` | auto | KV capacity override in pages / tokens (mutually exclusive; auto sizes from VRAM left after weights and MoE cache) | | `--page-size` | 1 | KV page size; DSV4 forces 128, the TRTLLM backend needs 16/32/64, SWA models require 1 | | `--cache-type` | radix | `radix` (prefix reuse; SWA/GDN-aware variants picked automatically) or `naive` | @@ -122,6 +124,65 @@ so a client can gate its attachment controls without reading the checkpoint conf | `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain | | `--allowed-local-media-path` | off | Directory `file://` image refs may be read from; unset rejects local files | +## ft info + +```bash +ft info [ft serve flags] [--json] [--gpu-memory-gib N | --gpu-free-gib N] +``` + +Forecasts what `ft serve [flags]` would put on the GPU, and whether it +fits, without reading a single weight: it resolves the configuration the way the +engine does, builds the model on PyTorch's meta device, and sizes every pool +with the engine's own sizing code. It runs in a few seconds, needs no CUDA, and +takes every `ft serve` flag, so "what if" combinations can be tried before a +multi-minute load: + +```bash +ft info ~/models/Qwen3.8-27B-NVFP4 --text-model-only +ft info Qwen/Qwen3-30B-A3B --gpu-memory-gib 24 --max-running-requests 1 +``` + +The report lists: + +- the architecture, the format each layer type is stored in on the GPU, and the + checkpoint's dtypes; +- the weights resident on the GPU per category (attention, linear attention, + MLP, routed experts, embeddings / lm_head, vision encoder, MTP heads, rope + tables) next to what the checkpoint stores, so a skipped vision tower + (`--text-model-only`), MTP heads the engine does not load, or offloaded experts + living in host RAM are visible; +- the budget: `--memory-ratio` x free memory, minus weights, the MoE expert slot + cache, the GDN state pool (slots x bytes), and the KV cache (bytes per token x + tokens); +- what lives outside the budget: the attention backend's workspace, the page + table, the CUDA graphs, the prefill activations of one `--max-prefill-length` + chunk, and `--vram-reserve-mb`; the "free after init" line is the number the + engine logs as `Free memory after initialization`; +- a verdict -- `fits`, `tight` or `does not fit` -- with the max context and the + number of concurrent requests, and tips: flag changes with their predicted + effect, plus the cheapest combination that fits. + +Free memory comes from NVML for the local GPU (`--gpu` picks the card), minus +the CUDA context the engine creates first. `--gpu-memory-gib N` plans for an +empty card of N GiB instead, `--gpu-free-gib N` for N GiB free as `nvidia-smi` +reports it. Without a GPU and without these flags the weights and per-token +costs are still reported. `--json` prints the same data for tools; the exit code +is 1 when the configuration does not fit, else 0. + +Weights come from the meta-built model, whose parameters are exactly what the +loader fills (tensors the family skips are absent, re-quantized ones already +have their stored size); the checkpoint column reads only the safetensors +headers -- from the local directory, the Hugging Face cache, or the Hub through +HTTP range requests. The load overhead (kernel modules, library handles), the +prefill activations and the CUDA graph memory are estimates and are labelled so. + +`ft serve` (and the offline `LLM` class) runs the same forecast right before it +loads the weights, against the free memory it has just measured. A configuration +whose KV cache cannot be sized even before those estimates is refused within +seconds, with the forecast and the tips in the log, instead of failing with +`Not enough memory for KV cache` after the load. A tight one only logs a warning +and continues. `--skip-preflight` loads regardless. + ## ft shell ```bash diff --git a/docs/quickstart.md b/docs/quickstart.md index bea8d2e5c4..66bf5a2134 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -13,6 +13,14 @@ and MoE backends, cache sizes, tool-call and reasoning parsers — resolves from the checkpoint and the GPU; see [cli.md](cli.md) for the flags. The server is ready when the log reaches `API server is ready to serve on 127.0.0.1:1919`. +To check first whether a model fits your GPU, and which flags to change if it +does not, run `ft info` with the same arguments; it takes seconds and reads no +weights: + +```bash +ft info ~/models/Qwen3.6-35B-A3B +``` + ## Send a request Check what is being served: diff --git a/python/freetoken/cli.py b/python/freetoken/cli.py index 4e6deff230..a0d9373630 100644 --- a/python/freetoken/cli.py +++ b/python/freetoken/cli.py @@ -11,6 +11,7 @@ def _print_help(file: TextIO) -> None: Commands: serve Start the FreeToken API server + info Forecast a model's GPU memory for a serve configuration, without loading it shell Chat with a FreeToken server in the terminal ctl Query and manage a running FreeToken server daemon Run the FreeToken supervisor (persistent engine service) @@ -31,6 +32,17 @@ def _run_serve(argv: list[str]) -> int: return 0 +def _run_info(argv: list[str]) -> int: + import os + + if "--verbose" not in argv: + # before the first freetoken logger exists: the level is read once + os.environ.setdefault("LOG_LEVEL", "WARNING") + from freetoken.server.info import main + + return main(argv, prog="ft info") + + def _run_shell(argv: list[str]) -> int: from freetoken.shell import main @@ -92,6 +104,7 @@ def _run_bench(argv: list[str]) -> int: COMMANDS = { "serve": "_run_serve", + "info": "_run_info", "shell": "_run_shell", "ctl": "_run_ctl", "daemon": "_run_daemon", diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 3e38cb2165..d324864da7 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -69,6 +69,8 @@ class EngineConfig: cuda_graph_max_bs: int | None = None page_size: int = 1 memory_ratio: float = 0.9 + # --skip-preflight: load the weights even when the pre-load memory forecast says the config cannot fit + skip_preflight: bool = False # Hybrid GDN models default to the HybridRadixCache (cross-request GDN-state prefix reuse); # `--cache-type naive` opts out. linear_state_cache_ratio sizes the GDN snapshot cache as # ceil(ratio * max_running_req) extra slots. diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index e855bcd32f..3413eac48b 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -76,6 +76,34 @@ def _startup_kv_budget(memory_ratio: float, init_free_memory: int, new_free_memo return int(memory_ratio * init_free_memory) - (init_free_memory - new_free_memory) +def plan_moe_cache_auto( + config, pool_cls, *, baseline_free: int, weights_bytes: int, per_expert_bytes: int, + max_slots: int | None, +) -> tuple[int, int, bool]: + """--moe-cache-auto's (moe_cache_size, num_pages, prefill_overlap) from the measured (or, for + the pre-load forecast, predicted) baseline and resident weights.""" + from freetoken.engine.cache_budget import resolve_moe_cache_auto + + cache_per_page, fixed_cache_size, page_tokens, min_reserve = pool_cls.kv_cost(config) + fixed_cache_size += state_pool_bytes(config) # sibling GDN state pool, engine-summed + num_experts = config.model_config.num_experts + total_experts = config.model_config.num_moe_layers * num_experts + return resolve_moe_cache_auto( + baseline_free=baseline_free, + weights_bytes=weights_bytes, + memory_ratio=config.memory_ratio, + cache_per_page=cache_per_page, + fixed_cache_size=fixed_cache_size, + per_expert_bytes=per_expert_bytes, + num_experts=num_experts, + total_experts=total_experts, + prefill_overlap=config.moe_prefill_overlap, + kv_reserve_tokens=max(config.kv_reserve_tokens, min_reserve), + page_size=page_tokens, + max_slots=max_slots, + ) + + def _page_table_width(max_seq_len: int, page_size: int) -> int: """Column count for the page table. ``_write_page_table`` writes WHOLE trailing pages, so the highest column touched is ``align_ceil(max_seq_len, page_size) - 1`` -- which the 32-alignment @@ -357,6 +385,10 @@ def __init__(self, config: EngineConfig): set_rope_device(self.device) with torch.device("meta"), torch_dtype(config.dtype): self.model = create_model(config.model_config) + if not config.skip_preflight: + from .forecast import run_preflight + + run_preflight(config, self.model, init_free_memory) self._load_weights(config) if config.active_encoders: from freetoken.models.blocks import SupportsMultimodal @@ -592,24 +624,14 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None) Pure glue over the Phase-1 budget policy; isolated here so it is unit-testable without a GPU. Reused by the Phase-2 runtime rebuild. """ - from freetoken.engine.cache_budget import expert_bytes_per_slot, resolve_moe_cache_auto + from freetoken.engine.cache_budget import expert_bytes_per_slot - cache_per_page, fixed_cache_size, page_tokens, min_reserve = self._pool_cls.kv_cost(config) - fixed_cache_size += state_pool_bytes(config) # sibling GDN state pool, engine-summed - num_experts = config.model_config.num_experts - total_experts = config.model_config.num_moe_layers * num_experts - return resolve_moe_cache_auto( + return plan_moe_cache_auto( + config, + self._pool_cls, baseline_free=self._baseline_free, weights_bytes=self._weights_bytes, - memory_ratio=config.memory_ratio, - cache_per_page=cache_per_page, - fixed_cache_size=fixed_cache_size, per_expert_bytes=expert_bytes_per_slot(banks.sources), - num_experts=num_experts, - total_experts=total_experts, - prefill_overlap=config.moe_prefill_overlap, - kv_reserve_tokens=max(config.kv_reserve_tokens, min_reserve), - page_size=page_tokens, max_slots=method.slot_limit() if method is not None else None, ) diff --git a/python/freetoken/engine/forecast.py b/python/freetoken/engine/forecast.py new file mode 100644 index 0000000000..a63a3b21e4 --- /dev/null +++ b/python/freetoken/engine/forecast.py @@ -0,0 +1,843 @@ +"""Pre-load GPU memory forecast for one serve configuration. + +Prices what Engine.__init__ will put on the device before any weight is read: the resident +weights come from the meta-device model the engine builds anyway (its parameters are exactly +the tensors the loader materializes), and every pool is sized by the engine's own functions +(``_startup_kv_budget``, the pool family's ``kv_cost`` / ``plan_num_pages``, ``state_pool_bytes``, +``plan_moe_cache_auto``). Only what the engine does not price up +front -- load overhead, prefill activations, CUDA graphs -- is a heuristic here, and the verdict +never refuses on a heuristic alone. +""" + +from __future__ import annotations + +import copy +import itertools +import re +from collections import Counter +from dataclasses import asdict, dataclass, field +from typing import Any, Callable, Iterable + +import torch + +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +GiB = 1 << 30 +MiB = 1 << 20 + +# CUDA context + distributed setup held before the engine's first free-memory reading: nvidia-smi +# showed 22.79 GiB free where the engine logged 22.39 GiB (RTX 4090, CUDA 13.1). +CUDA_CONTEXT_BYTES = 400 * MiB +# Device memory the load takes beyond the parameters (kernel modules, library handles, finalize +# scratch): measured weights minus meta parameters was 0.31 GiB for Qwen3.8-27B-NVFP4 on an RTX 4090. +LOAD_OVERHEAD_BYTES = 320 * MiB +# Free memory "fits" keeps at the forecast prefill peak; less is reported as tight. +PEAK_MARGIN_BYTES = 256 * MiB +# Below this context (or the model's own limit) a configuration is reported as tight, not fits. +USEFUL_CONTEXT = 8192 +# fla's chunked GDN prefill keeps one [K, V] state per value head and 64-token chunk. +_GDN_CHUNK = 64 + +VERDICTS = ("does not fit", "tight", "fits") + +CATEGORY_LABELS = { + "attention": "attention", + "linear_attention": "linear attention (GDN)", + "mlp": "MLP / shared experts / routers", + "experts": "MoE routed experts", + "embeddings": "embeddings / lm_head", + "vision": "vision encoder", + "mtp": "MTP / draft heads", + "other": "norms / other", + "rope": "rope tables", +} + +_LANGUAGE_ROOTS = ("model.language_model.", "language_model.model.", "language_model.") + + +def tensor_category(name: str) -> str: + """Bucket a checkpoint or model tensor name for the report.""" + from freetoken.models.config import VISION_KEY_PREFIXES + + n = name + for root in _LANGUAGE_ROOTS: + if n.startswith(root): + n = n[len(root):] + break + if n.startswith("model.") and n[len("model."):].startswith(VISION_KEY_PREFIXES): + n = n[len("model."):] + if n.startswith(VISION_KEY_PREFIXES + ("vision_model.", "multi_modal_projector.", "mm_projector.")): + return "vision" + parts = n.split(".") + if "mtp" in parts or any("nextn" in p for p in parts): + return "mtp" + keys = set(parts) + if keys & {"embed_tokens", "lm_head", "wte", "embed"}: + return "embeddings" + if "experts" in keys: + return "experts" + if keys & {"linear_attn", "mamba", "mixer"}: + return "linear_attention" + if keys & {"self_attn", "attn", "attention"}: + return "attention" + if keys & {"mlp", "feed_forward", "ffn", "shared_expert", "shared_experts", "router"}: + return "mlp" + return "other" + + +def _nbytes(t: torch.Tensor) -> int: + return t.numel() * t.element_size() + + +def _walk(obj: Any, path: str = ""): + """(path, object) over the model's own object graph (BaseOP trees, lists, dicts).""" + seen: set[int] = set() + stack = [(path, obj)] + while stack: + p, o = stack.pop() + if id(o) in seen: + continue + seen.add(id(o)) + yield p, o + if isinstance(o, torch.Tensor): + continue + if isinstance(o, (list, tuple)): + stack.extend((p, x) for x in o) + elif isinstance(o, dict): + stack.extend((f"{p}.{k}", x) for k, x in o.items()) + elif type(o).__module__.startswith("freetoken") and hasattr(o, "__dict__"): + stack.extend((f"{p}.{k}" if p else k, x) for k, x in vars(o).items()) + + +def _eager_tensor_bytes(model) -> int: + """Device tensors built eagerly during the meta build (rope cos/sin tables): resident too.""" + storages: dict[int, int] = {} + for _, o in _walk(model): + if isinstance(o, torch.Tensor) and o.device.type != "meta": + storages[o.untyped_storage().data_ptr()] = o.untyped_storage().nbytes() + return sum(storages.values()) + + +_BLOCK_STACK = re.compile(r"^(?P.+?\.(?:blocks|layers))\.(?P\d+)\.") + + +def _streamed_encoder_bytes(model, state: dict[str, torch.Tensor]) -> tuple[int, int]: + """(block bytes moved to pinned host banks, device staging bytes) under --mm-encoder-weights host: + the vision block stacks stream through BlockWeightStreamer, the rest of the tower stays resident.""" + from freetoken.layers.base import OPList + from freetoken.models.weight_stream import BlockWeightStreamer + + stacks: dict[str, int] = {} + for name, t in state.items(): + if tensor_category(name) != "vision": + continue + m = _BLOCK_STACK.match(name) + if m: + stacks[m["stack"]] = stacks.get(m["stack"], 0) + _nbytes(t) + streamed = staging = 0 + for stack, nbytes in stacks.items(): + obj = model + try: + for seg in stack.split("."): + obj = getattr(obj, seg) + except AttributeError: + continue + blocks = obj.op_list if isinstance(obj, OPList) else obj + if not blocks: + continue + streamed += nbytes + staging += BlockWeightStreamer.staging_bytes(blocks) + return streamed, staging + + +@dataclass +class WeightReport: + gpu: dict[str, int] + host: dict[str, int] = field(default_factory=dict) + notes: list[str] = field(default_factory=list) + # device bytes of the vision tower under --mm-encoder-weights host (None: nothing streams) + vision_gpu_if_host: int | None = None + + @property + def gpu_total(self) -> int: + return sum(self.gpu.values()) + + +# DF11 (lossless entropy-coded bf16) buffers are sized by the data at load and are empty in the meta +# build; GLM-4.5-Air's attention and embedding measured 10.7-11.2 bits per weight. +_DF11_BITS_PER_WEIGHT = 11.2 + + +def _df11_estimate(model) -> Counter[str]: + """Per category, the estimated bytes of DF11 weights whose buffers are still placeholders.""" + found: Counter[str] = Counter() + for path, o in _walk(model): + low8 = getattr(o, "low8", None) if hasattr(o, "_DF11_KEYS") else None + if isinstance(low8, torch.Tensor) and low8.numel() == 0: + found[tensor_category(f"{path}.weight")] += int(o.n * _DF11_BITS_PER_WEIGHT / 8) + return found + + +def resident_weights(model, config) -> WeightReport: + """Per-category bytes the loaded model keeps on the GPU, from its meta-device build: the loader + casts every checkpoint tensor into these parameters, so skipped tensors (a vision tower under + --text-model-only, MTP heads the family does not serve, offloaded experts) are already absent + and quantize/upcast-at-load is already priced.""" + state = model.state_dict() + gpu: Counter[str] = Counter() + for name, t in state.items(): + gpu[tensor_category(name)] += _nbytes(t) + df11 = _df11_estimate(model) + gpu.update(df11) + report = WeightReport(gpu=dict(gpu)) + if df11: + report.notes.append(f"DF11-compressed weights are estimated at {_DF11_BITS_PER_WEIGHT} bits per weight; " + "their real size depends on the data.") + rope = _eager_tensor_bytes(model) + if rope: + report.gpu["rope"] = rope + vision = report.gpu.get("vision", 0) + if vision and config.active_encoders: + streamed, staging = _streamed_encoder_bytes(model, state) + if streamed: + report.vision_gpu_if_host = vision - streamed + staging + if config.mm.encoder_weights == "host": + report.gpu["vision"] = report.vision_gpu_if_host + report.host["vision"] = streamed + return report + + +def quant_by_category(model) -> dict[str, list[str]]: + """Per category, the storage formats the model's projections were built with ("nvfp4 via triton", + "bfloat16") -- what the loader converts the checkpoint into.""" + found: dict[str, set[str]] = {} + for path, o in _walk(model): + if not isinstance(getattr(o, "in_features", None), int): + continue + method = getattr(o, "quant_method", None) + weight = getattr(o, "weight", None) + kind = getattr(getattr(method, "kind", None), "value", None) + if kind not in (None, "none"): + label = f"{kind} via {method.kernel.name}" + elif isinstance(weight, torch.Tensor): + label = str(weight.dtype).removeprefix("torch.") + else: + continue + found.setdefault(tensor_category(f"{path}.weight"), set()).add(label) + return {cat: sorted(v) for cat, v in found.items()} + + +def activation_bytes_per_token(model, config) -> int: + """Heuristic transient bytes per prefill token: a few residual-stream copies plus the widest + projection's input and output, the GDN chunk states, and the routed-expert intermediates.""" + from freetoken.kvcache.linear_state_pool import _linear_local_dims + + mc = config.model_config + item = config.dtype.itemsize + tp = config.tp_info.size + skip = ("visual", "vision_tower", "embed_vision", "vision_embedder", "lm_head", "embed_tokens", "mtp") + widest = 0 + for path, o in _walk(model): + if path.split(".")[0] in skip: + continue + i, out = getattr(o, "in_features", None), getattr(o, "out_features", None) + if isinstance(i, int) and isinstance(out, int): + widest = max(widest, i + out) + per_token = item * (4 * mc.hidden_size + widest) + group = mc.linear_attention_group() if hasattr(mc, "linear_attention_group") else None + if group is not None: + _, _, v_heads = _linear_local_dims(group, tp) + per_token += item * v_heads * group.key_head_dim * group.value_head_dim // _GDN_CHUNK + if getattr(mc, "is_moe", False): + per_token += item * mc.num_experts_per_tok * 3 * mc.moe_intermediate_size // tp + return int(per_token) + + +def _is_offload(config) -> bool: + from freetoken.moe import is_offload_moe_strategy + + return bool(getattr(config.model_config, "is_moe", False)) and is_offload_moe_strategy(config.moe_strategy) + + +@dataclass +class _Plan: + ok: bool + reason: str = "" + moe_slots: int = 0 + num_pages: int = 0 + prefill_overlap: bool = True + + +def _solve(config, pool_cls, free_before: int, weights_bytes: int, per_expert: int, + max_slots: int | None, state_bytes: int) -> _Plan: + """The engine's startup solve (offload cache, then KV pages) against predicted numbers.""" + from freetoken.engine.engine import _startup_kv_budget, plan_moe_cache_auto + + slots, overlap, planned_pages = 0, config.moe_prefill_overlap, None + if _is_offload(config): + if config.moe_cache_auto: + if per_expert <= 0: + return _Plan(False, "the expert slot size is unknown, cannot plan --moe-cache-auto") + try: + slots, pages, overlap = plan_moe_cache_auto( + config, pool_cls, baseline_free=free_before, weights_bytes=weights_bytes, + per_expert_bytes=per_expert, max_slots=max_slots, + ) + except AssertionError as exc: + return _Plan(False, f"--moe-cache-auto cannot split the budget ({exc})") + if config.num_page_override is None: + planned_pages = pages + else: + slots = config.moe_cache_size + resident = weights_bytes + slots * per_expert + available = _startup_kv_budget(config.memory_ratio, free_before, free_before - resident) - state_bytes + if planned_pages is not None: + pages = planned_pages + else: + try: + pages = pool_cls.plan_num_pages(config, available) + except ValueError as exc: + return _Plan(False, str(exc), slots, 0, overlap) + if pages <= 1: + return _Plan(False, f"no room for the KV cache ({_gib(available)} left for it after weights " + f"and fixed pools)", slots, pages, overlap) + if config.num_page_override is not None: + cache_per_page, fixed, _, _ = pool_cls.kv_cost(config) + need = pages * cache_per_page + fixed + physical = free_before - resident - state_bytes + if need > physical: + return _Plan(False, f"--num-pages {pages} needs {_gib(need)} of KV cache but only " + f"{_gib(physical)} is free after the weights", slots, pages, overlap) + return _Plan(True, "", slots, pages, overlap) + + +def _fixed_workspace_bytes(backend: str) -> int: + """Fixed (non-paged) attention workspace the backend allocates after the pools, summed over + comma-separated prefill,decode parts: FlashInfer's float workspace floor (256 MiB, + attention/fi.py) and TRT-LLM's 128 MiB buffer (attention/trtllm.py). The engine does not + charge it to the pool budget, so it comes out of the memory left after init.""" + sizes = {"fi": 256 * MiB, "trtllm": 128 * MiB} + return sum(sizes.get(p.strip(), 0) for p in backend.split(",")) + + +@dataclass +class Forecast: + free_before: int | None + memory_ratio: float + weights_bytes: int + load_overhead: int + kv_bytes_per_token: int + kv_fixed_bytes: int + page_size: int + state_slots: int + state_bytes_per_slot: int + per_expert_bytes: int + max_running_req: int + max_seq_len: int + moe_slots: int = 0 + num_pages: int = 0 + kv_tokens: int = 0 + max_context: int = 0 + attention_workspace: int = 0 + page_table: int = 0 + cuda_graphs: int = 0 + graph_batch_sizes: list[int] = field(default_factory=list) + prefill_chunk: int = 0 + prefill_activations: int = 0 + free_after_init: int | None = None + free_at_peak: int | None = None + verdict: str = "unknown" + reasons: list[str] = field(default_factory=list) + + @property + def state_bytes(self) -> int: + return self.state_slots * self.state_bytes_per_slot + + @property + def kv_bytes(self) -> int: + cache_per_page = self.kv_bytes_per_token * self.page_size + return self.num_pages * cache_per_page + (self.kv_fixed_bytes if self.num_pages else 0) + + @property + def moe_slot_bytes(self) -> int: + return self.moe_slots * self.per_expert_bytes + + @property + def budget(self) -> int | None: + return None if self.free_before is None else int(self.memory_ratio * self.free_before) + + def to_dict(self) -> dict: + out = asdict(self) + out.update(state_bytes=self.state_bytes, kv_bytes=self.kv_bytes, + moe_slot_bytes=self.moe_slot_bytes, budget=self.budget) + return out + + +@dataclass +class ForecastInputs: + """Everything the forecast needs that does not change with the sizing flags.""" + + config: Any + weights: WeightReport + free_before: int | None + per_expert_bytes: int = 0 + max_slots: int | None = None + activation_per_token: int = 0 + load_overhead: int = LOAD_OVERHEAD_BYTES + + def forecast(self, config=None, *, weight_delta: int = 0, per_expert: int | None = None) -> Forecast: + return forecast_memory( + self.config if config is None else config, + self.weights.gpu_total + weight_delta, + self.free_before, + per_expert_bytes=self.per_expert_bytes if per_expert is None else per_expert, + max_slots=self.max_slots, + activation_per_token=self.activation_per_token, + load_overhead=self.load_overhead, + ) + + +def forecast_memory(config, weights_bytes: int, free_before: int | None, *, per_expert_bytes: int = 0, + max_slots: int | None = None, activation_per_token: int = 0, + load_overhead: int = LOAD_OVERHEAD_BYTES) -> Forecast: + """Forecast the startup geometry of a resolved ``config`` (after ``_adjust_config``) with + ``weights_bytes`` resident and ``free_before`` device bytes free before the load. + + The verdict refuses ("does not fit") only when the engine's own solve fails with the + parameters alone -- the startup assert it would hit after loading. The load overhead and the + heuristic workspaces can only make it "tight".""" + from freetoken.engine.engine import _page_table_width + from freetoken.engine.graph import _determine_cuda_graph_bs + from freetoken.kvcache import resolve_pool_class + from freetoken.kvcache.linear_state_pool import _linear_pool_num_slots, state_pool_bytes + + mc = config.model_config + pool_cls = resolve_pool_class(mc) + cache_per_page, kv_fixed, _, _ = pool_cls.kv_cost(config) + has_state = mc.linear_attention_group() is not None + state_slots = _linear_pool_num_slots(config) if has_state else 0 + state_bytes = state_pool_bytes(config) + fc = Forecast( + free_before=free_before, + memory_ratio=config.memory_ratio, + weights_bytes=weights_bytes, + load_overhead=load_overhead, + kv_bytes_per_token=cache_per_page // config.page_size, + kv_fixed_bytes=kv_fixed, + page_size=config.page_size, + state_slots=state_slots, + state_bytes_per_slot=state_bytes // state_slots if state_slots else 0, + per_expert_bytes=per_expert_bytes if _is_offload(config) else 0, + max_running_req=config.max_running_req, + max_seq_len=config.max_seq_len, + attention_workspace=_fixed_workspace_bytes(config.attention_backend), + ) + if free_before is None: + return fc + + solve = lambda w: _solve(config, pool_cls, free_before, w, fc.per_expert_bytes, max_slots, state_bytes) # noqa: E731 + lower = solve(weights_bytes) + if not lower.ok: + fc.verdict = "does not fit" + fc.reasons.append(lower.reason) + fc.moe_slots = lower.moe_slots + return fc + plan = solve(weights_bytes + load_overhead) + tight: list[str] = [] + if not plan.ok: + tight.append(f"with the estimated {_gib(load_overhead)} load overhead: {plan.reason}") + plan = lower + fc.moe_slots, fc.num_pages = plan.moe_slots, plan.num_pages + fc.kv_tokens = plan.num_pages * config.page_size + fc.max_context = min(config.max_seq_len, fc.kv_tokens) + fc.page_table = (config.max_running_req + 1) * _page_table_width(fc.max_context, config.page_size) * 4 + bs_list = _determine_cuda_graph_bs(config.cuda_graph_bs, config.cuda_graph_max_bs, free_before) + fc.graph_batch_sizes = sorted(bs_list) + if bs_list: + width = _page_table_width(fc.max_context, config.page_size) + fc.cuda_graphs = max(bs_list) * (mc.vocab_size * 4 + activation_per_token + width * 4) + len(bs_list) * 2 * MiB + fc.prefill_chunk = min(config.max_forward_len, fc.max_context) + fc.prefill_activations = fc.prefill_chunk * activation_per_token + + resident = weights_bytes + load_overhead + fc.moe_slot_bytes + fc.state_bytes + fc.kv_bytes + fc.free_after_init = free_before - resident - fc.attention_workspace - fc.page_table + fc.free_at_peak = fc.free_after_init - fc.cuda_graphs - fc.prefill_activations + spare = fc.free_at_peak + if spare < 0: + tight.append( + f"prefill workspace not covered: a {fc.prefill_chunk}-token prefill needs ~{_gib(fc.prefill_activations)} and the " + f"CUDA graphs ~{_gib(fc.cuda_graphs)}, but only {_gib(fc.free_after_init)} is free after init" + ) + elif spare < PEAK_MARGIN_BYTES: + tight.append(f"only {_gib(spare)} spare at the prefill peak") + if fc.max_context < min(USEFUL_CONTEXT, config.max_seq_len): + tight.append(f"max context is only {fc.max_context} tokens") + override = getattr(config, "max_seq_len_override", None) + if override is not None and fc.kv_tokens < override: + tight.append(f"--max-seq-len-override {override} is above the {fc.kv_tokens} KV tokens; " + f"requests are capped at {fc.max_context}") + fc.reasons.extend(tight) + fc.verdict = "tight" if tight else "fits" + return fc + + +# --------------------------------------------------------------------------------------------- +# Tips: each candidate is one flag change, priced by re-running the forecast on a config copy. +# --------------------------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class Change: + flag: str + key: str # one change per key in a combination + cost: int # what the user gives up; the suggested combination minimizes the sum + apply: Callable[[Any], None] = lambda config: None + weight_delta: int = 0 + per_expert: int | None = None + note: str = "" + + +def _set(**values): + def apply(config): + for k, v in values.items(): + object.__setattr__(config, k, v) + return apply + + +def _set_running_req(n: int): + def apply(config): + if config.cuda_graph_max_bs == config.max_running_req: + object.__setattr__(config, "cuda_graph_max_bs", n) + if config.cuda_graph_bs is not None: + object.__setattr__(config, "cuda_graph_bs", [bs for bs in config.cuda_graph_bs if bs <= n] or [1]) + object.__setattr__(config, "max_running_req", n) + return apply + + +def _tips_concurrency(inp: ForecastInputs) -> Iterable[Change]: + mr = inp.config.max_running_req + for n in sorted({1, mr // 2} - {0}): + if n < mr: + yield Change(f"--max-running-requests {n}", "max_running_req", 3 if n == 1 else 2, _set_running_req(n)) + + +def _tips_memory_ratio(inp: ForecastInputs) -> Iterable[Change]: + current = inp.config.memory_ratio + for ratio, cost in ((0.95, 2), (0.97, 3)): + if ratio > current + 1e-9: + yield Change(f"--memory-ratio {ratio}", "memory_ratio", cost, _set(memory_ratio=ratio), + note="less headroom outside the budget") + for ratio in [r for r in (0.95, 0.9, 0.85, 0.8) if r < current - 1e-9][:2]: + yield Change(f"--memory-ratio {ratio}", "memory_ratio", 1, _set(memory_ratio=ratio), + note="smaller KV cache, more headroom") + + +def _tips_prefill(inp: ForecastInputs) -> Iterable[Change]: + config = inp.config + current = getattr(config, "max_extend_tokens", None) + if current is None or getattr(config.model_config, "dsv4_args", None) is not None: + return + for length, cost in ((4096, 1), (2048, 2)): + if length < current: + yield Change(f"--max-prefill-length {length}", "max_extend_tokens", cost, + _set(max_extend_tokens=length), note="longer prompts prefill in more chunks") + + +def _tips_cache_type(inp: ForecastInputs) -> Iterable[Change]: + if inp.config.cache_type == "hybrid_radix": + yield Change("--cache-type naive", "cache_type", 6, _set(cache_type="naive"), + note="no cross-request prefix reuse") + + +def _tips_encoders(inp: ForecastInputs) -> Iterable[Change]: + vision = inp.weights.gpu.get("vision", 0) + if vision: + yield Change("--text-model-only", "encoders", 4, weight_delta=-vision, note="no image input") + host = inp.weights.vision_gpu_if_host + if inp.config.mm.encoder_weights == "gpu" and host is not None and host < vision: + yield Change("--mm-encoder-weights host", "encoder_weights", 1, weight_delta=host - vision, + note="vision blocks stream from pinned host RAM") + + +def _tips_moe(inp: ForecastInputs) -> Iterable[Change]: + config = inp.config + mc = config.model_config + if not getattr(mc, "is_moe", False): + return + total = mc.num_moe_layers * mc.num_experts + if config.moe_strategy == "fused": + experts = inp.weights.gpu.get("experts", 0) + if experts and total: + yield Change("--moe-strategy offload", "moe_strategy", 3, + _set(moe_strategy="offload", moe_cache_auto=True, moe_cache_size=0), + weight_delta=-experts, per_expert=experts // total, + note="experts stream from pinned host RAM into a GPU slot cache") + return + if not _is_offload(config): + return + if not config.moe_cache_auto: + yield Change("--moe-cache-auto", "moe_cache", 1, _set(moe_cache_auto=True)) + half = max(2 * mc.num_experts, config.moe_cache_size // 2) + if half < config.moe_cache_size: + yield Change(f"--moe-cache-size {half}", "moe_cache", 2, _set(moe_cache_size=half)) + else: + tokens = max(2 * config.kv_reserve_tokens, 32768) + yield Change(f"--kv-reserve-tokens {tokens}", "kv_reserve", 1, _set(kv_reserve_tokens=tokens), + note="fewer expert slots, more KV") + if config.moe_strategy != "cpu": + yield Change("--moe-strategy cpu", "moe_strategy", 5, + _set(moe_strategy="cpu", moe_cache_auto=False, moe_cache_size=2 * mc.num_experts, + moe_prefill_overlap=True), + note="experts decode on the CPU") + + +# A new sizing option (e.g. a KV-cache dtype flag) adds one generator here; its bytes-per-token +# effect reaches the forecast through the pool family's kv_cost, so nothing else changes. +TIP_CANDIDATES: list[Callable[[ForecastInputs], Iterable[Change]]] = [ + _tips_concurrency, _tips_memory_ratio, _tips_prefill, _tips_cache_type, _tips_encoders, _tips_moe, +] + + +@dataclass +class Tip: + flags: list[str] + forecast: Forecast + effect: str + notes: list[str] = field(default_factory=list) + + def to_dict(self) -> dict: + return {"flags": self.flags, "verdict": self.forecast.verdict, "kv_tokens": self.forecast.kv_tokens, + "max_context": self.forecast.max_context, "max_running_req": self.forecast.max_running_req, + "effect": self.effect, "notes": self.notes} + + +def _evaluate(inp: ForecastInputs, changes: tuple[Change, ...]) -> Forecast: + config = copy.copy(inp.config) + delta, per_expert = 0, None + for c in changes: + c.apply(config) + delta += c.weight_delta + per_expert = c.per_expert if c.per_expert is not None else per_expert + return inp.forecast(config, weight_delta=delta, per_expert=per_expert) + + +def _rank(fc: Forecast) -> tuple: + return VERDICTS.index(fc.verdict) if fc.verdict in VERDICTS else -1, fc.kv_tokens + + +def _improves(base: Forecast, new: Forecast) -> bool: + """A better verdict; at the same verdict, more KV for a fitting config, else more room where + the config is short (prefill headroom, or a context below USEFUL_CONTEXT).""" + if _rank(new)[0] != _rank(base)[0]: + return _rank(new)[0] > _rank(base)[0] + if base.verdict == "fits": + return new.kv_tokens > base.kv_tokens + if base.verdict != "tight": + return False + more_peak = (new.free_at_peak or 0) - (base.free_at_peak or 0) >= 64 * MiB + short_context = base.max_context < min(USEFUL_CONTEXT, base.max_seq_len) + return more_peak or (short_context and new.max_context > base.max_context) + + +def _effect(base: Forecast, new: Forecast, changes: tuple[Change, ...]) -> str: + parts = [] + if new.verdict != base.verdict: + parts.append(f"{base.verdict} -> {new.verdict}") + if new.kv_tokens != base.kv_tokens: + parts.append(f"KV {base.kv_tokens} -> {new.kv_tokens} tokens ({new.kv_tokens - base.kv_tokens:+d})") + delta = sum(c.weight_delta for c in changes) + if delta: + parts.append(f"{'saves' if delta < 0 else 'adds'} {_gib(abs(delta))} of weights") + if new.state_slots != base.state_slots: + parts.append(f"GDN state {base.state_slots} -> {new.state_slots} slots " + f"({_gib(base.state_bytes - new.state_bytes)} freed)") + if new.moe_slots != base.moe_slots: + parts.append(f"expert slots {base.moe_slots} -> {new.moe_slots}") + if new.prefill_activations != base.prefill_activations: + parts.append(f"prefill workspace {_gib(base.prefill_activations)} -> {_gib(new.prefill_activations)}") + if (new.free_at_peak is not None and base.free_at_peak is not None + and abs(new.free_at_peak - base.free_at_peak) >= 64 * MiB and not parts[1:]): + parts.append(f"free at prefill peak {_gib(base.free_at_peak)} -> {_gib(new.free_at_peak)}") + return ", ".join(parts) + + +def suggest_tips(inp: ForecastInputs, base: Forecast, max_combo: int = 3) -> tuple[list[Tip], Tip | None]: + """Single-flag tips that improve the forecast, and the cheapest combination (fewest given up) + that reaches "fits" -- or the best reachable one when nothing fits.""" + if inp.free_before is None: + return [], None + changes = [c for gen in TIP_CANDIDATES for c in gen(inp)] + tips = [] + for c in changes: + fc = _evaluate(inp, (c,)) + if _improves(base, fc): + tips.append((c, fc)) + tips.sort(key=lambda t: (-_rank(t[1])[0], t[0].cost, -t[1].kv_tokens)) + single = [Tip([c.flag], fc, _effect(base, fc, (c,)), [c.note] if c.note else []) for c, fc in tips] + if base.verdict == "fits": + return single, None + best = None + for k in range(1, max_combo + 1): + for combo in itertools.combinations(changes, k): + if len({c.key for c in combo}) < k: + continue + fc = _evaluate(inp, combo) + cost = sum(c.cost for c in combo) + key = (fc.verdict == "fits", -cost if fc.verdict == "fits" else 0, _rank(fc), -cost) + if best is None or key > best[0]: + best = (key, combo, fc) + if best is None or _rank(best[2])[0] <= _rank(base)[0]: + return single, None + _, combo, fc = best + return single, Tip([c.flag for c in combo], fc, _effect(base, fc, combo), [c.note for c in combo if c.note]) + + +# --------------------------------------------------------------------------------------------- +# Shared entry points: the inputs from a meta-built model, and the engine's preflight. +# --------------------------------------------------------------------------------------------- + + +def forecast_inputs(config, model, free_before: int | None) -> ForecastInputs: + """Inputs from the meta-device model of a resolved config, as Engine.__init__ builds it.""" + weights = resident_weights(model, config) + per_expert, max_slots = 0, None + if _is_offload(config): + from freetoken.engine.engine import shared_offload_method + from freetoken.moe.expert_banks import bank_bytes_per_expert + + method = shared_offload_method(model) + per_expert = bank_bytes_per_expert(config.model_config, method) or 0 + max_slots = method.slot_limit() if method is not None else None + mc = config.model_config + if per_expert: + weights.host["experts"] = per_expert * mc.num_moe_layers * mc.num_experts + else: + weights.notes.append("expert bank layout unknown for this format: the GPU slot cache is not priced") + return ForecastInputs( + config=config, weights=weights, free_before=free_before, per_expert_bytes=per_expert, + max_slots=max_slots, activation_per_token=activation_bytes_per_token(model, config), + ) + + +class PreflightRefused(RuntimeError): + """The forecast proves the configuration cannot start; raised before any weight is read.""" + + +def run_preflight(config, model, free_before: int) -> Forecast | None: + """Engine.__init__'s pre-load check: refuse a configuration whose startup solve cannot succeed, + warn on a tight one. Any failure of the forecast itself only logs: it must never be stricter + than the engine.""" + try: + inp = forecast_inputs(config, model, free_before) + fc = inp.forecast() + tips, combo = suggest_tips(inp, fc) if fc.verdict != "fits" else ([], None) + except Exception as exc: # noqa: BLE001 -- the forecast is advisory + logger.warning_rank0(f"Memory preflight skipped: {type(exc).__name__}: {exc}") + return None + if fc.verdict == "fits": + logger.info_rank0( + f"Memory preflight: fits -- {fc.kv_tokens} KV tokens, max context {fc.max_context}, " + f"{_requests(fc.max_running_req)} (forecast; see `ft info`)" + ) + return fc + text = format_forecast(inp, fc, tips, combo) + if fc.verdict == "does not fit": + logger.critical_rank0(text) + raise PreflightRefused( + f"this configuration cannot fit the GPU ({'; '.join(fc.reasons)})" + + (f"; try {' '.join(combo.flags)}" if combo else "") + + ". Pass --skip-preflight to load anyway." + ) + logger.warning_rank0(text) + return fc + + +# --------------------------------------------------------------------------------------------- +# Formatting +# --------------------------------------------------------------------------------------------- + + +def _gib(n: int | None) -> str: + if n is None: + return "?" + return f"{n / GiB:.2f} GiB" + + +def _requests(n: int) -> str: + return f"{n} concurrent request{'' if n == 1 else 's'}" + + +def _row(label: str, value: str, note: str = "", indent: int = 0) -> str: + return f" {label:<{44 + indent}} {value:>12}" + (f" {note}" if note else "") + + +def format_forecast(inp: ForecastInputs, fc: Forecast, tips: list[Tip], combo: Tip | None, + checkpoint: dict[str, int] | None = None) -> str: + config = inp.config + w = inp.weights + lines = [] + head = f"{'Weights on the GPU':<46}" + lines.append(head + (f"{'checkpoint':>12} {'GPU':>12}" if checkpoint is not None else f"{'GPU':>12}")) + cats = [c for c in CATEGORY_LABELS if w.gpu.get(c) or w.host.get(c) or (checkpoint or {}).get(c)] + for cat in cats: + g, h = w.gpu.get(cat, 0), w.host.get(cat, 0) + note = "" + if h: + note = f"+ {_gib(h)} in host RAM" + elif checkpoint is not None and checkpoint.get(cat) and not g: + note = "not loaded" + if cat == "vision" and not config.active_encoders: + note = "skipped (--text-model-only / --mm-disable)" + elif checkpoint is not None and checkpoint.get(cat) and g and abs(g - checkpoint[cat]) > 0.05 * checkpoint[cat]: + note = "smaller than stored (tied, skipped or re-packed)" if g < checkpoint[cat] else "upcast at load" + ck = f"{_gib(checkpoint.get(cat, 0)):>12} " if checkpoint is not None else "" + lines.append(f" {CATEGORY_LABELS[cat]:<44}{ck}{_gib(g):>12}" + (f" {note}" if note else "")) + pad = 13 if checkpoint is not None else 0 + lines.append(_row("load overhead (estimate)", _gib(fc.load_overhead), "kernel modules, library handles", pad)) + lines.append(_row("total", _gib(fc.weights_bytes + fc.load_overhead), "", pad)) + lines.append("") + if fc.free_before is None: + lines.append("GPU free memory unknown: pass --gpu-memory-gib or --gpu-free-gib for a verdict.") + lines.append(_row("KV cache per token", f"{fc.kv_bytes_per_token} B")) + if fc.state_slots: + lines.append(_row(f"GDN state ({fc.state_slots} slots)", _gib(fc.state_bytes))) + return "\n".join(lines) + lines.append(f"Budget: memory_ratio {fc.memory_ratio:g} x {_gib(fc.free_before)} free = {_gib(fc.budget)}") + lines.append(_row("weights + load overhead", _gib(fc.weights_bytes + fc.load_overhead))) + if fc.moe_slots or fc.per_expert_bytes: + lines.append(_row(f"expert slot cache ({fc.moe_slots} x {fc.per_expert_bytes / MiB:.1f} MiB)", + _gib(fc.moe_slot_bytes))) + if fc.state_slots: + lines.append(_row(f"GDN state pool ({fc.state_slots} slots x {fc.state_bytes_per_slot / MiB:.1f} MiB)", + _gib(fc.state_bytes))) + kv_label = f"KV cache ({fc.kv_bytes_per_token} B/token x {fc.kv_tokens} tokens)" + lines.append(_row(kv_label, _gib(fc.kv_bytes))) + if fc.verdict != "does not fit": + lines.append("Outside the budget") + lines.append(_row(f"attention workspace ({config.attention_backend})", _gib(fc.attention_workspace))) + lines.append(_row("page table", _gib(fc.page_table))) + lines.append(_row("free after init (engine log line)", _gib(fc.free_after_init))) + bs = ",".join(map(str, fc.graph_batch_sizes)) or "off" + lines.append(_row(f"CUDA graphs (bs {bs}, estimate)", _gib(fc.cuda_graphs))) + lines.append(_row(f"prefill activations ({fc.prefill_chunk} tokens, estimate)", _gib(fc.prefill_activations))) + lines.append(_row("free at prefill peak", _gib(fc.free_at_peak))) + lines.append("") + verdict = fc.verdict.upper() + if fc.verdict != "does not fit": + verdict += (f": max context {fc.max_context} tokens, {_requests(fc.max_running_req)}" + f" ({fc.kv_tokens // max(1, fc.max_context)} at full context)") + lines.append(f"Verdict: {verdict}") + lines.extend(f" - {r}" for r in fc.reasons) + if tips or combo: + lines.append("") + lines.append("Tips") + for tip in tips: + note = f" [{'; '.join(tip.notes)}]" if tip.notes else "" + lines.append(f" {' '.join(tip.flags):<32} {tip.effect}{note}") + if combo: + fcc = combo.forecast + lines.append(f" Suggested: {' '.join(combo.flags)} -> {fcc.verdict}, {fcc.kv_tokens} KV tokens, " + f"max context {fcc.max_context}, {_requests(fcc.max_running_req)}") + return "\n".join(lines) diff --git a/python/freetoken/gpu_select.py b/python/freetoken/gpu_select.py index 2c7d012b23..1d5e382f7c 100644 --- a/python/freetoken/gpu_select.py +++ b/python/freetoken/gpu_select.py @@ -68,13 +68,8 @@ def single_gpu_arg(value: str) -> str: return entries[0] -def _nvml_uuids() -> "list[str] | None": - """Full GPU UUIDs in physical (nvidia-smi) order, or None when NVML is unavailable. - - Own ctypes loader instead of torch's _raw_device_uuid_nvml: that helper only knows the Linux library name, raises (not None) when the library is missing, and is private API. - NVML exports are cdecl on every platform, so CDLL is right on Windows too (same as nvidia-ml-py). - None on any failure -- no library, a stub library without the _v2 symbols, WSL, a dead device -- and callers fall back. - """ +def _load_nvml(): + """The NVML library handle, or None when it cannot be loaded.""" import ctypes if os.name == "nt": @@ -85,16 +80,58 @@ def _nvml_uuids() -> "list[str] | None": ] else: candidates = ["libnvidia-ml.so.1"] + for name in candidates: + try: + return ctypes.CDLL(name) + except OSError: + continue + return None + + +def nvml_memory_info(spec: str = "0") -> "tuple[str, int, int] | None": + """(name, total bytes, free bytes) of the GPU ``spec`` names (nvidia-smi index or UUID) without + initializing CUDA, or None when NVML is unavailable or the GPU is not found.""" + import ctypes + + class _Memory(ctypes.Structure): + _fields_ = [("total", ctypes.c_ulonglong), ("free", ctypes.c_ulonglong), ("used", ctypes.c_ulonglong)] + try: - for name in candidates: - try: - lib = ctypes.CDLL(name) - break - except OSError: - continue - else: + lib = _load_nvml() + if lib is None or lib.nvmlInit() != 0: return None - if lib.nvmlInit() != 0: + try: + handle = ctypes.c_void_p() + if is_gpu_uuid(spec): + rc = lib.nvmlDeviceGetHandleByUUID(spec.encode("ascii"), ctypes.byref(handle)) + else: + rc = lib.nvmlDeviceGetHandleByIndex_v2(int(spec), ctypes.byref(handle)) + if rc != 0: + return None + mem = _Memory() + if lib.nvmlDeviceGetMemoryInfo(handle, ctypes.byref(mem)) != 0: + return None + buf = ctypes.create_string_buffer(96) + name = buf.value.decode("utf-8", "replace") if lib.nvmlDeviceGetName(handle, buf, 96) == 0 else "GPU" + return name, int(mem.total), int(mem.free) + finally: + lib.nvmlShutdown() + except (OSError, AttributeError, ValueError): + return None + + +def _nvml_uuids() -> "list[str] | None": + """Full GPU UUIDs in physical (nvidia-smi) order, or None when NVML is unavailable. + + Own ctypes loader instead of torch's _raw_device_uuid_nvml: that helper only knows the Linux library name, raises (not None) when the library is missing, and is private API. + NVML exports are cdecl on every platform, so CDLL is right on Windows too (same as nvidia-ml-py). + None on any failure -- no library, a stub library without the _v2 symbols, WSL, a dead device -- and callers fall back. + """ + import ctypes + + try: + lib = _load_nvml() + if lib is None or lib.nvmlInit() != 0: return None try: count = ctypes.c_int() diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py index 95669e8c81..40a37a8333 100644 --- a/python/freetoken/kvcache/base.py +++ b/python/freetoken/kvcache/base.py @@ -61,15 +61,23 @@ def kv_cost(cls, config, **kwargs) -> tuple[int, int, int, int]: ``num_swa_pages=`` (a rebuild's target window, priced before the config override is written).""" + @classmethod + def plan_num_pages(cls, config, available_memory: int) -> int: + """The page count ``solve_num_pages`` would pick, without its checks or log line: may + be <= 1 when the budget is too small. The pre-load memory forecast calls this.""" + cache_per_page, fixed_cache_size, _, _ = cls.kv_cost(config) + num_pages = config.num_page_override + if num_pages is None: + num_pages = (available_memory - fixed_cache_size) // cache_per_page + return num_pages + @classmethod def solve_num_pages(cls, config, available_memory: int) -> int: """Largest usable page count fitting ``available_memory`` (bytes measured by the engine: memory_ratio x baseline free minus resident weights and any sibling pool's fixed cost), honoring ``config.num_page_override``.""" cache_per_page, fixed_cache_size, _, _ = cls.kv_cost(config) - num_pages = config.num_page_override - if num_pages is None: - num_pages = (available_memory - fixed_cache_size) // cache_per_page + num_pages = cls.plan_num_pages(config, available_memory) assert num_pages > 1, "Not enough memory for KV cache, try reducing --num-pages" real_kv_size = num_pages * cache_per_page + fixed_cache_size logger.info( diff --git a/python/freetoken/kvcache/dsv4_paged_pool.py b/python/freetoken/kvcache/dsv4_paged_pool.py index 3ce23eea7d..1be717d12a 100644 --- a/python/freetoken/kvcache/dsv4_paged_pool.py +++ b/python/freetoken/kvcache/dsv4_paged_pool.py @@ -327,16 +327,14 @@ def kv_cost(cls, config) -> tuple[int, int, int, int]: return per_page, fixed, config.page_size, min_reserve_tokens @classmethod - def solve_num_pages(cls, config, available_memory: int) -> int: + def _plan_sizes(cls, config, available_memory: int): # Solve the largest budget-respecting anchor with the exact per-tier byte model. # num_pages is in P (window) units and anchors full_token = num_pages*P (the FULL # cmp/idx tiers). The window working-set floor is honored in PAGES and the total is # byte-checked here. A budget too small for even the minimal working set raises a # graceful config error, not a late OOM. - from freetoken.utils import mem_GB - from .dsv4_cost_model import _dsv4_pool_sizes, _dsv4_swa_ratio, _dsv4_window_floor_pages - from .dsv4_cost_model import dsv4_pool_bytes, dsv4_solve_num_pages + from .dsv4_cost_model import dsv4_solve_num_pages dsv4_args = config.model_config.dsv4_args P = dsv4_args.window_size @@ -361,6 +359,21 @@ def solve_num_pages(cls, config, available_memory: int) -> int: f"--num-pages or lower max_running_req/max_seq_len" ) sizes = _dsv4_pool_sizes(config, num_pages + 1) # +1 for dummy page + return num_pages, sizes + + @classmethod + def plan_num_pages(cls, config, available_memory: int) -> int: + return cls._plan_sizes(config, available_memory)[0] + + @classmethod + def solve_num_pages(cls, config, available_memory: int) -> int: + from freetoken.utils import mem_GB + + from .dsv4_cost_model import dsv4_pool_bytes + + num_pages, sizes = cls._plan_sizes(config, available_memory) + dsv4_args = config.model_config.dsv4_args + P = dsv4_args.window_size assert num_pages > 1, "Not enough memory for KV cache, try reducing --num-pages" real = dsv4_pool_bytes(sizes, dsv4_args, config.max_running_req + 1) logger.info( diff --git a/python/freetoken/models/loader.py b/python/freetoken/models/loader.py index 0f310ccbb2..3d595583d6 100644 --- a/python/freetoken/models/loader.py +++ b/python/freetoken/models/loader.py @@ -54,6 +54,13 @@ def iter_weight_files(model_path: str) -> list[str]: return [f for f in files if not f.endswith("consolidated.safetensors")] or files +def read_safetensors_header(path: str) -> dict: + """A shard's JSON header (8-byte little-endian length, then JSON); no tensor data is read.""" + with open(path, "rb") as fh: + n = struct.unpack(" dict[str, str]: """Tensor name -> shard basename, from the index or from each shard's header when the checkpoint ships none.""" index = os.path.join(folder, "model.safetensors.index.json") @@ -62,9 +69,7 @@ def safetensors_weight_map(folder: str) -> dict[str, str]: return json.load(f)["weight_map"] weight_map: dict[str, str] = {} for path in sorted(iter_weight_files(folder)): - with open(path, "rb") as fh: - n = struct.unpack(" int: return self.staging.numel() + @staticmethod + def staging_bytes(blocks: Sequence[BaseOP]) -> int: + """``device_bytes`` a streamer over ``blocks`` would hold, from the (possibly meta) tensors alone.""" + return 2 * _slots(blocks[0])[1] + def _bind_all_to_host(self) -> None: # between forwards the attributes point at the pinned bank, so state_dict() stays complete and correct for b, (slots, _) in enumerate(self._layouts): diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index e3a768b6d0..6cb2808433 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -252,24 +252,33 @@ def bank_bytes_estimate(model_config, method=None) -> int | None: format-tag table sizes the GGUF format. ``None`` for unknown formats or missing dims (callers then skip the pre-load sizing).""" layers = getattr(model_config, "num_moe_layers", None) + per_expert = bank_bytes_per_expert(model_config, method) if method is not None and layers: - per_expert = sum( + return layers * method.cfg.num_experts * per_expert + experts = getattr(model_config, "num_experts", None) + if per_expert is None or not all((layers, experts)): + return None + return layers * experts * per_expert + + +def bank_bytes_per_expert(model_config, method=None) -> int | None: + """Host bank bytes of one expert, which is also what one GPU slot of the offload cache holds + (``cache_budget.expert_bytes_per_slot`` over the loaded banks). ``None`` when unknown.""" + if method is not None: + return sum( math.prod(spec.shape) * torch.empty((), dtype=spec.dtype).element_size() for spec in method.layout().values() if not spec.resident ) - return layers * method.cfg.num_experts * per_expert expert_quant = getattr(model_config, "expert_quant", "none") fmt = expert_quant if expert_quant != "none" else ( getattr(model_config, "moe_weight_format", None) or "bf16" ) per_expert = _BANK_BYTES_PER_EXPERT.get(fmt) - layers = getattr(model_config, "num_moe_layers", None) - experts = getattr(model_config, "num_experts", None) hidden = getattr(model_config, "hidden_size", None) inter = getattr(model_config, "moe_intermediate_size", None) - if per_expert is None or not all((layers, experts, hidden, inter)): + if per_expert is None or not all((hidden, inter)): return None - return layers * experts * per_expert(hidden, inter) + return per_expert(hidden, inter) def load_expert_banks( diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 23a62f8419..8d97fd1388 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -333,6 +333,14 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--skip-preflight", + action="store_true", + default=ServerArgs.skip_preflight, + help="Load the weights even when the pre-load memory forecast (see `ft info`) says this " + "configuration cannot fit the GPU.", + ) + assert ServerArgs.use_dummy_weight == False parser.add_argument( "--dummy-weight", diff --git a/python/freetoken/server/info.py b/python/freetoken/server/info.py new file mode 100644 index 0000000000..390e8dbb4e --- /dev/null +++ b/python/freetoken/server/info.py @@ -0,0 +1,344 @@ +"""ft info: forecast a model's GPU memory under a serve configuration without loading its weights.""" + +from __future__ import annotations + +import argparse +import glob +import json +import os +import sys +from collections import Counter +from dataclasses import dataclass, field +from typing import Any, Sequence + +from freetoken.engine.forecast import CATEGORY_LABELS + +GiB = 1 << 30 + +USAGE = """usage: {prog} [ft serve flags] [--json] [--gpu-memory-gib N | --gpu-free-gib N] [--verbose] + +Forecast what `ft serve [flags]` puts on the GPU -- weights by category, KV cache, +fixed pools, workspaces -- and whether it fits, without reading any weight. Accepts every +`ft serve` flag, so "what if" combinations can be tried before a long load. + +Options of ft info itself: + --json print the forecast as JSON + --gpu-memory-gib N plan for an empty card with N GiB of memory instead of the local GPU + --gpu-free-gib N plan for N GiB free (as nvidia-smi reports it) instead of the local GPU + --verbose keep the engine's INFO logs + +The `ft serve` flags follow:""" + + +def _info_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--json", action="store_true") + parser.add_argument("--verbose", action="store_true") + group = parser.add_mutually_exclusive_group() + group.add_argument("--gpu-memory-gib", type=float, default=None) + group.add_argument("--gpu-free-gib", type=float, default=None) + return parser + + +def split_argv(argv: Sequence[str]) -> tuple[argparse.Namespace, list[str]]: + """(ft info's own options, the argv for the serve parser with the model as --model).""" + opts, rest = _info_parser().parse_known_args(list(argv)) + if rest and not rest[0].startswith("-"): + rest = ["--model", rest[0], *rest[1:]] + return opts, rest + + +def checkpoint_tensors(model_path: str) -> tuple[dict[str, tuple[str, int]], str] | None: + """Tensor name -> (dtype, bytes) from the safetensors headers alone, and where they came from: + a local directory or the HF cache, else the Hub through HTTP range requests (headers only).""" + from freetoken.models.loader import read_safetensors_header + + folder = model_path if os.path.isdir(model_path) else None + if folder is None: + try: + from huggingface_hub import snapshot_download + + folder = snapshot_download(model_path, local_files_only=True) + except Exception: # noqa: BLE001 -- not cached: fall back to the Hub below + folder = None + if folder is not None: + index = os.path.join(folder, "model.safetensors.index.json") + if os.path.isfile(index): + with open(index, encoding="utf-8") as f: + shards = [os.path.join(folder, s) for s in sorted(set(json.load(f)["weight_map"].values()))] + else: + shards = [p for p in sorted(glob.glob(os.path.join(folder, "*.safetensors"))) + if not p.endswith("consolidated.safetensors")] + if shards and all(os.path.isfile(s) for s in shards): + out: dict[str, tuple[str, int]] = {} + for shard in shards: + for name, meta in read_safetensors_header(shard).items(): + if name != "__metadata__": + begin, end = meta["data_offsets"] + out[name] = (meta["dtype"], end - begin) + return out, "local safetensors headers" + if os.path.isdir(model_path): + return None + try: + from huggingface_hub import get_safetensors_metadata + + meta = get_safetensors_metadata(model_path) + except Exception: # noqa: BLE001 -- offline or not a safetensors repo: no checkpoint view + return None + out = {} + for file_meta in meta.files_metadata.values(): + for name, info in file_meta.tensors.items(): + begin, end = info.data_offsets + out[name] = (info.dtype, end - begin) + return out, "Hub safetensors headers (HTTP range requests)" + + +@dataclass +class GpuInfo: + free_before: int | None + source: str + name: str | None = None + total: int | None = None + free_now: int | None = None + + +def gpu_info(opts: argparse.Namespace, config) -> GpuInfo: + """Free memory the engine would measure before the load: the card's free memory minus the + CUDA context the engine process creates first. NVML only -- no CUDA context here.""" + from freetoken.engine.forecast import CUDA_CONTEXT_BYTES + from freetoken.gpu_select import nvml_memory_info + + if opts.gpu_memory_gib is not None: + total = int(opts.gpu_memory_gib * GiB) + return GpuInfo(total - CUDA_CONTEXT_BYTES, f"--gpu-memory-gib {opts.gpu_memory_gib:g} (empty card)", + total=total, free_now=total) + if opts.gpu_free_gib is not None: + free = int(opts.gpu_free_gib * GiB) + return GpuInfo(free - CUDA_CONTEXT_BYTES, f"--gpu-free-gib {opts.gpu_free_gib:g}", free_now=free) + spec = config.gpu[0] if config.gpu else None + if spec is None: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible is not None: + entries = [e.strip() for e in visible.split(",") if e.strip()] + if not entries: + return GpuInfo(None, "no GPU visible (CUDA_VISIBLE_DEVICES is empty)") + spec = entries[0] + info = nvml_memory_info(spec or "0") + if info is None: + return GpuInfo(None, "no GPU found (NVML unavailable)") + name, total, free = info + return GpuInfo(free - CUDA_CONTEXT_BYTES, "NVML", name=name, total=total, free_now=free) + + +@dataclass +class InfoReport: + config: Any + gpu: GpuInfo + inputs: Any + forecast: Any + tips: list + combo: Any + checkpoint: dict[str, int] | None = None + checkpoint_dtypes: dict[str, int] | None = None + checkpoint_source: str | None = None + quant: dict[str, list[str]] = field(default_factory=dict) + notes: list[str] = field(default_factory=list) + + +def _quant_summary(report: "InfoReport") -> dict[str, str]: + return {CATEGORY_LABELS.get(cat, cat): ", ".join(kinds) for cat, kinds in report.quant.items()} + + +def _layout_summary(mc) -> str: + kinds = sorted({spec.attn_type.value for spec in mc.kv_cache_group_specs() if spec.num_layers > 0}) + if mc.linear_attention_group() is not None: + kinds.append(f"linear ({mc.linear_attention_group().variant}, {len(mc.linear_attention_group().layer_ids)} layers)") + text = f"{mc.num_layers} layers, hidden {mc.hidden_size}, attention: {', '.join(kinds)}" + if getattr(mc, "is_moe", False): + text += f", MoE {mc.num_experts} experts top-{mc.num_experts_per_tok} on {mc.num_moe_layers} layers" + return text + + +def analyze(config, opts: argparse.Namespace) -> InfoReport: + """Resolve ``config`` as Engine.__init__ does, build the model on the meta device and + forecast; reads only config files and safetensors headers.""" + import torch + + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.engine.engine import _adjust_config, _adjust_ftw_quant_backend + from freetoken.engine.forecast import forecast_inputs, quant_by_category, suggest_tips, tensor_category + from freetoken.layers import set_rope_device + from freetoken.layers.quantization import QuantBackend, set_quant_backend + from freetoken.models import create_model + from freetoken.utils import torch_dtype + + tp = try_get_tp_info() + if tp is None: + set_tp_info(rank=0, size=config.tp_info.size) + elif tp.size != config.tp_info.size: + raise RuntimeError(f"TP size {tp.size} is already set in this process, not {config.tp_info.size}") + # Read free memory first: _adjust_config's arch probe creates this process's CUDA context, + # which the reading would otherwise include on top of the context allowance it subtracts. + gpu = gpu_info(opts, config) + set_quant_backend(_adjust_ftw_quant_backend(config.model_path, QuantBackend.parse(config.quant_backend))) + _adjust_config(config) + set_rope_device(torch.device("cpu")) + found = checkpoint_tensors(config.model_path) + try: + with torch.device("meta"), torch_dtype(config.dtype): + model = create_model(config.model_config) + except Exception as exc: # noqa: BLE001 -- fall back to the checkpoint headers, flagged as such + if found is None: + raise + model, meta_error = None, exc + if model is not None: + inputs = forecast_inputs(config, model, gpu.free_before) + quant = quant_by_category(model) + else: + inputs = header_inputs(config, found[0], gpu.free_before) + inputs.weights.notes.append( + f"The model could not be built on the meta device ({type(meta_error).__name__}: {meta_error}); " + "GPU weights are the checkpoint's stored bytes, so re-quantized or upcast tensors are not priced." + ) + quant = {} + fc = inputs.forecast() + tips, combo = suggest_tips(inputs, fc) + report = InfoReport(config, gpu, inputs, fc, tips, combo, quant=quant) + if found is not None: + tensors, report.checkpoint_source = found + by_cat: Counter[str] = Counter() + by_dtype: Counter[str] = Counter() + for name, (dtype, nbytes) in tensors.items(): + by_cat[tensor_category(name)] += nbytes + by_dtype[dtype] += nbytes + report.checkpoint, report.checkpoint_dtypes = dict(by_cat), dict(by_dtype) + report.notes = inputs.weights.notes + [ + "Kernels that repack weights in finalize_quant are assumed to keep the byte count.", + "Load overhead, prefill activations and CUDA graph memory are estimates; the verdict only " + "refuses when the engine's own sizing fails without them.", + ] + if gpu.free_before is not None and gpu.source == "NVML": + report.notes.append("Free memory is read now; other processes on the GPU change it.") + return report + + +def header_inputs(config, tensors: dict[str, tuple[str, int]], free_before: int | None): + """Forecast inputs from the safetensors headers alone, for a family the meta build cannot run: + the tensors the engine would not put on the GPU (an inactive vision tower, MTP heads, + offloaded experts) are dropped by category.""" + from freetoken.engine.forecast import ForecastInputs, WeightReport, _is_offload, tensor_category + + mc = config.model_config + gpu: Counter[str] = Counter() + host: Counter[str] = Counter() + offload = _is_offload(config) + for name, (_, nbytes) in tensors.items(): + cat = tensor_category(name) + if cat == "mtp" or (cat == "vision" and not config.active_encoders): + continue + (host if cat == "experts" and offload else gpu)[cat] += nbytes + total_experts = getattr(mc, "num_moe_layers", 0) * getattr(mc, "num_experts", 0) + per_expert = host["experts"] // total_experts if offload and total_experts else 0 + item = config.dtype.itemsize + return ForecastInputs( + config=config, weights=WeightReport(gpu=dict(gpu), host=dict(host)), free_before=free_before, + per_expert_bytes=per_expert, + activation_per_token=item * (4 * mc.hidden_size + 3 * mc.intermediate_size), + ) + + +def report_dict(r: InfoReport) -> dict: + config, mc = r.config, r.config.model_config + fc = r.forecast + return { + "model": config.model_path, + "architecture": config.hf_config.architectures[0], + "model_type": mc.model_type, + "layout": _layout_summary(mc), + "quantization": _quant_summary(r), + "resolved": { + "attention_backend": config.attention_backend, + "cache_type": config.cache_type, + "page_size": config.page_size, + "moe_strategy": config.moe_strategy if getattr(mc, "is_moe", False) else None, + "max_running_req": config.max_running_req, + "memory_ratio": config.memory_ratio, + "max_prefill_length": getattr(config, "max_extend_tokens", None), + "max_seq_len": config.max_seq_len, + "dtype": str(config.dtype).removeprefix("torch."), + }, + "gpu": vars(r.gpu), + "checkpoint": None if r.checkpoint is None else { + "source": r.checkpoint_source, "bytes_by_category": r.checkpoint, "bytes_by_dtype": r.checkpoint_dtypes, + }, + "weights": {"gpu": r.inputs.weights.gpu, "host": r.inputs.weights.host, "load_overhead": fc.load_overhead}, + "forecast": fc.to_dict(), + "tips": [t.to_dict() for t in r.tips], + "suggested": r.combo.to_dict() if r.combo else None, + "notes": r.notes, + } + + +def format_report(r: InfoReport) -> str: + from freetoken.engine.forecast import _gib, format_forecast + + config, mc = r.config, r.config.model_config + quant = "; ".join(f"{k}: {v}" for k, v in _quant_summary(r).items()) + lines = [ + f"Model {config.model_path}", + f"Architecture {config.hf_config.architectures[0]} ({mc.model_type}): {_layout_summary(mc)}", + f"Quantization {quant}", + ] + if r.checkpoint_dtypes: + dtypes = ", ".join(f"{k} {_gib(v)}" for k, v in sorted(r.checkpoint_dtypes.items(), key=lambda kv: -kv[1])) + lines.append(f"Checkpoint {dtypes} ({r.checkpoint_source})") + resolved = (f"attention {config.attention_backend}, cache {config.cache_type}, page size {config.page_size}, " + f"max_running_req {config.max_running_req}, memory_ratio {config.memory_ratio:g}, " + f"max prefill {getattr(config, 'max_extend_tokens', '-')}, max_seq_len {config.max_seq_len}") + if getattr(mc, "is_moe", False): + resolved = f"moe {config.moe_strategy}, " + resolved + lines.append(f"Resolved {resolved}") + g = r.gpu + if g.free_before is None: + lines.append(f"GPU {g.source}") + else: + card = f"{g.name}: " if g.name else "" + total = f"{_gib(g.total)} total, " if g.total else "" + lines.append(f"GPU {card}{total}{_gib(g.free_now)} free ({g.source}) -> " + f"{_gib(g.free_before)} free before loading (after the CUDA context)") + lines.append("") + lines.append(format_forecast(r.inputs, r.forecast, r.tips, r.combo, r.checkpoint)) + lines.append("") + lines.extend(f"Note: {n}" for n in r.notes) + return "\n".join(lines) + + +def main(argv: Sequence[str] | None = None, prog: str = "ft info") -> int: + argv = list(sys.argv[1:] if argv is None else argv) + if not argv or argv[0] in ("-h", "--help"): + print(USAGE.format(prog=prog), file=sys.stdout if argv else sys.stderr) + if not argv: + return 2 + from freetoken.server.args import parse_args + + try: + parse_args(["--help"], prog=prog) + except SystemExit: + pass + return 0 + opts, serve_argv = split_argv(argv) + if not opts.verbose: + os.environ.setdefault("LOG_LEVEL", "WARNING") + from freetoken.server.args import parse_args + + config, _ = parse_args(serve_argv, prog=prog) + try: + report = analyze(config, opts) + except (ValueError, RuntimeError, NotImplementedError) as exc: + print(f"{prog}: error: {type(exc).__name__}: {exc}", file=sys.stderr) + return 2 + if opts.json: + print(json.dumps(report_dict(report), indent=2, default=str)) + else: + print(format_report(report)) + return 1 if report.forecast.verdict == "does not fit" else 0 diff --git a/tests/engine/test_forecast.py b/tests/engine/test_forecast.py new file mode 100644 index 0000000000..8bbb6ca410 --- /dev/null +++ b/tests/engine/test_forecast.py @@ -0,0 +1,323 @@ +"""Pre-load memory forecast (engine/forecast.py): CPU-only, on synthetic header-only checkpoints. + +The forecast must reuse the engine's own sizing functions, so each check prices the same +configuration through those functions directly and compares.""" + +import json +import struct +import time + +import pytest +import torch + +from freetoken.engine import forecast as fc_mod +from freetoken.engine.forecast import GiB, MiB, PreflightRefused, tensor_category + +_DTYPE_BYTES = {"BF16": 2, "F32": 4, "F8_E4M3": 1, "U8": 1} + + +def write_checkpoint(path, config: dict, tensors: dict[str, tuple[str, list[int]]]) -> str: + """config.json plus one safetensors shard holding only its header (no tensor data).""" + path.mkdir(parents=True, exist_ok=True) + (path / "config.json").write_text(json.dumps(config)) + header, offset = {}, 0 + for name, (dtype, shape) in tensors.items(): + nbytes = _DTYPE_BYTES[dtype] + for d in shape: + nbytes *= d + header[name] = {"dtype": dtype, "shape": shape, "data_offsets": [offset, offset + nbytes]} + offset += nbytes + blob = json.dumps(header).encode() + (path / "model.safetensors").write_bytes(struct.pack(" dict[str, tuple[str, list[int]]]: + t = {"model.embed_tokens.weight": ("BF16", [256, 128]), "lm_head.weight": ("BF16", [256, 128]), + "model.norm.weight": ("BF16", [128])} + for i in range(2): + p = f"model.layers.{i}." + t.update({ + p + "self_attn.q_proj.weight": ("BF16", [128, 128]), p + "self_attn.k_proj.weight": ("BF16", [64, 128]), + p + "self_attn.v_proj.weight": ("BF16", [64, 128]), p + "self_attn.o_proj.weight": ("BF16", [128, 128]), + p + "self_attn.q_norm.weight": ("BF16", [64]), p + "self_attn.k_norm.weight": ("BF16", [64]), + p + "mlp.gate_proj.weight": ("BF16", [256, 128]), p + "mlp.up_proj.weight": ("BF16", [256, 128]), + p + "mlp.down_proj.weight": ("BF16", [128, 256]), + p + "input_layernorm.weight": ("BF16", [128]), p + "post_attention_layernorm.weight": ("BF16", [128]), + }) + return t + + +def _analyze(model_path, *flags, free_gib=8.0): + from freetoken.server.info import analyze, split_argv + from freetoken.server.args import parse_args + + opts, argv = split_argv([model_path, "--gpu-free-gib", str(free_gib), *flags]) + config, _ = parse_args(argv, prog="ft info") + return analyze(config, opts) + + +@pytest.fixture +def tiny_qwen3(tmp_path): + return write_checkpoint(tmp_path / "tiny-qwen3", TINY_QWEN3, tiny_qwen3_tensors()) + + +def test_tensor_category_buckets(): + cases = { + "model.layers.3.self_attn.qkv_proj.weight": "attention", + "model.language_model.layers.0.linear_attn.in_proj_qkv.weight": "linear_attention", + "model.layers.1.mlp.experts.7.down_proj.weight": "experts", + "model.layers.1.mlp.gate.weight": "mlp", + "model.layers.1.mlp.shared_expert.up_proj.weight": "mlp", + "model.embed_tokens.weight": "embeddings", + "lm_head.weight_scale": "embeddings", + "model.visual.blocks.0.attn.qkv.weight": "vision", + "visual.merger.linear_fc1.weight": "vision", + "mtp.layers.0.mlp.experts.0.gate_proj.weight": "mtp", + "model.layers.0.input_layernorm.weight": "other", + } + for name, cat in cases.items(): + assert tensor_category(name) == cat, name + + +def test_forecast_matches_engine_sizing(tiny_qwen3): + from freetoken.engine.engine import _startup_kv_budget + from freetoken.kvcache.mha_pool import MHAKVCache + + r = _analyze(tiny_qwen3) + config, f = r.config, r.forecast + # weights: the meta model's parameters (bf16, untied lm_head) + the eager rope table + params = sum(t.numel() * t.element_size() for t in _meta_model(config).state_dict().values()) + assert f.weights_bytes == params + r.inputs.weights.gpu["rope"] + # KV price is the pool family's own kv_cost + cache_per_page, fixed, _, _ = MHAKVCache.kv_cost(config) + assert f.kv_bytes_per_token * config.page_size == cache_per_page == KV_PER_TOKEN * config.page_size + # page count is the engine's startup solve over the forecast resident bytes + free = r.gpu.free_before + resident = f.weights_bytes + f.load_overhead + expected = MHAKVCache.plan_num_pages(config, _startup_kv_budget(config.memory_ratio, free, free - resident)) + assert f.num_pages == expected + assert f.kv_tokens == expected * config.page_size + assert f.max_context == min(config.max_seq_len, f.kv_tokens) + assert f.verdict == "fits", f.reasons + + +def _meta_model(config): + from freetoken.models import create_model + from freetoken.utils import torch_dtype + + with torch.device("meta"), torch_dtype(config.dtype): + return create_model(config.model_config) + + +def test_num_pages_override_is_priced(tiny_qwen3): + r = _analyze(tiny_qwen3, "--num-pages", "300") + assert r.forecast.num_pages == 300 + assert r.forecast.kv_bytes == 300 * r.forecast.kv_bytes_per_token + + +def test_hopeless_config_does_not_fit_and_tips_recover(tiny_qwen3): + # the budget cannot hold the load overhead estimate on top of the weights, but the weights + # alone fit: tight, never refused on an estimate + r = _analyze(tiny_qwen3, free_gib=0.65) + assert r.forecast.verdict == "tight" + assert "load overhead" in r.forecast.reasons[0] + # a --num-pages request larger than the whole card is refused + r = _analyze(tiny_qwen3, "--num-pages", str(10 * GiB // (KV_PER_TOKEN)), free_gib=1.0) + assert r.forecast.verdict == "does not fit" + assert "--num-pages" in r.forecast.reasons[0] + + +def test_tips_price_flag_changes(tiny_qwen3): + r = _analyze(tiny_qwen3, "--memory-ratio", "0.8", free_gib=32) + assert r.forecast.verdict == "fits" + tip = next(t for t in r.tips if t.flags == ["--memory-ratio 0.95"]) + assert tip.forecast.kv_tokens > r.forecast.kv_tokens + assert tip.forecast.memory_ratio == 0.95 + # the base config is left untouched by the what-if evaluation + assert r.config.memory_ratio == 0.8 + + +def test_tight_headroom_suggests_a_lower_memory_ratio(tiny_qwen3): + r = _analyze(tiny_qwen3, "--memory-ratio", "0.97", free_gib=4) + assert r.forecast.verdict == "tight" + lower = [t for t in r.tips if t.flags[0] in ("--memory-ratio 0.95", "--memory-ratio 0.9")] + assert lower and all(t.forecast.free_at_peak > r.forecast.free_at_peak for t in lower) + assert not any(t.flags[0] == "--memory-ratio 0.99" for t in r.tips) + + +# --------------------------------------------------------------------------------------------- +# Preflight (Engine.__init__ calls run_preflight between the meta build and the weight load) +# --------------------------------------------------------------------------------------------- + + +def _resolved(model_path, *flags): + from freetoken.distributed import set_tp_info + from freetoken.engine.engine import _adjust_config + from freetoken.layers import set_rope_device + from freetoken.server.args import parse_args + + from freetoken.distributed import try_get_tp_info + + config, _ = parse_args(["--model", model_path, *flags]) + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + _adjust_config(config) + set_rope_device(torch.device("cpu")) + return config, _meta_model(config) + + +def test_preflight_refuses_hopeless_config_fast(tiny_qwen3): + config, model = _resolved(tiny_qwen3) + start = time.monotonic() + with pytest.raises(PreflightRefused, match="--skip-preflight"): + fc_mod.run_preflight(config, model, free_before=MiB) + assert time.monotonic() - start < 10 + + +def test_preflight_lets_tight_config_through_with_warning(tiny_qwen3, monkeypatch): + config, model = _resolved(tiny_qwen3) + warnings = [] + monkeypatch.setattr(fc_mod.logger, "warning_rank0", lambda msg, *a, **k: warnings.append(msg)) + # weights fit, the load overhead estimate does not: the engine might still start + weights = fc_mod.resident_weights(model, config).gpu_total + fc = fc_mod.run_preflight(config, model, free_before=int((weights + 64 * MiB) / config.memory_ratio)) + assert fc is not None and fc.verdict == "tight" + assert warnings and "TIGHT" in warnings[0] + + +def test_preflight_never_fails_on_its_own_errors(tiny_qwen3, monkeypatch): + config, model = _resolved(tiny_qwen3) + monkeypatch.setattr(fc_mod, "forecast_inputs", lambda *a: 1 / 0) + assert fc_mod.run_preflight(config, model, free_before=GiB) is None + + +# --------------------------------------------------------------------------------------------- +# Hybrid GDN: the sibling state pool and the tips that shrink it +# --------------------------------------------------------------------------------------------- + +TINY_QWEN3_5 = { + "architectures": ["Qwen3_5ForCausalLM"], "model_type": "qwen3_5_text", "hidden_size": 128, + "intermediate_size": 256, "num_hidden_layers": 4, "num_attention_heads": 2, "num_key_value_heads": 1, + "head_dim": 64, "vocab_size": 256, "max_position_embeddings": 32768, "rms_norm_eps": 1e-6, + "hidden_act": "silu", "tie_word_embeddings": False, "dtype": "bfloat16", + "layer_types": ["linear_attention", "linear_attention", "linear_attention", "full_attention"], + "full_attention_interval": 4, "linear_conv_kernel_dim": 4, "linear_key_head_dim": 128, + "linear_num_key_heads": 2, "linear_num_value_heads": 16, "linear_value_head_dim": 128, + "attn_output_gate": True, "partial_rotary_factor": 0.25, + "rope_parameters": {"rope_theta": 10000000, "rope_type": "default", "partial_rotary_factor": 0.25}, +} + + +@pytest.fixture +def tiny_hybrid(tmp_path): + return write_checkpoint(tmp_path / "tiny-qwen3_5", TINY_QWEN3_5, + {"model.embed_tokens.weight": ("BF16", [256, 128])}) + + +def test_hybrid_state_pool_is_the_engines(tiny_hybrid): + from freetoken.kvcache.linear_state_pool import _linear_pool_num_slots, state_pool_bytes + + r = _analyze(tiny_hybrid) + config, f = r.config, r.forecast + assert config.cache_type == "hybrid_radix" + assert f.state_slots == _linear_pool_num_slots(config) == 4 * 4 + 8 + 1 + assert f.state_bytes == state_pool_bytes(config) + # only the one full-attention layer holds paged KV + assert f.kv_bytes_per_token == 2 * 1 * 64 * 2 * 1 + + +def test_hybrid_no_fit_suggests_a_smaller_state_pool(tiny_hybrid): + from freetoken.engine.forecast import forecast_inputs, suggest_tips + + config, model = _resolved(tiny_hybrid) + inp = forecast_inputs(config, model, None) + base = inp.forecast() + # the weights fit, the 25-slot state pool does not + inp.free_before = int((inp.weights.gpu_total + base.state_bytes - 4 * MiB) / config.memory_ratio) + base = inp.forecast() + assert base.verdict == "does not fit" + tips, combo = suggest_tips(inp, base) + by_flag = {t.flags[0]: t for t in tips} + assert by_flag["--max-running-requests 1"].forecast.state_slots == 4 + 4 + 1 + assert by_flag["--cache-type naive"].forecast.state_slots == 4 + 1 + assert combo is not None and combo.forecast.verdict != "does not fit" + assert config.max_running_req == 4 and config.cache_type == "hybrid_radix" + + +# --------------------------------------------------------------------------------------------- +# MoE offload: the expert slot cache comes from the engine's --moe-cache-auto plan +# --------------------------------------------------------------------------------------------- + +TINY_QWEN3_MOE = { + "architectures": ["Qwen3MoeForCausalLM"], "model_type": "qwen3_moe", "hidden_size": 128, + "intermediate_size": 256, "moe_intermediate_size": 64, "num_experts": 8, "num_experts_per_tok": 2, + "num_hidden_layers": 2, "num_attention_heads": 2, "num_key_value_heads": 1, "head_dim": 64, + "vocab_size": 256, "max_position_embeddings": 4096, "rms_norm_eps": 1e-6, "rope_theta": 10000, + "hidden_act": "silu", "tie_word_embeddings": False, "torch_dtype": "bfloat16", "norm_topk_prob": True, + "decoder_sparse_step": 1, "mlp_only_layers": [], +} + + +def test_moe_offload_slots_follow_the_auto_plan(tmp_path): + from freetoken.engine.engine import plan_moe_cache_auto + from freetoken.kvcache import resolve_pool_class + + path = write_checkpoint(tmp_path / "tiny-moe", TINY_QWEN3_MOE, + {"model.layers.0.mlp.experts.0.gate_proj.weight": ("BF16", [64, 128])}) + r = _analyze(path) + config, f, inp = r.config, r.forecast, r.inputs + assert config.moe_strategy == "offload" and config.moe_cache_auto + # bf16 experts: gate+up+down rows of one expert + assert f.per_expert_bytes == 3 * 64 * 128 * 2 + assert "experts" not in inp.weights.gpu and inp.weights.host["experts"] == 2 * 8 * f.per_expert_bytes + size, pages, _ = plan_moe_cache_auto( + config, resolve_pool_class(config.model_config), baseline_free=r.gpu.free_before, + weights_bytes=f.weights_bytes + f.load_overhead, per_expert_bytes=f.per_expert_bytes, + max_slots=inp.max_slots, + ) + assert (f.moe_slots, f.num_pages) == (size, pages) + + +def test_free_memory_is_read_before_this_process_creates_a_cuda_context(tiny_qwen3, monkeypatch): + """The arch probe in _adjust_config initializes CUDA on a GPU machine; a reading taken after + it would count this process's own context twice.""" + from freetoken.engine import engine + from freetoken.server import info + + calls = [] + real_gpu, real_adjust = info.gpu_info, engine._adjust_config + monkeypatch.setattr(info, "gpu_info", lambda *a, **k: calls.append("gpu_info") or real_gpu(*a, **k)) + monkeypatch.setattr(engine, "_adjust_config", lambda *a, **k: calls.append("adjust") or real_adjust(*a, **k)) + _analyze(tiny_qwen3) + assert calls[:2] == ["gpu_info", "adjust"] + + +def test_df11_placeholders_are_priced_from_the_weight_count(): + """DF11 buffers are empty until load, so the forecast estimates them instead of counting 0.""" + from freetoken.layers.base import OPList + from freetoken.models.glm4_moe.df11_embedding import EmbeddingDF11 + from freetoken.models.glm4_moe.df11_linear import LinearDF11 + + holder = OPList([]) + holder.self_attn = LinearDF11(128, 256, has_bias=False) + holder.embed_tokens = EmbeddingDF11(512, 128) + loaded = LinearDF11(64, 64, has_bias=False) + loaded.low8 = torch.empty(64 * 64, dtype=torch.uint8) + holder.mlp = loaded + est = fc_mod._df11_estimate(holder) + bits = fc_mod._DF11_BITS_PER_WEIGHT + assert est == {"attention": int(128 * 256 * bits / 8), "embeddings": int(512 * 128 * bits / 8)} diff --git a/tests/kvcache/test_pool_sizing_surface.py b/tests/kvcache/test_pool_sizing_surface.py index 516090b099..239ba64032 100644 --- a/tests/kvcache/test_pool_sizing_surface.py +++ b/tests/kvcache/test_pool_sizing_surface.py @@ -112,6 +112,37 @@ def test_dsv4_kv_cost_and_floor_parity(): DSV4PagedKVCache.solve_num_pages(_dsv4_config(num_page_override=1), 0) +def test_plan_num_pages_is_solve_without_the_assert(): + # the pre-load forecast reads the same page count the startup solve allocates, and a + # too-small budget as a number instead of an AssertionError + from freetoken.kvcache.mha_pool import MHAKVCache + + config = _generic_config() + per_page = MHAKVCache.kv_cost(config)[0] + for available in (per_page * 2, per_page * 100 + 5, 10**9): + assert MHAKVCache.plan_num_pages(config, available) == MHAKVCache.solve_num_pages(config, available) + assert MHAKVCache.plan_num_pages(config, per_page) == 1 + assert MHAKVCache.plan_num_pages(config, -per_page) == -1 + with pytest.raises(AssertionError): + MHAKVCache.solve_num_pages(config, per_page) + assert MHAKVCache.plan_num_pages(_generic_config(num_page_override=7), 0) == 7 + + +def test_dsv4_plan_num_pages_matches_solve(): + pytest.importorskip("freetoken.kvcache.dsv4_cost_model") + from freetoken.kvcache.dsv4_cost_model import _dsv4_window_floor_pages + from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + + floor = _dsv4_window_floor_pages(_dsv4_config(), 128) + config = _dsv4_config(num_page_override=floor + 3) + assert DSV4PagedKVCache.plan_num_pages(config, 0) == DSV4PagedKVCache.solve_num_pages(config, 0) == floor + 3 + budget = 64 << 30 + assert DSV4PagedKVCache.plan_num_pages(_dsv4_config(), budget) == DSV4PagedKVCache.solve_num_pages( + _dsv4_config(), budget) + with pytest.raises(ValueError, match="window working-set floor"): + DSV4PagedKVCache.plan_num_pages(_dsv4_config(num_page_override=1), 0) + + def test_startup_kv_budget_composition(): # The engine-side budget must stay ratio*old - (old - new); a sign/order slip here # mis-sizes every model's startup KV pool with the whole CPU suite green. diff --git a/tests/server/test_info.py b/tests/server/test_info.py new file mode 100644 index 0000000000..60ccbe197f --- /dev/null +++ b/tests/server/test_info.py @@ -0,0 +1,138 @@ +"""ft info (server/info.py): argument handling, the safetensors header view, JSON and exit codes.""" + +import json +import math + +import pytest + +from freetoken.engine.forecast import GiB, tensor_category +from tests.engine.test_forecast import ( + _DTYPE_BYTES, KV_PER_TOKEN, TINY_QWEN3, _analyze, tiny_qwen3_tensors, write_checkpoint, +) + + +@pytest.fixture +def tiny_qwen3(tmp_path): + return write_checkpoint(tmp_path / "tiny-qwen3", TINY_QWEN3, tiny_qwen3_tensors()) + + +def test_split_argv_takes_the_model_positionally(): + from freetoken.server.info import split_argv + + opts, rest = split_argv(["Qwen/Qwen3-0.6B", "--json", "--max-running-requests", "1", "--gpu-memory-gib", "24"]) + assert opts.json and opts.gpu_memory_gib == 24 and opts.gpu_free_gib is None + assert rest == ["--model", "Qwen/Qwen3-0.6B", "--max-running-requests", "1"] + _, rest = split_argv(["--model", "m", "--text-model-only"]) + assert rest == ["--model", "m", "--text-model-only"] + with pytest.raises(SystemExit): + split_argv(["m", "--gpu-memory-gib", "24", "--gpu-free-gib", "20"]) + + +def test_hub_headers_without_download(monkeypatch): + import huggingface_hub + from types import SimpleNamespace + + from freetoken.server.info import checkpoint_tensors + + def not_cached(*a, **k): + raise FileNotFoundError + + meta = SimpleNamespace(files_metadata={"model.safetensors": SimpleNamespace(tensors={ + "model.embed_tokens.weight": SimpleNamespace(dtype="BF16", data_offsets=(0, 4096)), + "model.layers.0.self_attn.q_proj.weight": SimpleNamespace(dtype="F8_E4M3", data_offsets=(4096, 6144)), + })}) + monkeypatch.setattr(huggingface_hub, "snapshot_download", not_cached) + monkeypatch.setattr(huggingface_hub, "get_safetensors_metadata", lambda repo: meta) + tensors, source = checkpoint_tensors("org/not-cached") + assert tensors == {"model.embed_tokens.weight": ("BF16", 4096), + "model.layers.0.self_attn.q_proj.weight": ("F8_E4M3", 2048)} + assert "HTTP range" in source + + +def test_checkpoint_view_reads_headers_only(tiny_qwen3): + r = _analyze(tiny_qwen3) + expected: dict[str, int] = {} + for name, (dtype, shape) in tiny_qwen3_tensors().items(): + n = _DTYPE_BYTES[dtype] + for d in shape: + n *= d + expected[tensor_category(name)] = expected.get(tensor_category(name), 0) + n + assert r.checkpoint == expected + assert r.checkpoint_source == "local safetensors headers" + + +def test_json_report_round_trips(tiny_qwen3, capsys): + from freetoken.server.info import main + + rc = main([tiny_qwen3, "--gpu-free-gib", "1", "--json", "--max-running-requests", "2"]) + doc = json.loads(capsys.readouterr().out) + assert rc == 0 + assert doc["architecture"] == "Qwen3ForCausalLM" + assert doc["resolved"]["max_running_req"] == 2 + assert doc["forecast"]["verdict"] in ("fits", "tight") + assert doc["forecast"]["kv_bytes_per_token"] == KV_PER_TOKEN + assert doc["checkpoint"]["bytes_by_dtype"] == {"BF16": sum( + 2 * math.prod(shape) for _, shape in tiny_qwen3_tensors().values())} + + +def test_cli_exit_code_flags_no_fit(tiny_qwen3, capsys): + from freetoken.server.info import main + + assert main([tiny_qwen3, "--gpu-free-gib", "0.01"]) == 1 + assert "DOES NOT FIT" in capsys.readouterr().out + + +def test_gpu_memory_option_plans_for_another_card(tiny_qwen3, monkeypatch): + from freetoken.engine.forecast import CUDA_CONTEXT_BYTES + from freetoken.server.args import parse_args + from freetoken.server.info import gpu_info, split_argv + + opts, argv = split_argv([tiny_qwen3, "--gpu-memory-gib", "24"]) + config, _ = parse_args(argv, prog="ft info") + info = gpu_info(opts, config) + assert info.free_before == 24 * GiB - CUDA_CONTEXT_BYTES + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "") + opts, _ = split_argv([tiny_qwen3]) + assert gpu_info(opts, config).free_before is None + + +def test_no_gpu_still_reports_weights(tiny_qwen3, monkeypatch, capsys): + from freetoken.server.info import main + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "") + assert main([tiny_qwen3]) == 0 + out = capsys.readouterr().out + assert "GPU free memory unknown" in out and "attention" in out + + +def test_skip_preflight_flag_parses(tiny_qwen3): + from freetoken.server.args import parse_args + + assert parse_args(["--model", tiny_qwen3])[0].skip_preflight is False + assert parse_args(["--model", tiny_qwen3, "--skip-preflight"])[0].skip_preflight is True + + +def test_header_fallback_when_the_meta_build_fails(tiny_qwen3, monkeypatch): + import freetoken.models + + def broken(model_config): + raise NotImplementedError("no meta build") + + monkeypatch.setattr(freetoken.models, "create_model", broken) + r = _analyze(tiny_qwen3) + stored = sum(2 * math.prod(shape) for _, shape in tiny_qwen3_tensors().values()) + assert r.inputs.weights.gpu_total == stored + assert r.forecast.kv_bytes_per_token == KV_PER_TOKEN + assert any("meta device" in n for n in r.notes) + + +def test_config_errors_are_reported_without_a_traceback(tiny_qwen3, capsys, monkeypatch): + import freetoken.engine.engine + from freetoken.server.info import main + + def bad_config(config): + raise ValueError("--moe-strategy fused cannot hold nvfp4 experts") + + monkeypatch.setattr(freetoken.engine.engine, "_adjust_config", bad_config) + assert main([tiny_qwen3, "--gpu-free-gib", "8"]) == 2 + assert "nvfp4 experts" in capsys.readouterr().err