[None][fix] Keep ADP ranks in collective lockstep on request errors and fail fast on desync by roborluo · Pull Request #16687 · NVIDIA/TensorRT-LLM

roborluo · 2026-07-21T21:01:52Z

Summary

Under attention DP, _enqueue_responses performs a tp_gather; collectives must be entered by every rank in the same order. A request-scoped error can be observed by only one rank, so directly publishing its response can cross the peers' next collective (usually tp_allgather(batch_size)).

This change:

Buffers non-fatal ADP error responses and flushes them at shared executor-loop safepoints.
Uses one end-of-loop flush rather than a second flush in can_queue iterations.
Flushes before the retry and clean-scheduler-exit paths, so buffered client errors are not silently stranded.
Fails fast on malformed gathered response payloads, giving a clear collective-desynchronization error.
Avoids publishing a rank-local fatal error through a TP response gather; it shuts down locally instead of risking another rank-divergent collective.

Test Coverage

Added focused unit coverage for buffered-response delivery, empty ADP participation, retry/shutdown exits, the single-flush behavior, and rank-local fatal handling.

Local macOS validation:

python3 -m py_compile — passed
git diff --check — passed
Pytest cannot import the CUDA-dependent TensorRT-LLM runtime on macOS (cuda-python has no supported macOS bindings); full test execution needs CUDA/Linux CI.

PR Checklist

PR title follows the required format
PR description clearly explains the issue and solution
Commits are signed off (DCO)
Coding guidelines followed (yapf clean)

Dev Engineer Review

Attention-DP synchronization: Buffers non-fatal transfer responses and flushes them at synchronized executor points, including idle non-overlap passes.
Error handling: Handles rank-local fatal errors locally and avoids divergent TP response gathers. Flushes pending responses during shutdown and benchmark retries.
Collective validation: Raises a bounded RuntimeError for malformed gathered payloads and skips None contributions.
Scope: Changes are limited to executor error handling and unit tests. No public API, configuration, or test-list changes are present.
Regression review: The additional collective remains synchronized because can_queue is TP-uniform under Attention-DP.

QA Engineer Review

Added coverage: TestPendingTransferResponseFlush covers rank-local fatal errors, empty Attention-DP contributions, buffered-response delivery and clearing, and executor-loop flushes during normal exit, benchmark retry, idle passes, and overlap-loop paths.
Test-list coverage: No tests/integration/test_lists/ files changed. The added unit tests are not listed in test-db or QA files.
Verdict: sufficient.

coderabbitai · 2026-07-21T21:06:01Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7bec13ad-4085-470a-8072-8db79e4c949f

📥 Commits

Reviewing files that changed from the base of the PR and between 0fe5df7 and f15f0b2.

📒 Files selected for processing (2)

tensorrt_llm/_torch/pyexecutor/py_executor.py
tests/unittest/_torch/executor/test_py_executor.py

🚧 Files skipped from review as they are similar to previous changes (2)

tensorrt_llm/_torch/pyexecutor/py_executor.py
tests/unittest/_torch/executor/test_py_executor.py

Walkthrough

PyExecutor now buffers transfer responses until synchronized publication. It flushes responses during executor-loop passes, shutdown, benchmark retries, and error handling. Attention-DP gathering handles None and malformed collective payloads explicitly.

Changes

Attention-DP Response Coordination

