From 5f21a5754d718881b372e8a5e88f3c38a1a5da33 Mon Sep 17 00:00:00 2001
From: yipin <1259478761@qq.com>
Date: Fri, 13 Nov 2026 11:14:10 -0500
Subject: [PATCH] fix(engine): drop prefill overlap when moe-cache-auto cannot
 fund its slot floor

---
 python/freetoken/engine/cache_budget.py |  6 ++++--
 tests/engine/test_cache_budget.py       | 11 +++++++++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py
index 891ac2ac11..8683d92c25 100644
--- a/python/freetoken/engine/cache_budget.py
+++ b/python/freetoken/engine/cache_budget.py
@@ -83,9 +83,11 @@ def plan_cache_budget(
     kv_reserve_bytes = kv_reserve_pages * cache_per_page
     # MoE-priority: reserve KV first, then experts greedily take the remaining budget.
     raw = (budget_bytes - kv_reserve_bytes) // per_expert_bytes
+    # Overlap is optional: when the budget cannot fund its 2*num_experts floor beside the
+    # KV reserve, plan without it instead of clamping up past the budget.
+    if overlap and raw < 2 * num_experts:
+        overlap, lo = False, num_experts
     moe_cache_size = max(lo, min(raw, hi))
-    # A tiny budget may have forced moe_cache_size below 2*num_experts even with overlap on.
-    overlap = overlap and moe_cache_size >= 2 * num_experts
 
     remaining = budget_bytes - moe_cache_size * per_expert_bytes
     num_pages = max(remaining // cache_per_page, kv_reserve_pages)
diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py
index 5c1f93b348..288269f912 100644
--- a/tests/engine/test_cache_budget.py
+++ b/tests/engine/test_cache_budget.py
@@ -60,6 +60,17 @@ def test_small_cache_disables_prefill_overlap():
     assert size == 12
 
 
+def test_budget_below_the_overlap_floor_disables_prefill_overlap():
+    # 2*num_experts slots (800 B) + the KV reserve (100 B) exceed the 700 B budget, but
+    # num_experts slots fit beside the reserve: plan without overlap instead of asserting.
+    size, pages, overlap = plan_cache_budget(
+        budget_bytes=700, per_expert_bytes=100, cache_per_page=10,
+        num_experts=4, total_experts=50, prefill_overlap=True,
+        kv_reserve_pages=10, max_slots=50,
+    )
+    assert (size, pages, overlap) == (6, 10, False)
+
+
 def test_insufficient_kv_memory_raises():
     with pytest.raises(AssertionError, match="not enough memory"):
         plan_cache_budget(