[https://nvbugs/6114140][fix] avoid indefinite block in disagg gen-side KV transfer - #14414
Closed
Shixiaowei02 wants to merge 1 commit into
Closed
[https://nvbugs/6114140][fix] avoid indefinite block in disagg gen-side KV transfer#14414Shixiaowei02 wants to merge 1 commit into
Shixiaowei02 wants to merge 1 commit into
Conversation
…de KV transfer CacheTransceiver::checkGenTransferStatus used a bare std::future::get() to wait on receiver futures, which blocks forever when the corresponding sender never completes (e.g. UCX path unhealthy on gb300 NVL72 nodes). The sender-side checkContextTransferStatus already mirrors this with std::future::wait_for(kvTransferSenderFutureTimeoutMs) so the executor loop can periodically check py_kv_transfer_timed_out and cancel a stuck request. Mirror the same pattern on the gen side: * cacheTransceiver.cpp: when atLeastRequestNum has a value (i.e. the caller did not request unbounded blocking) read kvTransferSenderFutureTimeoutMs from the config and use it as the receiver-side wait_for budget. Timed-out futures stay in mRequesterFutures for the next iteration; only ready or caller-permitted-unbounded futures are completed/erased. * py_executor.py: after the C++ call returns, walk the active generation requests and, for any that have crossed kv_transfer_timeout_ms, call cancel_request and surface DISAGG_TRANS_ERROR -- symmetric with the existing logic in _check_disagg_ctx_cache_transfer_status. Verified that test_disaggregated_ctxtp2_genpp2[TinyLlama-1.1B-Chat-v1.0] passes on gb300 (theia0288, ~114s) with this fix; also passes the unfixed paths used by other ctxtp/genpp combinations because the new timeout only activates when atLeastRequestNum is set. The waive in waives.txt is removed so CI exercises this path again. Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Contributor
📝 WalkthroughWalkthroughThe PR adds bounded-wait timeout handling for disaggregated generation KV-cache transfers. The C++ receiver derives a timeout from config and replaces blocking future.get() with time-bounded wait_for, logging and re-polling on timeout rather than blocking indefinitely. The Python executor adds a cancellation path for timed-out in-progress transfers. A previously waived test is re-enabled. ChangesDisaggregated Generation KV-Transfer Timeout Handling
Sequence DiagramsequenceDiagram
participant PyExecutor
participant CacheTransceiver
participant Request as Request Future
PyExecutor->>CacheTransceiver: checkGenTransferStatus()
CacheTransceiver->>Request: wait_for(timeout)
alt Timeout expires
Request-->>CacheTransceiver: timeout
CacheTransceiver->>CacheTransceiver: log warning
CacheTransceiver-->>PyExecutor: keep pending for re-poll
PyExecutor->>PyExecutor: detect py_kv_transfer_timed_out
PyExecutor->>CacheTransceiver: cancel_request(req)
CacheTransceiver->>Request: mark kDISAGG_TRANS_ERROR
CacheTransceiver-->>PyExecutor: cancellation done
PyExecutor->>PyExecutor: mark DISAGG_TRANS_ERROR
else Future ready
Request-->>CacheTransceiver: ready
CacheTransceiver->>CacheTransceiver: update KV bandwidth
CacheTransceiver->>CacheTransceiver: log completion
CacheTransceiver-->>PyExecutor: request complete
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes The changes span C++ and Python with moderate logic density around timeout configuration, wait_for semantics, and request cancellation coordination. The timeout flow is straightforward but requires understanding the interaction between the receiver-side bounded wait and the executor-side cancellation path. Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 3835-3842: The loop handling timed-out generation transfers sets
req.state = LlmRequestState.DISAGG_TRANS_ERROR after
kv_cache_transceiver.cancel_request(req), but
_check_cache_transfer_errors("generation requests") may remove the Python
request from active_requests before the matching C++ future in
CacheTransceiver::mRequesterFutures is reaped; fix by either (A) updating
kv_cache_transceiver.cancel_request to also erase the corresponding future from
mRequesterFutures (ensure CacheTransceiver::cancelRequest removes the requester
future on successful cancel), or (B) change the Python-side flow so cancelled
generation transfers remain pollable until check_gen_transfer_status() has
reaped them (e.g., do not clear py_kv_transfer_start_time or remove the request
from active_requests inside _check_cache_transfer_errors while
is_disagg_generation_transmission_in_progress is true); reference symbols:
is_disagg_generation_transmission_in_progress, py_kv_transfer_timed_out,
kv_cache_transceiver.cancel_request, py_kv_transfer_start_time,
LlmRequestState.DISAGG_TRANS_ERROR, _check_cache_transfer_errors,
check_gen_transfer_status, CacheTransceiver::cancelRequest, mRequesterFutures.
🪄 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: 3ecc16c3-dcdf-4d6a-b7f4-2654395f8e8b
📒 Files selected for processing (3)
cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpptensorrt_llm/_torch/pyexecutor/py_executor.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
Comment on lines
+3835
to
3842
| for req in self.active_requests: | ||
| if (req.is_disagg_generation_transmission_in_progress | ||
| and req.py_kv_transfer_timed_out): | ||
| is_cancelled = self.kv_cache_transceiver.cancel_request(req) | ||
| if is_cancelled: | ||
| req.py_kv_transfer_start_time = None | ||
| req.state = LlmRequestState.DISAGG_TRANS_ERROR | ||
| self._check_cache_transfer_errors("generation requests") |
Contributor
There was a problem hiding this comment.
Don't remove timed-out gen requests before the C++ future is reaped.
On a successful timeout cancel, this path flips the request to DISAGG_TRANS_ERROR, and _check_cache_transfer_errors() can drop it from active_requests immediately afterward. But CacheTransceiver::cancelRequest() does not remove the matching entry from mRequesterFutures, so once the Python request is gone the executor may stop polling and leave the cancelled receiver future stranded in C++ state indefinitely. Please either erase the future as part of cancel_request() or keep cancelled gen transfers pollable until check_gen_transfer_status() has reaped them.
🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 3835 - 3842, The
loop handling timed-out generation transfers sets req.state =
LlmRequestState.DISAGG_TRANS_ERROR after
kv_cache_transceiver.cancel_request(req), but
_check_cache_transfer_errors("generation requests") may remove the Python
request from active_requests before the matching C++ future in
CacheTransceiver::mRequesterFutures is reaped; fix by either (A) updating
kv_cache_transceiver.cancel_request to also erase the corresponding future from
mRequesterFutures (ensure CacheTransceiver::cancelRequest removes the requester
future on successful cancel), or (B) change the Python-side flow so cancelled
generation transfers remain pollable until check_gen_transfer_status() has
reaped them (e.g., do not clear py_kv_transfer_start_time or remove the request
from active_requests inside _check_cache_transfer_errors while
is_disagg_generation_transmission_in_progress is true); reference symbols:
is_disagg_generation_transmission_in_progress, py_kv_transfer_timed_out,
kv_cache_transceiver.cancel_request, py_kv_transfer_start_time,
LlmRequestState.DISAGG_TRANS_ERROR, _check_cache_transfer_errors,
check_gen_transfer_status, CacheTransceiver::cancelRequest, mRequesterFutures.
Collaborator
|
PR_Github #49700 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49700 [ run ] completed with state
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
CacheTransceiver::checkGenTransferStatusused a barestd::future::get()on each receiver future, which blocks forever when the corresponding sender never completes. On gb300 NVL72 the receiver-side disagg KV transfer forctxtp2_genpp2has been hitting this on the CI cluster, leaving the PP executor loop stuck inside_recv_disagg_gen_cache → _check_disagg_gen_cache_transfer_status → kv_cache_transceiver.check_gen_transfer_statusuntil the 300 s server-start-timeout fires the client off.The sender side (
CacheTransceiver::checkContextTransferStatus) already deals with this by readingkv_transfer_sender_future_timeout_msfrom the config and usingfuture.wait_for(...)so the executor loop can fall back to Python, observepy_kv_transfer_timed_out, and cancel the stuck request. This PR mirrors the same pattern on the gen side.Changes
cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp— incheckGenTransferStatus, whenatLeastRequestNumhas a value, readgetKvTransferSenderFutureTimeoutMs()and use it as the receiver-sidewait_forbudget. Timed-out futures are left inmRequesterFuturesfor the next iteration; only ready or unbounded-blocking futures are completed/erased. Exception paths still markkDISAGG_TRANS_ERRORand erase.tensorrt_llm/_torch/pyexecutor/py_executor.py—_check_disagg_gen_cache_transfer_statusnow walks active gen requests and, for any whosepy_kv_transfer_timed_outis set, callskv_cache_transceiver.cancel_request(req)and marksDISAGG_TRANS_ERROR— symmetric with the existing logic in_check_disagg_ctx_cache_transfer_status.tests/integration/test_lists/waives.txt— remove thetest_disaggregated_ctxtp2_genpp2[TinyLlama-1.1B-Chat-v1.0]SKIP entry that pointed at https://nvbugs/6114140, since the test now runs to completion with the C++ side no longer blocking indefinitely.Test plan
test_disaggregated_ctxtp2_genpp2[TinyLlama-1.1B-Chat-v1.0]passes on gb300 (theia0288, ~114 s) with this fix.wait_forbudget only activates whenatLeastRequestNumis set (i.e., the timeout is opt-in via the existingkv_transfer_sender_future_timeout_msconfig, default 1000 ms)./bot runon full CI for confirmation.🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Tests