Bound prefill delayer all-branch delay and decay the max_prefill_bs high-watermark (#32880)
This commit is contained in:
@@ -244,6 +244,11 @@ class PrefillDelayer:
|
|||||||
self.skip_first_delayer = False
|
self.skip_first_delayer = False
|
||||||
pass
|
pass
|
||||||
else:
|
else:
|
||||||
|
# Bound the wait like the "mixed" branch: on a saturated
|
||||||
|
# engine slot_condition may never turn false, so cap the
|
||||||
|
# delay by max_delay_passes.
|
||||||
|
prev_delayed_count = prev_state.delayed_count if prev_state else 0
|
||||||
|
if prev_delayed_count < self._max_delay_passes - 1:
|
||||||
next_state = prev_state or _State()
|
next_state = prev_state or _State()
|
||||||
next_state = next_state.bump_delayed_count()
|
next_state = next_state.bump_delayed_count()
|
||||||
return _NegotiateOutput(
|
return _NegotiateOutput(
|
||||||
@@ -252,6 +257,13 @@ class PrefillDelayer:
|
|||||||
output_reason="delay",
|
output_reason="delay",
|
||||||
**debug_info,
|
**debug_info,
|
||||||
)
|
)
|
||||||
|
return _NegotiateOutput(
|
||||||
|
next_state=None,
|
||||||
|
output_allow=True,
|
||||||
|
output_reason="wait_timeout",
|
||||||
|
**debug_info,
|
||||||
|
**wait_info,
|
||||||
|
)
|
||||||
exist_previous_wait = prev_state is not None
|
exist_previous_wait = prev_state is not None
|
||||||
return _NegotiateOutput(
|
return _NegotiateOutput(
|
||||||
next_state=None,
|
next_state=None,
|
||||||
|
|||||||
@@ -1141,7 +1141,7 @@ class Scheduler(
|
|||||||
self.schedule_low_priority_values_first,
|
self.schedule_low_priority_values_first,
|
||||||
)
|
)
|
||||||
self.prefill_delayer: Optional[PrefillDelayer] = None
|
self.prefill_delayer: Optional[PrefillDelayer] = None
|
||||||
self.max_prefill_bs: int = 0
|
self.max_prefill_bs: float = 0.0
|
||||||
if get_schedule().enable_prefill_delayer:
|
if get_schedule().enable_prefill_delayer:
|
||||||
if get_disagg().disaggregation_mode == "decode":
|
if get_disagg().disaggregation_mode == "decode":
|
||||||
logger.info(
|
logger.info(
|
||||||
@@ -3010,6 +3010,11 @@ class Scheduler(
|
|||||||
def get_new_batch_prefill(self, running_batch: ScheduleBatch) -> NextBatchPlan:
|
def get_new_batch_prefill(self, running_batch: ScheduleBatch) -> NextBatchPlan:
|
||||||
prefill_delayer_single_pass = None
|
prefill_delayer_single_pass = None
|
||||||
if self.prefill_delayer:
|
if self.prefill_delayer:
|
||||||
|
# Decay the max-prefill-bs high-watermark once per pass so one
|
||||||
|
# unusually large admission burst does not permanently raise the
|
||||||
|
# slot_condition bar in the delayer (0.998/pass ~= half-life of
|
||||||
|
# ~350 forward passes).
|
||||||
|
self.max_prefill_bs *= 0.998
|
||||||
# Get max usage across all pools for prefill delay decision
|
# Get max usage across all pools for prefill delay decision
|
||||||
max_pool_usage = (
|
max_pool_usage = (
|
||||||
self.pool_stats_observer.get_pool_stats().get_max_pool_usage()
|
self.pool_stats_observer.get_pool_stats().get_max_pool_usage()
|
||||||
@@ -3104,7 +3109,7 @@ class Scheduler(
|
|||||||
chunked_prefill_size,
|
chunked_prefill_size,
|
||||||
running_bs if self.is_mixed_chunk else 0,
|
running_bs if self.is_mixed_chunk else 0,
|
||||||
self.priority_scheduling_preemption_threshold,
|
self.priority_scheduling_preemption_threshold,
|
||||||
max_prefill_bs=self.max_prefill_bs,
|
max_prefill_bs=int(self.max_prefill_bs),
|
||||||
max_running_requests=self.max_running_requests,
|
max_running_requests=self.max_running_requests,
|
||||||
prefill_max_requests=get_schedule().prefill_max_requests,
|
prefill_max_requests=get_schedule().prefill_max_requests,
|
||||||
prefill_delayer_single_pass=prefill_delayer_single_pass,
|
prefill_delayer_single_pass=prefill_delayer_single_pass,
|
||||||
|
|||||||
@@ -371,6 +371,53 @@ _NEGOTIATE_TEST_CASES = [
|
|||||||
# One queue-trigger delay was recorded before the wall-clock release.
|
# One queue-trigger delay was recorded before the wall-clock release.
|
||||||
expected_wait_forward_passes=1,
|
expected_wait_forward_passes=1,
|
||||||
),
|
),
|
||||||
|
# slot_condition (all-branch) must not delay forever: with 128-100=28
|
||||||
|
# free slots < max_prefill_bs=80 the delay holds, but it must release
|
||||||
|
# with wait_timeout after max_delay_passes, like the mixed branch.
|
||||||
|
NegotiateTestCase(
|
||||||
|
name="slot_condition_pass_cap_timeout",
|
||||||
|
max_delay_passes=3,
|
||||||
|
token_usage_low_watermark=0.8,
|
||||||
|
calls=[
|
||||||
|
# skip_first_delayer consumes the first would-be delay.
|
||||||
|
NegotiateCall(
|
||||||
|
prefillable=[True, True, True, True],
|
||||||
|
token_usage=[0.9, 0.9, 0.9, 0.9],
|
||||||
|
running_batch=[100, 100, 100, 100],
|
||||||
|
max_prefill_bs=[80, 80, 80, 80],
|
||||||
|
waiting_queue_len=[10, 10, 10, 10],
|
||||||
|
max_running_requests=128,
|
||||||
|
),
|
||||||
|
NegotiateCall(
|
||||||
|
prefillable=[True, True, True, True],
|
||||||
|
token_usage=[0.9, 0.9, 0.9, 0.9],
|
||||||
|
running_batch=[100, 100, 100, 100],
|
||||||
|
max_prefill_bs=[80, 80, 80, 80],
|
||||||
|
waiting_queue_len=[10, 10, 10, 10],
|
||||||
|
max_running_requests=128,
|
||||||
|
),
|
||||||
|
NegotiateCall(
|
||||||
|
prefillable=[True, True, True, True],
|
||||||
|
token_usage=[0.9, 0.9, 0.9, 0.9],
|
||||||
|
running_batch=[100, 100, 100, 100],
|
||||||
|
max_prefill_bs=[80, 80, 80, 80],
|
||||||
|
waiting_queue_len=[10, 10, 10, 10],
|
||||||
|
max_running_requests=128,
|
||||||
|
),
|
||||||
|
NegotiateCall(
|
||||||
|
prefillable=[True, True, True, True],
|
||||||
|
token_usage=[0.9, 0.9, 0.9, 0.9],
|
||||||
|
running_batch=[100, 100, 100, 100],
|
||||||
|
max_prefill_bs=[80, 80, 80, 80],
|
||||||
|
waiting_queue_len=[10, 10, 10, 10],
|
||||||
|
max_running_requests=128,
|
||||||
|
),
|
||||||
|
],
|
||||||
|
expected_allow=True,
|
||||||
|
expected_reason="wait_timeout",
|
||||||
|
# Two slot-condition delays accumulated after the skip-first pass.
|
||||||
|
expected_wait_forward_passes=2,
|
||||||
|
),
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user