[TRTLLM-12721][feat] Add disagg transfer state consensus by chienchunhung · Pull Request #15139 · NVIDIA/TensorRT-LLM

chienchunhung · 2026-06-09T05:06:02Z

Description

Adds a V1 CacheTransceiver request-state consensus step for disaggregated KV transfer terminal transitions.

The implementation records local completed/failed transfer outcomes, packs {request_id, state} entries, gathers them through the existing communication backend, and applies a deterministic consensus result before mutating request state. Context-side consensus uses the existing TP/CP sync group and then the PP group; generation-side consensus uses the existing generation sync group.

This keeps PR#1 scoped to the consensus foundation for safe in-flight cancellation. It does not yet add deadline enforcement, deferred cleanup/quarantine, richer cancellation states, or hard mid-transfer abort semantics; those remain for the follow-up cancellation/lifetime PR.

Dependency graph

Arrows point from prerequisite to dependent. PR numbers in graph nodes are clickable.

graph TD
    PR15139["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15139'>#15139</a>: transfer state consensus (this PR, merged)"]
    PR15181["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15181'>#15181</a>: bounded C++ transfer status polling (inflight)"]
    PR15356["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15356'>#15356</a>: bounded V2 context transfer polling (inflight)"]
    PR15238["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15238'>#15238</a>: in-flight cancellation + buffer poison (draft)"]
    WORK_BLOCKALL["blockAll / wait-all cancellation (planned)"]
    WORK_BUFFER["multi-slot buffers + unpoison recovery (planned)"]

    PR15139 -->|satisfied| PR15238
    PR15181 -->|blocking| PR15356
    PR15181 -->|blocking| PR15238
    PR15356 -->|blocking| PR15238
    PR15238 -.->|planned| WORK_BLOCKALL
    PR15238 -.->|planned| WORK_BUFFER

    classDef merged fill:#dcfce7,stroke:#16a34a,color:#14532d;
    classDef inflight fill:#dbeafe,stroke:#2563eb,color:#1e3a8a;
    classDef draft fill:#ffedd5,stroke:#f97316,color:#7c2d12;
    classDef current fill:#ede9fe,stroke:#7c3aed,color:#3b0764,stroke-width:3px;
    classDef downstream fill:#f3f4f6,stroke:#6b7280,color:#374151,stroke-dasharray:5 5;
    linkStyle 0 stroke:#16a34a,stroke-width:2px;
    linkStyle 1,2,3 stroke:#ea580c,stroke-width:3px;
    linkStyle 4,5 stroke:#6b7280,stroke-width:2px,stroke-dasharray:5 5;

    class PR15139 current;
    class PR15181,PR15356 inflight;
    class PR15238 draft;
    class WORK_BLOCKALL,WORK_BUFFER downstream;

Notes

Failed terminal state wins over completed once all participating ranks report a terminal state.
Generation transfer no longer overwrites a receiver-side kDISAGG_TRANS_ERROR with kDISAGG_GENERATION_TRANS_COMPLETE.
The existing ready-request-ID collectives are intentionally left in place for this first slice.

Follow-ups

Audit the existing ready-request-ID gather path after this consensus foundation lands:
- checkContextTransferStatus: gatherRequestIds(syncComm, contextCompleteRequestIds)
- checkGenTransferStatus: gatherRequestIds(syncComm, genTransferReadyRequestIds)
If the new {request_id, completed/failed} transfer-state consensus fully covers the same cross-rank agreement semantics, remove or fold those ready-ID gathers into the consensus path.
Keep this as a separate cleanup/performance PR so this PR stays focused on correctness of terminal-state consensus.

Testing

pre-commit run clang-format --files cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
pre-commit run codespell --files cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
git diff --check upstream/main..HEAD

Attempted full file-scoped pre-commit as well, but local hooks fail in scripts/check_test_list.py with TypeError: unsupported operand type(s) for |: 'type' and 'NoneType' under the local hook environment; unrelated C++ hooks passed.

Summary by CodeRabbit

Bug Fixes
- Enhanced cross-rank consensus mechanism for KV cache transfer outcomes in distributed configurations
Refactor
- Improved cache transceiver with optimized pipeline-parallel communication setup and deferred state transitions for increased reliability

chienchunhung · 2026-06-09T05:09:27Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-09T05:15:11Z