Layer / File(s)	Summary
Buffer and flush transfer responses `tensorrt_llm/_torch/pyexecutor/py_executor.py`	Request termination waits until buffered transfer responses are synchronously published.
Flush responses across executor-loop exits and passes `tensorrt_llm/_torch/pyexecutor/py_executor.py`, `tests/unittest/_torch/executor/test_py_executor.py`	Executor passes flush pending responses after processing and before shutdown, idle exits, and benchmark retries. Tests cover normal and overlap-loop cleanup without duplicate flushes.
Coordinate Attention-DP errors and validate gathers `tensorrt_llm/_torch/pyexecutor/py_executor.py`, `tests/unittest/_torch/executor/test_py_executor.py`	Attention-DP paths buffer non-fatal errors and avoid rank-local fatal gathers. Collective gathering ignores `None`, accepts list or tuple payloads, and raises a diagnostic `RuntimeError` for malformed payloads.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant ExecutorLoop
  participant ErrorHandler
  participant AttentionDPGather
  participant PendingTransferResponses
  ExecutorLoop->>ErrorHandler: process current executor errors
  ErrorHandler->>AttentionDPGather: gather synchronized responses
  AttentionDPGather-->>ErrorHandler: return validated contributions
  ErrorHandler->>PendingTransferResponses: buffer or flush responses
  ExecutorLoop->>PendingTransferResponses: flush after each loop pass

Possibly related PRs

NVIDIA/TensorRT-LLM#17082: Both PRs modify PyExecutor response handling. This PR changes deferred response flushing, while the related PR changes speculative-decoding metric accumulation.

Suggested labels: Release Blocker

Suggested reviewers: tabrizian

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 52.38% 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 follows the required format and clearly summarizes the ADP collective desynchronization fix.
Description check	✅ Passed	The description explains the issue, solution, test coverage, platform limitation, and checklist status in sufficient detail.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

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/pyexecutor/py_executor.py`:
- Around line 6644-6651: Replace the unbounded resp!r interpolation in
_enqueue_responses with a bounded, safe summary that cannot expose full response
contents or depend on an unsafe custom __repr__. Preserve the existing type
information and collective-desynchronization context in the RuntimeError
message.
- Around line 6634-6652: Update the response-merging logic in _enqueue_responses
to validate every item in each gathered list or tuple before calling
gather_responses.extend. Require each entry to be a two-field response
structure, rejecting malformed values such as scalar entries or entries with the
wrong length via the same clear RuntimeError path, then extend only validated
entries.

🪄 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: f54097f4-1650-45d1-9543-e3eb9780cce6

📥 Commits

Reviewing files that changed from the base of the PR and between 6fabfdc and afaf86c.

📒 Files selected for processing (1)

tensorrt_llm/_torch/pyexecutor/py_executor.py

…t errors Under attention DP, _enqueue_responses gathers per-rank response lists via tp_gather, and TP collectives pair by call order, not by type. _handle_errors is reached only by the rank(s) that observed an error: for request-scoped failures such as a failed disagg generation KV transfer, a single rank enqueues its error responses solo while the other ranks proceed to the per-step tp_allgather(batch_size). The collectives cross: the erroring rank deserializes a response-list payload inside safe_allgather (UnpicklingError: pickle data was truncated) and rank 0's response gather receives a stray batch-size int (TypeError: 'int' object is not iterable in gather_responses.extend), while the remaining ranks hang until the hang detector fires. Observed on a GB300 TP8 ADP disagg gen server: rank 2 hit 'Error occurred during generation transfer' (dataTransceiver.cpp), entered _handle_errors, and its solo gather crossed with peers' _can_queue allgather in the same iteration. Fix: on the non-fatal path, buffer error responses in _pending_transfer_responses instead of enqueueing directly. The buffer is flushed at _flush_pending_transfer_responses, which every rank executes unconditionally at synchronized points in the executor loops, so the gather is always entered in lockstep. Fatal errors keep the direct enqueue since the executor is shutting down. This covers every per-rank-divergent caller of _handle_errors in one place. Same failure family as the per-rank-divergent collective entry fixed in NVIDIA#14020; complements the fail-fast diagnostic in NVIDIA#16687. Signed-off-by: Bofeng Luo <bofengl@nvidia.com>

roborluo · 2026-07-21T21:25:58Z

/bot run --disable-fail-fast

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/pyexecutor/py_executor.py`:
- Around line 6527-6539: Move the _flush_pending_transfer_responses() call out
of the can_queue-guarded block in the non-overlap scheduling loop so it executes
at an unconditional, rank-symmetric point. Ensure it also runs when no batch can
be queued, while preserving collective ordering across ranks and existing
behavior when a batch is available.
- Around line 6527-6539: The non-fatal attention-DP path must retain each
request’s result queue until its buffered error response is flushed. Update the
termination flow around _do_terminate_request and
_flush_pending_transfer_responses so result_wait_queues entries are not removed
before _enqueue_responses can deliver the response; perform cleanup after
synchronized flushing, while preserving immediate cleanup for paths that do not
buffer responses.

