[https://nvbugs/6194812][fix] Restore non-blocking PP send-handle progress in broadcast loop

[tensorrt-cicd]

Summary

Restores the per-iteration MPI-progress driver in _broadcast_sample_state_loop that was removed by PR #12888 (which fixed a separate PP shutdown hang, NVBugs/6050489).

The pipeline-parallel sample-state ring broadcast forwards the sampled token/state between PP ranks with a non-blocking isend_object (mpi4py pkl5 / UCX rendezvous). For a rendezvous transfer the bulk payload is only pushed when the sender drives the progress engine (ucp_worker_progress), and in the default build that happens only when the application calls into MPI — there is no background progress thread.

• Root cause: PR [https://nvbugs/6050489][fix] fix agg pp4 hang issue #12888 (commit d1731cb5) removed the if self.executed_batch_queue.empty(): self.wait_on_pp_send_handles(...) flush from _broadcast_sample_state_loop in tensorrt_llm/_torch/pyexecutor/py_executor.py. That flush was also the broadcast thread's only MPI-progress driver during the idle gap between decode steps. Without it, after posting the isend the loop spins in Python and makes no MPI calls, so the sender's UCX worker is never pumped, the payload sits parked, and the next PP stage's recv_sample_state stalls until the sender incidentally re-enters MPI on the following step. Same kernels and same kernel time — the regression is a host-side PP-bubble dilation (NVBugs/6194812).
• Approach: Restore the empty-queue progress driver, but route it through a new progress_pp_send_handles helper that uses a non-blocking, bounded Request.test() poll (max 64) instead of the original blocking wait(). test() pumps the sender's MPI/UCX progress so the parked send flushes immediately, while the bounded non-blocking form cannot block forever if the receiver has already exited — so it stays safe against the shutdown hang from NVBugs/6050489. Surgical change, single file.

Test case

test_perf[qwen3_0.6b-bench-pytorch-bfloat16-maxbs:512-maxnt:2048-input_output_len:8000,1000-reqs:256-con:1-pp:4-gpus:4]

The combination of properties is what makes this case a sensitive detector:

• pp:4 — after the last stage samples, the sampled state is ring-broadcast around all 4 ranks (rank3 → rank0 → rank1 → rank2) so every rank updates its request/KV bookkeeping and rank0 gets the next input token. This ring sits on the per-token critical path, and the 3 hops are a strict sequential chain (each rank forwards the bytes it just received).
• con:1 — a single request in flight; no other microbatch work to overlap, so any broadcast stall is a fully-exposed GPU idle bubble instead of being hidden.
• input_output_len 8000,1000 — 1000 output tokens = 1000 serial decode steps, each paying the stall once. The small qwen3 0.6b bfloat16 model makes per-step compute cheap, so the host-side comm stall dominates step time.

Small model + single request + 4-way PP + long decode ⇒ the sample-state broadcast latency is the bottleneck, multiplied over 1000 steps.

What handle.test() does

handle is the MPI.Request returned by isend_object. handle.test() returns (flag, status); [0] is the boolean completion flag. Each call does two things:

1. Non-blockingly pumps the sender's MPI/UCX progress engine — the side effect that actually pushes the parked rendezvous payload forward.
2. Reports whether the send completed, so the handle can be cleared.

def progress_pp_send_handles(self, send_handles, microbatch_id, max_polls: int = 64):
    handle = send_handles[microbatch_id]
    if handle is None:
        return
    for _ in range(max_polls):
        if handle.test()[0]:
            send_handles[microbatch_id] = None
            return

It is non-blocking and bounded, so it restores the per-step progress cadence (the perf win) without re-introducing the blocking wait() that #12888 removed to fix the shutdown hang.

Why it's effective

The ring is a 3-hop sequential chain — this fix does not parallelize the hops. What it removes is the per-hop dead wait: in the regressed build each hop's stall is almost entirely "waiting for the sender to re-enter MPI", not data movement (the payload is small; real transfer is ~1 ms). nsys shows recv_sample_state going from ~3 ms → 269 ms average (1.85 s worst case) in the regressed build, with GPU kernels byte-for-byte identical (same FMHA, GEMM, NCCL kernels and durations).

Legend: # GPU forward · S sample · . GPU idle / transfer parked (no progress) · > transfer moving (pumped) · R* recv completes

Before (with #12888) — each ring hop parks:

        FORWARD (activations r0->r1->r2->r3)   SAMPLE-STATE RING (each hop parks ~269ms)    NEXT ITER
r0   |# F0 |.....|.....|....|.....................................................R*|# F0(N+1)
r1   |.....|# F1 |.....|....|.......................................................R*|# F1
r2   |.....|.....|# F2 |....|.........................................................R*|
r3   |.....|.....|.....|#F3 | S* .....parked.....|.....parked.....|.....parked.....|
                            hop1 r3->r0 (~269ms)  hop2 r0->r1 (~269ms) hop3 r1->r2 (~269ms)
                            +------------ all 4 GPUs idle the whole time -----------+

After (this PR) — test() pumps each sender, ~1 ms/hop:

        FORWARD (activations r0->r1->r2->r3)   RING (pumped, ~1ms/hop)   NEXT ITER
r0   |# F0 |.....|.....|....|R*|# F0(N+1) ...
r1   |.....|# F1 |.....|....|  |R*|# F1 ...
r2   |.....|.....|# F2 |....|  |  R*|# F2 ...
r3   |.....|.....|.....|#F3 |S*>|
                            hop1> hop2> hop3>  (~1ms each, back-to-back)
                            +- rank3 keeps pumping via test() -> payload pushed now -+

The GPU work (# widths) is identical in both; only the ring row differs — parked vs moving. Removing the per-hop dead wait collapses the broadcast from ~3×269 ms of bubble to ~3 ms, recovering essentially all of the throughput back to the good-commit baseline.

Perf

	Bad commit	Threshold	After fix	Good commit
qwen3_0.6b-bench-pytorch-bfloat16-pp:4-gpus:4 (GB300, lyris)	177.08 tok/s	223.19	270.55 tok/s	269.29

total_output_throughput, higher-is-better. Fix recovers full perf (+52.8% over bad commit, +0.5% over good commit, within run-to-run noise).

Independent reproduction on GB300 (lyris), total_token_throughput, plain run (no profiler), building the wheel from source for each variant:

Variant	tok/s
TOT with #12888 (regressed)	1560
TOT + this PR (#14850)	2428 (+55.6%)
TOT with #12888's removed lines manually restored (control)	2464

Wall-clock for the benchmark drops 28:26 → 19:05. nsys confirms identical GPU kernels across all three; the only delta is host-side recv_sample_state latency.

Test plan

• Verify fix on the same GPU type as the original failure (GB300, lyris)
• Check for regressions in related tests (nsys profile compare vs bad/good — kernel set unchanged, host gap restored)

Links

• Bug: https://nvbugs/6194812
• Culprit PR: [https://nvbugs/6050489][fix] fix agg pp4 hang issue #12888
• Auto-bisect/fix session: /lustre/fsw/coreai_comparch_trtllm/chenfeiz/agent/20260525-120524-bug-triage-6194812/

🤖 Generated with repair-bot

Summary by CodeRabbit

• Bug Fixes
  • Fixed potential deadlock and stall issues during shutdown in parallel execution scenarios
  • Enhanced communication progress handling to prevent hangs during application termination

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR prevents MPI stalls and deadlocks during PyExecutor shutdown by adding bounded polling of outstanding pipeline-parallel (PP) send handles. When the broadcast queue is empty and about to block, the thread now polls pending MPI sends instead of relying solely on the main loop to advance them.

Changes

MPI Progress Polling During Shutdown

Layer / File(s)	Summary
New MPI send handle polling helper tensorrt_llm/_torch/pyexecutor/py_executor.py	Added progress_pp_send_handles(...) method that non-blockingly tests MPI send request handles for a given microbatch ID using bounded polling (up to 64 iterations by default) and clears handles upon completion.
Broadcast loop integration with bounded send polling tensorrt_llm/_torch/pyexecutor/py_executor.py	Updated _broadcast_sample_state_loop to check queue emptiness and invoke bounded send polling before blocking, keeping MPI progress moving safely during shutdown scenarios.

🎯 2 (Simple) | ⏱️ ~8 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.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 PR title clearly identifies the main change (restoring non-blocking PP send-handle progress in broadcast loop) and includes the NVBugs reference with the appropriate [fix] type.
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 PR description is comprehensive and well-structured, covering root cause, approach, test cases, performance metrics, and links.

_{✏️ 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 and usage tips.}

[achartier]

Looks like that should be under some kind of lock

[achartier]

Why 64?

[achartier]

@bo-nv @pcastonguay @Tabrizian FYI since this is a follow-up to #12888

[chenfeiz0326]

● All three runs complete. Final results vs. baseline data:

  ┌────────────────┬──────────────────────┬──────────┬──────────┬───────┬────────────────────────────────┐
  │      Run       │        Commit        │ Measured │ Expected │ Match │       Threshold (223.19)       │
  ├────────────────┼──────────────────────┼──────────┼──────────┼───────┼────────────────────────────────┤
  │ GOOD           │ b9ce4b69             │ 264.134  │ 269.29   │ 98.1% │ ✓ above                        │
  ├────────────────┼──────────────────────┼──────────┼──────────┼───────┼────────────────────────────────┤
  │ BAD (main TOT) │ 66577ac9             │ 172.541  │ 177.08   │ 97.4% │ ✗ below — regression confirmed │
  ├────────────────┼──────────────────────┼──────────┼──────────┼───────┼────────────────────────────────┤
  │ FIX            │ 0005c08e (PR #14850) │ 267.012  │ 270.55   │ 98.7% │ ✓ above                        │
  └────────────────┴──────────────────────┴──────────┴──────────┴───────┴────────────────────────────────┘

  Verdict — fix verified. Main TOT still reproduces the regression (172.5, ~21% below threshold), and the fix branch fully recovers perf to ~267 (~20% above threshold, within 1.1% of the recorded post-fix number).

[bo-nv]

I don't understand why #12888 causes the regression, pr12888 just reverted #12528, which I wrongly introduced. Maybe the baseline is wrong, we need perf data before karen-sy@9a24d1a CC @chenfeiz0326

[chenfeiz0326]

/bot run --disable-fail-fast --stage-list "*GB200*PerfSanity*,*GB300*PerfSanity*"

[tensorrt-cicd]

PR_Github #52498 [ run ] triggered by Bot. Commit: 0005c08 Link to invocation

[tensorrt-cicd]

PR_Github #52498 [ run ] completed with state FAILURE. Commit: 0005c08
/LLM/main/L0_MergeRequest_PR pipeline #41790 (Partly Tested) 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