diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index fcee2c1daef..c9e14f824e2 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -1204,8 +1204,6 @@ def post_process_requests( # so gate the update to keep the metric a truthful sparse # sample instead of polluting it with zeros. if step_time > 0: - if request.tpot is None: - request.tpot = [] per_token_step_time = step_time / len(tokens) request.tpot.extend([per_token_step_time] * len(tokens)) diff --git a/megatron/core/inference/inference_request.py b/megatron/core/inference/inference_request.py index 33fbcdf6518..d6e7c67a959 100644 --- a/megatron/core/inference/inference_request.py +++ b/megatron/core/inference/inference_request.py @@ -155,7 +155,7 @@ class InferenceRequest: prompt_top_n_logprobs: Optional[List[Dict[str, float]]] = None generated_top_n_logprobs: Optional[List[Dict[str, float]]] = None generated_length: Optional[int] = None - tpot: Optional[List[int]] = None + tpot: List[float] = field(default_factory=list) def __post_init__(self): if self.sampling_params is None and self.inference_parameters is not None: diff --git a/tests/unit_tests/inference/engines/test_dynamic_events.py b/tests/unit_tests/inference/engines/test_dynamic_events.py index a991dd5329c..79445f0255f 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_events.py +++ b/tests/unit_tests/inference/engines/test_dynamic_events.py @@ -416,6 +416,53 @@ def test_merge_preserves_precomputed_block_hashes(): mock_compute.assert_not_called() +# ============================================================================ +# Test 5c: Merge tolerates sparse tpot across checkpoints +# ============================================================================ + + +def test_merge_with_unpopulated_tpot_in_later_segment(): + """Regression test: merge() must not crash when a post-checkpoint sub-request + never observed a logging step and therefore has an empty tpot list while the + pre-checkpoint sub-request has populated entries. + + The dynamic engine populates `tpot` lazily, only on logging steps + (`step_time > 0`). A request that gets evicted then finishes within fewer + than `logging_step_interval` decode steps after recompute will produce a + record whose later sub-request never accumulated any tpot samples, while + earlier sub-requests did. `merge()` must concatenate what's there without + erroring on the empty segment. + """ + # Pre-checkpoint sub-request: lived through a logging step, so tpot has data. + req1 = DynamicInferenceRequest( + request_id=1, + prompt_tokens=torch.tensor([1, 2, 3], dtype=torch.int64), + sampling_params=SamplingParams(num_tokens_to_generate=10), + ) + req1.generated_tokens.extend([100, 101]) + req1.tpot = [0.10, 0.11] + + # Post-checkpoint sub-request: recomputed but never saw a logging step + # before finishing, so tpot is left at its default. + req2 = DynamicInferenceRequest( + request_id=1, + prompt_tokens=torch.tensor([1, 2, 3, 100, 101], dtype=torch.int64), + sampling_params=SamplingParams(num_tokens_to_generate=8), + ) + req2.generated_tokens.extend([200]) + + record = DynamicInferenceRequestRecord() + record.requests.append(req1) + record.requests.append(req2) + + merged = record.merge() + + # The pre-checkpoint segment's measurements must survive the merge, + # and the empty post-checkpoint segment must contribute nothing. + assert merged.tpot == [0.10, 0.11] + assert merged.generated_tokens == [100, 101, 200] + + # ============================================================================ # Test 6: TTFT Calculation from Event Timestamps # ============================================================================