🪄 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: 2f30393e-ae51-46b8-a77a-16dcd2986907

📥 Commits

Reviewing files that changed from the base of the PR and between afaf86c and 33bd50a.

📒 Files selected for processing (1)

tensorrt_llm/_torch/pyexecutor/py_executor.py

roborluo · 2026-07-21T22:08:14Z

/bot run --disable-fail-fast

…t errors Under attention DP, _enqueue_responses gathers per-rank response lists via tp_gather, and TP collectives pair by call order, not by type. _handle_errors is reached only by the rank(s) that observed an error: for request-scoped failures such as a failed disagg generation KV transfer, a single rank enqueues its error responses solo while the other ranks proceed to the per-step tp_allgather(batch_size). The collectives cross: the erroring rank deserializes a response-list payload inside safe_allgather (UnpicklingError: pickle data was truncated) and rank 0's response gather receives a stray batch-size int (TypeError: 'int' object is not iterable in gather_responses.extend), while the remaining ranks hang until the hang detector fires. Observed on a GB300 TP8 ADP disagg gen server: rank 2 hit 'Error occurred during generation transfer' (dataTransceiver.cpp), entered _handle_errors, and its solo gather crossed with peers' _can_queue allgather in the same iteration. Fix: on the non-fatal path, buffer error responses in _pending_transfer_responses instead of enqueueing directly. The buffer is flushed at _flush_pending_transfer_responses, which every rank executes unconditionally at synchronized points in the executor loops, so the gather is always entered in lockstep. Fatal errors keep the direct enqueue since the executor is shutting down. This covers every per-rank-divergent caller of _handle_errors in one place. Same failure family as the per-rank-divergent collective entry fixed in NVIDIA#14020; complements the fail-fast diagnostic in NVIDIA#16687. Signed-off-by: Bofeng Luo <bofengl@nvidia.com>

roborluo · 2026-07-21T22:26:07Z

/bot run --disable-fail-fast

…t errors Under attention DP, _enqueue_responses gathers per-rank response lists via tp_gather, and TP collectives pair by call order, not by type. _handle_errors is reached only by the rank(s) that observed an error: for request-scoped failures such as a failed disagg generation KV transfer, a single rank enqueues its error responses solo while the other ranks proceed to the per-step tp_allgather(batch_size). The collectives cross: the erroring rank deserializes a response-list payload inside safe_allgather (UnpicklingError: pickle data was truncated) and rank 0's response gather receives a stray batch-size int (TypeError: 'int' object is not iterable in gather_responses.extend), while the remaining ranks hang until the hang detector fires. Observed on a GB300 TP8 ADP disagg gen server: rank 2 hit 'Error occurred during generation transfer' (dataTransceiver.cpp), entered _handle_errors, and its solo gather crossed with peers' _can_queue allgather in the same iteration. Fix: on the non-fatal path, buffer error responses in _pending_transfer_responses instead of enqueueing directly. The buffer is flushed at _flush_pending_transfer_responses, which every rank executes unconditionally at synchronized points in the executor loops, so the gather is always entered in lockstep. Fatal errors keep the direct enqueue since the executor is shutting down. This covers every per-rank-divergent caller of _handle_errors in one place. Same failure family as the per-rank-divergent collective entry fixed in NVIDIA#14020; complements the fail-fast diagnostic in NVIDIA#16687. Signed-off-by: Bofeng Luo <bofengl@nvidia.com>

