feat: Atn-dp fpm upstream integration

[indrajit96]

Overview:

Enable TRT-LLM ForwardPassMetrics under Attention-DP now that NVIDIA/TensorRT-LLM#13649 emits one IterationStats row per Attention-DP rank.
DO not merge until NVIDIA/TensorRT-LLM#13649 merges

Details:

• Remove the Dynamo-side gate that disabled FPM when attention_dp_size > 1.
• Initialize FpmDirectPublisher with dp_size=attention_dp_size so dp_rank=1..N-1 publishes are valid.
• Map TRT-LLM attentionDpRank to Dynamo FPM dp_rank.
• Preserve rank-local scheduled metrics from each TRT-LLM row.
• Preserve queued metrics only where TRT-LLM reports them, which is rank 0 for Attention-DP.
• Update comments and tests to reflect per-rank Attention-DP FPM fanout.

No FPM schema change. Planner still reads normal ForwardPassMetrics events from the forward-pass-metrics event plane.

Before this change, Dynamo disabled TRT-LLM FPM when attention_dp_size > 1, so planner did not receive TRT-LLM ForwardPassMetrics for Attention-DP workers.

After this change, planner sees one FPM stream per Attention-DP rank, keyed by the existing (worker_id, dp_rank) fields:

• scheduled fields are rank-local, so planner can detect Attention-DP load imbalance.
• queued fields are present only on dp_rank=0, because TRT-LLM's executor request queue is rank-0/global and must not be double-counted.
• wall_time is repeated on each rank row for the same forward pass.

Example decoded FPM events from H100x2 E2E validation:

{
  "version": 1,
  "worker_id": "7587894670641032969",
  "dp_rank": 0,
  "counter_id": 273,
  "wall_time": 0.03973197937011719,
  "scheduled_requests": {
    "num_prefill_requests": 2,
    "sum_prefill_tokens": 2,
    "var_prefill_length": 0.0,
    "sum_prefill_kv_tokens": 53,
    "num_decode_requests": 6,
    "sum_decode_kv_tokens": 198,
    "var_decode_kv_tokens": 0.0
  },
  "queued_requests": {
    "num_prefill_requests": 4,
    "sum_prefill_tokens": 111,
    "var_prefill_length": 0.0,
    "num_decode_requests": 0,
    "sum_decode_kv_tokens": 0,
    "var_decode_kv_tokens": 0.0
  }
}

{
  "version": 1,
  "worker_id": "7587894670641032969",
  "dp_rank": 1,
  "counter_id": 273,
  "wall_time": 0.03973197937011719,
  "scheduled_requests": {
    "num_prefill_requests": 2,
    "sum_prefill_tokens": 20,
    "var_prefill_length": 0.0,
    "sum_prefill_kv_tokens": 36,
    "num_decode_requests": 6,
    "sum_decode_kv_tokens": 193,
    "var_decode_kv_tokens": 0.0
  },
  "queued_requests": {
    "num_prefill_requests": 0,
    "sum_prefill_tokens": 0,
    "var_prefill_length": 0.0,
    "num_decode_requests": 0,
    "sum_decode_kv_tokens": 0,
    "var_decode_kv_tokens": 0.0
  }
}

Where should the reviewer start?

• components/src/dynamo/trtllm/publisher.py
• components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

• Relates to [None][feat] Emit per-rank Attention-DP iteration stats NVIDIA/TensorRT-LLM#13649

Summary by CodeRabbit

• Bug Fixes

  • Enhanced FPM publisher initialization to consistently handle all distributed parallelism configurations.

• Tests

  • Updated test suite to reflect new publisher initialization behavior.

[github-actions]

🌿 Fern Docs Preview: https://nvidia-preview-c484bc4d-69a9-4caa-9ec2-6b8cc3a8cd3e.docs.buildwithfern.com/dynamo/dev

[coderabbitai]

Walkthrough

Publisher changes from conditionally initializing FpmDirectPublisher based on attention_dp_size to unconditionally initializing it with dp_size=1 for all configurations, falling back to disabling FPM on runtime errors. The corresponding test was updated to verify this new behavior, and stat-publishing documentation was clarified.

Changes

FPM Publisher Initialization Refactoring

Layer / File(s)	Summary
Gating Logic Removal components/src/dynamo/trtllm/publisher.py	Removed self.fpm_enabled = (self.attention_dp_size == 1) conditional gating in __init__; replaced with comments describing FPM as a single rank-0-emitted logical channel under attention-DP.
Core Initialization components/src/dynamo/trtllm/publisher.py	Publisher.initialize() now unconditionally attempts FpmDirectPublisher construction with dp_size=1 instead of checking self.fpm_enabled first; on RuntimeError, logs warning and sets self.fpm_publisher to None.
Stats Documentation components/src/dynamo/trtllm/publisher.py	Updated _publish_stats_task comments to clarify that attentionDpRank is top-level in serialized stats and FPM source fields are nested under stat["inflightBatchingStats"].
Test Coverage components/src/dynamo/trtllm/tests/test_trtllm_fpm_publisher.py	Replaced test_publisher_does_not_init_fpm_publisher_under_attention_dp with test_publisher_initializes_fpm_publisher_under_attention_dp, asserting unconditional FPM initialization with dp_size == 1 even when attention_dp_size > 1.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 57.14% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'feat: Atn-dp fpm upstream integration' directly references the main change of enabling ForwardPassMetrics under Attention-DP and integrating upstream TRT-LLM improvements.
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.
Description check	✅ Passed	The pull request description fully covers all required template sections with comprehensive details about changes, objectives, and implementation rationale.

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

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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