diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 2e4bf1f06667..bac770103a45 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"); @@ -339,6 +339,10 @@ void CacheTransceiver::respondAndSendAsync(LlmRequest* llmRequest) } setContextState(llmRequest); auto future = mCacheSender->sendAsync(*llmRequest); + TLLM_LOG_DEBUG("respondAndSendAsync: adding request %ld to mSenderFutures (ptr=%p, transferStart=%ld, size=%zu)", + llmRequest->mRequestId, static_cast(llmRequest), + static_cast(llmRequest->getKvCacheTransferStart().time_since_epoch().count()), + mSenderFutures.size() + 1); mSenderFutures.emplace_back(llmRequest, std::move(future)); } @@ -485,10 +489,35 @@ 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()) + std::optional kvTransferTimeoutMs = std::nullopt; + // 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(); + kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + } + { + senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); + kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + } + + // Log mSenderFutures state for diagnosing dangling pointer issues. + // Each entry's pointer address and request ID are logged so we can detect + // when a pointer's underlying memory is freed (reqId changes to 0). + if (!mSenderFutures.empty()) + { + TLLM_LOG_DEBUG("checkContextTransferStatus: mSenderFutures.size()=%zu, blockAll=%d, " + "kvTransferTimeoutMs=%d", + mSenderFutures.size(), blockAll ? 1 : 0, + kvTransferTimeoutMs.value_or(-1)); + for (size_t i = 0; i < mSenderFutures.size(); ++i) + { + auto& [req, fut] = mSenderFutures[i]; + auto startTs = req->getKvCacheTransferStart().time_since_epoch().count(); + TLLM_LOG_DEBUG(" [%zu] ptr=%p reqId=%ld startTs=%ld", + i, static_cast(req), req->mRequestId, static_cast(startTs)); + } } auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupTPInDPComm : mGroupTensorParaComm; @@ -551,8 +580,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 +594,57 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } else if (status == std::future_status::timeout) { + // Check if total elapsed time exceeds kv_transfer_timeout_ms. + // Without this, stuck transfers retry the per-iteration timeout forever, + // holding KV blocks indefinitely and exhausting the cache pool. + if (kvTransferTimeoutMs.has_value()) + { + auto transferStart = request->getKvCacheTransferStart(); + // Guard: if transfer start was never set (TimePoint epoch), + // the request pointer may be stale or the start time was not recorded. + // Treat as timed out immediately to avoid infinite retry. + bool startTimeValid = transferStart.time_since_epoch().count() > 0; + bool shouldTimeout = !startTimeValid; + long elapsedMs = 0; + if (startTimeValid) + { + auto elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - transferStart); + elapsedMs = static_cast(elapsed.count()); + shouldTimeout = elapsedMs > kvTransferTimeoutMs.value(); + } + if (shouldTimeout) + { + // IMPORTANT: Do NOT dereference request->setState() or call + // mCacheSender->cancelRequest(*request) here. The LlmRequest* + // in mSenderFutures is a raw pointer with no ownership — the + // Python layer may have already freed the request. Dereferencing + // a dangling pointer causes heap corruption (free(): invalid + // next size). Just erase the stale entry from mSenderFutures. + // The Python layer handles state transitions via + // _end_transfer_and_maybe_terminate when it processes the + // error/completion from check_context_transfer_status. + if (startTimeValid) + { + TLLM_LOG_WARNING( + "Context KV cache transfer timed out: elapsed %ld ms > limit %d ms. " + "Removing entry from mSenderFutures (ptr=%p).", + elapsedMs, kvTransferTimeoutMs.value(), + static_cast(request)); + } + else + { + TLLM_LOG_WARNING( + "Removing stale entry from mSenderFutures: transfer start time is " + "uninitialized (request pointer %p may be dangling).", + static_cast(request)); + } + it = mSenderFutures.erase(it); + continue; + } + } TLLM_LOG_WARNING("Timed out waiting for context KV cache transfer after %d milliseconds.", - senderFutureTimeoutMs.value()); + timeoutMs); ++it; } else @@ -580,10 +659,14 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } catch (std::exception const& e) { - TLLM_LOG_ERROR( - "Error occurred during context transfer for request %ld: %s", request->mRequestId, e.what()); - request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); - requestsStatus.errorRequestIds.insert(request->mRequestId); + // Do NOT dereference request here — the pointer may be dangling. + // The future threw an exception (e.g. Broken promise from sender + // thread crash), and the request may have been freed concurrently. + // Just log the error and erase the entry. + TLLM_LOG_WARNING( + "Error during context transfer (ptr=%p): %s. " + "Removing entry from mSenderFutures.", + static_cast(request), e.what()); it = mSenderFutures.erase(it); } } @@ -593,12 +676,29 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } } + if (!requestsStatus.completedRequestIds.empty() || !requestsStatus.errorRequestIds.empty()) + { + TLLM_LOG_DEBUG("checkContextTransferStatus done: completed=%zu, errors=%zu, " + "mSenderFutures.size()=%zu", + requestsStatus.completedRequestIds.size(), + requestsStatus.errorRequestIds.size(), + mSenderFutures.size()); + } + return requestsStatus; } void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastRequestNum) { bool blockAll = !atLeastRequestNum.has_value(); + std::optional receiverFutureTimeoutMs = std::nullopt; + // 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(); + } + std::vector genTransferReadyRequestIds; for (auto&& [request, future] : mRequesterFutures) { @@ -709,20 +809,59 @@ 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()) { try { - 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()) + // Wait for up to a specified timeout + 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); + + // 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 const freqIt = frequencyMap.find(it->first->mRequestId); + if (blockAll || (freqIt != frequencyMap.end() && freqIt->second == syncSize)) + { + 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) + { + TLLM_LOG_WARNING( + "Timed out waiting for generation KV cache transfer after %d milliseconds.", timeoutMs); + ++it; + } + else { - auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; - updateKVCacheTransferBW(syncComm, it->first); + 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 +869,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 { diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 7db8a320dbda..38db4be50e84 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1686,6 +1686,14 @@ std::pair> WindowBlockManager::sto } if (pinBlocks) { + // Claim block from free queue before pinning so that + // unpinBlocksById can safely releaseBlock it back later. + // Without this, the block stays in the free queue while pinned, + // and the subsequent releaseBlock creates a duplicate entry. + if (!searchRoot->hasRefs()) + { + mEvictionPolicy->claimBlock(searchRoot, searchRoot->getPriority(), searchRoot->getDurationMs()); + } searchRoot->incRefCount(); pinnedBlockIds.push_back(searchRoot->getBlockId()); } diff --git a/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py b/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py index 17eee1d339c1..1e3a3666e7a4 100644 --- a/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py @@ -163,14 +163,21 @@ def compute_batch_gpu_times(self, requests): # Compute once per batch, reuse for all requests if batch_gpu_forward_time is None: - batch_gpu_forward_time = perf.gpu_forward_start_event.elapsed_time( - perf.gpu_forward_end_event - ) - batch_gpu_sample_time = ( - perf.gpu_forward_end_event.elapsed_time(perf.gpu_sample_end_event) - if perf.gpu_sample_end_event - else 0.0 - ) + try: + batch_gpu_forward_time = perf.gpu_forward_start_event.elapsed_time( + perf.gpu_forward_end_event + ) + batch_gpu_sample_time = ( + perf.gpu_forward_end_event.elapsed_time(perf.gpu_sample_end_event) + if perf.gpu_sample_end_event + else 0.0 + ) + except RuntimeError: + # CUDA event timing can fail if events were not recorded + # on the current stream. Skip metrics for this batch rather + # than crashing the executor thread. + batch_gpu_forward_time = 0.0 + batch_gpu_sample_time = 0.0 target["gpu_forward_time"] = batch_gpu_forward_time target["gpu_sample_time"] = batch_gpu_sample_time