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
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,