diff --git a/python/sglang/srt/managers/prefill_delayer.py b/python/sglang/srt/managers/prefill_delayer.py index 1eb21ae64..d775a35c9 100644 --- a/python/sglang/srt/managers/prefill_delayer.py +++ b/python/sglang/srt/managers/prefill_delayer.py @@ -33,6 +33,11 @@ class _NegotiateOutput(NamedTuple): output_reason: str num_prefillable: int num_token_watermark_force_allow: int + # Accumulated wait of the prefill being released on this pass. Carried + # explicitly because `next_state` is None on every release path and thus + # cannot convey it to the metrics observation. + wait_forward_passes: int = 0 + wait_seconds: float = 0.0 class PrefillDelayer: @@ -175,6 +180,16 @@ class PrefillDelayer: num_token_watermark_force_allow=global_token_watermark_force_allow.sum().item(), ) + # Wait accumulated so far, taken from prev_state. Release paths attach + # this so the wait histograms observe the real value; delay paths leave + # the defaults (0) since the wait isn't finished and isn't observed. + wait_info = dict( + wait_forward_passes=prev_state.delayed_count if prev_state else 0, + wait_seconds=( + (time.perf_counter() - prev_state.start_time) if prev_state else 0.0 + ), + ) + # Compute outputs if prefillable_status == "all": # Safety valve: low KV usage means GPU is underutilized, skip @@ -185,6 +200,7 @@ class PrefillDelayer: output_allow=True, output_reason="token_watermark", **debug_info, + **wait_info, ) if not self.enable_dp_attention: @@ -242,6 +258,7 @@ class PrefillDelayer: output_allow=True, output_reason="wait_success" if exist_previous_wait else "no_wait", **debug_info, + **wait_info, ) elif prefillable_status == "none": return _NegotiateOutput( @@ -250,6 +267,7 @@ class PrefillDelayer: output_allow=True, output_reason="", **debug_info, + **wait_info, ) elif prefillable_status == "mixed": if global_exists_token_watermark_force_allow: @@ -258,6 +276,7 @@ class PrefillDelayer: output_allow=True, output_reason="token_watermark", **debug_info, + **wait_info, ) prev_delayed_count = prev_state.delayed_count if prev_state else 0 @@ -276,6 +295,7 @@ class PrefillDelayer: output_allow=True, output_reason="wait_timeout", **debug_info, + **wait_info, ) else: raise NotImplementedError @@ -376,14 +396,9 @@ def _record_single_pass_result( } if metrics_collector is not None: - if (s := output.next_state) is not None: - wait_seconds = time.perf_counter() - s.start_time - forward_passes = s.delayed_count - else: - wait_seconds = forward_passes = 0 metrics_collector.observe_prefill_delayer_outcome( - forward_passes=forward_passes, - wait_seconds=wait_seconds, + forward_passes=output.wait_forward_passes, + wait_seconds=output.wait_seconds, input_estimation=output.input_estimation, output_allow=output.output_allow, output_reason=output.output_reason, diff --git a/test/registered/scheduler/test_prefill_delayer.py b/test/registered/scheduler/test_prefill_delayer.py index c4ba97e0f..1bc2b9cb9 100644 --- a/test/registered/scheduler/test_prefill_delayer.py +++ b/test/registered/scheduler/test_prefill_delayer.py @@ -66,6 +66,9 @@ class NegotiateTestCase: # to exercise the legacy slot-only code paths. queue_min_ratio: Optional[float] = None max_delay_ms: Optional[float] = None + # Expected accumulated wait surfaced on the final (release) outcome. When + # set, asserts the wait histograms would observe this value instead of 0. + expected_wait_forward_passes: Optional[int] = None def _run_negotiate_test(rank, test_cases): @@ -113,6 +116,17 @@ def _run_negotiate_test(rank, test_cases): case.expected_reason, ), f"Case {case.name} rank {rank}" + if case.expected_wait_forward_passes is not None: + assert result.wait_forward_passes == case.expected_wait_forward_passes, ( + f"Case {case.name} rank {rank}: wait_forward_passes " + f"{result.wait_forward_passes} != {case.expected_wait_forward_passes}" + ) + # On a release after a real wait, seconds must be observed too. + if case.expected_wait_forward_passes > 0: + assert ( + result.wait_seconds > 0.0 + ), f"Case {case.name} rank {rank}: wait_seconds not surfaced" + _NEGOTIATE_TEST_CASES = [ NegotiateTestCase( @@ -127,6 +141,8 @@ _NEGOTIATE_TEST_CASES = [ ], expected_allow=True, expected_reason="no_wait", + # No prior wait, so the histograms legitimately observe 0. + expected_wait_forward_passes=0, ), NegotiateTestCase( name="all_prefillable_with_previous_wait", @@ -144,6 +160,9 @@ _NEGOTIATE_TEST_CASES = [ ], expected_allow=True, expected_reason="wait_success", + # One mixed delay preceded the release, so the wait histograms must + # observe 1 forward pass (regression guard for #25949). + expected_wait_forward_passes=1, ), NegotiateTestCase( name="none_prefillable", @@ -230,6 +249,9 @@ _NEGOTIATE_TEST_CASES = [ ], expected_allow=True, expected_reason="wait_timeout", + # Two delays accumulated before timing out; the timeout release must + # still surface that wait to the histograms. + expected_wait_forward_passes=2, ), # Queue-based trigger: waiting queue below queue_min = min(running * R, # max_prefill_bs) should defer prefill. With R=0.5, running=100 and @@ -346,6 +368,8 @@ _NEGOTIATE_TEST_CASES = [ ], expected_allow=True, expected_reason="wait_success", + # One queue-trigger delay was recorded before the wall-clock release. + expected_wait_forward_passes=1, ), ]