Skip to content

Records time to first token in the dynamic inference coordinator #2081

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
sidsingh-nvidia wants to merge 3 commits into from
Closed

Records time to first token in the dynamic inference coordinator #2081

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
426f16a
record first token generation timestamp
sidsingh-nvidia Oct 16, 2025
0ed2d5e
return tpots from engine
sidsingh-nvidia Oct 19, 2025
f09bd99
format
sidsingh-nvidia Oct 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion megatron/core/inference/data_parallel_inference_coordinator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

import logging
import time
from collections import deque
from itertools import cycle, repeat
from typing import List, Tuple
Expand Down Expand Up @@ -179,6 +180,7 @@ def postprocess(
request: DynamicInferenceRequest = self.requests[request_id]
# Handle chunked prefill similar to the engine logic
if chunked_prefill_request_id == -1 or request_id != chunked_prefill_request_id:
request.generation_timestamps.append(time.perf_counter())
request.generated_tokens.append(token)

if request_log_probs is not None:
Expand Down Expand Up @@ -226,7 +228,12 @@ def postprocess(
request = self.requests.pop(fid)
request.generated_length = len(request.generated_tokens)
request.generated_text = self.tokenizer.detokenize(request.generated_tokens)

request.tpot = []
for i, gen_timestamp in enumerate(request.generation_timestamps):
prev_timestamp = (
request.arrival_time if i == 0 else request.generation_timestamps[i - 1]
)
request.tpot.append(gen_timestamp - prev_timestamp)
client_identity = self.request_id_to_client_id[fid]
client_request_identity = self.request_id_to_client_request_id[fid]
del self.request_id_to_client_id[fid]
Expand Down
2 changes: 2 additions & 0 deletions megatron/core/inference/inference_request.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,7 @@ class DynamicInferenceRequest(InferenceRequest):

request_id: int
generated_tokens: List[int] = field(default_factory=list)
generation_timestamps: List[float] = field(default_factory=list)
prompt: Optional[str] = None
prompt_tokens: Optional[torch.Tensor] = None
# remaining prompt tokens are used for chunked prefill
Expand All @@ -141,6 +142,7 @@ class DynamicInferenceRequest(InferenceRequest):
def __post_init__(self):
if self.prompt_tokens is not None:
self.remaining_prompt_tokens = copy.deepcopy(self.prompt_tokens)
self.arrival_time = time.perf_counter()

@property
def remaining_prompt_length(self):
Expand Down
Loading