Fix prefill delayer wait histograms always observing 0 (#25975)

Co-authored-by: kingjameschan <170807154+kingjameschan@users.noreply.github.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Humphrey <181440142+humphreysun98@users.noreply.github.com>
This commit is contained in:
Kingjameschan
2026-06-14 00:35:29 -07:00
committed by GitHub
co-authored by kingjameschan Claude Opus 4.7 Humphrey
parent 8c5320b37e
commit b796338271
2 changed files with 46 additions and 7 deletions
+22 -7
View File
@@ -33,6 +33,11 @@ class _NegotiateOutput(NamedTuple):
output_reason: str output_reason: str
num_prefillable: int num_prefillable: int
num_token_watermark_force_allow: 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: class PrefillDelayer:
@@ -175,6 +180,16 @@ class PrefillDelayer:
num_token_watermark_force_allow=global_token_watermark_force_allow.sum().item(), 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 # Compute outputs
if prefillable_status == "all": if prefillable_status == "all":
# Safety valve: low KV usage means GPU is underutilized, skip # Safety valve: low KV usage means GPU is underutilized, skip
@@ -185,6 +200,7 @@ class PrefillDelayer:
output_allow=True, output_allow=True,
output_reason="token_watermark", output_reason="token_watermark",
**debug_info, **debug_info,
**wait_info,
) )
if not self.enable_dp_attention: if not self.enable_dp_attention:
@@ -242,6 +258,7 @@ class PrefillDelayer:
output_allow=True, output_allow=True,
output_reason="wait_success" if exist_previous_wait else "no_wait", output_reason="wait_success" if exist_previous_wait else "no_wait",
**debug_info, **debug_info,
**wait_info,
) )
elif prefillable_status == "none": elif prefillable_status == "none":
return _NegotiateOutput( return _NegotiateOutput(
@@ -250,6 +267,7 @@ class PrefillDelayer:
output_allow=True, output_allow=True,
output_reason="", output_reason="",
**debug_info, **debug_info,
**wait_info,
) )
elif prefillable_status == "mixed": elif prefillable_status == "mixed":
if global_exists_token_watermark_force_allow: if global_exists_token_watermark_force_allow:
@@ -258,6 +276,7 @@ class PrefillDelayer:
output_allow=True, output_allow=True,
output_reason="token_watermark", output_reason="token_watermark",
**debug_info, **debug_info,
**wait_info,
) )
prev_delayed_count = prev_state.delayed_count if prev_state else 0 prev_delayed_count = prev_state.delayed_count if prev_state else 0
@@ -276,6 +295,7 @@ class PrefillDelayer:
output_allow=True, output_allow=True,
output_reason="wait_timeout", output_reason="wait_timeout",
**debug_info, **debug_info,
**wait_info,
) )
else: else:
raise NotImplementedError raise NotImplementedError
@@ -376,14 +396,9 @@ def _record_single_pass_result(
} }
if metrics_collector is not None: 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( metrics_collector.observe_prefill_delayer_outcome(
forward_passes=forward_passes, forward_passes=output.wait_forward_passes,
wait_seconds=wait_seconds, wait_seconds=output.wait_seconds,
input_estimation=output.input_estimation, input_estimation=output.input_estimation,
output_allow=output.output_allow, output_allow=output.output_allow,
output_reason=output.output_reason, output_reason=output.output_reason,
@@ -66,6 +66,9 @@ class NegotiateTestCase:
# to exercise the legacy slot-only code paths. # to exercise the legacy slot-only code paths.
queue_min_ratio: Optional[float] = None queue_min_ratio: Optional[float] = None
max_delay_ms: 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): def _run_negotiate_test(rank, test_cases):
@@ -113,6 +116,17 @@ def _run_negotiate_test(rank, test_cases):
case.expected_reason, case.expected_reason,
), f"Case {case.name} rank {rank}" ), 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 = [ _NEGOTIATE_TEST_CASES = [
NegotiateTestCase( NegotiateTestCase(
@@ -127,6 +141,8 @@ _NEGOTIATE_TEST_CASES = [
], ],
expected_allow=True, expected_allow=True,
expected_reason="no_wait", expected_reason="no_wait",
# No prior wait, so the histograms legitimately observe 0.
expected_wait_forward_passes=0,
), ),
NegotiateTestCase( NegotiateTestCase(
name="all_prefillable_with_previous_wait", name="all_prefillable_with_previous_wait",
@@ -144,6 +160,9 @@ _NEGOTIATE_TEST_CASES = [
], ],
expected_allow=True, expected_allow=True,
expected_reason="wait_success", 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( NegotiateTestCase(
name="none_prefillable", name="none_prefillable",
@@ -230,6 +249,9 @@ _NEGOTIATE_TEST_CASES = [
], ],
expected_allow=True, expected_allow=True,
expected_reason="wait_timeout", 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, # Queue-based trigger: waiting queue below queue_min = min(running * R,
# max_prefill_bs) should defer prefill. With R=0.5, running=100 and # max_prefill_bs) should defer prefill. With R=0.5, running=100 and
@@ -346,6 +368,8 @@ _NEGOTIATE_TEST_CASES = [
], ],
expected_allow=True, expected_allow=True,
expected_reason="wait_success", expected_reason="wait_success",
# One queue-trigger delay was recorded before the wall-clock release.
expected_wait_forward_passes=1,
), ),
] ]