Add enhanced event tracking with TTFT measurement and compact serialization. (#3253)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

Author: lmcafee-nvidia

examples/inference/gpt/gpt_dynamic_inference.py

@@ -195,6 +195,8 @@ def _add_request():
                 request.request_id = finished_request.request_id
                 request.events = finished_request.events
 
+                request.ttft = finished_request.ttft
+
                 # Update prompt, in case engine has been suspended and resumed.
                 request.prompt_tokens = finished_request.prompt_tokens.tolist()
                 request.prompt_text = finished_request.prompt
@@ -409,6 +411,7 @@ def escape_str(s):
                         "generated_text": req.output_text,
                         "generated_tokens": req.output_tokens,
                         "latency": req.time_end - req.time_start,
+                        "ttft": req.ttft,  # Time-to-first-token in seconds
                         "cuda_graph_request_count_map": result["cuda_graph_request_count_map"],
                         "step_count": engine.step_count,
                         "top_n_logprobs": getattr(req, 'generated_top_n_logprobs', None),

examples/inference/gpt/utils.py

@@ -69,6 +69,7 @@ def __init__(
         self.time_arrival = None
         self.time_start = None
         self.time_end = None
+        self.ttft = None  # Time-to-first-token in seconds
         self.state = "not-started"
         self.sampling_params: SamplingParams = (
             sampling_params

megatron/core/inference/config.py

@@ -170,6 +170,13 @@ class InferenceConfig:
     requests when they are paused during bookkeeping.
     """
 
+    track_generated_token_events: bool = False
+    """
+    Whether to track per-token events with timestamps for each generated token.
+    When enabled, each generated token creates a GENERATED_TOKEN event with a
+    timestamp, useful for per-token latency analysis.
+    """
+
     metrics_writer: Optional["WandbModule"] = None
     """Wandb module for writing metrics."""
 

megatron/core/inference/engines/dynamic_engine.py

@@ -30,6 +30,8 @@
 from megatron.core.inference.engines.abstract_engine import AbstractEngine
 from megatron.core.inference.headers import Headers, UnknownHeaderError
 from megatron.core.inference.inference_request import (
+    DynamicInferenceEvent,
+    DynamicInferenceEventType,
     DynamicInferenceRequest,
     DynamicInferenceRequestRecord,
     Status,
@@ -173,6 +175,7 @@ def __init__(self, controller: TextGenerationController, context: DynamicInferen
         self.controller = controller
         self.context = context
         self.track_paused_request_events = inference_config.track_paused_request_events
+        self.track_generated_token_events = inference_config.track_generated_token_events
         self.enable_chunked_prefill = inference_config.enable_chunked_prefill
         self.metrics_writer = inference_config.metrics_writer
         self.logging_step_interval = inference_config.logging_step_interval
@@ -710,6 +713,7 @@ def _add_request(
                 record=DynamicInferenceRequestRecord.from_request(request),
                 future=self._loop.create_future(),
             )
+            request.add_event_add_engine()  # Record when request enters engine
 
         if request.status is None:
             request.status = Status.ACTIVE_AND_GENERATING_TOKENS
@@ -882,7 +886,21 @@ def post_process_requests(
                 # Skip appending token for requests being finished due to stop words
                 # (they already have their final token from the previous step)
                 if request_id not in self.stop_word_being_finished_ids:
+                    is_first_token = len(request.generated_tokens) == 0
                     request.generated_tokens.append(token)
+                    if self.track_generated_token_events:
+                        event_generated_token = request.add_event_generated_token(token)
+                    if is_first_token:
+                        if self.track_generated_token_events:
+                            first_token_event = event_generated_token
+                        else:
+                            first_token_event = DynamicInferenceEvent(
+                                type=DynamicInferenceEventType.GENERATED_TOKEN,
+                                payload={"token_id": token},
+                            )
+                        request.ttft = (
+                            first_token_event.timestamp - request.event_add_engine.timestamp
+                        )
                     if request.tpot is None:
                         request.tpot = []
                     request.tpot.append(step_time)
@@ -894,6 +912,7 @@ def post_process_requests(
                     # Request finished by normal means (termination_id, max_length, or stop word from previous step)
                     request.generated_length = len(request.generated_tokens)
                     request.status = Status.COMPLETED
+                    request.add_event_finish()
                     finished_entry = self.requests.pop(request_id)
                     finished_request = finished_entry.record[-1]
                     finished_request.generated_length = len(finished_request.generated_tokens)
@@ -1102,7 +1121,7 @@ def schedule_non_chunked_prefill(self):
                     self._loop.create_task, self._notify_cond_for_new_request()
                 )
                 req.remaining_prompt_tokens = req.remaining_prompt_tokens.new_empty(0)
-                req.add_event_add()
+                req.add_event_add_context()
                 self.waiting_request_ids.popleft()
             else:
                 break
@@ -1148,7 +1167,7 @@ def schedule_chunked_prefill(self):
                         self._loop.create_task, self._notify_cond_for_new_request()
                     )
                     req.remaining_prompt_tokens = req.remaining_prompt_tokens.new_empty(0)
-                    req.add_event_add()
+                    req.add_event_add_context()
                     # Fully scheduled, so we remove from waiting pool
                     self.waiting_request_ids.popleft()
                     # Only this case we keep checking the rest of the waiting queue
@@ -1274,9 +1293,7 @@ async def async_bookkeep(
                 newly_paused_request_ids = newly_paused_request_ids.tolist()
                 [self.get_request(i).add_event_pause() for i in newly_paused_request_ids]
 
-            # Mark requests finished.
-            [self.get_request(i).add_event_finish() for i in finished_request_ids.tolist()]
-            # Add finished events.
+            # Process finished requests (adds FINISH events and returns records).
             (active_request_ids, finished_request_records) = self.post_process_requests(
                 active_request_ids,
                 finished_request_ids,

megatron/core/inference/inference_request.py

@@ -159,7 +159,9 @@ def _post_deserialize(self, obj: dict):
 class DynamicInferenceEventType(Enum):
     """Dynamic inference event type."""
 
-    ADD = auto()
+    ADD_ENGINE = auto()  # When request is added to engine via _add_request()
+    ADD_CONTEXT = auto()  # When request is added to context (scheduled for prefill)
+    GENERATED_TOKEN = auto()  # When an output token is generated (payload = {"token_id": int})
     PAUSE = auto()
     EVICT = auto()
     FINISH = auto()
@@ -202,33 +204,46 @@ def __post_init__(self):
             DynamicInferenceEventType.ERROR_NONTRANSIENT,
         ):
             assert self.payload is not None
+        elif self.type == DynamicInferenceEventType.GENERATED_TOKEN:
+            assert (
+                self.payload is not None
+                and isinstance(self.payload, dict)
+                and "token_id" in self.payload
+            )
         else:
             assert self.payload is None
 
     def __str__(self):
-        payload_str = "" if self.payload is None else f", {type(self.payload).__name__}"
+        if self.type == DynamicInferenceEventType.GENERATED_TOKEN:
+            payload_str = f", token={self.payload['token_id']}"
+        elif self.payload is None:
+            payload_str = ""
+        else:
+            payload_str = f", {type(self.payload).__name__}"
         return f"[{self.timestamp:.3f}] {self.type.name}{payload_str}"
 
     def serialize(self) -> dict:
         """Converts the instance into a serializable dictionary.
 
         Returns:
-            (dict) A dictionary representation of the instance suitable for
-                serialization.
+            dict: Full event dict.
         """
-
-        # Dataclass to dict.
         torch.cuda.nvtx.range_push("DynamicInferenceEvent.serialize")
         # do not use asdict(self) - it has very high CPU overheads
         # and if there are tensors, it will try to deepcopy them
         obj = self.__dict__.copy()
         obj["type"] = self.type.name
 
         # Serialize payload.
-        if self.payload:
-            from .contexts.dynamic_context import ContextErrorFactory  # avoid circular import.
+        if self.payload is not None:
+            if self.type in (
+                DynamicInferenceEventType.ERROR_TRANSIENT,
+                DynamicInferenceEventType.ERROR_NONTRANSIENT,
+            ):
+                from .contexts.dynamic_context import ContextErrorFactory  # avoid circular import.
+
+                obj["payload"] = ContextErrorFactory.serialize(self.payload)
 
-            obj["payload"] = ContextErrorFactory.serialize(self.payload)
         torch.cuda.nvtx.range_pop()
         return obj
 
@@ -237,22 +252,25 @@ def deserialize(cls, obj: dict) -> "DynamicInferenceEvent":
         """Deserialize event.
 
         Args:
-            obj (dict): Serialized event data.
+            obj: Serialized event data dict.
 
         Returns:
             (DynamicInferenceEvent) Deserialized event.
         """
+        event_type = DynamicInferenceEventType[obj["type"]]
 
-        # Initialize event.
-        event = cls(**{**obj, "type": DynamicInferenceEventType[obj["type"]]})
+        # Pre-process payload before construction (since __post_init__ validates types).
+        init_obj = {**obj, "type": event_type}
+        if obj["payload"] is not None:
+            if event_type in (
+                DynamicInferenceEventType.ERROR_TRANSIENT,
+                DynamicInferenceEventType.ERROR_NONTRANSIENT,
+            ):
+                from .contexts.dynamic_context import ContextErrorFactory  # avoid circular import.
 
-        # Deserialize payload.
-        if obj["payload"]:
-            from .contexts.dynamic_context import ContextErrorFactory  # avoid circular import.
+                init_obj["payload"] = ContextErrorFactory.deserialize(obj["payload"])
 
-            event.payload = ContextErrorFactory.deserialize(obj["payload"])
-
-        return event
+        return cls(**init_obj)
 
 
 @experimental_api
@@ -265,7 +283,6 @@ class DynamicInferenceRequest(InferenceRequest):
     """
 
     request_id: int
-    generated_tokens: List[int] = field(default_factory=list)
     prompt: Optional[str] = None
     prompt_tokens: Optional[torch.Tensor] = None
     # remaining prompt tokens are used for chunked prefill
@@ -289,7 +306,10 @@ def remaining_prompt_length(self):
         """
         return len(self.remaining_prompt_tokens)
 
+    ttft: Optional[float] = None
     events: List[DynamicInferenceEvent] = field(default_factory=list)
+    event_add_engine: Optional[DynamicInferenceEvent] = field(default=None, repr=False)
+    generated_tokens: List[int] = field(default_factory=list)
 
     def __str__(self):
         return ", ".join(
@@ -302,7 +322,7 @@ def __str__(self):
             )
         )
 
-    def serialize(self) -> dict:
+    def serialize(self):
         """Converts the instance into a serializable dictionary.
 
         Returns:
@@ -312,6 +332,7 @@ def serialize(self) -> dict:
         torch.cuda.nvtx.range_push("DynamicInferenceRequest.serialize")
         obj = super().serialize()
         obj["events"] = [e.serialize() for e in self.events]
+        obj.pop("event_add_engine", None)
 
         # Sanity check routing_indices: Tensor [total_tokens - 1, num_layers, topk]
         if self.routing_indices is not None:
@@ -328,7 +349,7 @@ def serialize(self) -> dict:
 
     def _post_deserialize(self, obj):
         super()._post_deserialize(obj)
-        self.events = [DynamicInferenceEvent.deserialize(e) for e in obj["events"]]
+        self.events = [DynamicInferenceEvent.deserialize(e) for e in obj.get("events", [])]
 
     @property
     def tracked_metadata(self) -> List[Any]:
@@ -370,13 +391,30 @@ def get_metadata_types() -> List[Tuple[str, torch.dtype, bool]]:
             ("top_n_logprobs", torch.int32, False),  # CPU for torch sampling
         ]
 
-    def add_event(self, type: DynamicInferenceEventType, payload: Optional[Any] = None) -> None:
+    def add_event(
+        self, type: DynamicInferenceEventType, payload: Optional[Any] = None
+    ) -> DynamicInferenceEvent:
         """Add event."""
-        self.events.append(DynamicInferenceEvent(type=type, payload=payload))
+        event = DynamicInferenceEvent(type=type, payload=payload)
+        self.events.append(event)
+        return event
+
+    def add_event_add_engine(self):
+        """Add 'add_engine' event - called when request enters the engine queue."""
+        self.event_add_engine = self.add_event(DynamicInferenceEventType.ADD_ENGINE)
+        return self.event_add_engine
 
-    def add_event_add(self):
-        """Add 'add' event."""
-        return self.add_event(DynamicInferenceEventType.ADD)
+    def add_event_add_context(self):
+        """Add 'add_context' event - called when request is added to context for prefill."""
+        return self.add_event(DynamicInferenceEventType.ADD_CONTEXT)
+
+    def add_event_generated_token(self, token: int):
+        """Add 'generated_token' event - records each generated token.
+
+        Args:
+            token (int): The token ID that was generated.
+        """
+        return self.add_event(DynamicInferenceEventType.GENERATED_TOKEN, {"token_id": token})
 
     def add_event_pause(self):
         """Add 'pause' event."""
@@ -535,6 +573,7 @@ def merge_lists(key):
             generated_log_probs=merge_lists("generated_log_probs"),
             generated_top_n_logprobs=merge_lists("generated_top_n_logprobs"),
             sampling_params=self.requests[0].sampling_params,
+            ttft=self.requests[0].ttft,
             tpot=merge_lists("tpot"),
             status=self.requests[-1].status,
             latency=self.latency,

megatron/training/arguments.py

@@ -1593,6 +1593,11 @@ def _add_inference_args(parser):
                        help='Track paused request ids by adding \'paused\' events '
                        'to each request\'s event history. This has a very minor '
                        'impact on latency.')
+    group.add_argument('--inference-dynamic-batching-track-generated-token-events',
+                       action='store_true',
+                       help='Track per-token events with timestamps for each generated token. '
+                       'When enabled, each generated token creates a GENERATED_TOKEN event '
+                       'with a timestamp, useful for per-token latency analysis.')
     group.add_argument('--decode-only-cuda-graphs',
                        action='store_true', default=False,
                        help='Only use cuda graphs for decode-only steps, not prefill and mixed steps.')