diff --git a/docs/cli.md b/docs/cli.md index 5a759374a..8c79202c1 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -70,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) | +| `--vram-reserve-mb` | 0 | VRAM (MiB) left free for the rest of the machine at the largest prefill: startup runs one max-length chunk and shrinks the MoE expert cache until the reserve stays free; 0 skips the check | | `--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` | diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py index 891ac2ac1..6c5c0c5ad 100644 --- a/python/freetoken/engine/cache_budget.py +++ b/python/freetoken/engine/cache_budget.py @@ -45,6 +45,21 @@ def net_cache_budget_bytes( return int(memory_ratio * baseline_free) - weights_bytes - fixed_cache_size +# A rebuild re-captures the CUDA graphs and returns less than its slots held (1.24 of 1.26 GiB +# measured), so a drop sized to the shortfall alone creeps toward the reserve without reaching it. +_RESERVE_FIT_MARGIN = 128 << 20 + + +def slots_to_free_for_reserve(free_at_peak: int, reserve: int, per_expert_bytes: int) -> int: + """Expert slots to drop so the engine's measured peak leaves ``reserve`` bytes of the device + free, rounded up to whole slots past a margin for what the rebuild itself costs.""" + assert per_expert_bytes > 0, "per_expert_bytes must be positive" + shortfall = reserve - free_at_peak + if shortfall <= 0: + return 0 + return div_ceil(shortfall + max(shortfall // 16, _RESERVE_FIT_MARGIN), per_expert_bytes) + + def required_bytes( moe_cache_size: int, num_pages: int, per_expert_bytes: int, cache_per_page: int ) -> int: diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 23afd3573..99c00923f 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -86,6 +86,8 @@ class EngineConfig: cuda_graph_max_bs: int | None = None page_size: int = 1 memory_ratio: float = 0.9 + # --vram-reserve-mb: device memory left free for the rest of the machine at the largest prefill + vram_reserve_mb: int = 0 # 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 7982095ac..74350d48b 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -358,6 +358,11 @@ def _materialize_loaded_weight_state_dict( return state_dict +# Shrinking the expert cache re-captures the CUDA graphs, which moves the peak it was sized +# against; a second measurement settles it, and a third that still falls short never will. +_PEAK_FIT_ATTEMPTS = 3 + + class ForwardOutput(NamedTuple): next_tokens_gpu: torch.Tensor next_tokens_cpu: torch.Tensor @@ -536,6 +541,8 @@ def __init__(self, config: EngineConfig): if config.attention_backend.split(",")[0] == "triton": # Prefill runs on the first comma part; warm its autotune cache. self._warmup_prefill() + if config.vram_reserve_mb: + self._fit_prefill_peak(config.vram_reserve_mb << 20) def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: if config.tp_info.size == 1 or config.use_pynccl: @@ -1008,17 +1015,11 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: self.cpu_moe_executor = executor def _sync_get_memory(self) -> Tuple[int, int]: - """Get the min and max free memory across TP ranks.""" + """Get the min and max free memory across TP ranks, after emptying the allocator cache.""" torch.cuda.synchronize(self.device) torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats(self.device) - free_memory = get_free_memory(self.device) - free_mem_tensor = torch.tensor([free_memory, -free_memory], device="cpu", dtype=torch.int64) - torch.distributed.all_reduce( - free_mem_tensor, op=torch.distributed.ReduceOp.MIN, group=self.tp_cpu_group - ) - min_free_memory = int(free_mem_tensor[0].item()) - max_free_memory = -int(free_mem_tensor[1].item()) + min_free_memory, max_free_memory = self._free_memory_across_ranks() if max_free_memory - min_free_memory > 2 * 1024 * 1024 * 1024: logger.error( f"Memory across TP ranks are imbalanced:" @@ -1028,6 +1029,16 @@ def _sync_get_memory(self) -> Tuple[int, int]: return min_free_memory, max_free_memory + def _free_memory_across_ranks(self) -> Tuple[int, int]: + """The min and max free device memory across TP ranks, with what the allocator caches still + held: after a forward, that is the peak the forward reached.""" + free_memory = get_free_memory(self.device) + free_mem_tensor = torch.tensor([free_memory, -free_memory], device="cpu", dtype=torch.int64) + torch.distributed.all_reduce( + free_mem_tensor, op=torch.distributed.ReduceOp.MIN, group=self.tp_cpu_group + ) + return int(free_mem_tensor[0].item()), -int(free_mem_tensor[1].item()) + def _target_moe_and_expert_bytes(self, moe_cache_size: int | None, pool_sizes: list[int] | None) -> tuple[int, int]: cache = self.moe_offload_cache if cache is None: @@ -1247,9 +1258,7 @@ def _warmup_prefill(self) -> None: """Compile the Triton prefill path before the first real request. Decode CUDA graph capture warms the decode path, but the first prefill - can still pay Triton/cublas setup costs. Use the dummy request row and - restore it afterwards so padded decode graph replay keeps using the - dedicated dummy KV slot. + can still pay Triton/cublas setup costs. """ if self.max_seq_len < 2: return @@ -1260,19 +1269,66 @@ def _warmup_prefill(self) -> None: warmup_lens = sorted({length for length in warmup_lens if length >= 2}) if not warmup_lens: return + seconds = self._prefill_dummy(warmup_lens, input_ids=None) + logger.info_rank0(f"Prefill warmup complete for lengths {warmup_lens} in {seconds:.3f} s") + + @torch.inference_mode() + def _fit_prefill_peak(self, reserve: int) -> None: + """Shrink the expert cache until a max-length prefill leaves ``reserve`` bytes free: the + budget does not price that chunk's activations, which the allocator keeps once reached.""" + from freetoken.engine.cache_budget import slots_to_free_for_reserve + + length = min(self.config.max_forward_len, self.max_seq_len) + # Real-looking ids, so routing spreads the chunk over the experts as a prompt does. + ids = torch.randint( + self.config.model_config.vocab_size, (length,), + generator=torch.Generator().manual_seed(0), dtype=torch.int32, + ) + for attempt in range(_PEAK_FIT_ATTEMPTS): + seconds = self._prefill_dummy([length], input_ids=ids) + free, _ = self._free_memory_across_ranks() + logger.info_rank0( + f"A {length}-token prefill ({seconds:.1f} s) leaves {mem_GB(free)} free against " + f"a {mem_GB(reserve)} reserve" + ) + if free >= reserve: + return + if self.moe_offload_cache is None or attempt == _PEAK_FIT_ATTEMPTS - 1: + break + size, per_expert = self._target_moe_and_expert_bytes(None, None) + drop = slots_to_free_for_reserve(free, reserve, per_expert) + logger.info_rank0(f"Shrinking the expert cache from {size} to {size - drop} slots") + try: + self.rebuild_runtime_cache(moe_cache_size=size - drop) + except CacheRebuildRejected as exc: + raise RuntimeError( + f"--vram-reserve-mb {reserve >> 20} cannot be kept free: {exc}. Lower " + "--memory-ratio or --max-prefill-length, or the reserve." + ) from exc + raise RuntimeError( + f"a {length}-token prefill still leaves less than --vram-reserve-mb {reserve >> 20} " + "free, and the expert cache cannot shrink further; lower --memory-ratio or " + "--max-prefill-length" + ) + def _prefill_dummy(self, lengths: list[int], input_ids: torch.Tensor | None) -> float: + """Prefill each length on the dummy row, restored afterwards so padded decode replay keeps + its dedicated KV slot; returns the GPU seconds spent.""" dummy_row = self.page_table[self.dummy_req.table_idx] dummy_slot = int(dummy_row[0].item()) started = torch.cuda.Event(enable_timing=True) ended = torch.cuda.Event(enable_timing=True) started.record(self.stream) try: - for length in warmup_lens: + for length in lengths: dummy_row[:length] = torch.arange( length, dtype=torch.int32, device=self.device ) + token_ids = ( + torch.zeros(length, dtype=torch.int32) if input_ids is None else input_ids[:length] + ) warm_req = Req( - input_ids=torch.zeros(length, dtype=torch.int32, device="cpu"), + input_ids=token_ids, table_idx=self.dummy_req.table_idx, cached_len=0, output_len=1, @@ -1282,7 +1338,7 @@ def _warmup_prefill(self) -> None: ) batch = Batch(reqs=[warm_req], phase="prefill") batch.padded_reqs = batch.reqs - batch.input_ids = torch.zeros(length, dtype=torch.int32, device=self.device) + batch.input_ids = token_ids.to(self.device) batch.positions = torch.arange(length, dtype=torch.int32, device=self.device) if self.config.model_config.model_is_mrope: batch.mrope_positions = ( @@ -1298,10 +1354,7 @@ def _warmup_prefill(self) -> None: self.moe_offload_cache.reset() ended.record(self.stream) torch.cuda.synchronize(self.device) - logger.info_rank0( - f"Prefill warmup complete for lengths {warmup_lens} " - f"in {started.elapsed_time(ended) / 1000.0:.3f} s" - ) + return started.elapsed_time(ended) / 1000.0 def shutdown(self) -> None: self.graph_runner.destroy_cuda_graphs() diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index cca78e632..f1a7d06a4 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -360,6 +360,16 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--vram-reserve-mb", + type=int, + default=ServerArgs.vram_reserve_mb, + help="VRAM (MiB) to leave free for the rest of the machine at this engine's largest " + "prefill. The budget prices weights and caches but not a max-length chunk's activations, " + "which the allocator keeps once reached; with a reserve set, startup runs one such chunk " + "and shrinks the MoE expert cache until the reserve stays free. 0 skips the check.", + ) + assert ServerArgs.use_dummy_weight == False parser.add_argument( "--dummy-weight", diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index 44130d354..f00dbe646 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -573,3 +573,74 @@ def test_moe_pool_flags_and_decode_batch_clamp(): ) _clamp_decode_batch_to_pools(config, [12, 16, 9]) assert (config.max_running_req, config.cuda_graph_max_bs, config.cuda_graph_bs) == (2, 2, [1, 2]) + + +# ---- the startup fit to a max-length prefill: slots_to_free_for_reserve, Engine._fit_prefill_peak ---- + +MIB = 1 << 20 +EXPERT_BYTES = 2_772_480 # one NVFP4 expert of Qwen3.8-Flash-Next + + +def test_the_reserve_shortfall_rounds_up_to_whole_slots_past_a_margin(): + from freetoken.engine.cache_budget import slots_to_free_for_reserve + + def drop(free_at_peak: int) -> int: + return slots_to_free_for_reserve(free_at_peak=free_at_peak, reserve=2000 * MIB, per_expert_bytes=EXPERT_BYTES) + + assert drop(2000 * MIB) == 0 and drop(3000 * MIB) == 0 + # one byte short still drops the whole margin, so the rebuild's own cost cannot swallow it + assert drop(2000 * MIB - 1) * EXPERT_BYTES >= 128 * MIB + # the measured case: 430 MiB left at the peak against a 2000 MiB reserve + assert drop(430 * MIB) * EXPERT_BYTES >= 1570 * MIB + 128 * MIB + + +class _ReserveFitEngine: + """What _fit_prefill_peak drives: a max-length prefill that leaves ``free`` bytes, and a cache + rebuild that hands back ``returned`` of what the dropped slots held, less ``rebuild_cost``.""" + + def __init__(self, free: int, *, returned: float = 1.0, rebuild_cost: int = 0, refuses: bool = False): + self.config = SimpleNamespace(max_forward_len=8192, model_config=SimpleNamespace(vocab_size=248320)) + self.max_seq_len = 1 << 20 + self.moe_offload_cache = SimpleNamespace(cache_size=2900) + self.free, self.returned, self.rebuild_cost, self.refuses = free, returned, rebuild_cost, refuses + + def _prefill_dummy(self, lengths, input_ids) -> float: + return 1.0 + + def _free_memory_across_ranks(self) -> tuple[int, int]: + return self.free, self.free + + def _target_moe_and_expert_bytes(self, moe_cache_size, pool_sizes): + return self.moe_offload_cache.cache_size, EXPERT_BYTES + + def rebuild_runtime_cache(self, *, moe_cache_size: int) -> None: + from freetoken.kvcache.base import CacheRebuildRejected + + if self.refuses: + raise CacheRebuildRejected(f"moe_cache_size={moe_cache_size} is below the minimum") + dropped = (self.moe_offload_cache.cache_size - moe_cache_size) * EXPERT_BYTES + self.free += int(dropped * self.returned) - self.rebuild_cost + self.moe_offload_cache.cache_size = moe_cache_size + + +def test_the_fit_reaches_the_reserve_when_a_rebuild_costs_memory_of_its_own(): + # A start that failed: 0.69 GiB free against 2000 MiB, and 1.26 GiB of slots returned 1.24 GiB. + # Drops sized to the shortfall alone crept to 1.95 GiB and never arrived in three measurements. + from freetoken.engine.engine import Engine + + engine = _ReserveFitEngine(706 * MIB, rebuild_cost=20 * MIB) + Engine._fit_prefill_peak(engine, 2000 * MIB) + assert engine.free >= 2000 * MIB and engine.moe_offload_cache.cache_size < 2900 + + +@pytest.mark.parametrize( + ("returned", "refuses", "message"), + [(0.0, False, "still leaves less than --vram-reserve-mb 2000"), (1.0, True, "--vram-reserve-mb 2000 cannot be kept free")], + ids=["frees-nothing", "refused"], +) +def test_a_reserve_the_expert_cache_cannot_free_fails_the_start(returned, refuses, message): + from freetoken.engine.engine import Engine + + engine = _ReserveFitEngine(706 * MIB, returned=returned, refuses=refuses) + with pytest.raises(RuntimeError, match=message): + Engine._fit_prefill_peak(engine, 2000 * MIB)