feat(engine): add --vram-reserve-mb to keep VRAM free at the largest prefill

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` |

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:

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.

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()

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",

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)