Add enhanced event tracking with TTFT measurement and compact serialization. by lmcafee-nvidia · Pull Request #3253 · NVIDIA/Megatron-LM

lmcafee-nvidia · 2026-02-04T19:43:05Z

What does this PR do ?

Summary

This PR enhances the inference event tracking system with precise time-to-first-token (TTFT)
measurement, per-token event tracking, and configurable serialization modes.

Key improvements:

Precise TTFT tracking: Split request lifecycle into ADD_ENGINE and ADD_CONTEXT events to
enable accurate TTFT measurement from request submission to first token generation
Per-token event tracking: New GENERATED_TOKEN events capture each token as it's generated with
individual timestamps
Compact serialization mode: Add --inference-dynamic-batching-track-generated-token-events
flag to control whether token events serialize as full dicts (with timestamps) or just token IDs,
reducing overhead when timestamps aren't needed
Comprehensive test suite: Redesigned event tracking tests with 10 comprehensive test classes
covering event ordering, timestamps, serialization modes, and edge cases

Changes:

Event types: ADD_ENGINE, ADD_CONTEXT, GENERATED_TOKEN, FINISH
TTFT calculation: first_GENERATED_TOKEN.timestamp - ADD_ENGINE.timestamp
CLI flags: --inference-dynamic-batching-track-generated-token-events for verbose token event
tracking
Test coverage: New integration tests for event timestamps and serialization modes

Bug fixes:

Fixed FINISH event ordering in dynamic engine
Fixed DynamicInferenceRequest instantiation bug

Test plan

All existing inference functional tests pass
New unit tests for event tracking (10 comprehensive test classes)
Integration tests for event timestamps
TTFT measurement validated in gpt_dynamic_inference.py output
Compact vs. full serialization modes tested

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

Replace the single ADD event with three separate events to enable precise time-to-first-token measurement: - ADD_ENGINE: When request is added to engine via _add_request() - ADD_CONTEXT: When request is scheduled for prefill - FIRST_TOKEN: When first output token is about to be generated TTFT is now calculated as FIRST_TOKEN - ADD_ENGINE and included in the JSON output from gpt_dynamic_inference.py. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

- Add tests/unit_tests/inference/test_utils.py with TestPriority enum for selective test execution based on priority levels - Add tests/unit_tests/inference/engines/test_dynamic_events.py with comprehensive tests for DynamicInferenceEvent and event lifecycle - Use consistent convention: CRITICAL=1, LOW=4 with skipif pattern TEST_PRIORITY < TestPriority.LEVEL Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

copy-pr-bot · 2026-02-04T19:43:09Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

- Update test_events to use new event types (ADD_ENGINE, ADD_CONTEXT, FIRST_TOKEN) instead of old ADD type - Add test_event_timestamps integration test that verifies: - Completed requests have expected event sequence - Event timestamps are monotonically increasing - TTFT (FIRST_TOKEN - ADD_ENGINE) is positive - Total request time >= TTFT Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

examples/inference/gpt/gpt_dynamic_inference.py

tests/unit_tests/inference/test_utils.py

tests/unit_tests/inference/engines/test_dynamic_events.py

