Add request-level timestamp for when prefill finishes (#14860)

This commit is contained in:
Scott Lee
2025-12-17 13:34:46 -08:00
committed by GitHub
parent 4128d4f5cb
commit 169a75dfb8
7 changed files with 42 additions and 6 deletions

View File

@@ -425,6 +425,9 @@ class SchedulerDisaggregationPrefillMixin:
zip(batch.reqs, next_token_ids, strict=True)
):
if req.is_chunked <= 0:
if req.time_stats.prefill_finished_ts == 0.0:
req.time_stats.prefill_finished_ts = time.time()
# There is no output_ids for prefill
req.output_ids.append(next_token_id)
self.tree_cache.cache_unfinished_req(req) # update the tree and lock

View File

@@ -302,6 +302,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
forward_entry_time=recv_obj.forward_entry_time,
prefill_launch_delay=recv_obj.prefill_launch_delay,
prefill_launch_latency=recv_obj.prefill_launch_latency,
prefill_finished_ts=recv_obj.prefill_finished_ts,
)
def handle_multimodal_decode_req(self, recv_obj: BatchMultimodalDecodeReq):

View File

@@ -93,6 +93,10 @@ class RequestTimingMetricsMixin:
# Calculated as: prefill_end_time_host - prefill_start_time_host
prefill_launch_latency: Optional[List[Optional[float]]]
# Prefill finished time: timestamp when prefill phase completes (wall clock time).
# This marks when the prefill computation finishes.
prefill_finished_ts: Optional[List[Optional[float]]]
@dataclass
class SpeculativeDecodingMetricsMixin:

View File

@@ -128,6 +128,9 @@ def _handle_output_by_index(output, i):
prefill_launch_latency=_extract_field_by_index(
output, "prefill_launch_latency", i
),
prefill_finished_ts=_extract_field_by_index(
output, "prefill_finished_ts", i
),
finished_reasons=_extract_field_by_index(output, "finished_reasons", i),
decoded_texts=_extract_field_by_index(output, "decoded_texts", i),
decode_ids=_extract_field_by_index(output, "decode_ids", i),
@@ -216,6 +219,9 @@ def _handle_output_by_index(output, i):
prefill_launch_latency=_extract_field_by_index(
output, "prefill_launch_latency", i
),
prefill_finished_ts=_extract_field_by_index(
output, "prefill_finished_ts", i
),
finished_reasons=_extract_field_by_index(output, "finished_reasons", i),
output_strs=_extract_field_by_index(output, "output_strs", i),
output_ids=_extract_field_by_index(output, "output_ids", i),

View File

@@ -108,6 +108,9 @@ class SchedulerOutputProcessorMixin:
continue
if req.is_chunked <= 0:
if req.time_stats.prefill_finished_ts == 0.0:
req.time_stats.prefill_finished_ts = time.time()
# req output_ids are set here
req.output_ids.append(next_token_id)
req.check_finished()
@@ -793,6 +796,7 @@ class SchedulerOutputProcessorMixin:
forward_entry_times = []
prefill_launch_delays = []
prefill_launch_latencies = []
prefill_finished_timestamps = []
if return_logprob:
input_token_logprobs_val = []
@@ -901,6 +905,9 @@ class SchedulerOutputProcessorMixin:
prefill_launch_latencies.append(
req.time_stats.get_prefill_launch_latency()
)
prefill_finished_timestamps.append(
req.time_stats.get_prefill_finished_ts()
)
if not self.spec_algorithm.is_none():
spec_verify_ct.append(req.spec_verify_ct)
@@ -1001,6 +1008,7 @@ class SchedulerOutputProcessorMixin:
forward_entry_time=forward_entry_times,
prefill_launch_delay=prefill_launch_delays,
prefill_launch_latency=prefill_launch_latencies,
prefill_finished_ts=prefill_finished_timestamps,
finished_reasons=finished_reasons,
decoded_texts=decoded_texts,
decode_ids=decode_ids_list,
@@ -1044,6 +1052,7 @@ class SchedulerOutputProcessorMixin:
forward_entry_times = []
prefill_launch_delays = []
prefill_launch_latencies = []
prefill_finished_timestamps = []
retraction_counts = []
for req in reqs:
if req.finished():
@@ -1061,6 +1070,9 @@ class SchedulerOutputProcessorMixin:
prefill_launch_latencies.append(
req.time_stats.get_prefill_launch_latency()
)
prefill_finished_timestamps.append(
req.time_stats.get_prefill_finished_ts()
)
retraction_counts.append(req.retraction_count)
self.send_to_detokenizer.send_output(
BatchEmbeddingOutput(
@@ -1070,6 +1082,7 @@ class SchedulerOutputProcessorMixin:
forward_entry_time=forward_entry_times,
prefill_launch_delay=prefill_launch_delays,
prefill_launch_latency=prefill_launch_latencies,
prefill_finished_ts=prefill_finished_timestamps,
finished_reasons=finished_reasons,
embeddings=embeddings,
prompt_tokens=prompt_tokens,

View File

@@ -1538,6 +1538,9 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
self._add_metric_if_present(
recv_obj, "prefill_launch_latency", meta_info, i
)
self._add_metric_if_present(
recv_obj, "prefill_finished_ts", meta_info, i
)
if getattr(state.obj, "return_logprob", False):
self.convert_logprob_style(
@@ -1836,12 +1839,6 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
meta_info["request_sent_to_scheduler_ts"] = (
state.request_sent_to_scheduler_ts
)
# For embeddings, there's no separate prefill phase, so omit `prefill_finished_ts`.
if (
not isinstance(recv_obj, BatchEmbeddingOutput)
and state.first_token_time > 0
):
meta_info["prefill_finished_ts"] = state.first_token_time
if state.response_sent_to_client_ts > 0:
meta_info["response_sent_to_client_ts"] = state.response_sent_to_client_ts
if state.finished_time > 0:

View File

@@ -64,6 +64,13 @@ class TimeStats:
prefill_start_time_host: float = 0.0
prefill_end_time_host: float = 0.0
# Timestamp when prefill phase finishes, obtained from `time.time()`.
# Note that this differs from the other `_time` fields tracked by the
# `TimeStats` class, which are obtained from `time.perf_counter()`.
# We use `time.time()` instead of `time.perf_counter()` here in order to
# maintain unit consistency with other timestamp fields tracked by the `ReqState` class.
prefill_finished_ts: float = 0.0
def get_queueing_time(self) -> float:
return self.forward_entry_time - self.wait_queue_entry_time
@@ -77,6 +84,11 @@ class TimeStats:
return self.prefill_end_time_host - self.prefill_start_time_host
return None
def get_prefill_finished_ts(self) -> Optional[float]:
if self.prefill_finished_ts > 0.0:
return self.prefill_finished_ts
return None
def convert_to_duration(self) -> str:
if self.disagg_mode == DisaggregationMode.NULL:
queue_duration = self.forward_entry_time - self.wait_queue_entry_time