PR_Github #52963 [ run ] triggered by Bot. Commit: 5115366 Link to invocation

tensorrt-cicd · 2026-06-09T11:50:18Z

PR_Github #52963 [ run ] completed with state SUCCESS. Commit: 5115366
/LLM/main/L0_MergeRequest_PR pipeline #42204 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 · 2026-06-09T17:02:42Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-09T17:10:52Z

PR_Github #53107 [ run ] triggered by Bot. Commit: 27b39f2 Link to invocation

chienchunhung · 2026-06-10T00:13:52Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-10T00:19:44Z

PR_Github #53177 [ run ] triggered by Bot. Commit: 61dc211 Link to invocation

coderabbitai · 2026-06-10T00:22:26Z

📝 Walkthrough

Walkthrough

CacheTransceiver now implements cross-rank consensus for KV cache transfer outcomes. Local transfer completions and failures are recorded per-request, gathered across MPI/ProcessGroup ranks via new utilities, reduced into consensus sets, and then committed to request states. Both context and generation transfer paths now defer state transitions until after consensus.

Changes

KV cache transfer consensus mechanism

Layer / File(s)	Summary
State tracking members and includes `cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h`	Header adds `<unordered_map>` include and new private containers: sender/requester completed/failed request ID sets, and unordered_maps tracking requests awaiting consensus decision.
Consensus utility infrastructure `cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp`	Anonymous-namespace helpers pack per-request terminal states (completed/failed), gather outcomes across the sync group via MPI/ProcessGroup, and reduce them into global consensus sets for later commit.
Constructor pipeline-parallel communication setup `cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp`	Constructor initializes `mGroupPipeParaComm` when pipeline parallelism is enabled, using a split color computed from tensor/context parallel ranks.
Context-side transfer status refactoring `cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp`	Context transfer checking records local terminal outcomes per-request instead of mutating request state directly; exceptions and unexpected future statuses are recorded as failures for consensus. After gathering and reducing, final request states are set (complete vs error) and local tracking maps are cleaned up.
Generation-side transfer status refactoring `cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp`	Generation transfer checking similarly records local completion/error outcomes, defers state transitions until after consensus reduction, updates request states based on consensus results, gathers KV transfer bandwidth/time for completed requests, and cleans up local tracking.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

NVIDIA/TensorRT-LLM#14768: Both PRs refactor checkContextTransferStatus and checkGenTransferStatus in cacheTransceiver.cpp, but this PR adds full cross-rank consensus mechanism while the other focuses on shared-ptr lifecycle and observe-only timeout/dedup behavior.

Suggested reviewers

reasonsolo
bo-nv
dongxuy04
pcastonguay
chuangz0

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: adding disaggregated transfer state consensus to CacheTransceiver, which aligns with the core functionality of the changeset.
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.
Description check	✅ Passed	PR description covers the core change (consensus mechanism), testing steps, and known limitations with sufficient clarity.

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

coderabbitai

🧹 Nitpick comments (1)

cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp (1)
71-75: 💤 Low value

Consider adding an explicit kPending or sentinel value to prevent accidental zero-initialization bugs.

TransferConsensusState starts at 1, leaving 0 unassigned. If any packed buffer is zero-initialized or truncated, parsing could yield undefined behavior when casting to this enum. Adding kPending = 0 or a sentinel would make such cases detectable.
♻️ Optional: Add sentinel value
 enum class TransferConsensusState : std::uint64_t
 {
+    kPending = 0,   // Sentinel for uninitialized/invalid state
     kCompleted = 1,
     kFailed = 2,
 };
And add a default case in the switch at line 158-162 to catch unknown values.
🤖 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 71 - 75,
Add an explicit zero-value sentinel to the enum by defining
TransferConsensusState::kPending = 0 so zero-initialized or truncated buffers
map to a known state, and update the switch that handles TransferConsensusState
(the switch covering the consensus handling logic) to include a default branch
that treats unknown/invalid enum values robustly (e.g., log error and
handle/return/fail gracefully) to make such cases detectable and safe.

