[TRTLLM-12721][fix] Bound disagg transfer status polling by chienchunhung · Pull Request #15181 · NVIDIA/TensorRT-LLM

chienchunhung · 2026-06-09T23:04:41Z

Summary

Reviewer callout: without this fix, a request that has exceeded the PyExecutor timeout may still not be marked timed out because the PyExecutor loop can be stuck inside the C++ transceiver waiting on the KV transfer future (future.get() / future wait). In that state, Python timeout handling cannot make progress because control has not returned from check_*_transfer_status().
Bound finite disagg context/generation transfer status polling to short wait slices in the C++ transceiver so the PyExecutor loop can return to scheduling, observe request timeouts, and continue processing other requests.
Keep poll-slice timeout non-terminal: if the transfer future is not ready, the request remains in progress and is checked again by later executor polls.
Preserve drain-all behavior for nullopt/None status checks.
Update direct transceiver tests so finite polls are treated as scheduler polls rather than completion barriers.
Align the test helper completion deadline with the default CacheTransceiverConfig.kv_transfer_timeout_ms plus a small margin.

Dependency graph

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

#15422 is a standalone main-based teardown hardening PR. It is independent of the bounded-polling stack, but should merge before #15238 if possible so the cancellation PR can build on the safer shutdown/lifetime behavior after rebasing.

