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:
co-authored by
kingjameschan
Claude Opus 4.7
Humphrey
parent
8c5320b37e
commit
b796338271
@@ -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,
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user