[Scheduler] Honor explicit min-free-slots thresholds (#33403)

Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
This commit is contained in:
paulzhang-tm
2026-08-05 01:44:18 -07:00
committed by GitHub
co-authored by hnyls2002 Liangsheng Yin
parent c0d5ebd6c4
commit a6e5fa7081
3 changed files with 35 additions and 34 deletions
@@ -8,21 +8,18 @@ def resolve_min_free_slots(
) -> Optional[int]: ) -> Optional[int]:
"""Resolve the min-free-slots threshold (None = disabled). """Resolve the min-free-slots threshold (None = disabled).
A user value (>1) is capped to the DFlash formula so the trigger never An explicit user value always wins, capped by max_running_requests
delays more aggressively than the legacy heuristic. When unset, DFlash (<= 1 disables). When unset, DFlash workloads fall back to the legacy
workloads fall back to the formula (preserving the always-on behavior); formula (preserving the always-on behavior, disabled when
other workloads stay disabled. Also disabled when max_running_requests < 8. max_running_requests < 8); other workloads stay disabled.
""" """
max_running_requests = max(0, int(max_running_requests)) max_running_requests = max(0, int(max_running_requests))
formula = min(4, max(2, (max_running_requests + 5) // 6)) if user_value is not None:
if user_value is None: threshold = min(user_value, max_running_requests)
user_value = formula if is_dflash_family else None return threshold if threshold > 1 else None
if is_dflash_family and max_running_requests >= 8:
if user_value is None or user_value <= 1: return min(4, max(2, (max_running_requests + 5) // 6))
return None return None
if max_running_requests < 8:
return None
return min(user_value, formula)
class MinFreeSlotsDelayer: class MinFreeSlotsDelayer:
+10 -4
View File
@@ -3234,10 +3234,11 @@ class ServerArgs:
"Hold new prefills until at least N running-request slots have freed " "Hold new prefills until at least N running-request slots have freed "
"up, so they are admitted in one batch instead of one at a time. " "up, so they are admitted in one batch instead of one at a time. "
"Useful when each admission is disproportionately expensive, e.g. " "Useful when each admission is disproportionately expensive, e.g. "
"speculative decoding with a separate draft prefill pass. Capped to " "speculative decoding with a separate draft prefill pass. An "
"the DFlash formula (disabled when max-running-requests < 8; " "explicit value always wins, capped by max-running-requests "
"min(4, max(2, (max-run + 5) // 6))). DFlash workloads auto-enable " "(1 disables). When unset, DFlash workloads auto-enable the "
"this with the formula when unset; other workloads stay disabled." "formula; other workloads stay disabled. Not supported with "
"pipeline parallelism."
), ),
NS("schedule"), NS("schedule"),
] = None ] = None
@@ -8722,6 +8723,11 @@ class ServerArgs:
assert ( assert (
self.disable_overlap_schedule and self.speculative_algorithm is None self.disable_overlap_schedule and self.speculative_algorithm is None
), "Pipeline parallelism is not compatible with overlap schedule, speculative decoding" ), "Pipeline parallelism is not compatible with overlap schedule, speculative decoding"
assert self.min_free_slots_delay is None, (
"--min-free-slots-delay is not supported with pipeline "
"parallelism: allocatable slots per microbatch are bounded by "
"pp-max-micro-batch-size, so the threshold may never be reached"
)
assert not ( assert not (
self.dp_size > 1 and self.nnodes != 1 and not self.enable_dp_attention self.dp_size > 1 and self.nnodes != 1 and not self.enable_dp_attention
@@ -10,19 +10,15 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
class TestResolveMinFreeSlots(unittest.TestCase): class TestResolveMinFreeSlots(unittest.TestCase):
"""Unit tests for resolve_min_free_slots threshold resolution."""
def test_unset_non_dflash_disables(self): def test_unset_non_dflash_disables(self):
# Unset + not DFlash -> trigger stays disabled.
self.assertIsNone(resolve_min_free_slots(None, 512, is_dflash_family=False)) self.assertIsNone(resolve_min_free_slots(None, 512, is_dflash_family=False))
def test_unset_dflash_auto_enables(self): def test_unset_dflash_auto_enables(self):
# Unset + DFlash -> falls back to the legacy formula (full mapping).
self.assertEqual(resolve_min_free_slots(None, 512, is_dflash_family=True), 4) self.assertEqual(resolve_min_free_slots(None, 512, is_dflash_family=True), 4)
self.assertEqual(resolve_min_free_slots(None, 8, is_dflash_family=True), 2) self.assertEqual(resolve_min_free_slots(None, 8, is_dflash_family=True), 2)
def test_unset_dflash_small_cluster_disables(self): def test_unset_dflash_small_cluster_disables(self):
# DFlash auto-default still respects the < 8 guard.
self.assertIsNone(resolve_min_free_slots(None, 7, is_dflash_family=True)) self.assertIsNone(resolve_min_free_slots(None, 7, is_dflash_family=True))
self.assertIsNone(resolve_min_free_slots(None, 0, is_dflash_family=True)) self.assertIsNone(resolve_min_free_slots(None, 0, is_dflash_family=True))
@@ -31,27 +27,29 @@ class TestResolveMinFreeSlots(unittest.TestCase):
self.assertIsNone(resolve_min_free_slots(1, 512)) self.assertIsNone(resolve_min_free_slots(1, 512))
self.assertIsNone(resolve_min_free_slots(0, 512)) self.assertIsNone(resolve_min_free_slots(0, 512))
def test_small_cluster_disables(self): def test_explicit_value_survives_small_cluster(self):
# max_running_requests < 8 disables, matching DFlash. # The < 8 guard belongs to the DFlash auto-default, not explicit values.
self.assertIsNone(resolve_min_free_slots(4, 7)) self.assertEqual(resolve_min_free_slots(4, 7), 4)
self.assertEqual(resolve_min_free_slots(4, 7, is_dflash_family=True), 4)
def test_caps_to_formula(self): def test_non_dflash_uses_explicit_value(self):
# Capped down so it never delays more aggressively than DFlash.
self.assertEqual(resolve_min_free_slots(10, 512), 4)
self.assertEqual(resolve_min_free_slots(10, 8), 2) # (8 + 5) // 6 = 2
def test_respects_smaller_user_value(self):
# Below the formula cap is taken as-is.
self.assertEqual(resolve_min_free_slots(3, 512), 3)
self.assertEqual(resolve_min_free_slots(2, 8), 2) self.assertEqual(resolve_min_free_slots(2, 8), 2)
self.assertEqual(resolve_min_free_slots(3, 512), 3)
self.assertEqual(resolve_min_free_slots(8, 512), 8)
self.assertEqual(resolve_min_free_slots(16, 512), 16)
def test_explicit_value_is_capped_to_max_running_requests(self):
self.assertEqual(resolve_min_free_slots(16, 8), 8)
def test_user_value_overrides_dflash_default(self): def test_user_value_overrides_dflash_default(self):
# An explicit user value wins over the DFlash auto-default.
self.assertEqual(resolve_min_free_slots(3, 512, is_dflash_family=True), 3) self.assertEqual(resolve_min_free_slots(3, 512, is_dflash_family=True), 3)
self.assertEqual(resolve_min_free_slots(16, 512, is_dflash_family=True), 16)
def test_explicit_one_disables_dflash_default(self):
self.assertIsNone(resolve_min_free_slots(1, 512, is_dflash_family=True))
class TestMinFreeSlotsDelayer(unittest.TestCase): class TestMinFreeSlotsDelayer(unittest.TestCase):
"""Unit tests for the per-rank local should_delay decision."""
def test_delays_below_threshold(self): def test_delays_below_threshold(self):
delayer = MinFreeSlotsDelayer(min_free_slots=4) delayer = MinFreeSlotsDelayer(min_free_slots=4)