Skip to content

[https://nvbugs/6388153][fix] Modify the broadcast of sample state in pp to default to synchronous mode. - #16170

Merged
WeiHaocheng merged 5 commits into
NVIDIA:mainfrom
WeiHaocheng:fix/nvbug-6388153-pp-sync-sample-state
Aug 6, 2026
Merged

[https://nvbugs/6388153][fix] Modify the broadcast of sample state in pp to default to synchronous mode.#16170
WeiHaocheng merged 5 commits into
NVIDIA:mainfrom
WeiHaocheng:fix/nvbug-6388153-pp-sync-sample-state

Conversation

@WeiHaocheng

@WeiHaocheng WeiHaocheng commented Jul 9, 2026

Copy link
Copy Markdown
Collaborator

Under pipeline parallelism, sample states are relayed between PP ranks
(last rank -> 0 -> 1 -> ...) by a dedicated background thread. The relay
thread needs the GIL to run: when the executor thread blocks inside a
GIL-holding native call that waits on GPU progress (e.g. a DeepGEMM JIT
kernel cold load whose lazy loading requires context synchronization),
the relay starves, the downstream rank stops advancing, its unlaunched
forward never pairs the in-flight NCCL p2p kernel the blocked rank's GPU
is waiting on, and the ranks deadlock.

Make the executor thread relay each sample state inline in its
pre-defined iteration instead (default; set
TLLM_PP_ASYNC_BROADCAST_SAMPLE_STATE=1 to restore the asynchronous
relay thread):

- The relay recv/apply/isend runs at stage 2 of the PP executor loop, so
  delivery is ordered with the loop and immune to GIL starvation. The
  background thread, its duplicated MPI communicator, and the enqueue
  queue are only created in asynchronous mode.
- Before entering a forward pass, drain all pending relay isends. Large
  sample-state messages complete via the MPI rendezvous protocol, which
  needs the sender to keep entering MPI calls; a forward that blocks in
  a native call without MPI progress would otherwise starve the
  downstream receiver and close the same deadlock ring through a
  different edge (verified: without this, the test hangs 5/5 with all
  four ranks in a circular MPI wait).
- On ranks other than 0, a missing relayed batch now fails loudly
  instead of blocking forever, since the synchronous relay deposits
  batches in the same iteration they are handled.

Verified on GB300 PP4 (TestDeepSeekV3Lite::test_fp8_block_scales_4gpus,
20 runs): 0 hangs with 102 mid-inference DeepGEMM cold loads per run
still present (DG_JIT_PRINT_LOAD_TIME audit) - the deadlock trigger
remains exercised and no longer hangs - vs ~25% hang rate before.

Summary

  • Fixes a pipeline-parallel deadlock in sample-state relay.
  • Uses synchronous relay by default.
  • Preserves asynchronous relay with TLLM_PP_ASYNC_BROADCAST_SAMPLE_STATE=1.
  • Creates asynchronous relay resources only when asynchronous mode is enabled.
  • Drains pending relay sends before forward execution.
  • Reports an explicit error when a nonzero PP rank receives no batch.
  • Removes eleven obsolete waiver entries from tests/integration/test_lists/waives.txt.

Dev Engineer Review

  • The synchronous relay path improves MPI progress and prevents indefinite waits.
  • The asynchronous path remains available through the environment variable.
  • The default behavior changes from asynchronous to synchronous.
  • The synchronous path changes num_micro_batches from max(pp_size, 1024) to pp_size.
  • The release notes should document the correctness and performance trade-off.
  • PP throughput measurements should quantify the impact of synchronous relay and send draining.
  • The waiver changes require verification of the final entry set and test-list formatting.
  • The unrelated TestGPTOSS::test_w4_chunked_prefill[trtllm-fp8] waiver for nvbug 6529692 should remain in a separate PR.

QA Engineer Review

  • Only tests/integration/test_lists/waives.txt changed.
  • Eleven obsolete waiver entries were removed.
  • No test-db/ or qa/ files were modified.
  • CBTS coverage data is unavailable for this test-list-only change.
  • Verdict: needs follow-up.

@WeiHaocheng WeiHaocheng self-assigned this Jul 9, 2026
@WeiHaocheng
WeiHaocheng marked this pull request as ready for review July 9, 2026 03:46
@WeiHaocheng
WeiHaocheng requested a review from a team as a code owner July 9, 2026 03:46
@WeiHaocheng
WeiHaocheng requested a review from byshiue July 9, 2026 03:46
@coderabbitai

coderabbitai Bot commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

Walkthrough

PyExecutor initializes PP sample-state relay before microbatch allocation and defaults to synchronous relay. Synchronous execution drains pending sends and broadcasts inline. The waiver list removes obsolete DeepSeekV3Lite entries.

Changes

PP Synchronous Sample-State Relay

Layer / File(s) Summary
Relay mode configuration and execution
tensorrt_llm/_torch/pyexecutor/py_executor.py
PyExecutor initializes relay mode before microbatch allocation and removes duplicate initialization. Synchronous mode uses pipeline-size storage, drains pending sends, and broadcasts batches inline. Asynchronous mode queues batches for background relay.

DeepSeekV3Lite Test Waiver Updates

Layer / File(s) Summary
Remove obsolete DeepSeekV3Lite waivers
tests/integration/test_lists/waives.txt
The waiver list removes obsolete BF16, FP8, NVFP4, auto-dtype, and guided-decoding entries. Remaining entries are preserved.

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

Sequence Diagram(s)

sequenceDiagram
  participant PyExecutor
  participant ExecutedBatchQueue
  participant RingBroadcast

  alt synchronous relay
    PyExecutor->>PyExecutor: drain pending sample-state sends
    PyExecutor->>RingBroadcast: broadcast executed batch inline
  else asynchronous relay
    PyExecutor->>ExecutedBatchQueue: enqueue executed batch
    ExecutedBatchQueue->>RingBroadcast: relay sample state in background
  end
Loading

Possibly related PRs

Suggested reviewers: bowenfu, qijune, reasonsolo

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 0.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 follows the required format and clearly describes the synchronous PP sample-state broadcast change.
Description check ✅ Passed The description clearly explains the deadlock, solution, configuration, safeguards, and validation results, but it omits the template headings and checklist.
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 💡 1
⚔️ Resolve merge conflicts 💡
  • Resolve merge conflict in branch fix/nvbug-6388153-pp-sync-sample-state
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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

@WeiHaocheng
WeiHaocheng force-pushed the fix/nvbug-6388153-pp-sync-sample-state branch from 354e901 to 5c7af1d Compare July 9, 2026 03:57
@litaotju

litaotju commented Jul 9, 2026

Copy link
Copy Markdown
Collaborator

LGTM. Can merge after @yuxianq take a look.

@litaotju
litaotju requested a review from yuxianq July 9, 2026 04:01
Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py
@WeiHaocheng

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58597 [ run ] triggered by Bot. Commit: b1f5cda Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@WeiHaocheng

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@WeiHaocheng
WeiHaocheng enabled auto-merge (squash) July 10, 2026 15:17
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58670 [ run ] triggered by Bot. Commit: 42e4067 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58670 [ run ] completed with state SUCCESS. Commit: 42e4067
/LLM/main/L0_MergeRequest_PR pipeline #47256 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

@WeiHaocheng

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58757 [ run ] triggered by Bot. Commit: 42e4067 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58757 [ run ] completed with state FAILURE. Commit: 42e4067
/LLM/main/L0_MergeRequest_PR pipeline #47339 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

@WeiHaocheng

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58923 [ run ] triggered by Bot. Commit: 42e4067 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58923 [ run ] completed with state SUCCESS. Commit: 42e4067
/LLM/main/L0_MergeRequest_PR pipeline #47458 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

@WeiHaocheng
WeiHaocheng force-pushed the fix/nvbug-6388153-pp-sync-sample-state branch from 42e4067 to 7255092 Compare July 13, 2026 09:50
@WeiHaocheng

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58956 [ run ] triggered by Bot. Commit: 7255092 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #58956 [ run ] completed with state SUCCESS. Commit: 7255092
/LLM/main/L0_MergeRequest_PR pipeline #47488 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

@WeiHaocheng

Copy link
Copy Markdown
Collaborator Author

/bot run

1 similar comment
@WeiHaocheng

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63986 [ run ] completed with state SUCCESS. Commit: c068f0a
/LLM/main/L0_MergeRequest_PR pipeline #51921 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ 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

@BowenFu

BowenFu commented Aug 5, 2026

Copy link
Copy Markdown

The deadlock analysis reads right, and relaying inline removes the GIL-starvation window cleanly. Two things I'd want said out loud before this merges, since it's already at five approvals.

1. There's an unrelated new waiver in here. The diff removes the two nvbugs/6388153 entries — correct, that's the bug being fixed — but also adds:

full:GB300/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_chunked_prefill[trtllm-fp8] SKIP (https://nvbugs/6529692)

Different bug, and test_w4_chunked_prefill is single-GPU, so it never touches the PP relay path this PR changes. That belongs in its own PR — folding a new skip into a [fix] is how waivers stop being traceable to a decision.

2. The title says "relay synchronously"; the change is a default flip. TLLM_PP_ASYNC_BROADCAST_SAMPLE_STATE goes from defaulting to "1" to "0", and the deleted comment described it as "only for debugging purposes / some tests can disable it to get a deterministic behavior." So every PP deployment moves to the inline relay, and the old default becomes the opt-in. That's the right call if async can deadlock — but it swaps a documented perf path for a correctness one, and the description frames it as a fix rather than a behavior change. Worth being explicit in the PR body and release notes.

Related: num_micro_batches goes from max(pp_size, 1024) to pp_size on the new default path (py_executor.py:740-751). I think that's sound — with inline retirement each microbatch is retired in its own iteration, so the 1024-slot ring only ever existed to absorb the async thread's lag — but it is the kind of thing worth a sentence, since it's a 1024 → pp_size change to a structure the whole PP loop indexes.

Has anyone measured PP throughput before/after? The relay now serializes with forward instead of overlapping it, plus there's a new pre-forward wait_on_pp_send_handles drain. If there's a number on 6388153's repro config, that would settle it.

Not blocking — deferring to the PP/runtime owners on the tradeoff. But I'd split the waiver out.

@WeiHaocheng

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64151 [ run ] triggered by Bot. Commit: c068f0a Link to invocation

@yuxianq

yuxianq commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64158 [ run ] triggered by Bot. Commit: c068f0a Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64151 [ run ] completed with state ABORTED. Commit: c068f0a

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@WeiHaocheng

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64272 [ run ] triggered by Bot. Commit: c068f0a Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@WeiHaocheng
WeiHaocheng force-pushed the fix/nvbug-6388153-pp-sync-sample-state branch from c068f0a to 9559a9e Compare August 6, 2026 09:25
@coderabbitai

coderabbitai Bot commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

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.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/py_executor.py (1)

740-754: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Create relay-only resources only in async mode.

Line 744 selects synchronous relay by default. However, start_worker() still creates the relay queue, duplicates the MPI communicator, and starts broadcast_sample_state_handler for every PP executor. This adds an unnecessary MPI collective, daemon thread, and CUDA thread initialization on the default path.

Keep executed_batch_response_queue for both modes. Guard executed_batch_queue, communicator duplication, thread startup, and the matching shutdown sentinel and join with self.pp_async_broadcast_sample_state.

Proposed fix
 if self.dist.pp_size > 1:
-    self.executed_batch_queue = Queue(maxsize=self.num_micro_batches)
     self.executed_batch_response_queue = Queue(maxsize=-1)
+    if self.pp_async_broadcast_sample_state:
+        self.executed_batch_queue = Queue(maxsize=self.num_micro_batches)
+        self._broadcast_mpi_comm = mpi_comm().Dup()
+        self.broadcast_sample_state_handler = threading.Thread(...)
+        self.broadcast_sample_state_handler.start()
 
-if self.dist.pp_size > 1:
+if self.dist.pp_size > 1 and self.pp_async_broadcast_sample_state:
     self.executed_batch_queue.put(None)
     self.broadcast_sample_state_handler.join()
🤖 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 740 - 754, Update
start_worker() so relay-only resources are created only when
self.pp_async_broadcast_sample_state is true: conditionally initialize
executed_batch_queue, duplicate the MPI communicator, and start
broadcast_sample_state_handler, while retaining executed_batch_response_queue
for both modes. Apply the same condition to the matching shutdown sentinel and
thread join, leaving synchronous relay behavior free of those resources.
🤖 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.

Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 740-754: Update start_worker() so relay-only resources are created
only when self.pp_async_broadcast_sample_state is true: conditionally initialize
executed_batch_queue, duplicate the MPI communicator, and start
broadcast_sample_state_handler, while retaining executed_batch_response_queue
for both modes. Apply the same condition to the matching shutdown sentinel and
thread join, leaving synchronous relay behavior free of those resources.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 373f7743-18cd-4535-846f-b3e0f62cfe53

📥 Commits

Reviewing files that changed from the base of the PR and between e3e5cca and 9559a9e.

📒 Files selected for processing (2)
  • tensorrt_llm/_torch/pyexecutor/py_executor.py
  • tests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
  • tests/integration/test_lists/waives.txt

@WeiHaocheng WeiHaocheng changed the title [https://nvbugs/6388153][fix] Relay PP sample states synchronously on… [https://nvbugs/6388153][fix] Modify the broadcast of sample state in pp to default to synchronous mode. Aug 6, 2026
@WeiHaocheng
WeiHaocheng force-pushed the fix/nvbug-6388153-pp-sync-sample-state branch from 9559a9e to 62e9a01 Compare August 6, 2026 09:38
@coderabbitai

coderabbitai Bot commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

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.

1 similar comment
@coderabbitai

coderabbitai Bot commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

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.

@yuxianq

yuxianq commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator

/bot skip --comment "failed tests are unrelated to this PR, skip CI"

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64298 [ skip ] triggered by Bot. Commit: 33861f1 Link to invocation

Signed-off-by: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com>
@WeiHaocheng
WeiHaocheng force-pushed the fix/nvbug-6388153-pp-sync-sample-state branch from 33861f1 to 3de1064 Compare August 6, 2026 10:05
@coderabbitai

coderabbitai Bot commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

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: WeiHaocheng <469657736@qq.com>
@WeiHaocheng

Copy link
Copy Markdown
Collaborator Author

/bot skip --comment "failed tests are unrelated to this PR, skip CI"

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64298 [ skip ] completed with state SUCCESS. Commit: 33861f1
Skipping testing for commit 33861f1

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64302 [ skip ] triggered by Bot. Commit: 68e0150 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64302 [ skip ] completed with state SUCCESS. Commit: 68e0150
Skipping testing for commit 68e0150

Link to invocation

@WeiHaocheng
WeiHaocheng merged commit 02540d6 into NVIDIA:main Aug 6, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

9 participants