[https://nvbugs/6388153][fix] Relay PP sample states synchronously on…

[WeiHaocheng]

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 by CodeRabbit

• Bug Fixes
  • Improved pipeline-parallel execution reliability to prevent hangs and deadlocks in sample-state synchronization.
  • Made message handling more robust when asynchronous relay is disabled, helping ensure pending updates are processed before forward execution.
  • Added clearer failure handling when execution states fall out of sync, making scheduling issues surface as explicit errors instead of stalling.

[coderabbitai]

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: 887a8738-be2f-425c-ae5b-ed92dfcce8f7

📥 Commits

Reviewing files that changed from the base of the PR and between 60583d2 and 354e901.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/pyexecutor/py_executor.py

---

📝 Walkthrough

Walkthrough

This change modifies PyExecutor's pipeline-parallel sample-state relay behavior. The TLLM_PP_ASYNC_BROADCAST_SAMPLE_STATE default flips from "1" to "0", enabling synchronous inline relay by default. Queue setup, shutdown, forward-pass drainage, and executed-batch dequeue logic are adjusted to support both relay modes without deadlocks.

Changes

PP Synchronous Sample-State Relay

Layer / File(s)	Summary
Relay mode default and queue setup tensorrt_llm/_torch/pyexecutor/py_executor.py	Adds Empty import, flips the async broadcast flag default to "0", updates comments, and conditionally initializes/tears down executed_batch_queue and the relay thread based on relay mode and PP size.
Synchronous relay execution path tensorrt_llm/_torch/pyexecutor/py_executor.py	Drains pending send_handles before forward passes in synchronous mode, switches executed-batch relay to an immediate inline call instead of queuing, and updates dequeue logic to use get_nowait() with a RuntimeError on missing sample states.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant ExecutorLoopPP
  participant RingBroadcast
  participant ExecutedBatchQueue
  participant NonZeroRank

  alt async relay enabled
    ExecutorLoopPP->>ExecutedBatchQueue: enqueue executed_batch
    ExecutedBatchQueue->>RingBroadcast: background thread relay
    NonZeroRank->>NonZeroRank: get() blocking on executed_batch_response_queue
  else synchronous relay (default)
    ExecutorLoopPP->>RingBroadcast: _ring_broadcast_sample_state(executed_batch)
    ExecutorLoopPP->>ExecutorLoopPP: drain send_handles before forward
    NonZeroRank->>NonZeroRank: get_nowait() and raise RuntimeError if missing
  end

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 inconclusive)

Check name	Status	Explanation	Resolution
Title check	❓ Inconclusive	The title is specific and related, but the trailing ellipsis leaves the change summary incomplete.	Use a complete one-line title that states synchronous PP sample-state relay without truncation.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description check	✅ Passed	It explains the bug, fix, and validation, so the core template content is present despite missing headings.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.}

[litaotju]

LGTM. Can merge after @yuxianq take a look.

[WeiHaocheng]

/bot run --disable-fail-fast

[tensorrt-cicd]

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

[tensorrt-cicd]

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]

/bot run --disable-fail-fast

[tensorrt-cicd]

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