[https://nvbugs/6448152][perf] make C++ context-transfer consensus asynchronous

[chienchunhung]

Summary

Replace the scheduler-hot-path blocking pipeline-parallel terminal-outcome collective in the C++ context cache transceiver with asynchronous ordered agreement.

For the qualified topology, every context pipeline rank publishes one immutable local terminal vote without waiting. The final pipeline rank reduces those votes and publishes one authoritative ordered commit. Request state transitions, retained bookkeeping, and global resource reclamation remain gated on that commit, preserving cross-rank consistency while removing the blocking collective from a naturally staggered scheduler path.

Important

The implementation is functionally validated on this pull request's current head, but current-head token throughput is not fully recovered. The exact-head GB300 run completed 512/512 requests at 794.61 output tokens/s. The 1548.84 output tokens/s result below is a matched result from the earlier frozen validation tree and must not be presented as current-head performance.

Motivation

The blocking context-side pipeline-parallel collective introduced by the original transfer-state-consensus change—not the cross-rank consistency requirement itself—caused the original PP4 throughput regression. That change fixed a real correctness issue: pipeline ranks could otherwise apply different terminal KV-transfer outcomes. Its synchronous execution shape, however, placed the collective at a staggered pipeline scheduler point and converted rank-entry skew into idle time.

The causal comparison used adjacent frozen trees with the same sender wake-up fix and CI compatibility change on both sides. The only product-tree difference between the matched control and treatment was the original two-file consistency patch; its stable patch ID (df1d9fe1964957dd0a50112a3e46c820a01b7f4f) exactly matched the merged change. The post-change treatment timed out, while transition-only tracing localized thousands of accumulated seconds of asymmetric waiting to the context pipeline consensus rendezvous.

The asynchronous protocol retains global agreement and ordered reclamation without forcing every rank to enter the terminal-outcome operation simultaneously.

Protocol and safety invariants

1. Local terminal votes are immutable for a request epoch.
2. COMPLETED is committed only when every participant votes completed.
3. Failure or cancellation dominates completed.
4. Timeout is a sticky, repeatable, nonterminal proposal and never authorizes resource release by itself.
5. With in-flight cancellation enabled, local cancellation is retried until quiescence permits a terminal vote.
6. Timeout events and terminal votes share one ordered stream to the coordinator; timeout updates and final commits share one ordered stream to followers.
7. No externally visible terminal transition or globally retained resource release occurs before the authoritative commit.
8. Protocol mode and version are checked once across the pipeline group at startup.
9. Shutdown exchanges close markers and drains outstanding MPI sends before communicator-backed storage is destroyed.

Qualified scope and fallback

The asynchronous integration is currently qualified only for the context/sender-side C++ cache transceiver:

Dimension	Qualified value
Transfer direction	Context/sender side
Data path	NIXL with UCX
Control plane	MPI
Tensor parallelism	TP1
Context parallelism	CP1
Pipeline parallelism	PP greater than one
Attention data parallelism	Disabled

All other configurations retain the existing reducer and state-transition path. This is an explicitly qualified fast path, not a correctness-semantics change for unqualified topologies.

The generation-side status path is not optimized here. Eligible workers construct the coordinator symmetrically at startup, but the generation status path does not publish to or poll it; there is no per-request generation-side coordinator traffic. Extending this to TP, CP, attention-DP, or generation-side domains requires explicit request-owner membership and hierarchical shard-to-stage-to-pipeline agreement. Merely relaxing the TP1/CP1 guard could allow different lanes for the same request to commit different outcomes.

Verification

Targeted performance validation uses the exact GB300 PP4 disaggregated workload:

disagg_upload-e2e-gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL

Historical causal A/B and repair evidence

Arm	CI pipeline	Functional result	Output-token throughput	Interpretation
Matched pre-change control	GB300 pipeline	512/512 successful	1557.83 tok/s	Valid matched control
Matched post-change treatment	GB300 pipeline	481/512 successful; timed out	598.42 tok/s raw	Censored; verifies the original consistency patch still causes the stall after removing the sender-race confound
Transition-only lifecycle trace	GB300 pipeline	480/512 successful; timed out	approximately 818 tok/s before teardown	Censored; localizes the steady-state loss to asymmetric context PP consensus entry
Unsafe local-only causal bypass	GB300 pipeline	512/512 successful	1564.71 tok/s	Causal upper bound; recovers throughput but removes required global consistency
Asynchronous ordered commit	GB300 pipeline	512/512 successful, 0 failed	1548.84 tok/s	Safe repair; 99.42% of matched control

The lifecycle trace measured negligible TP consensus time but approximately 4240, 3204, and 2126 accumulated seconds of PP waiting on context ranks 0, 1, and 2, versus 0.082 seconds on the last pipeline rank. Removing only context-side global consensus recovered throughput, establishing the causal phase, but was unsafe. The asynchronous run then recovered matched-control throughput while keeping coordinator mode active on all four context pipeline ranks. Every rank published a local vote and observed the same authoritative global commit; recovery was not obtained by bypassing agreement or globally gated resource retention.

Together, these runs verify both claims: the blocking rendezvous introduced by the original consistency change caused a valid performance regression, and asynchronous ordered commit fixes that regression on the controlled tree without reverting its correctness guarantee.

Current review head

The same targeted workload was run on exact head a2fca7808c94f90ca396fb047fd3d0887e293706 in the current-head GB300 pipeline:

• 512/512 requests succeeded with zero request failures.
• Official valid output-token throughput was 794.61 tok/s, approximately 51.0% of the historical matched control.
• Coordinator mode was active on all four context pipeline ranks, with authoritative commits and clean shutdown.
• The targeted stage was red only because of the expected performance threshold; it was not a functional, transport, consensus, or infrastructure failure.

Therefore, this head validates protocol correctness and lifecycle safety, but it does not demonstrate recovery of current-head absolute throughput.

Current-head regression isolation

Separate TEST ONLY diagnostics preserve the production coordinator while changing one factor at a time:

• A bounded phase trace completed 512/512 at 795.57 tok/s. Fetch/activation, sender submission, transfer completion, coordinator vote/commit, and cleanup were fast; the first inflated interval was the context pipeline forward/device dependency loop.
• A CuTe BF16 automatic-selection discriminator completed 512/512 at 814.02 tok/s, only 2.44% above the exact review head. Automatic CuTe BMM/GEMM selection is not the dominant regression.
• A generation-transfer scheduler-accounting discriminator completed 512/512 at 798.99 tok/s, only 0.55% above the exact review head. Logical accounting for transfer-in-progress requests is not the dominant regression.
• Admission telemetry recorded 512 candidates admitted, zero deferred, and no active transfer-block pressure at the decisions. Transfer p95 was approximately 156 ms. Transport-budget admission did not create the long interval before the next transfer candidate.
• A historical-source/current-runtime crossover was inconclusive before benchmark execution: historical DeepEP failed to link against the current NVSHMEM symbols, so no GB300 test child, request count, or performance metric exists.
• A current-source/historical-runtime reverse crossover is now running the exact PP4 stage on the current production source with only the five runtime-image tags reverted to the known-fast historical bundle. Its result is pending and does not change this production change's verified correctness claim.

The evidence currently places the unresolved slowdown upstream of transfer and asynchronous consensus, in the context pipeline forward/device path or its runtime stack. That investigation remains separate from this production change.

Why this change is still worth merging

The lower current-head number does not invalidate the historical causal result:

• the original blocking-rendezvous regression was isolated with a matched adjacent-tree A/B;
• this implementation removes that rendezvous while preserving authoritative cross-rank agreement and global resource-retention rules;
• the current review head completes the exact workload without request, transport, consensus, or shutdown failures; and
• one-factor current-head diagnostics place the remaining throughput loss upstream of transfer consensus.

This change therefore fixes a real, independently verified issue and is worth merging after normal review and full-CI gates. The merge claim is deliberately scoped: it fixes the consensus-rendezvous regression, but it does not claim to restore the current branch to approximately 1558 output tokens/s. The additional current-head factor remains under investigation.

Additional coverage

• Companion targeted coverage on the byte-identical product implementation completed with 3430 passed, 0 failed, and 1447 skipped.
• Staggered success/failure coordinator coverage passed with both two and eight MPI ranks and shut down without abandoned work or pending sends.
• Flag-off compatibility retained the existing path.

Full CI on the official review head is in progress through the current full-CI helper. This pull request should remain draft until required review and full-CI gates are complete.

Follow-up qualification

Future topology work should use hierarchical agreement:

1. aggregate shard outcomes within each pipeline stage over the request-owning TP/CP ranks;
2. publish one immutable stage vote to the pipeline coordinator;
3. compute one authoritative ordered request-and-epoch commit across pipeline stages;
4. fan that commit back to every participating shard; and
5. transition state or reclaim resources only after authoritative commit and local quiescence.

Each extension must cover success, failure, cancellation, timeout, duplicate and out-of-order messages, request-ID reuse with epochs, clean shutdown/drain, and a matched throughput comparison. Independent DP replicas must remain separate agreement domains unless a request explicitly spans them.

Summary by CodeRabbit

• New Features

  • Added distributed coordination for context-transfer completion, failure, and timeout outcomes.
  • Added asynchronous MPI communication support, including nonblocking sends, message checks, receives, and request completion polling.
  • Added dedicated communication tags for context-transfer coordination.

• Bug Fixes

  • Improved consistency of context-transfer status and timeout handling across participating processes.

• Tests

  • Added unit and multi-device coverage for consensus, timeout behavior, asynchronous communication, and shutdown scenarios.

[chienchunhung]

/bot run --disable-fail-fast --stage-list "GB300-12_GPUs-3_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU4-GEN1-NODE2-GPU8-Post-Merge-1"

[chienchunhung]

/bot run --disable-fail-fast --stage-list "GB300-12_GPUs-3_Nodes-PyTorch-Disagg-PerfSanity-CTX1-NODE1-GPU4-GEN1-NODE2-GPU8-Post-Merge-1"

[tensorrt-cicd]

PR_Github #60750 [ run ] triggered by Bot. Commit: a2fca78 Link to invocation

[tensorrt-cicd]

PR_Github #60750 [ run ] completed with state FAILURE. Commit: a2fca78
/LLM/main/L0_MergeRequest_PR pipeline #49035 (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

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61052 [ run ] triggered by Bot. Commit: a2fca78 Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes

The PR adds an MPI-based ContextTransferCoordinator with vote reduction, timeout handling, asynchronous peer messaging, and ordered shutdown. CacheTransceiver conditionally uses it for eligible multi-rank topologies, with new transport APIs, MPI tags, build wiring, and unit or multi-GPU tests.

Context-transfer consensus

Layer / File(s)	Summary
Coordination contracts and MPI transport cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h, cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h, cpp/include/tensorrt_llm/runtime/utils/mpi*.h	Defines consensus result types and coordinator APIs, adds MPI point-to-point helpers, context-transfer tags, and MpiRequest::isCompleted().
Reducer and asynchronous coordinator implementation cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp, cpp/tensorrt_llm/batch_manager/CMakeLists.txt	Implements vote validation, consensus reduction, asynchronous event processing, peer updates, shutdown signaling, and static-library build registration.
Cache-transceiver consensus integration cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp, cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h	Enables coordination for eligible topologies, publishes terminal outcomes and timeouts, polls consensus results, and persists timed-out request IDs for cancellation.
Reducer and MPI integration tests cpp/tests/unit_tests/batch_manager/*, cpp/tests/unit_tests/multi_gpu/*	Tests reducer validation and timeout behavior, staggered MPI consensus, and polling completion of asynchronous MPI sends.

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

Suggested labels: api-compatible

Suggested reviewers: pcastonguay, shixiaowei02, chuangz0

Sequence Diagram(s)

sequenceDiagram
  participant CacheTransceiver
  participant ContextTransferCoordinator
  participant PeerRanks
  participant MPI
  CacheTransceiver->>ContextTransferCoordinator: publish local outcome or timeout
  ContextTransferCoordinator->>MPI: send asynchronous event
  MPI->>PeerRanks: deliver vote event
  PeerRanks->>MPI: send coordinator update
  MPI->>ContextTransferCoordinator: receive votes and updates
  ContextTransferCoordinator-->>CacheTransceiver: poll consensus result

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 2.70% 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 is specific, concise, and accurately describes the main change to asynchronous context-transfer consensus.
Description check	✅ Passed	The description follows the template well and includes the required summary, motivation, invariants, scope, verification, and follow-up sections.
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]

🧹 Nitpick comments (1)

cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp (1)

1025-1036: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Blocking consensus loop is a pure busy-wait.

In blockAll mode this spins on poll() + std::this_thread::yield() until every awaiting request commits. Local transfers are already drained (the terminal loop calls future.get()), so this loop only waits for MPI vote propagation — but if a peer is slow to finish its own transfers, this rank pins a core at ~100% for the entire wait. Consider a short backoff sleep (e.g. kTransferFuturePollIntervalMs) instead of a tight yield.

♻️ Suggested backoff

             if (blockAll
                 && consensusOutcome.completedRequestIds.size() + consensusOutcome.failedRequestIds.size()
                     < mSenderRequestsAwaitingConsensus.size())
             {
-                std::this_thread::yield();
+                std::this_thread::sleep_for(std::chrono::milliseconds(kTransferFuturePollIntervalMs));
             }

🤖 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 `@cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp` around lines 1025 -
1036, Replace the tight std::this_thread::yield() polling in the blockAll
consensus loop around mergeCoordinatorOutcome() with a short sleep using the
existing kTransferFuturePollIntervalMs interval. Preserve the loop’s completion
condition and continue merging outcomes before each wait.

🤖 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 `@cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp`:
- Around line 1025-1036: Replace the tight std::this_thread::yield() polling in
the blockAll consensus loop around mergeCoordinatorOutcome() with a short sleep
using the existing kTransferFuturePollIntervalMs interval. Preserve the loop’s
completion condition and continue merging outcomes before each wait.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 1d8480c1-a9f6-4c54-8a78-aa332505d6c8

📥 Commits

Reviewing files that changed from the base of the PR and between 6df5134 and a2fca78.

📒 Files selected for processing (11)

• cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h
• cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h
• cpp/include/tensorrt_llm/runtime/utils/mpiTags.h
• cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h
• cpp/tensorrt_llm/batch_manager/CMakeLists.txt
• cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
• cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp
• cpp/tests/unit_tests/batch_manager/CMakeLists.txt
• cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp
• cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp
• cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp

[tensorrt-cicd]

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

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61120 [ run ] triggered by Bot. Commit: a2fca78 Link to invocation

[tensorrt-cicd]

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

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61202 [ run ] triggered by Bot. Commit: a2fca78 Link to invocation

[tensorrt-cicd]

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

[chienchunhung]

/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-PyTorch-3"

[tensorrt-cicd]

PR_Github #61310 [ run ] triggered by Bot. Commit: a2fca78 Link to invocation

[tensorrt-cicd]

PR_Github #61310 [ run ] completed with state SUCCESS. Commit: a2fca78
/LLM/main/L0_MergeRequest_PR pipeline #49539 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61345 [ run ] triggered by Bot. Commit: a2fca78 Link to invocation

[tensorrt-cicd]

PR_Github #61345 [ run ] completed with state SUCCESS. Commit: a2fca78
/LLM/main/L0_MergeRequest_PR pipeline #49572 completed with status: 'SUCCESS'

CI Report

Link to invocation