From 426f16a0b6f9804039178c001ca47884cd794d44 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Thu, 16 Oct 2025 15:47:49 -0700 Subject: [PATCH 1/3] record first token generation timestamp --- megatron/core/inference/data_parallel_inference_coordinator.py | 3 +++ megatron/core/inference/inference_request.py | 1 + 2 files changed, 4 insertions(+) diff --git a/megatron/core/inference/data_parallel_inference_coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator.py index 101acb31e1f..ca583d40810 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator.py +++ b/megatron/core/inference/data_parallel_inference_coordinator.py @@ -6,6 +6,7 @@ from typing import List, Tuple import torch +import time from megatron.core.inference.headers import Headers from megatron.core.inference.inference_request import DynamicInferenceRequest @@ -179,6 +180,8 @@ 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: + if request.timestamp_of_first_token is None: + request.timestamp_of_first_token = time.perf_counter() request.generated_tokens.append(token) if request_log_probs is not None: diff --git a/megatron/core/inference/inference_request.py b/megatron/core/inference/inference_request.py index d4b956e58e2..4cf30732eb0 100644 --- a/megatron/core/inference/inference_request.py +++ b/megatron/core/inference/inference_request.py @@ -136,6 +136,7 @@ class DynamicInferenceRequest(InferenceRequest): # remaining prompt tokens are used for chunked prefill remaining_prompt_tokens: Optional[torch.Tensor] = None latency: Optional[float] = None + timestamp_of_first_token: Optional[float] = None finished_chunk_token_count = 0 def __post_init__(self): From 0ed2d5e117c0a2ab0f3e484a51f4bc3335495145 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Sun, 19 Oct 2025 05:46:59 -0700 Subject: [PATCH 2/3] return tpots from engine --- .../core/inference/data_parallel_inference_coordinator.py | 8 +++++--- megatron/core/inference/inference_request.py | 3 ++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/megatron/core/inference/data_parallel_inference_coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator.py index ca583d40810..6df05f7dc6c 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator.py +++ b/megatron/core/inference/data_parallel_inference_coordinator.py @@ -180,8 +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: - if request.timestamp_of_first_token is None: - request.timestamp_of_first_token = time.perf_counter() + request.generation_timestamps.append(time.perf_counter()) request.generated_tokens.append(token) if request_log_probs is not None: @@ -229,7 +228,10 @@ 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] diff --git a/megatron/core/inference/inference_request.py b/megatron/core/inference/inference_request.py index 4cf30732eb0..0d4641bc2a9 100644 --- a/megatron/core/inference/inference_request.py +++ b/megatron/core/inference/inference_request.py @@ -131,17 +131,18 @@ 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 remaining_prompt_tokens: Optional[torch.Tensor] = None latency: Optional[float] = None - timestamp_of_first_token: Optional[float] = None finished_chunk_token_count = 0 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): From f09bd99828ca1a891f010de94c35adffa151f2b0 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Sun, 19 Oct 2025 05:48:38 -0700 Subject: [PATCH 3/3] format --- .../core/inference/data_parallel_inference_coordinator.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/megatron/core/inference/data_parallel_inference_coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator.py index 6df05f7dc6c..efd93f9d293 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator.py +++ b/megatron/core/inference/data_parallel_inference_coordinator.py @@ -1,12 +1,12 @@ # 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 import torch -import time from megatron.core.inference.headers import Headers from megatron.core.inference.inference_request import DynamicInferenceRequest @@ -229,8 +229,10 @@ def postprocess( 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] + 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]