roborluo · 2026-07-21T23:54:51Z

Addressed review feedback in the amended commit:

Unconditional flush in the non-overlap loop — added _flush_pending_transfer_responses() outside the if can_queue block (next to _handle_kv_transfer_timeouts_synced(), which exists for the same reason), mirroring the overlap loop's unconditional flush. Buffered error responses are now delivered even when no rank can queue a batch. can_queue is TP-uniform under attention DP (it is the result of the _can_queue consensus allgather), so the additional collective stays in lockstep.
Bounded desync error message — dropped the unbounded resp!r from the RuntimeError; the type name is the diagnostic signal (the known mismatch partner is an int batch size).

Not taken: per-entry (two-field tuple) validation of gathered items. The guard targets crossed collectives, where the stray payload is a non-list scalar; contributions that are lists of malformed entries would indicate a different bug and still fail loudly at the unpack in the consumer loop.

roborluo · 2026-07-22T00:47:28Z

/bot run --disable-fail-fast

chienchunhung

Thanks for the PR!

roborluo · 2026-08-04T00:30:50Z

Thanks for the comments, Will check

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 `@tests/unittest/_torch/executor/test_py_executor.py`:
- Around line 1937-1949: Initialize the missing is_benchmark_disagg attribute in
_make_executor_loop_stub, using the default non-benchmark value expected by
_executor_loop before entering its loop. Keep the existing stub setup unchanged
otherwise.
- Around line 1887-2019: Add a focused test for PyExecutor._enqueue_responses
that mocks tp_gather() to return [1] and verifies the method raises RuntimeError
matching "TP collective desync". Initialize only the executor state and
dependencies required by _enqueue_responses, keeping the test alongside the
existing response-flush coverage.

🪄 Autofix

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: a1d4bc35-8fda-40b6-ae2d-f74ef7140ccc

📥 Commits

Reviewing files that changed from the base of the PR and between dcc4dc6 and 12b5e6f.

📒 Files selected for processing (2)

tensorrt_llm/_torch/pyexecutor/py_executor.py
tests/unittest/_torch/executor/test_py_executor.py

roborluo · 2026-08-09T01:16:05Z

@coderabbitai resume

coderabbitai · 2026-08-09T01:16:10Z

✅ Action performed

Reviews resumed.

coderabbitai

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
3982-4003: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Flush buffered responses on overlap-loop early exits.

This adds the required flushes only to _executor_loop. _executor_loop_overlap can also buffer a non-fatal ADP response in _handle_disagg_cache_errors_synced(), but its clean-exit path at Line 4476 and retry path at Lines 4481-4487 break or continue without flushing. The queued client can then wait until timeout.

Add the same flush before both overlap-loop early exits. Add focused overlap-loop tests for clean shutdown and benchmark retry.
Proposed fix
                 if scheduled_batch is None:
+                    self._flush_pending_transfer_responses()
                     break

                 can_forward, should_retry = self._check_benchmark_disagg_gate(
                     scheduled_batch, can_forward)
                 if should_retry:
                     if self._is_kv_manager_v2:
                         for req in scheduled_batch.generation_requests:
                             self.kv_cache_manager.revert_allocate_generation(
                                 req)
                     self._finalize_adp_dummy_allocation(False)
+                    self._flush_pending_transfer_responses()
                     continue
🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 3982 - 4003,
Update _executor_loop_overlap to call _flush_pending_transfer_responses()
immediately before both its clean-exit break and benchmark-retry continue paths,
matching the flush behavior in _executor_loop. Add focused tests covering
buffered-response delivery during overlap-loop clean shutdown and benchmark
retry.

🤖 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.

Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 3982-4003: Update _executor_loop_overlap to call
_flush_pending_transfer_responses() immediately before both its clean-exit break
and benchmark-retry continue paths, matching the flush behavior in
_executor_loop. Add focused tests covering buffered-response delivery during
overlap-loop clean shutdown and benchmark retry.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a95abea6-7666-408d-b78d-b245c9c4ac11

