From d51a8e2f9146d0a3d2ad0e6806273ddbc5e55b00 Mon Sep 17 00:00:00 2001 From: dejay2 <218806300+dejay2@users.noreply.github.com> Date: Wed, 2 Sep 2026 05:04:04 +0100 Subject: [PATCH] fix(engine): reserve the pool's dummy page in the --moe-cache-auto KV floor plan_cache_budget works in USABLE pages, while create_kv_pool allocates num_pages + 1: every pool family keeps page 0 as an unreachable dummy/sentinel. The plan prices the KV half as num_pages * cache_per_page and hands experts everything the reserve does not protect (the fill is greedy and deliberately leaves no headroom), so it under-counted the KV pool by exactly one page's worth of bytes, and that page was then allocated out of memory the split had already given to expert slots. On a model with a large cache_per_page that page is not small, and the failure mode is the one --moe-cache-auto exists to prevent: an arithmetic plan that says it fits, followed by a CUDA OOM in the KV allocation at boot. The same off-by-one made --kv-reserve-tokens quietly under-deliver: it is documented as a KV-cache token floor, but ceil(N / page_size) usable pages minus the dummy is less than N tokens of reachable capacity. Reserve the sentinel page on top of the user-visible floor and say so in the --kv-reserve-tokens help text. Tested (Windows, RTX 5090): python -m pytest -q tests/engine tests/server before: 643 passed, 2 failed after: 644 passed, 2 failed The two failures are pre-existing on this box and unrelated: both tests/engine/test_cache_budget.py::test_adjust_config_* need flashinfer, which is not installed here. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01RG8BXfsSZi1nh4wMZnhJQK --- python/freetoken/engine/cache_budget.py | 4 +++- python/freetoken/server/args.py | 5 ++++- tests/engine/test_cache_budget.py | 20 ++++++++++++++++++++ 3 files changed, 27 insertions(+), 2 deletions(-) diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py index ab7c0a9f1..8f85cf44d 100644 --- a/python/freetoken/engine/cache_budget.py +++ b/python/freetoken/engine/cache_budget.py @@ -118,7 +118,9 @@ def resolve_moe_cache_auto( """ budget_bytes = net_cache_budget_bytes(memory_ratio, baseline_free, weights_bytes, fixed_cache_size) max_slots = 992 if quant_format == "nvfp4_marlin" else total_experts - kv_reserve_pages = div_ceil(kv_reserve_tokens, page_size) + # Every pool keeps page 0 as an unreachable dummy/sentinel. The CLI floor is expressed in + # usable tokens, so reserve that internal page in addition to the user-visible capacity. + kv_reserve_pages = div_ceil(kv_reserve_tokens, page_size) + 1 return plan_cache_budget( budget_bytes=budget_bytes, per_expert_bytes=per_expert_bytes, diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 5b4db587d..ab0e0cd63 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -537,7 +537,10 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "--kv-reserve-tokens", type=int, default=ServerArgs.kv_reserve_tokens, - help="KV-cache token floor reserved before --moe-cache-auto fills experts.", + help=( + "Usable KV-cache token floor reserved before --moe-cache-auto fills experts " + "(the internal dummy page is additional)." + ), ) parser.add_argument( diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index 9ac2a4f4c..40e1055b2 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -112,6 +112,26 @@ def test_resolve_auto_applies_ratio_once_and_marlin_cap(): assert size == 8 and pages == 40 and overlap is True +def test_resolve_auto_reserves_usable_tokens_beyond_the_dummy_page(): + size, pages, overlap = resolve_moe_cache_auto( + baseline_free=940, + weights_bytes=0, + memory_ratio=1.0, + cache_per_page=10, + fixed_cache_size=0, + per_expert_bytes=100, + num_experts=2, + total_experts=50, + prefill_overlap=False, + kv_reserve_tokens=256, + page_size=64, + quant_format="bf16", + ) + assert overlap is False + assert (pages - 1) * 64 >= 256 + assert size == 8 and pages == 14 + + def test_resolve_auto_marlin_caps_slots(): size, _, _ = resolve_moe_cache_auto( baseline_free=10_000_000, weights_bytes=0, memory_ratio=1.0,