Fix PyExecutor FPM iteration timing

[tedzhouhk]

Summary

• Add batch-matched gpuForwardTimeMS for PyExecutor iteration stats using CUDA events around the batch _forward_step.
• Snapshot scheduler/FPM counters before forward mutates request state, so inflightBatchingStats stays matched to the exported iteration.
• Pool FPM-owned CUDA timing events for PP and iter-stats-only modes; when PerfMetricsManager provides pooled events, this code reuses those instead.
• Fix the iter-stats regression helper so explicit pre-forward iter_states snapshots are actually passed into _update_iter_stats.

Validation

• python -m py_compile tensorrt_llm/_torch/pyexecutor/py_executor.py tests/unittest/pyexecutor/test_iter_stats_populate.py tensorrt_llm/_torch/pyexecutor/adp_iter_stats.py tensorrt_llm/executor/base_worker.py
• Downstream Dynamo FPM publisher test: 29 passed in 0.11s
• Local A6000 load test with Dynamo + TRT-LLM PyTorch backend, Qwen/Qwen3-0.6B, FPM enabled:
  • aiperf concurrency=50, ISL=5000, OSL=500, requests=200
  • active FPM records after trimming idle margins: 2049
  • idle records inside trimmed window: 0
  • active gpuForwardTimeMS wall-time p50=40.8 ms, p95=71.7 ms, max=159.2 ms

TRT-LLM pytest was not runnable in my local checkout because the source tree lacks generated tensorrt_llm.bindings and the available TRT-LLM venv does not include pytest; the modified files were syntax-checked instead.

Summary by CodeRabbit

• New Features

  • Added GPU forward pass timing metric to iteration statistics, providing granular performance insights alongside existing host and device timing measurements.

• Bug Fixes

  • Enhanced iteration-state snapshots to ensure accurate scheduler metrics and timing data capture across distributed execution.

• Tests

  • Updated test coverage to validate new timing metric propagation and scheduler state snapshot behavior.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR extends iteration statistics collection to capture GPU forward timing via CUDA events and snapshot per-iteration scheduler state before request mutation. The metric propagates through ADP fanout and direct serialization paths across all executor modes (pipeline-parallelism, non-PP, overlap), with corresponding test updates validating the new timing field and iter_states-driven stat override behavior.

Changes

GPU Forward Timing Metrics in Iteration Statistics

Layer / File(s)	Summary
ADP Iteration Stats Record Extension tensorrt_llm/_torch/pyexecutor/adp_iter_stats.py	ADPIterStatsRecord gains optional gpu_forward_time_ms field. ADPIterStatsBuffer stores rank-0 values in a per-iteration dictionary, accepts the metric in queue(...), removes it in _discard(...), and includes it when finalizing records sent to all ranks.
PyExecutor GPU Timing Infrastructure tensorrt_llm/_torch/pyexecutor/py_executor.py	BatchState adds CUDA event handles (start/end), iteration-state snapshots dict, and pooled-event ownership flag. PyExecutor introduces event pooling for reuse, helper methods to detect dummy requests, snapshot scheduler state before forward mutation, and compute GPU forward elapsed time from CUDA events.
Iteration Stats Update with iter_states Override tensorrt_llm/_torch/pyexecutor/py_executor.py	_update_iter_stats accepts optional iter_states dict to override inflight batching counters instead of reading live state. _populate_req_stats refactors dummy-request filtering and sources context/gen/paused/token counters and KV amounts from pre-forward snapshots when available.
GPU Timing Computation and Event Release tensorrt_llm/_torch/pyexecutor/py_executor.py	_append_iter_stats signature extends to include gpu_forward_time_ms. _process_iter_stats computes elapsed time from CUDA events with exception handling and releases pooled events in a try-finally block to prevent leaks.
Stats Serialization and Export Paths tensorrt_llm/_torch/pyexecutor/py_executor.py, tensorrt_llm/executor/base_worker.py	Metric plumbing routes gpu_forward_time_ms into both ADP queue and legacy export. ADP path conditionally injects gpuForwardTimeMS JSON field. Legacy path includes metric in 8-tuple. BaseWorker._stats_serializer extracts and conditionally emits the new field.
Executor Loop Integration Across Variants tensorrt_llm/_torch/pyexecutor/py_executor.py	Pipeline-parallelism, non-PP, and overlap executor loops each snapshot iter_states, create/pool CUDA timing events, pass events through forward (with perf recording), construct BatchState/BatchStatePP with timing and snapshots, and invoke _process_iter_stats. _forward_step_inter_pp accepts optional timing events and wraps forward with perf recording.
Test Coverage for Timing Metrics and iter_states tests/unittest/executor/test_stats_serializer.py, tests/unittest/pyexecutor/test_iter_stats_populate.py	Stats serializer test validates 8-tuple handling and gpuForwardTimeMS JSON serialization. Iter stats test harness extends to accept explicit iter_states snapshots. New regression tests verify num_ctx_tokens and num_gen_kv_tokens source from batch snapshots. ADP fanout tests assert timing fields propagate across all ranks.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is missing required template sections: no PR title following template format, no explicit 'Test Coverage' section, and the PR checklist is incomplete with no items checked.	Add PR title with ticket/issue format (e.g., [TRTLLM-XXXX][fix]), complete 'Test Coverage' section listing relevant tests, and check all applicable items in the PR Checklist.
Docstring Coverage	⚠️ Warning	Docstring coverage is 40.63% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Fix PyExecutor FPM iteration timing' clearly and concisely describes the primary change: fixing iteration timing measurement for PyExecutor's FPM (ForwardPassMetrics) system.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[tedzhouhk]

