From db0b59a8ff892825094280c32272e5e6ced92744 Mon Sep 17 00:00:00 2001 From: yifjiang <19356972+yifjiang@users.noreply.github.com> Date: Mon, 23 Mar 2026 20:35:34 -0700 Subject: [PATCH 1/4] [None][fix] Add bounded timeout to gen-side KV transfer in C++ CacheTransceiver CacheTransceiver::checkGenTransferStatus called future.get() without a timeout, causing an unbounded block when the KV transfer never completes. This leads to decode worker hangs in disaggregated serving. The context (send) path in checkContextTransferStatus already uses future.wait_for() with kv_transfer_sender_future_timeout_ms. Apply the same pattern to the generation (receive) path: use wait_for() with a bounded timeout, log a warning on timeout, and skip to the next iteration instead of blocking forever. Signed-off-by: yifjiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.cpp | 68 +++++++++++++------ 1 file changed, 48 insertions(+), 20 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 2e4bf1f06667..7266048899a1 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -599,6 +599,13 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastRequestNum) { bool blockAll = !atLeastRequestNum.has_value(); + std::optional receiverFutureTimeoutMs = std::nullopt; + // If blockAll is true, we want to block and not use a timeout + if (!blockAll && mCacheTransceiverConfig.has_value()) + { + receiverFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); + } + std::vector genTransferReadyRequestIds; for (auto&& [request, future] : mRequesterFutures) { @@ -715,14 +722,47 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR { try { - it->second.get(); - it->first->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); + // Wait for up to a specified timeout + auto status = it->second.wait_for( + std::chrono::milliseconds(receiverFutureTimeoutMs.value_or(0))); + if (status == std::future_status::ready || !receiverFutureTimeoutMs.has_value()) + { + it->second.get(); + it->first->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); - // Gather the kv cache transfer time from all workers and update to leader rank - if (!common::getEnvKVCacheTimeOutputPath().empty()) + // Gather the kv cache transfer time from all workers and update to leader rank + if (!common::getEnvKVCacheTimeOutputPath().empty()) + { + auto syncComm + = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; + updateKVCacheTransferBW(syncComm, it->first); + } + if (useMPI()) + { + TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), + "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", + it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); + } + else + { + TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), + "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", + it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); + } + it = mRequesterFutures.erase(it); + } + else if (status == std::future_status::timeout) { - auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; - updateKVCacheTransferBW(syncComm, it->first); + TLLM_LOG_WARNING("Timed out waiting for generation KV cache transfer after %d milliseconds.", + receiverFutureTimeoutMs.value()); + ++it; + } + else + { + TLLM_LOG_ERROR("Future returned unexpected status for request %ld. Marking as error", + it->first->mRequestId); + it->first->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + it = mRequesterFutures.erase(it); } } catch (std::exception const& e) @@ -730,20 +770,8 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR TLLM_LOG_ERROR( "Error occurred during generation transfer for request %ld: %s", it->first->mRequestId, e.what()); it->first->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + it = mRequesterFutures.erase(it); } - if (useMPI()) - { - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", - it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); - } - else - { - TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), - "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", - it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); - } - it = mRequesterFutures.erase(it); } else { From 4921fbeca88ff7d08b46a5664a78f190b88691b9 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Tue, 7 Apr 2026 00:30:23 +0000 Subject: [PATCH 2/4] fix: apply bounded timeout unconditionally in checkGenTransferStatus and checkContextTransferStatus Remove the !blockAll guard so the timeout applies in all code paths, including when the scheduler calls with atLeastRequestNum=nullopt. Previously, blockAll=true caused the timeout to be skipped, falling through to future.get() which blocks indefinitely on stalled transfers. Also use value_or(1000) instead of value_or(0) as the default timeout when no config is set, ensuring there is always a bounded wait. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.cpp | 33 ++++++++++--------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 7266048899a1..d133bd7f3c5e 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -485,8 +485,9 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { bool blockAll = !atLeastRequestNum.has_value(); std::optional senderFutureTimeoutMs = std::nullopt; - // If blockAll is true, we want to block and not use a timeout - if (!blockAll && mCacheTransceiverConfig.has_value()) + // Always use a bounded timeout to prevent unbounded blocking. + // The caller (scheduler) loops, so timed-out transfers retry on next iteration. + if (mCacheTransceiverConfig.has_value()) { senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); } @@ -551,8 +552,9 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( try { // Wait for up to a specified timeout - auto status = future.wait_for(std::chrono::milliseconds(senderFutureTimeoutMs.value_or(0))); - if (status == std::future_status::ready || !senderFutureTimeoutMs.has_value()) + auto const timeoutMs = senderFutureTimeoutMs.value_or(1000); + auto status = future.wait_for(std::chrono::milliseconds(timeoutMs)); + if (status == std::future_status::ready) { future.get(); requestsStatus.completedRequestIds.insert(request->mRequestId); @@ -564,8 +566,8 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } else if (status == std::future_status::timeout) { - TLLM_LOG_WARNING("Timed out waiting for context KV cache transfer after %d milliseconds.", - senderFutureTimeoutMs.value()); + TLLM_LOG_WARNING( + "Timed out waiting for context KV cache transfer after %d milliseconds.", timeoutMs); ++it; } else @@ -600,8 +602,9 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR { bool blockAll = !atLeastRequestNum.has_value(); std::optional receiverFutureTimeoutMs = std::nullopt; - // If blockAll is true, we want to block and not use a timeout - if (!blockAll && mCacheTransceiverConfig.has_value()) + // Always use a bounded timeout to prevent unbounded blocking. + // The caller (scheduler) loops, so timed-out transfers retry on next iteration. + if (mCacheTransceiverConfig.has_value()) { receiverFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); } @@ -723,9 +726,9 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR try { // Wait for up to a specified timeout - auto status = it->second.wait_for( - std::chrono::milliseconds(receiverFutureTimeoutMs.value_or(0))); - if (status == std::future_status::ready || !receiverFutureTimeoutMs.has_value()) + auto const timeoutMs = receiverFutureTimeoutMs.value_or(1000); + auto status = it->second.wait_for(std::chrono::milliseconds(timeoutMs)); + if (status == std::future_status::ready) { it->second.get(); it->first->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); @@ -753,14 +756,14 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR } else if (status == std::future_status::timeout) { - TLLM_LOG_WARNING("Timed out waiting for generation KV cache transfer after %d milliseconds.", - receiverFutureTimeoutMs.value()); + TLLM_LOG_WARNING( + "Timed out waiting for generation KV cache transfer after %d milliseconds.", timeoutMs); ++it; } else { - TLLM_LOG_ERROR("Future returned unexpected status for request %ld. Marking as error", - it->first->mRequestId); + TLLM_LOG_ERROR( + "Future returned unexpected status for request %ld. Marking as error", it->first->mRequestId); it->first->setState(LlmRequestState::kDISAGG_TRANS_ERROR); it = mRequesterFutures.erase(it); } From 37063ba56f55db867ca8c5f43a98148474135e8f Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Wed, 8 Apr 2026 17:59:51 +0000 Subject: [PATCH 3/4] fix: guard updateKVCacheTransferBW collective and apply bounded timeout unconditionally Address review feedback: - Guard updateKVCacheTransferBW timing collective so it only runs when all ranks block together (blockAll) or the request was confirmed ready on every rank in the initial poll (freqIt->second == syncSize). This prevents hangs in allgather when a peer timed out and skipped the request. - Keep bounded timeout on both context and gen sides unconditionally (remove !blockAll guard) with value_or(1000) default, so the scheduler loop is never blocked indefinitely on a single future. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.cpp | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index d133bd7f3c5e..b9d2b860fd06 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -719,6 +719,7 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(), atLeastRequestNum.value_or(0)); } + auto const syncSize = (syncComm != nullptr) ? syncComm->getSize() : 1; for (auto it = mRequesterFutures.begin(); it != mRequesterFutures.end();) { if (blockAll || toCompleteIdSet.find(it->first->mRequestId) != toCompleteIdSet.end()) @@ -733,12 +734,17 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR it->second.get(); it->first->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); - // Gather the kv cache transfer time from all workers and update to leader rank + // Gather the kv cache transfer time from all workers and update to leader rank. + // Only call the timing collective when either all ranks block together (blockAll) + // or the request was confirmed ready on every rank in the initial poll, to avoid + // hanging in allgather when a peer timed out and skipped this request. if (!common::getEnvKVCacheTimeOutputPath().empty()) { - auto syncComm - = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; - updateKVCacheTransferBW(syncComm, it->first); + auto const freqIt = frequencyMap.find(it->first->mRequestId); + if (blockAll || (freqIt != frequencyMap.end() && freqIt->second == syncSize)) + { + updateKVCacheTransferBW(syncComm, it->first); + } } if (useMPI()) { @@ -762,8 +768,8 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR } else { - TLLM_LOG_ERROR( - "Future returned unexpected status for request %ld. Marking as error", it->first->mRequestId); + TLLM_LOG_ERROR("Future returned unexpected status for request %ld. Marking as error", + it->first->mRequestId); it->first->setState(LlmRequestState::kDISAGG_TRANS_ERROR); it = mRequesterFutures.erase(it); } From e97c9369d6f9214036c775a089a2d153c38cb2c0 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Wed, 8 Apr 2026 13:19:50 -0700 Subject: [PATCH 4/4] [None][fix] Fix clang-format on TLLM_LOG_ERROR call Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index b9d2b860fd06..95028e2791ff 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -768,8 +768,8 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR } else { - TLLM_LOG_ERROR("Future returned unexpected status for request %ld. Marking as error", - it->first->mRequestId); + TLLM_LOG_ERROR( + "Future returned unexpected status for request %ld. Marking as error", it->first->mRequestId); it->first->setState(LlmRequestState::kDISAGG_TRANS_ERROR); it = mRequesterFutures.erase(it); }