graph TD
    PR15139["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15139'>#15139</a>: transfer state consensus (merged)"]
    PR15422["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15422'>#15422</a>: teardown hardening (open for review, independent)"]
    PR15181["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15181'>#15181</a>: bounded C++ transfer status polling (this PR, 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
    PR15422 -.->|preferred hardening| 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 stroke:#64748b,stroke-width:2px,stroke-dasharray:3 3;
    linkStyle 5,6 stroke:#6b7280,stroke-width:2px,stroke-dasharray:5 5;

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

Scope

This change applies only to the C++ transceiver path.
This intentionally does not implement deadline enforcement, request cancellation, deferred cleanup, or poison/unpoison buffer handling.
Transfer timeout warnings remain observe-only in this PR; actual in-flight cancellation/deadline enforcement should be handled separately.

Testing

git diff --check upstream/main..HEAD
python -m py_compile tests/unittest/others/test_kv_cache_transceiver.py

Targeted pytest could not be run locally because this environment is missing transformers during repo import.

Note: the final commit was created with --no-verify because local pre-commit hooks invoke system python3 3.9, and the existing repo test-list hook script uses newer str | None typing syntax.

Summary by CodeRabbit

Refactor
- Optimized cache transfer polling intervals for improved scheduler efficiency
- Enhanced timeout handling and status tracking with improved logging during transfers
Tests
- Expanded test coverage for cache transfer operations and timeout behaviors

chienchunhung · 2026-06-10T00:14:38Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-10T00:20:39Z

PR_Github #53178 [ run ] triggered by Bot. Commit: 09121fb Link to invocation

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

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-10T05:09:01Z

PR_Github #53236 [ run ] triggered by Bot. Commit: 7e90bda Link to invocation

tensorrt-cicd · 2026-06-10T05:14:37Z

PR_Github #53178 [ run ] completed with state ABORTED. Commit: 09121fb

Link to invocation

tensorrt-cicd · 2026-06-10T11:40:02Z

PR_Github #53236 [ run ] completed with state SUCCESS. Commit: 7e90bda
/LLM/main/L0_MergeRequest_PR pipeline #42430 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-10T23:02:53Z

/bot run --disable-fail-fast

coderabbitai · 2026-06-10T23:05:21Z

📝 Walkthrough

Walkthrough

This PR converts KV cache transfer status checks from blocking single wait_for calls to scheduler-friendly short polling intervals. Both context and generation transfer paths now compute bounded poll slices from configurable timeouts, log observe-only warnings when per-request timeouts are exceeded, and track timed-out request IDs. Tests are refactored to poll status repeatedly and a new test verifies bounded polling keeps in-flight transfers non-terminal.

Changes

KV Cache Transfer Polling and Timeout Tracking

Layer / File(s)	Summary
Context transfer polling interval and loop refactor `cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp`	Adds `<chrono>` header, introduces `getTransferFuturePollInterval()` helper that computes bounded polling intervals from optional timeouts, extracts sender-future timeout from config, and refactors the context sender future loop to emit observe-only timeout warnings, repeatedly call `wait_for(futurePollInterval)`, and update completion/error state with timed-out ID tracking.
Generation transfer polling interval and loop refactor `cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp`	Updates `blockAll` to `bool const`, extracts requester-future timeout from config using the same polling helper, and refactors the generation requester future loop to mirror context-side polling behavior with per-request timeout warnings, short polling waits, and timed-out ID erasure on completion or error.
Test polling constants and helper function `tests/unittest/others/test_kv_cache_transceiver.py`	Adds module-level constants for KV transfer timeout based on `CacheTransceiverConfig` defaults, and introduces `wait_for_transfer_completion()` helper that polls provided callables until completion or timeout.
Existing test refactoring to use polling `tests/unittest/others/test_kv_cache_transceiver.py`	Refactors `test_kv_cache_transceiver_single_process`, `test_async_transfer_keeps_llm_request_alive`, and `test_hybrid_cache_transceiver_single_process` to poll `check_context_transfer_status` and `check_gen_transfer_status` repeatedly, accumulate completed context request IDs, and wait for generation request terminal state instead of performing single immediate checks.
Bounded polling behavior test `tests/unittest/others/test_kv_cache_transceiver.py`	Adds `test_context_transfer_bounded_poll_keeps_request_in_progress` to verify that polling with short bounded intervals keeps in-flight context transfers in `DISAGG_CONTEXT_TRANS_IN_PROGRESS` with no completed/failed IDs, then drives completion using the standard polling helper.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Suggested reviewers

Shixiaowei02
pcastonguay

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 26.09% 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 clearly and specifically describes the main change: bounding disagg transfer status polling in the transceiver to prevent blocking the PyExecutor loop.
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 detailed and well-structured, explaining the problem, solution, testing, and scope with clear callouts and dependency graphs.

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

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/unittest/others/test_kv_cache_transceiver.py (1)

61-73: 💤 Low value

Consider adding a docstring to the polling helper.

The wait_for_transfer_completion helper is used by multiple tests and would benefit from a brief docstring explaining its parameters (poll_fn, is_done_fn, timeout_s) and behavior (polls repeatedly until completion or timeout).

📝 Suggested docstring

 def wait_for_transfer_completion(poll_fn,
                                  is_done_fn,
                                  timeout_s=DEFAULT_KV_TRANSFER_TIMEOUT_S +
                                  KV_TRANSFER_COMPLETION_MARGIN_S):
+    """Poll until KV cache transfer completes or timeout expires.
+
+    Args:
+        poll_fn: Callable that polls transfer status (e.g., check_context_transfer_status).
+        is_done_fn: Callable that returns True when transfer is complete.
+        timeout_s: Maximum time to wait before asserting failure.
+    """
     deadline = time.monotonic() + timeout_s

🤖 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 `@tests/unittest/others/test_kv_cache_transceiver.py` around lines 61 - 73, Add
a concise docstring to wait_for_transfer_completion describing its purpose (poll
until a KV transfer completes or timeout), parameters (poll_fn: callable invoked
each loop to drive progress, is_done_fn: callable returning True when transfer
finished, timeout_s: seconds to wait before failing), default behavior (uses
DEFAULT_KV_TRANSFER_TIMEOUT_S + KV_TRANSFER_COMPLETION_MARGIN_S and polls every
0.01s), and the failure behavior (raises an assertion on timeout). Place this
docstring immediately below the def wait_for_transfer_completion(...) signature.

🤖 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.

Inline comments:
In `@tests/unittest/others/test_kv_cache_transceiver.py`:
- Line 477: The pytest timeout on the test is too short and may kill the test
before wait_for_transfer_completion can assert; change the decorator on this
test from `@pytest.mark.timeout`(60) to `@pytest.mark.timeout`(120) so it matches
other transfer-completion tests (e.g., test_kv_cache_transceiver_single_process,
test_async_transfer_keeps_llm_request_alive,
test_hybrid_cache_transceiver_single_process) and allows
wait_for_transfer_completion (the helper used around line 544) to finish.

---

Nitpick comments:
In `@tests/unittest/others/test_kv_cache_transceiver.py`:
- Around line 61-73: Add a concise docstring to wait_for_transfer_completion
describing its purpose (poll until a KV transfer completes or timeout),
parameters (poll_fn: callable invoked each loop to drive progress, is_done_fn:
callable returning True when transfer finished, timeout_s: seconds to wait
before failing), default behavior (uses DEFAULT_KV_TRANSFER_TIMEOUT_S +
KV_TRANSFER_COMPLETION_MARGIN_S and polls every 0.01s), and the failure behavior
(raises an assertion on timeout). Place this docstring immediately below the def
wait_for_transfer_completion(...) signature.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 891d89d2-3362-41d2-89f5-db6f40845a91

📥 Commits

Reviewing files that changed from the base of the PR and between 7301075 and 7d71ebd.

📒 Files selected for processing (2)

cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
tests/unittest/others/test_kv_cache_transceiver.py

tensorrt-cicd · 2026-06-10T23:08:36Z

PR_Github #53393 [ run ] triggered by Bot. Commit: 7d71ebd Link to invocation

chienchunhung · 2026-06-10T23:59:47Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-11T00:06:22Z

PR_Github #53403 [ run ] triggered by Bot. Commit: 4654ab7 Link to invocation

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

PR_Github #53393 [ run ] completed with state ABORTED. Commit: 7d71ebd

Link to invocation

tensorrt-cicd · 2026-06-11T11:22:55Z

PR_Github #53403 [ run ] completed with state SUCCESS. Commit: 4654ab7
/LLM/main/L0_MergeRequest_PR pipeline #42576 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:37:52Z

/bot run --disable-fail-fast

chienchunhung · 2026-06-24T05:51:53Z

/bot run

chienchunhung · 2026-06-24T05:52:54Z

/bot run

chienchunhung · 2026-06-24T05:56:30Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-24T05:58:53Z

PR_Github #55413 [ run ] triggered by Bot. Commit: a179842 Link to invocation

tensorrt-cicd · 2026-06-24T06:00:23Z

PR_Github #55416 [ run ] triggered by Bot. Commit: a179842 Link to invocation

tensorrt-cicd · 2026-06-24T06:00:26Z

PR_Github #55413 [ run ] completed with state ABORTED. Commit: a179842

Link to invocation

tensorrt-cicd · 2026-06-24T06:03:11Z

PR_Github #55420 [ run ] triggered by Bot. Commit: a179842 Link to invocation

tensorrt-cicd · 2026-06-24T13:26:31Z

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

Link to invocation

chienchunhung · 2026-06-24T16:16:05Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-24T16:23:58Z

PR_Github #55539 [ run ] triggered by Bot. Commit: a179842 Link to invocation

tensorrt-cicd · 2026-06-24T19:02:39Z

PR_Github #55539 [ run ] completed with state SUCCESS. Commit: a179842
/LLM/main/L0_MergeRequest_PR pipeline #44462 completed with status: 'SUCCESS'

CI Report

Link to invocation

chienchunhung · 2026-06-24T20:59:10Z

Superseded by rebase; removing stale pre-rebase bot command.

chienchunhung · 2026-06-24T21:05:14Z

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

tensorrt-cicd · 2026-06-24T21:05:15Z

PR_Github #55583 [ run ] triggered by Bot. Commit: e8f3802 Link to invocation

tensorrt-cicd · 2026-06-24T21:11:11Z

PR_Github #55588 [ run ] triggered by Bot. Commit: e8f3802 Link to invocation

tensorrt-cicd · 2026-06-24T21:15:04Z

PR_Github #55583 [ run ] completed with state ABORTED. Commit: e8f3802

Link to invocation

tensorrt-cicd · 2026-06-24T23:51:52Z

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

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

chienchunhung · 2026-06-26T20:35:55Z

Closing this PR; PR#15356 will carry the changes made by this PR.

github-actions Bot assigned chienchunhung Jun 9, 2026

chienchunhung mentioned this pull request Jun 9, 2026

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

Merged

chienchunhung force-pushed the codex/disagg-bounded-transfer-poll branch from b5e32ac to 09121fb Compare June 10, 2026 00:05

chienchunhung force-pushed the codex/disagg-bounded-transfer-poll branch from 09121fb to 7e90bda Compare June 10, 2026 05:01

chienchunhung force-pushed the codex/disagg-bounded-transfer-poll branch 3 times, most recently from 2517dff to 7d71ebd Compare June 10, 2026 22:52

chienchunhung requested review from Shixiaowei02 and pcastonguay June 10, 2026 22:58

chienchunhung marked this pull request as ready for review June 10, 2026 22:59

chienchunhung requested a review from a team as a code owner June 10, 2026 22:59

coderabbitai Bot reviewed Jun 10, 2026

View reviewed changes

Comment thread tests/unittest/others/test_kv_cache_transceiver.py Outdated

chienchunhung force-pushed the codex/disagg-bounded-transfer-poll branch from 7d71ebd to 4654ab7 Compare June 10, 2026 23:59

chienchunhung mentioned this pull request Jun 11, 2026

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

Merged

Shixiaowei02 approved these changes Jun 11, 2026

View reviewed changes

chienchunhung force-pushed the codex/disagg-bounded-transfer-poll branch from 4654ab7 to d2006d7 Compare June 11, 2026 16:37

chienchunhung force-pushed the codex/disagg-bounded-transfer-poll branch from 9e50e98 to a179842 Compare June 24, 2026 05:50

chienchunhung force-pushed the codex/disagg-bounded-transfer-poll branch from 13c0eb2 to e8f3802 Compare June 24, 2026 21:03

xinhe-nv approved these changes Jun 25, 2026

View reviewed changes

chienchunhung added 6 commits June 25, 2026 21:53

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

b288d60

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

[TRTLLM-12721][test] Fix Python transceiver completion id check

0a85598

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

[TRTLLM-12721][fix] Initialize generation transfer timing

9ffc576

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

[TRTLLM-12721][fix] Restore bounded transfer backpressure

e273a42

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

[TRTLLM-12721][test] extend Qwen HELIX transfer timeout

34c3c96

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

[TRTLLM-12721][fix] Account active KV transfers in scheduler

47cf0af

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

chienchunhung force-pushed the codex/disagg-bounded-transfer-poll branch from e8f3802 to 47cf0af Compare June 26, 2026 04:58

chienchunhung closed this Jun 26, 2026

auto-merge was automatically disabled June 26, 2026 20:35
Pull request was closed

Uh oh!

Conversation

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

Uh oh!

Summary

Dependency graph

Scope

Testing

Summary by CodeRabbit

Uh oh!

chienchunhung commented Jun 10, 2026

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!

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

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

chienchunhung commented Jun 10, 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!

chienchunhung commented Jun 24, 2026

Uh oh!

chienchunhung commented Jun 24, 2026

Uh oh!

chienchunhung commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

chienchunhung commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

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

Uh oh!

Uh oh!

chienchunhung commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

Uh oh!

tensorrt-cicd commented Jun 24, 2026

chienchunhung commented Jun 9, 2026 •

edited

Loading

coderabbitai Bot commented Jun 10, 2026 •

edited

Loading

chienchunhung commented Jun 24, 2026 •

edited

Loading