Refactor event tracking to use a unified GENERATED_TOKEN event that captures each token as it's generated, replacing the separate FIRST_TOKEN event and generated_tokens list. Changes: - Replace FIRST_TOKEN enum with GENERATED_TOKEN (payload = token id) - Remove generated_tokens field, add property that extracts from events - Replace add_event_first_token() with add_event_generated_token(token) - Update TTFT calculation to use first GENERATED_TOKEN timestamp - Update serialization to handle GENERATED_TOKEN payload - Update tests for new event structure Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Complete rewrite of test_dynamic_events.py with 81 tests across 12 test classes covering all 9 event types, payload validation, serialization, and lifecycle sequences. Test classes: - TestDynamicInferenceEventType: Verifies all 9 event types exist - TestEventCreation: Tests creating events with correct payloads - TestPayloadValidation: Tests payload validation rules - TestEventSerialization: Tests serialize/deserialize roundtrip - TestErrorPayloads: Tests error types as event payloads - TestRequestEventMethods: Tests all add_event_*() helper methods - TestGeneratedTokensProperty: Tests generated_tokens property - TestRequestSerialization: Tests request serialization with events - TestRequestRecordMerge: Tests DynamicInferenceRequestRecord - TestEventLifecycleSequences: Tests common event sequences - TestEventTimestamps: Tests timestamp behavior and TTFT - TestEventEdgeCases: Tests __str__, rapid events, many tokens Note: Includes workaround for DynamicInferenceRequest instantiation bug where generated_tokens property conflicts with inherited field. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Add no-op setter for generated_tokens property to allow dataclass __init__ to work. The parent class InferenceRequest has generated_tokens as a field, which means the dataclass-generated __init__ tries to set it. Without a setter, this fails with "property has no setter" error. The setter ignores any value passed since generated_tokens is computed dynamically from GENERATED_TOKEN events. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Remove workaround helper function and use DynamicInferenceRequest directly now that the generated_tokens property bug is fixed. Also restore the original merge() tests that now work correctly. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Move add_event_finish() call inside post_process_requests() so it happens after the final GENERATED_TOKEN event is added, not before. This ensures the event sequence is: ADD_ENGINE -> ADD_CONTEXT -> GENERATED_TOKEN(s) -> FINISH Previously, FINISH was added before post_process_requests() was called, which meant FINISH came before the last generated token. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Add --inference-dynamic-batching-track-generated-token-events CLI flag to control whether GENERATED_TOKEN events serialize as full dicts with timestamps or just as integer token IDs. When disabled (default), reduces serialization overhead by omitting timestamp data for token events. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Replace 89-test class-based suite with 10 substantial standalone tests that maximize coverage while reducing complexity. Remove TestPriority system. Each test covers multiple related behaviors in meaningful scenarios including all 9 event types, serialization, TTFT calculation, error handling, and complex multi-checkpoint lifecycles. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Resolve conflicts keeping event tracking features (track_generated_token_events, compact serialization, TTFT) while adopting main's improvements (NVTX profiling, _post_deserialize pattern, dict copy optimization). Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

No test files import TestPriority, so the module is dead code. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

megatron/core/inference/inference_request.py

examples/inference/gpt/gpt_dynamic_inference.py

megatron/training/arguments.py

tests/unit_tests/inference/engines/test_dynamic_engine.py

Address PR review: generated_tokens as a property computed from events was O(n) on every access (called every forward step). Restore it as a plain List[int] field, keep events as supplementary timestamp data, and have the track_generated_token_events flag control whether GENERATED_TOKEN events are created at all (rather than controlling serialization format). Also removes compact dict|int serialization mode entirely and fixes a missing nvtx range_pop() in DynamicInferenceEvent.serialize(). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

tdene · 2026-02-12T00:55:37Z

megatron/core/inference/engines/dynamic_engine.py

+        self.track_generated_token_events = getattr(
+            inference_config, 'track_generated_token_events', False
+        )


nit: I don't believe the getattr part is needed, since it's being added to inference_config.py

lmcafee-nvidia · 2026-02-12T01:29:58Z

@shanmugamr1992 @Phlip79 , please review when you have time. thank you!

Phlip79

Great tests!

Also address PR review: replace getattr with direct attribute access for track_generated_token_events.

lmcafee-nvidia · 2026-02-12T17:28:43Z

/ok to test 76d49cb

ko3n1g · 2026-02-12T18:15:14Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/21958742266

…zation. (NVIDIA#3253) Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

lmcafee-nvidia and others added 2 commits February 4, 2026 11:22

lmcafee-nvidia added this to the Core 0.15 milestone Feb 4, 2026

lmcafee-nvidia self-assigned this Feb 4, 2026

lmcafee-nvidia requested a review from a team as a code owner February 4, 2026 19:43

lmcafee-nvidia added the Expert Review Apply this label to indicate that your PR is ready for expert review. label Feb 4, 2026

lmcafee-nvidia requested review from a team as code owners February 4, 2026 19:43

ko3n1g requested a review from a team February 4, 2026 19:43

kvareddy approved these changes Feb 5, 2026

View reviewed changes