/bot run

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (2)

tests/unittest/pyexecutor/test_iter_stats_populate.py (1)

550-576: Coverage assessment: sufficient for this PR’s iter-state/timing regressions.

tests/unittest/pyexecutor/test_iter_stats_populate.py now covers both snapshot-sourced counters (num_ctx_tokens, num_gen_kv_tokens) and ADP fanout timing-field propagation. For this PR objective, coverage is sufficient.
Follow-up outside this PR (optional): add one negative-path test where timing fields are absent end-to-end in fanout output.

Also applies to: 693-720

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/pyexecutor/test_iter_stats_populate.py` around lines 550 -
576, Add a negative-path unit test in
tests/unittest/pyexecutor/test_iter_stats_populate.py that verifies ADP fanout
timing fields are absent when not present in the inputs: use
_invoke_update_iter_stats with a _StubScheduledBatch (e.g., gen_reqs or
context_reqs as appropriate) and provide iter_states that omit timing fields,
then inspect the returned stats.inflight_batching_stats (or the fanout output)
and assert the timing fields are not present or are None/zero as expected; place
the test alongside test_num_ctx_tokens_... and test_num_gen_kv_tokens_... and
name it to reflect the "timing fields absent" negative path.

tests/unittest/executor/test_stats_serializer.py (1)

216-240: Coverage assessment: sufficient for serializer tuple-shape evolution.

tests/unittest/executor/test_stats_serializer.py now explicitly validates the 8-tuple contract (hostStepTimeMS, omitted prevDeviceStepTimeMS on None, schedulerMode, gpuForwardTimeMS). For this PR scope, coverage in this file is sufficient.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/executor/test_stats_serializer.py` around lines 216 - 240, The
test `test_serializer_8_tuple_emits_new_timing_and_scheduler_mode` already
covers the new 8-tuple serializer contract for BaseWorker._stats_serializer by
asserting hostStepTimeMS == 12.5, that prevDeviceStepTimeMS is omitted when
None, schedulerMode == "overlap", and gpuForwardTimeMS == 4.25; no code changes
are required—keep the existing test as-is and ensure
BaseWorker._stats_serializer continues to produce the documented 8-tuple JSON
output.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 3022-3033: In the branch where self.enable_iter_perf_stats is true
but sample_state is None, ensure any timing events borrowed from
_iter_timing_event_pool are returned to the pool instead of being leaked;
specifically, when skipping the call to _process_iter_stats (and after a
_forward_step() that may have returned None), release/return
gpu_forward_start_event, gpu_forward_end_event and any entries in
gpu_forward_events_from_iter_pool back into _iter_timing_event_pool (or call the
pool's release method) so borrowed CUDA events are not leaked.

---

Nitpick comments:
In `@tests/unittest/executor/test_stats_serializer.py`:
- Around line 216-240: The test
`test_serializer_8_tuple_emits_new_timing_and_scheduler_mode` already covers the
new 8-tuple serializer contract for BaseWorker._stats_serializer by asserting
hostStepTimeMS == 12.5, that prevDeviceStepTimeMS is omitted when None,
schedulerMode == "overlap", and gpuForwardTimeMS == 4.25; no code changes are
required—keep the existing test as-is and ensure BaseWorker._stats_serializer
continues to produce the documented 8-tuple JSON output.

In `@tests/unittest/pyexecutor/test_iter_stats_populate.py`:
- Around line 550-576: Add a negative-path unit test in
tests/unittest/pyexecutor/test_iter_stats_populate.py that verifies ADP fanout
timing fields are absent when not present in the inputs: use
_invoke_update_iter_stats with a _StubScheduledBatch (e.g., gen_reqs or
context_reqs as appropriate) and provide iter_states that omit timing fields,
then inspect the returned stats.inflight_batching_stats (or the fanout output)
and assert the timing fields are not present or are None/zero as expected; place
the test alongside test_num_ctx_tokens_... and test_num_gen_kv_tokens_... and
name it to reflect the "timing fields absent" negative path.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: cf2d0457-0914-4c9e-a8dc-71c25195f112

📥 Commits

Reviewing files that changed from the base of the PR and between 5f64e7d and 163eb86.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/pyexecutor/adp_iter_stats.py
• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tensorrt_llm/executor/base_worker.py
• tests/unittest/executor/test_stats_serializer.py
• tests/unittest/pyexecutor/test_iter_stats_populate.py

[hchings]

TRT-LLM pytest was not runnable in my local checkout because the source tree lacks generated tensorrt_llm.bindings and the available TRT-LLM venv does not include pytest; the modified files were syntax-checked instead.

Hi @tedzhouhk, thanks for the contribution! In the future, it'd still be better to run new tests locally first, as it'll take TRTLLM's CI quite a while to get to the new tests which might delay time-sensitive MRs. For tensorrt_llm.bindings, it'll need to compile the TRTLLM source following https://nvidia.github.io/TensorRT-LLM/installation/build-from-source.html#build-from-source.

[tedzhouhk]

/bot run

[achartier]

/bot run

[tensorrt-cicd]

PR_Github #52450 [ run ] triggered by Bot. Commit: e5f9d41 Link to invocation

[tensorrt-cicd]

PR_Github #52450 [ run ] completed with state FAILURE. Commit: e5f9d41
/LLM/main/L0_MergeRequest_PR pipeline #41743 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[achartier]

/bot run

[tensorrt-cicd]

PR_Github #52453 [ run ] triggered by Bot. Commit: 00ec987 Link to invocation

[tensorrt-cicd]

PR_Github #52453 [ run ] completed with state FAILURE. Commit: 00ec987
/LLM/main/L0_MergeRequest_PR pipeline #41746 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[hchings]

/bot run

[tensorrt-cicd]

PR_Github #52819 [ run ] triggered by Bot. Commit: 6acf289 Link to invocation

[tensorrt-cicd]

PR_Github #52819 [ run ] completed with state SUCCESS. Commit: 6acf289
/LLM/main/L0_MergeRequest_PR pipeline #42073 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[hchings]

/bot run

[tensorrt-cicd]

PR_Github #53126 [ run ] triggered by Bot. Commit: 6acf289 Link to invocation

[tensorrt-cicd]

PR_Github #53126 [ run ] completed with state SUCCESS. Commit: 6acf289
/LLM/main/L0_MergeRequest_PR pipeline #42332 completed with status: 'SUCCESS'

CI Report

Link to invocation