[None][fix] Fix disagg gen-consensus deadlock under pipeline parallelism

[Shixiaowei02]

Description

This pull request introduces important changes to the pipeline and parallelism (PP/TP) consensus and scheduling logic for disaggregated generation in the distributed execution engine. The main focus is on improving synchronization and preventing deadlocks in collective operations by ensuring all pipeline parallel (PP) ranks make consistent decisions about entering or skipping group collectives, particularly during disaggregated generation KV cache transfers.

The most important changes are:

Consensus and Synchronization Improvements:

• Modified _gen_consensus and _gen_consensus_outcome in transceiver.py to use a two-stage tensor-parallel (TP) then pipeline-parallel (PP) consensus instead of a single world allgather when attention data parallelism is disabled, preventing deadlocks caused by cross-communicator collectives. [1] [2]

KV Cache Transfer Coordination:

• Added a disagg_gen_transfer_in_progress flag to SerializableSchedulerOutput in scheduler.py to propagate the scheduling rank's transfer status to all PP ranks via the existing schedule broadcast, ensuring all ranks enter or skip the transfer-status allgather together. [1] [2]
• Updated _pp_schedule_and_propagate in py_executor.py to set and propagate the disagg_gen_transfer_in_progress flag, and to adopt this flag on every PP rank for consistent collective participation. [1] [2]
• Modified _check_disagg_gen_transfer_status in py_executor.py to gate the group allgather on the propagated flag, avoiding divergent per-rank decisions that could lead to deadlock.

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Summary by CodeRabbit

• Bug Fixes
  • Improved coordination during distributed generation so all pipeline stages stay in sync when KV transfer is active.
  • Reduced unnecessary synchronization overhead when attention data parallelism is disabled.
  • Made generation consensus handling more consistent across different distributed execution paths.

[Shixiaowei02]

/bot run --add-multi-gpu-test --disable-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

Fixes two disaggregated-generation collective communication issues: adds a disagg_gen_transfer_in_progress boolean to SerializableSchedulerOutput, propagates it from the scheduling rank to all PP ranks, and uses it to gate the disagg-gen transfer status allgather uniformly. Separately, KvCacheTransceiverV2 consensus and outcome methods now branch on enable_attention_dp, routing to TP-then-PP staged allgathers instead of the WORLD-based gen-stage path when attention DP is disabled.

Changes

Disaggregated generation deadlock and consensus fixes

Layer / File(s)	Summary
PP disagg-gen transfer flag propagation tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py, tensorrt_llm/_torch/pyexecutor/py_executor.py	SerializableSchedulerOutput gains disagg_gen_transfer_in_progress: bool = False; from_scheduler_result accepts and forwards it. The scheduling rank computes this flag from active requests with is_disagg_generation_transmission_in_progress and passes it into the broadcast output. All PP ranks read the propagated value into self._pp_disagg_gen_transfer_in_progress, which then overrides the local need_check in _check_disagg_gen_transfer_status to ensure all ranks enter or skip the status allgather in lockstep.
Non-ADP consensus path in KvCacheTransceiverV2 tensorrt_llm/_torch/disaggregation/transceiver.py	_gen_consensus returns the existing _ctx_consensus (TP/PP two-stage) result directly when enable_attention_dp is false, using pp_size as the sync size when true. _gen_consensus_outcome adds a TP-allgather-then-optional-PP-allgather path for non-ADP mode, bypassing the gen-stage allgather and gen_need_sync path.

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 20.00% 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 matches the template and clearly summarizes the main fix to disagg gen-consensus deadlock under PP.
Description check	✅ Passed	The description follows the template and explains the issue and fix well, but the Test Coverage section is left empty.
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.}

[tensorrt-cicd]

PR_Github #55516 [ run ] triggered by Bot. Commit: c493c61 Link to invocation

[coderabbitai]

Actionable comments posted: 2

🤖 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/disaggregation/transceiver.py`:
- Around line 363-365: In transceiver.py, the call site in the consensus flow is
using getattr(self._dist, "pp_allgather"), which triggers Ruff B009 and is
unnecessary. Update the _consensus_outcome invocation to access
self._dist.pp_allgather directly, keeping the same arguments and behavior while
removing the dynamic attribute lookup.

In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 1934-1940: The propagated disagg transfer flag in PyExecutor’s
scheduling path is missing newly started gen-init transfers, because it is
computed before _prepare_disagg_gen_init() can begin receives and before the
next poll loop reuses it. Update the logic around
disagg_gen_transfer_in_progress and
SerializableSchedulerOutput.from_scheduler_result in py_executor.py so the flag
also reflects any fitting_disagg_gen_init_requests that just started transfer,
ensuring the next status check is not skipped.

🪄 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: 115e5aba-1b55-4544-be00-7bf9a880410e

📥 Commits

Reviewing files that changed from the base of the PR and between eb0cbdb and c493c61.

📒 Files selected for processing (3)

• tensorrt_llm/_torch/disaggregation/transceiver.py
• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py

[tensorrt-cicd]

PR_Github #55516 [ run ] completed with state FAILURE. Commit: c493c61
/LLM/main/L0_MergeRequest_PR pipeline #44440 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

Link to invocation

[Shixiaowei02]

/bot run --add-multi-gpu-test --disable-fail-fast

[tensorrt-cicd]

PR_Github #55731 [ run ] triggered by Bot. Commit: fe65a28 Link to invocation

[Shixiaowei02]

/bot run --add-multi-gpu-test --disable-fail-fast

[tensorrt-cicd]

PR_Github #55786 [ run ] triggered by Bot. Commit: 9b3e53e Link to invocation

[tensorrt-cicd]

PR_Github #55731 [ run ] completed with state ABORTED. Commit: fe65a28

Link to invocation

[tensorrt-cicd]

PR_Github #55786 [ run ] completed with state FAILURE. Commit: 9b3e53e
/LLM/main/L0_MergeRequest_PR pipeline #44682 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

[Shixiaowei02]

duplicated with #15356