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