📥 Commits

Reviewing files that changed from the base of the PR and between dcc4dc6 and 713d49d.

📒 Files selected for processing (2)

tensorrt_llm/_torch/pyexecutor/py_executor.py
tests/unittest/_torch/executor/test_py_executor.py

🚧 Files skipped from review as they are similar to previous changes (1)

tests/unittest/_torch/executor/test_py_executor.py

coderabbitai · 2026-08-09T01:23:15Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

Signed-off-by: Bofeng Luo <bofengl@nvidia.com>

roborluo · 2026-08-09T02:09:27Z

/bot run --disable-fail-fast

roborluo requested review from a team as code owners July 21, 2026 21:01

roborluo requested review from Tabrizian, YihuiLu512, chienchunhung and zhaoyangwang-nvidia July 21, 2026 21:01

coderabbitai Bot reviewed Jul 21, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py

github-actions Bot assigned roborluo Jul 21, 2026

roborluo changed the title ~~[None][fix] Fail fast on ADP TP collective desync in _enqueue_responses~~ [None][fix] Keep ADP ranks in collective lockstep on request errors and fail fast on desync Jul 21, 2026

coderabbitai Bot reviewed Jul 21, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py Outdated

roborluo closed this Jul 21, 2026

roborluo reopened this Jul 21, 2026

roborluo force-pushed the bofengl/fix-adp-enqueue-responses-desync branch from 33bd50a to 8f1ea81 Compare July 21, 2026 22:15

roborluo force-pushed the bofengl/fix-adp-enqueue-responses-desync branch from 8f1ea81 to dcc4dc6 Compare July 21, 2026 23:54

chienchunhung reviewed Jul 28, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py Outdated

zhaoyangwang-nvidia requested changes Jul 30, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py Outdated

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py Outdated

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py

roborluo force-pushed the bofengl/fix-adp-enqueue-responses-desync branch from dcc4dc6 to 12b5e6f Compare August 9, 2026 00:30

coderabbitai Bot reviewed Aug 9, 2026

View reviewed changes

Comment thread tests/unittest/_torch/executor/test_py_executor.py

Comment thread tests/unittest/_torch/executor/test_py_executor.py

roborluo force-pushed the bofengl/fix-adp-enqueue-responses-desync branch 2 times, most recently from 577def0 to 7e8c912 Compare August 9, 2026 00:49

roborluo force-pushed the bofengl/fix-adp-enqueue-responses-desync branch 4 times, most recently from 1b457e5 to 713d49d Compare August 9, 2026 01:00

coderabbitai Bot reviewed Aug 9, 2026

View reviewed changes

roborluo force-pushed the bofengl/fix-adp-enqueue-responses-desync branch from 713d49d to 0fe5df7 Compare August 9, 2026 01:23

[None][fix] Keep ADP response gathers in lockstep

f15f0b2

Signed-off-by: Bofeng Luo <bofengl@nvidia.com>

roborluo force-pushed the bofengl/fix-adp-enqueue-responses-desync branch from 0fe5df7 to f15f0b2 Compare August 9, 2026 01:31

Uh oh!

Conversation

roborluo commented Jul 21, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test Coverage

PR Checklist

Dev Engineer Review

QA Engineer Review

Uh oh!

coderabbitai Bot commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

roborluo commented Jul 21, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

roborluo commented Jul 21, 2026

Uh oh!

roborluo commented Jul 21, 2026

Uh oh!

roborluo commented Jul 21, 2026

Uh oh!

roborluo commented Jul 22, 2026

Uh oh!

chienchunhung left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

roborluo commented Aug 4, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

roborluo commented Aug 9, 2026

Uh oh!

coderabbitai Bot commented Aug 9, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot commented Aug 9, 2026

Uh oh!

roborluo commented Aug 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

roborluo commented Jul 21, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 21, 2026 •

edited

Loading