🤖 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 71-75: Add an explicit zero-value sentinel to the enum by defining
TransferConsensusState::kPending = 0 so zero-initialized or truncated buffers
map to a known state, and update the switch that handles TransferConsensusState
(the switch covering the consensus handling logic) to include a default branch
that treats unknown/invalid enum values robustly (e.g., log error and
handle/return/fail gracefully) to make such cases detectable and safe.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 477aa235-789b-4c35-acea-207c25e5b630

📥 Commits

Reviewing files that changed from the base of the PR and between 9a7f76f and 61dc211.

📒 Files selected for processing (2)

cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h
cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp

tensorrt-cicd · 2026-06-10T00:24:23Z

PR_Github #53107 [ run ] completed with state ABORTED. Commit: 27b39f2

Link to invocation

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

chienchunhung · 2026-06-10T05:02:15Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-10T05:08:23Z

PR_Github #53235 [ run ] triggered by Bot. Commit: 019efd9 Link to invocation

tensorrt-cicd · 2026-06-10T05:13:12Z

PR_Github #53177 [ run ] completed with state ABORTED. Commit: 61dc211

Link to invocation

tensorrt-cicd · 2026-06-10T11:39:34Z

PR_Github #53235 [ run ] completed with state SUCCESS. Commit: 019efd9
/LLM/main/L0_MergeRequest_PR pipeline #42429 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 · 2026-06-10T17:50:30Z

/bot run --stage-list "DGX_H100-PyTorch-2,DGX_B200-8_GPUs-PyTorch-1,DGX_H100-4_GPUs-PyTorch-Others-2,GB10-PyTorch-1" --disable-fail-fast

chienchunhung · 2026-06-10T21:15:53Z

/bot run --disable-fail-fast --stage-list "DGX_H100-PyTorch-2,DGX_B200-8_GPUs-PyTorch-1,DGX_H100-4_GPUs-PyTorch-Others-2,GB10-PyTorch-1"

tensorrt-cicd · 2026-06-10T21:22:11Z

PR_Github #53377 [ run ] triggered by Bot. Commit: 019efd9 Link to invocation

tensorrt-cicd · 2026-06-11T02:41:30Z

PR_Github #53377 [ run ] completed with state FAILURE. Commit: 019efd9
/LLM/main/L0_MergeRequest_PR pipeline #42554 (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 · 2026-06-11T16:31:51Z

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

chienchunhung · 2026-06-11T16:34:48Z

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

tensorrt-cicd · 2026-06-11T16:39:07Z

PR_Github #53613 [ run ] triggered by Bot. Commit: 019efd9 Link to invocation

tensorrt-cicd · 2026-06-11T16:41:26Z

PR_Github #53615 [ run ] triggered by Bot. Commit: 019efd9 Link to invocation

tensorrt-cicd · 2026-06-11T18:42:08Z

PR_Github #53615 [ run ] completed with state SUCCESS. Commit: 019efd9
/LLM/main/L0_MergeRequest_PR pipeline #42761 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

chienchunhung · 2026-06-11T20:49:49Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-11T20:56:33Z

PR_Github #53674 [ run ] triggered by Bot. Commit: 019efd9 Link to invocation

tensorrt-cicd · 2026-06-11T21:22:46Z

PR_Github #53674 [ run ] completed with state SUCCESS. Commit: 019efd9
/LLM/main/L0_MergeRequest_PR pipeline #42812 completed with status: 'SUCCESS'

CI Report

Link to invocation

Tabrizian

There is a consensus gathering logic already in cache transceiver. Should we remove it now given what you are adding in this PR?

  if ((syncComm) && syncComm->getSize() > 1)
    {
        auto gatherRequestIdVec = gatherRequestIds(syncComm, contextCompleteRequestIds);
        for (auto&& requestId : gatherRequestIdVec)
        {
            frequencyMap[requestId]++;
        }
    }

chienchunhung · 2026-06-12T22:26:52Z

There is a consensus gathering logic already in cache transceiver. Should we remove it now given what you are adding in this PR?
  if ((syncComm) && syncComm->getSize() > 1)
    {
        auto gatherRequestIdVec = gatherRequestIds(syncComm, contextCompleteRequestIds);
        for (auto&& requestId : gatherRequestIdVec)
        {
            frequencyMap[requestId]++;
        }
    }

Thanks @Tabrizian !

Yes, this is the cleanup candidate I called out in the PR description:

The existing gatherRequestIds(...) path is still part of the readiness/selection logic for building toCompleteIdSet before we call future.get(), including the atLeastRequestNum behavior and ordering fallback. The new consensus added in this PR is a terminal-state consensus after a rank has observed local completion/failure; it prevents inconsistent final request-state transitions across ranks.

So they may overlap, but they are not obviously identical yet. The Follow-ups mentioned in the PR description is to explicitly track auditing/removing this path, including this context-side gather and the analogous generation-side gather:

checkContextTransferStatus: gatherRequestIds(syncComm, contextCompleteRequestIds)
checkGenTransferStatus: gatherRequestIds(syncComm, genTransferReadyRequestIds)

I’d like to keep that as a separate cleanup/performance PR so this one stays focused on correctness of terminal-state consensus.

github-actions Bot assigned chienchunhung Jun 9, 2026

chienchunhung force-pushed the codex/disagg-inflight-cancel-impl branch from 5115366 to 27b39f2 Compare June 9, 2026 17:02

chienchunhung mentioned this pull request Jun 9, 2026

[TRTLLM-12721][fix] Bound disagg transfer status polling #15181

Closed

chienchunhung force-pushed the codex/disagg-inflight-cancel-impl branch from 6a495f4 to 61dc211 Compare June 9, 2026 23:57

chienchunhung requested a review from Shixiaowei02 June 10, 2026 00:14

chienchunhung marked this pull request as ready for review June 10, 2026 00:16

chienchunhung requested a review from a team as a code owner June 10, 2026 00:16

chienchunhung requested a review from bo-nv June 10, 2026 00:16

coderabbitai Bot reviewed Jun 10, 2026

View reviewed changes

chienchunhung added 2 commits June 9, 2026 22:00

[TRTLLM-12721][feat] Add disagg transfer state consensus

bbf7ccf

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

[TRTLLM-12721][fix] Defer disagg transfer errors until consensus

019efd9

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

chienchunhung force-pushed the codex/disagg-inflight-cancel-impl branch from 61dc211 to 019efd9 Compare June 10, 2026 05:01

chienchunhung mentioned this pull request Jun 11, 2026

[TRTLLM-12721][fix] Add gated C++ NIXL in-flight cancellation and safe cleanup #15238

Merged

chienchunhung requested a review from Tabrizian June 12, 2026 21:55

Tabrizian reviewed Jun 12, 2026

View reviewed changes

Tabrizian approved these changes Jun 12, 2026

View reviewed changes

chienchunhung mentioned this pull request Jun 12, 2026

[TRTLLM-12721][perf] Remove ready-ID transfer gathers #15324

Closed

chienchunhung merged commit 57bb6ee into NVIDIA:main Jun 12, 2026
13 checks passed

coderabbitai Bot mentioned this pull request Jun 15, 2026

[https://nvbugs/6179661][fix] Fix disagg generation-side KV transfer timeout deadlocks and teardown crashes #15363

Closed

1 task

chienchunhung mentioned this pull request Jun 15, 2026

[TRTLLM-12721][fix] Bound disagg transfer polling and admission #15356

Merged

Uh oh!

Conversation

chienchunhung commented Jun 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Dependency graph

Notes

Follow-ups

Testing

Summary by CodeRabbit

Uh oh!

chienchunhung commented Jun 9, 2026

Uh oh!

tensorrt-cicd commented Jun 9, 2026

Uh oh!

tensorrt-cicd commented Jun 9, 2026

Uh oh!

chienchunhung commented Jun 9, 2026

Uh oh!

tensorrt-cicd commented Jun 9, 2026

Uh oh!

chienchunhung commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

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

Uh oh!

Walkthrough

Changes

Estimated code review effort

Possibly related PRs

Suggested reviewers

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

chienchunhung commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

chienchunhung commented Jun 10, 2026

Uh oh!

chienchunhung commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

chienchunhung commented Jun 11, 2026

Uh oh!

chienchunhung commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

chienchunhung commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

Tabrizian left a comment

Choose a reason for hiding this comment

Uh oh!

chienchunhung commented Jun 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

chienchunhung commented Jun 9, 2026 •

edited

Loading

coderabbitai Bot commented Jun 10, 2026 •

edited

Loading

chienchunhung commented Jun 12, 2026 •

edited

Loading