diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 9d28fa26c4ee..1ec047e551b6 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -27,6 +27,8 @@ #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/utils/pgUtils.h" +#include +#include #include #include #include @@ -55,6 +57,7 @@ class BaseKVCacheManager; class CacheSender; class CacheReceiver; +class ContextTransferCoordinator; class CacheTransceiverComm { @@ -149,6 +152,25 @@ class CacheTransceiverComm TLLM_THROW("Input arguments only supported in mpi"); } + [[nodiscard]] std::unique_ptr sendAsync( + void const* buffer, std::size_t size, mpi::MpiType dtype, int dest, mpi::MpiTag tag) const + { + TLLM_CHECK_WITH_INFO(isMpi(), "Point-to-point cache-transceiver status messages require MPI."); + return mMpiComm->sendAsync(buffer, size, dtype, dest, tag); + } + + [[nodiscard]] bool iprobe(int source, mpi::MpiTag tag, MPI_Status* status) const + { + TLLM_CHECK_WITH_INFO(isMpi(), "Point-to-point cache-transceiver status messages require MPI."); + return mMpiComm->iprobe(source, tag, status); + } + + void recv(void* buffer, std::size_t size, mpi::MpiType dtype, int source, mpi::MpiTag tag) const + { + TLLM_CHECK_WITH_INFO(isMpi(), "Point-to-point cache-transceiver status messages require MPI."); + static_cast(mMpiComm->recv(buffer, size, dtype, source, tag)); + } + CacheTransceiverComm split(int color, int key) { if (isMpi()) @@ -284,6 +306,9 @@ class CacheTransceiver : public BaseCacheTransceiver void setContextState(LlmRequest* llmRequest); + void logPhaseTrace(char const* event, std::optional requestId = std::nullopt, + char const* outcome = "none") const; + std::unique_ptr mCacheSender; std::unique_ptr mCacheReceiver; // shared_ptr (not raw LlmRequest*) so the futures hold a strong reference for @@ -307,11 +332,17 @@ class CacheTransceiver : public BaseCacheTransceiver std::shared_ptr mGroupComm; std::shared_ptr mGroupTensorParaComm, mGroupPipeParaComm, mGroupDataComm, mGroupTPInDPComm; + std::unique_ptr mContextTransferCoordinator; executor::kv_cache::CommState const* mCommState; std::unique_ptr mCacheState; std::unique_ptr mManager; std::optional mCacheTransceiverConfig; + bool mPhaseTraceEnabled{false}; + int mPhaseTraceWorldRank{-1}; + int mPhaseTracePpRank{-1}; + std::uint64_t mPhaseTraceStatusPolls{0}; + std::uint64_t mPhaseTraceFutureNotReady{0}; std::vector> mCacheTransBufferManagers; std::vector mCacheTransBufferManagerPtrs; diff --git a/cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h b/cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h new file mode 100644 index 000000000000..9b156c8887a0 --- /dev/null +++ b/cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h @@ -0,0 +1,109 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager +{ + +class CacheTransceiverComm; + +enum class ContextTransferVote : std::uint64_t +{ + kCompleted = 1, + kFailed = 2, +}; + +struct ContextTransferConsensusResult +{ + std::unordered_set completedRequestIds; + std::unordered_set failedRequestIds; + std::unordered_set timedOutRequestIds; +}; + +//! Accumulates one immutable terminal vote per participant and request. +class ContextTransferVoteReducer +{ +public: + explicit ContextTransferVoteReducer(int participantCount); + + void recordVote(int participantRank, std::uint64_t requestId, ContextTransferVote vote); + + //! Record a sticky nonterminal timeout proposal. The final outcome still waits for every terminal vote. + void recordTimeout(std::uint64_t requestId); + + [[nodiscard]] ContextTransferConsensusResult takeReady(); + + void clear() noexcept; + +private: + struct RequestVotes + { + explicit RequestVotes(int participantCount) + : votes(static_cast(participantCount), 0) + { + } + + std::vector votes; + int terminalCount{0}; + bool failed{false}; + bool timedOut{false}; + bool timeoutPending{false}; + }; + + int mParticipantCount; + std::unordered_map mRequestVotes; +}; + +//! Coordinates asynchronous context-transfer outcomes across a rank group. +//! +//! Timeout and terminal events share one ordered stream toward the coordinator. Timeout updates and final commits +//! share another ordered stream toward followers, so a request that times out can never commit success first. +class ContextTransferCoordinator +{ +public: + explicit ContextTransferCoordinator(std::shared_ptr comm); + ~ContextTransferCoordinator(); + + ContextTransferCoordinator(ContextTransferCoordinator const&) = delete; + ContextTransferCoordinator& operator=(ContextTransferCoordinator const&) = delete; + + //! Publish this rank's immutable local terminal outcome without waiting for peers. + void publishLocalOutcome(std::uint64_t requestId, bool failed); + + //! Publish an idempotent, sticky, nonterminal timeout proposal without waiting for peers. + void publishTimeout(std::uint64_t requestId); + + //! Make nonblocking protocol progress and return newly committed global outcomes. + [[nodiscard]] ContextTransferConsensusResult poll(); + + //! Exchange ordered close markers so no active MPI request outlives its backing buffer. Failure aborts closed. + void shutdown() noexcept; + +private: + class Impl; + std::unique_ptr mImpl; +}; + +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h b/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h index 32c086c84ee9..49c50e49b2e2 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h +++ b/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -71,7 +71,11 @@ enum class MpiTag : int // KvCacheEventManager kKvCacheEventSize = 1026, - kKvCacheEvent = 1027 + kKvCacheEvent = 1027, + + // Asynchronous context-transfer coordination. + kContextTransferEvent = 1028, + kContextTransferUpdate = 1029 }; } // namespace tensorrt_llm::mpi diff --git a/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h b/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h index 75ec7a534815..b474b4b12cbf 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h +++ b/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -235,6 +235,17 @@ class MpiRequest #endif } + [[nodiscard]] bool isCompleted() + { +#if ENABLE_MULTI_DEVICE + int completed = 0; + TLLM_MPI_CHECK(MPI_Test(&mRequest, &completed, MPI_STATUS_IGNORE)); + return completed != 0; +#else + TLLM_THROW("Multi device support is disabled."); +#endif + } + void cancel() { #if ENABLE_MULTI_DEVICE diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index d893d12fa3e3..c7fea62c023a 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & # AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); you may not @@ -30,6 +30,7 @@ set(SRCS capacityScheduler.cpp createNewDecoderRequests.cpp contextProgress.cpp + contextTransferCoordinator.cpp dataTransceiver.cpp decoderBuffers.cpp kvCacheManager.cpp diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index f730fb2aaf1e..ac5974e607fb 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -17,6 +17,7 @@ #include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h" #include "tensorrt_llm/executor/types.h" +#include #include #include #include @@ -36,6 +37,7 @@ #include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" #include "tensorrt_llm/batch_manager/contextProgress.h" +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" #include "tensorrt_llm/batch_manager/dataTransceiver.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/kvCacheType.h" @@ -73,6 +75,13 @@ using RequestIdType = LlmRequest::RequestIdType; constexpr int kTransferFuturePollIntervalMs = 10; +std::uint64_t monotonicNs() +{ + return static_cast( + std::chrono::duration_cast(std::chrono::steady_clock::now().time_since_epoch()) + .count()); +} + // Finite status checks are scheduler polls, not terminal deadlines. Pure polls // use short slices; calls that ask for at least one completion keep bounded // backpressure by waiting up to the configured future timeout. @@ -351,6 +360,9 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa std::vector const& rnnLayerNumPerPP) : mCacheTransceiverConfig{cacheTransceiverConfig} { + mPhaseTraceEnabled = common::getBoolEnv("TRTLLM_NVBUG_6448152_PHASE_TRACE"); + mPhaseTraceWorldRank = worldConfig.getRank(); + mPhaseTracePpRank = worldConfig.getPipelineParallelRank(); using tensorrt_llm::batch_manager::kv_cache_manager::CacheFormatter; TLLM_CHECK_WITH_INFO(mCacheTransceiverConfig.has_value(), "CacheTransceiverConfig is not set."); auto const backendType = mCacheTransceiverConfig.value().getBackendType(); @@ -569,15 +581,57 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa mCacheSender = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); mCacheReceiver = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); + // Keep automatic enablement within the currently qualified C++ NIXL/UCX TP1/CP1 pipeline topology. + bool const coordinatorTopologyEligible = worldConfig.getPipelineParallelism() > 1 && useMPI() + && backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL + && common::getEnvNixlBackend() == "UCX" && worldConfig.getTensorParallelism() == 1 + && worldConfig.getContextParallelism() == 1 && !mCacheState->getParallelConfig().mEnableAttentionDP; + if (worldConfig.getPipelineParallelism() > 1 && useMPI()) + { + TLLM_CHECK(mGroupPipeParaComm != nullptr); + constexpr std::uint64_t kCoordinatorProtocolVersion = 1; + std::uint64_t const localVersion = coordinatorTopologyEligible ? kCoordinatorProtocolVersion : 0; + bool const cancellationEnabled = common::getEnvDisaggEnableInflightCancel(); + std::uint64_t const localProtocolMode = (localVersion << 1) | static_cast(cancellationEnabled); + std::vector protocolModes(static_cast(mGroupPipeParaComm->getSize())); + mGroupPipeParaComm->allgather(&localProtocolMode, protocolModes.data(), 1, mpi::MpiType::kUINT64); + TLLM_CHECK_WITH_INFO(std::all_of(protocolModes.begin(), protocolModes.end(), + [&](std::uint64_t const mode) { return mode == localProtocolMode; }), + "Context-transfer consensus protocol version or cancellation mode differs across PP ranks."); + if (localVersion != 0) + { + mContextTransferCoordinator = std::make_unique(mGroupPipeParaComm); + TLLM_LOG_INFO( + "Enable asynchronous context-transfer consensus version %llu for PP group of size %d; in-flight " + "cancellation=%s.", + static_cast(kCoordinatorProtocolVersion), mGroupPipeParaComm->getSize(), + cancellationEnabled ? "enabled" : "disabled"); + } + } + initializeCommState(); + logPhaseTrace( + "cpp_mode", std::nullopt, mContextTransferCoordinator ? "coordinator_active" : "coordinator_inactive"); } CacheTransceiver::~CacheTransceiver() { + if (mPhaseTraceEnabled) + { + TLLM_LOG_INFO( + "NVBUG6448152_PHASE event=cpp_summary world_rank=%d pp_rank=%d request_id=0 request_id_valid=0 " + "monotonic_ns=%llu outcome=none status_polls=%llu future_not_ready=%llu sender_futures=%zu " + "awaiting_consensus=%zu", + mPhaseTraceWorldRank, mPhaseTracePpRank, static_cast(monotonicNs()), + static_cast(mPhaseTraceStatusPolls), + static_cast(mPhaseTraceFutureNotReady), mSenderFutures.size(), + mSenderRequestsAwaitingConsensus.size()); + } // Stop sender/receiver workers while the connection manager and transfer // plugin are still alive. The workers can access both during termination. mCacheSender.reset(); mCacheReceiver.reset(); + mContextTransferCoordinator.reset(); if (mWrapperLibHandle) { @@ -586,6 +640,20 @@ CacheTransceiver::~CacheTransceiver() } } +void CacheTransceiver::logPhaseTrace( + char const* event, std::optional requestId, char const* outcome) const +{ + if (!mPhaseTraceEnabled) + { + return; + } + TLLM_LOG_INFO("NVBUG6448152_PHASE event=%s world_rank=%d pp_rank=%d request_id=%" PRIu64 + " request_id_valid=%d monotonic_ns=%llu outcome=%s", + event, mPhaseTraceWorldRank, mPhaseTracePpRank, + requestId.has_value() ? static_cast(requestId.value()) : std::uint64_t{0}, + requestId.has_value() ? 1 : 0, static_cast(monotonicNs()), outcome); +} + void CacheTransceiver::initializeCommState() { mCommState = std::addressof(mCacheSender->getCommState()); @@ -624,7 +692,9 @@ void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmReques return; } setContextState(llmRequest.get()); + logPhaseTrace("cpp_send_async_submit", llmRequest->mRequestId); auto future = mCacheSender->sendAsync(llmRequest); + logPhaseTrace("cpp_send_async_return", llmRequest->mRequestId); mSenderFutures.emplace_back(std::move(llmRequest), std::move(future)); } @@ -640,7 +710,9 @@ void CacheTransceiver::respondAndSendLayerWise( llmRequest->setState(LlmRequestState::kDISAGG_CONTEXT_INIT_AND_TRANS); setContextState(llmRequest.get()); + logPhaseTrace("cpp_send_async_submit", llmRequest->mRequestId); auto future = mCacheSender->sendAsync(llmRequest); + logPhaseTrace("cpp_send_async_return", llmRequest->mRequestId); mSenderFutures.emplace_back(llmRequest, std::move(future)); } } @@ -799,6 +871,20 @@ void updateKVCacheTransferBW(std::shared_ptr const& mComm, RequestStatuses CacheTransceiver::checkContextTransferStatus( std::optional const& atLeastRequestNum, bool markComplete) { + if (mPhaseTraceEnabled) + { + ++mPhaseTraceStatusPolls; + if ((mPhaseTraceStatusPolls & (mPhaseTraceStatusPolls - 1)) == 0) + { + TLLM_LOG_INFO( + "NVBUG6448152_PHASE event=status_poll_checkpoint world_rank=%d pp_rank=%d request_id=0 " + "request_id_valid=0 monotonic_ns=%llu outcome=none status_polls=%llu sender_futures=%zu " + "awaiting_consensus=%zu", + mPhaseTraceWorldRank, mPhaseTracePpRank, static_cast(monotonicNs()), + static_cast(mPhaseTraceStatusPolls), mSenderFutures.size(), + mSenderRequestsAwaitingConsensus.size()); + } + } bool const blockAll = !atLeastRequestNum.has_value(); bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); TLLM_CHECK_WITH_INFO(!inflightCancelEnabled || !blockAll, @@ -823,10 +909,25 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( std::vector contextCompleteRequestIds; for (auto&& [request, future] : mSenderFutures) { - if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + auto const status = future.wait_for(std::chrono::milliseconds(0)); + if (status == std::future_status::ready) { contextCompleteRequestIds.push_back(request->mRequestId); } + else if (mPhaseTraceEnabled) + { + ++mPhaseTraceFutureNotReady; + if ((mPhaseTraceFutureNotReady & (mPhaseTraceFutureNotReady - 1)) == 0) + { + TLLM_LOG_INFO( + "NVBUG6448152_PHASE event=future_not_ready_checkpoint world_rank=%d pp_rank=%d request_id=0 " + "request_id_valid=0 monotonic_ns=%llu outcome=none future_not_ready=%llu sender_futures=%zu " + "awaiting_consensus=%zu", + mPhaseTraceWorldRank, mPhaseTracePpRank, static_cast(monotonicNs()), + static_cast(mPhaseTraceFutureNotReady), mSenderFutures.size(), + mSenderRequestsAwaitingConsensus.size()); + } + } } std::unordered_map frequencyMap; @@ -868,6 +969,28 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( toCompleteIdSet.insert(request->mRequestId); } + auto recordTimeout = [&](RequestIdType const requestId) + { + bool const inserted = mTimedOutSenderIds.insert(requestId).second; + if (inserted && inflightCancelEnabled && mContextTransferCoordinator) + { + mContextTransferCoordinator->publishTimeout(requestId); + } + return inserted; + }; + auto recordOutcome + = [&](RequestIdType const requestId, std::shared_ptr const& request, bool const failed) + { + logPhaseTrace("local_future_terminal", requestId, failed ? "failed" : "completed"); + recordLocalTransferOutcome(requestId, request, failed, mCompletedSenderRequestIds, mFailedSenderRequestIds, + mSenderRequestsAwaitingConsensus); + if (mContextTransferCoordinator) + { + logPhaseTrace("coordinator_vote", requestId, failed ? "failed" : "completed"); + mContextTransferCoordinator->publishLocalOutcome(requestId, failed); + } + }; + // Record local terminal outcomes for requests selected this round. The // request is reported only after all ranks in the sync group agree that the // request reached a terminal state. @@ -881,7 +1004,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( auto const elapsedMs = getTransferElapsedMs(request, LlmRequest::getSteadyClockNow()); if (elapsedMs > kvTransferTimeoutMs.value()) { - if (mTimedOutSenderIds.insert(requestId).second) + if (recordTimeout(requestId)) { TLLM_LOG_WARNING( "Context KV cache transfer for request %ld exceeded configured timeout: " @@ -893,6 +1016,8 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } if (blockAll || (toCompleteIdSet.find(requestId) != toCompleteIdSet.end())) { + bool terminal = false; + bool failed = false; try { auto const status = blockAll ? std::future_status::ready : future.wait_for(futureWaitInterval); @@ -902,7 +1027,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( if (kvTransferTimeoutMs.has_value()) { auto const elapsedMs = getTransferElapsedMs(request, request->getKvCacheTransferEnd()); - if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutSenderIds.insert(requestId).second) + if (elapsedMs > kvTransferTimeoutMs.value() && recordTimeout(requestId)) { TLLM_LOG_WARNING( "Context KV cache transfer for request %ld completed after its deadline: " @@ -911,9 +1036,8 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( inflightCancelEnabled ? "failing request" : "observe-only"); } } - recordLocalTransferOutcome(requestId, request, /*failed=*/false, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - it = mSenderFutures.erase(it); + failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR; + terminal = true; } else if (status == std::future_status::timeout) { @@ -927,25 +1051,28 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { TLLM_LOG_ERROR( "Future returned unexpected status for request %ld. Recording as failed.", requestId); - - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - it = mSenderFutures.erase(it); + failed = true; + terminal = true; } } catch (std::exception const& e) { TLLM_LOG_ERROR("Error occurred during context transfer for request %ld: %s", requestId, e.what()); - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - it = mSenderFutures.erase(it); + failed = true; + terminal = true; } catch (...) { TLLM_LOG_ERROR("Unknown error occurred during context transfer for request %ld", requestId); - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); + failed = true; + terminal = true; + } + if (terminal) + { + auto terminalRequest = request; it = mSenderFutures.erase(it); + // Publish outside the transfer-future try/catch. A protocol error must not rewrite an immutable vote. + recordOutcome(requestId, terminalRequest, failed); } } else @@ -955,12 +1082,51 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } RequestStatuses requestsStatus{}; - auto const consensusOutcome = reduceTransferStates(syncComm, mGroupPipeParaComm, mCompletedSenderRequestIds, - mFailedSenderRequestIds, inflightCancelEnabled ? mTimedOutSenderIds : std::unordered_set{}); + TransferConsensusOutcome consensusOutcome; + if (mContextTransferCoordinator) + { + auto mergeCoordinatorOutcome = [&]() + { + auto coordinatorOutcome = mContextTransferCoordinator->poll(); + consensusOutcome.completedRequestIds.insert( + coordinatorOutcome.completedRequestIds.begin(), coordinatorOutcome.completedRequestIds.end()); + consensusOutcome.failedRequestIds.insert( + coordinatorOutcome.failedRequestIds.begin(), coordinatorOutcome.failedRequestIds.end()); + consensusOutcome.timedOutRequestIds.insert( + coordinatorOutcome.timedOutRequestIds.begin(), coordinatorOutcome.timedOutRequestIds.end()); + }; + do + { + mergeCoordinatorOutcome(); + if (blockAll + && consensusOutcome.completedRequestIds.size() + consensusOutcome.failedRequestIds.size() + < mSenderRequestsAwaitingConsensus.size()) + { + std::this_thread::yield(); + } + } while (blockAll + && consensusOutcome.completedRequestIds.size() + consensusOutcome.failedRequestIds.size() + < mSenderRequestsAwaitingConsensus.size()); + + if (inflightCancelEnabled) + { + // A timeout update is transmitted once, but cancellation may be declined transiently. Keep the globally + // observed timeout active so rank-local cancellation is retried on every poll until terminal commit. + consensusOutcome.timedOutRequestIds.insert(mTimedOutSenderIds.begin(), mTimedOutSenderIds.end()); + } + } + else + { + consensusOutcome = reduceTransferStates(syncComm, mGroupPipeParaComm, mCompletedSenderRequestIds, + mFailedSenderRequestIds, inflightCancelEnabled ? mTimedOutSenderIds : std::unordered_set{}); + } if (inflightCancelEnabled) { for (auto const requestId : consensusOutcome.timedOutRequestIds) { + // Persist the global timeout even if this rank has not registered its local future yet. The one-shot + // coordinator update must remain actionable when that future appears on a later scheduler poll. + mTimedOutSenderIds.insert(requestId); auto const futureIt = std::find_if(mSenderFutures.begin(), mSenderFutures.end(), [requestId](auto const& entry) { return entry.first->mRequestId == requestId; }); if (futureIt == mSenderFutures.end() @@ -969,7 +1135,6 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { continue; } - mTimedOutSenderIds.insert(requestId); if (requestCancellationNoThrow( requestId, "Context", [&]() { return mCacheSender->cancelRequest(*futureIt->first); })) { @@ -989,6 +1154,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( continue; } requestIt->second->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + logPhaseTrace("coordinator_commit", requestId, "failed"); requestsStatus.errorRequestIds.insert(requestId); mTimedOutSenderIds.erase(requestId); mCancelRequestedSenderIds.erase(requestId); @@ -1003,6 +1169,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( continue; } requestsStatus.completedRequestIds.insert(requestId); + logPhaseTrace("coordinator_commit", requestId, "completed"); if (markComplete) { requestIt->second->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); diff --git a/cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp b/cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp new file mode 100644 index 000000000000..909b7858c574 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp @@ -0,0 +1,475 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" + +#include "tensorrt_llm/batch_manager/cacheTransceiver.h" +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/logger.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager +{ + +ContextTransferVoteReducer::ContextTransferVoteReducer(int const participantCount) + : mParticipantCount(participantCount) +{ + TLLM_CHECK_WITH_INFO(participantCount > 0, "Context-transfer consensus requires at least one participant."); +} + +void ContextTransferVoteReducer::recordVote( + int const participantRank, std::uint64_t const requestId, ContextTransferVote const vote) +{ + TLLM_CHECK_WITH_INFO(participantRank >= 0 && participantRank < mParticipantCount, + "Context-transfer consensus participant rank is out of range."); + TLLM_CHECK_WITH_INFO(vote == ContextTransferVote::kCompleted || vote == ContextTransferVote::kFailed, + "Context-transfer consensus received an invalid vote."); + + auto [requestIt, inserted] = mRequestVotes.try_emplace(requestId, mParticipantCount); + static_cast(inserted); + auto& requestVotes = requestIt->second; + auto& recordedVote = requestVotes.votes.at(static_cast(participantRank)); + auto const packedVote = static_cast(vote); + if (recordedVote != 0) + { + TLLM_CHECK_WITH_INFO( + recordedVote == packedVote, "Context-transfer participant changed its terminal vote for a request."); + return; + } + + recordedVote = packedVote; + ++requestVotes.terminalCount; + requestVotes.failed = requestVotes.failed || vote == ContextTransferVote::kFailed; +} + +void ContextTransferVoteReducer::recordTimeout(std::uint64_t const requestId) +{ + auto [requestIt, inserted] = mRequestVotes.try_emplace(requestId, mParticipantCount); + static_cast(inserted); + auto& requestVotes = requestIt->second; + if (!requestVotes.timedOut) + { + requestVotes.timedOut = true; + requestVotes.timeoutPending = true; + } +} + +ContextTransferConsensusResult ContextTransferVoteReducer::takeReady() +{ + ContextTransferConsensusResult result; + for (auto requestIt = mRequestVotes.begin(); requestIt != mRequestVotes.end();) + { + auto& requestVotes = requestIt->second; + if (requestVotes.timeoutPending) + { + result.timedOutRequestIds.insert(requestIt->first); + requestVotes.timeoutPending = false; + } + if (requestVotes.terminalCount != mParticipantCount) + { + ++requestIt; + continue; + } + + auto& terminalRequestIds + = (requestVotes.failed || requestVotes.timedOut) ? result.failedRequestIds : result.completedRequestIds; + terminalRequestIds.insert(requestIt->first); + requestIt = mRequestVotes.erase(requestIt); + } + return result; +} + +void ContextTransferVoteReducer::clear() noexcept +{ + mRequestVotes.clear(); +} + +class ContextTransferCoordinator::Impl +{ +public: + explicit Impl(std::shared_ptr comm) + : mComm(std::move(comm)) + , mCoordinatorRank(0) + , mReducer(mComm ? mComm->getSize() : 1) + { + TLLM_CHECK_WITH_INFO(mComm != nullptr, "Context-transfer coordination requires a communicator."); + TLLM_CHECK_WITH_INFO(mComm->isMpi(), "Asynchronous context-transfer coordination requires MPI."); + TLLM_CHECK_WITH_INFO( + mComm->getSize() > 1, "Asynchronous context-transfer coordination requires multiple participants."); + mCoordinatorRank = mComm->getSize() - 1; + } + + void publishLocalOutcome(std::uint64_t const requestId, bool const failed) + { + TLLM_CHECK_WITH_INFO(!mShutdown, "Cannot publish a context-transfer vote after coordinator shutdown."); + auto const vote = failed ? ContextTransferVote::kFailed : ContextTransferVote::kCompleted; + auto const [voteIt, inserted] = mPublishedLocalVotes.emplace(requestId, vote); + TLLM_CHECK_WITH_INFO( + inserted || voteIt->second == vote, "This rank changed its terminal vote for a context transfer."); + if (!inserted) + { + return; + } + + try + { + if (isCoordinator()) + { + mReducer.recordVote(mComm->getRank(), requestId, vote); + } + else + { + queuePacket( + requestId, static_cast(vote), mCoordinatorRank, mpi::MpiTag::kContextTransferEvent); + } + } + catch (...) + { + mPublishedLocalVotes.erase(voteIt); + throw; + } + } + + void publishTimeout(std::uint64_t const requestId) + { + TLLM_CHECK_WITH_INFO(!mShutdown, "Cannot publish a context-transfer timeout after coordinator shutdown."); + TLLM_CHECK_WITH_INFO(mPublishedLocalVotes.find(requestId) == mPublishedLocalVotes.end(), + "A context-transfer timeout must be published before its immutable terminal vote."); + auto const [timeoutIt, inserted] = mPublishedTimeouts.insert(requestId); + if (!inserted) + { + return; + } + + try + { + if (isCoordinator()) + { + mReducer.recordTimeout(requestId); + } + else + { + queuePacket(requestId, kTimedOutSignal, mCoordinatorRank, mpi::MpiTag::kContextTransferEvent); + } + } + catch (...) + { + mPublishedTimeouts.erase(timeoutIt); + throw; + } + } + + [[nodiscard]] ContextTransferConsensusResult poll() + { + TLLM_CHECK_WITH_INFO(!mShutdown, "Cannot poll context-transfer coordination after shutdown."); + return progress(); + } + + void shutdown() noexcept + { + if (mShutdown) + { + return; + } + mShutdown = true; + + try + { + if (!isCoordinator()) + { + queuePacket(/*requestId=*/0, kCloseMarker, mCoordinatorRank, mpi::MpiTag::kContextTransferEvent); + } + + auto const deadline = std::chrono::steady_clock::now() + kShutdownTimeout; + while (!shutdownComplete()) + { + static_cast(progress()); + if (!shutdownComplete()) + { + if (std::chrono::steady_clock::now() >= deadline) + { + TLLM_LOG_ERROR( + "Timed out shutting down asynchronous context-transfer coordinator; rank=%d " + "peer_closes=%zu/%d ack=%d pending_sends=%zu. Aborting to avoid freeing active MPI " + "requests.", + mComm->getRank(), mClosedPeers.size(), mComm->getSize() - 1, mCloseAcknowledged, + mPendingSends.size()); + std::abort(); + } + std::this_thread::yield(); + } + } + } + catch (std::exception const& error) + { + TLLM_LOG_ERROR("Failed to shut down asynchronous context-transfer coordinator: %s", error.what()); + std::abort(); + } + catch (...) + { + TLLM_LOG_ERROR("Failed to shut down asynchronous context-transfer coordinator with an unknown error."); + std::abort(); + } + } + +private: + static constexpr std::size_t kPacketFieldCount = 2; + static constexpr std::uint64_t kTimedOutSignal = 3; + static constexpr std::uint64_t kCloseMarker = 4; + static constexpr auto kShutdownTimeout = std::chrono::seconds(30); + + struct PendingSend + { + std::array packet{}; + std::unique_ptr request; + }; + + [[nodiscard]] bool isCoordinator() const + { + return mComm->getRank() == mCoordinatorRank; + } + + void queuePacket(std::uint64_t const requestId, std::uint64_t const value, int const peer, mpi::MpiTag const tag) + { + mPendingSends.emplace_back(); + auto& pendingSend = mPendingSends.back(); + pendingSend.packet = {requestId, value}; + try + { + pendingSend.request = mComm->sendAsync( + pendingSend.packet.data(), pendingSend.packet.size(), mpi::MpiType::kUINT64, peer, tag); + } + catch (...) + { + mPendingSends.pop_back(); + throw; + } + } + + void queueUpdateForPeers(std::uint64_t const requestId, std::uint64_t const update) + { + for (int peer = 0; peer < mComm->getSize(); ++peer) + { + if (peer != mCoordinatorRank) + { + queuePacket(requestId, update, peer, mpi::MpiTag::kContextTransferUpdate); + } + } + } + + void reapCompletedSends() + { + for (auto sendIt = mPendingSends.begin(); sendIt != mPendingSends.end();) + { + TLLM_CHECK(sendIt->request); + if (sendIt->request->isCompleted()) + { + sendIt = mPendingSends.erase(sendIt); + } + else + { + ++sendIt; + } + } + } + + void drainVotes() + { + for (int peer = 0; peer < mComm->getSize(); ++peer) + { + if (peer == mCoordinatorRank) + { + continue; + } + + MPI_Status status{}; + while (mComm->iprobe(peer, mpi::MpiTag::kContextTransferEvent, &status)) + { + std::array packet{}; + mComm->recv( + packet.data(), packet.size(), mpi::MpiType::kUINT64, peer, mpi::MpiTag::kContextTransferEvent); + if (packet.back() == kCloseMarker) + { + TLLM_CHECK_WITH_INFO( + mClosedPeers.insert(peer).second, "Received a duplicate context-transfer close marker."); + continue; + } + TLLM_CHECK_WITH_INFO(mClosedPeers.find(peer) == mClosedPeers.end(), + "Received a context-transfer vote after its peer close marker."); + if (packet.back() == kTimedOutSignal) + { + mReducer.recordTimeout(packet.front()); + continue; + } + mReducer.recordVote(peer, packet.front(), static_cast(packet.back())); + } + } + } + + ContextTransferConsensusResult completeCoordinatorUpdates() + { + auto result = mReducer.takeReady(); + for (auto const requestId : result.timedOutRequestIds) + { + queueUpdateForPeers(requestId, kTimedOutSignal); + } + for (auto const requestId : result.failedRequestIds) + { + queueUpdateForPeers(requestId, static_cast(ContextTransferVote::kFailed)); + mPublishedLocalVotes.erase(requestId); + mPublishedTimeouts.erase(requestId); + } + for (auto const requestId : result.completedRequestIds) + { + queueUpdateForPeers(requestId, static_cast(ContextTransferVote::kCompleted)); + mPublishedLocalVotes.erase(requestId); + mPublishedTimeouts.erase(requestId); + } + return result; + } + + ContextTransferConsensusResult drainUpdates() + { + ContextTransferConsensusResult result; + MPI_Status status{}; + while (mComm->iprobe(mCoordinatorRank, mpi::MpiTag::kContextTransferUpdate, &status)) + { + std::array packet{}; + mComm->recv(packet.data(), packet.size(), mpi::MpiType::kUINT64, mCoordinatorRank, + mpi::MpiTag::kContextTransferUpdate); + if (packet.back() == kCloseMarker) + { + TLLM_CHECK_WITH_INFO(!mCloseAcknowledged, "Received a duplicate coordinator close marker."); + mCloseAcknowledged = true; + mPublishedLocalVotes.clear(); + mPublishedTimeouts.clear(); + continue; + } + + if (packet.back() == kTimedOutSignal) + { + result.timedOutRequestIds.insert(packet.front()); + continue; + } + + auto const outcome = static_cast(packet.back()); + TLLM_CHECK_WITH_INFO(outcome == ContextTransferVote::kCompleted || outcome == ContextTransferVote::kFailed, + "Received an invalid context-transfer commit outcome."); + auto const localVoteIt = mPublishedLocalVotes.find(packet.front()); + TLLM_CHECK_WITH_INFO(localVoteIt != mPublishedLocalVotes.end(), + "Received a context-transfer commit before publishing the local terminal vote."); + if (outcome == ContextTransferVote::kFailed) + { + result.failedRequestIds.insert(packet.front()); + } + else + { + result.completedRequestIds.insert(packet.front()); + } + mPublishedLocalVotes.erase(localVoteIt); + mPublishedTimeouts.erase(packet.front()); + } + return result; + } + + ContextTransferConsensusResult progress() + { + reapCompletedSends(); + if (isCoordinator()) + { + drainVotes(); + auto result = completeCoordinatorUpdates(); + if (mShutdown && !mCloseSent && mClosedPeers.size() == static_cast(mComm->getSize() - 1)) + { + // Shutdown is an explicit abort epoch. A peer close proves that no more votes will arrive from that + // peer, so incomplete requests cannot reach a global decision and are intentionally abandoned. + mReducer.clear(); + mPublishedLocalVotes.clear(); + mPublishedTimeouts.clear(); + for (int peer = 0; peer < mCoordinatorRank; ++peer) + { + queuePacket(/*requestId=*/0, kCloseMarker, peer, mpi::MpiTag::kContextTransferUpdate); + } + mCloseSent = true; + } + return result; + } + return drainUpdates(); + } + + [[nodiscard]] bool shutdownComplete() const + { + if (isCoordinator()) + { + return mCloseSent && mPendingSends.empty(); + } + return mCloseAcknowledged && mPendingSends.empty(); + } + + std::shared_ptr mComm; + int mCoordinatorRank; + ContextTransferVoteReducer mReducer; + std::unordered_map mPublishedLocalVotes; + std::unordered_set mPublishedTimeouts; + std::unordered_set mClosedPeers; + std::list mPendingSends; + bool mShutdown{false}; + bool mCloseSent{false}; + bool mCloseAcknowledged{false}; +}; + +ContextTransferCoordinator::ContextTransferCoordinator(std::shared_ptr comm) + : mImpl(std::make_unique(std::move(comm))) +{ +} + +ContextTransferCoordinator::~ContextTransferCoordinator() +{ + shutdown(); +} + +void ContextTransferCoordinator::publishLocalOutcome(std::uint64_t const requestId, bool const failed) +{ + mImpl->publishLocalOutcome(requestId, failed); +} + +void ContextTransferCoordinator::publishTimeout(std::uint64_t const requestId) +{ + mImpl->publishTimeout(requestId); +} + +ContextTransferConsensusResult ContextTransferCoordinator::poll() +{ + return mImpl->poll(); +} + +void ContextTransferCoordinator::shutdown() noexcept +{ + if (mImpl) + { + mImpl->shutdown(); + } +} + +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index a8f6c7dd4e8f..75329d192a8c 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -20,6 +20,7 @@ add_gtest(cacheTransBufferTest cacheTransBufferTest.cpp) add_gtest(bufferIndexHolderTest bufferIndexHolderTest.cpp) add_gtest(capacitySchedulerTest capacitySchedulerTest.cpp) add_gtest(contextProgressTest contextProgressTest.cu) +add_gtest(contextTransferCoordinatorTest contextTransferCoordinatorTest.cpp) add_gtest(evictionPolicyTest evictionPolicyTest.cpp) add_gtest(kvCacheManagerTest kvCacheManagerTest.cpp) add_gtest(kvCacheManagerFabricMemoryTest kvCacheManagerFabricMemoryTest.cpp) diff --git a/cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp b/cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp new file mode 100644 index 000000000000..7b0f7aba3fc4 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp @@ -0,0 +1,143 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" + +#include + +namespace tensorrt_llm::batch_manager +{ +namespace +{ + +TEST(ContextTransferVoteReducerTest, WaitsForEveryParticipant) +{ + ContextTransferVoteReducer reducer(4); + reducer.recordVote(0, 17, ContextTransferVote::kCompleted); + reducer.recordVote(1, 17, ContextTransferVote::kCompleted); + reducer.recordVote(2, 17, ContextTransferVote::kCompleted); + EXPECT_TRUE(reducer.takeReady().completedRequestIds.empty()); + + reducer.recordVote(3, 17, ContextTransferVote::kCompleted); + auto const result = reducer.takeReady(); + EXPECT_EQ(result.completedRequestIds, std::unordered_set{17}); + EXPECT_TRUE(result.failedRequestIds.empty()); + EXPECT_TRUE(result.timedOutRequestIds.empty()); +} + +TEST(ContextTransferVoteReducerTest, FailureWinsAfterEveryParticipantVotes) +{ + ContextTransferVoteReducer reducer(4); + reducer.recordVote(0, 23, ContextTransferVote::kCompleted); + reducer.recordVote(1, 23, ContextTransferVote::kFailed); + reducer.recordVote(2, 23, ContextTransferVote::kCompleted); + reducer.recordVote(3, 23, ContextTransferVote::kCompleted); + + auto const result = reducer.takeReady(); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{23}); +} + +TEST(ContextTransferVoteReducerTest, RejectsChangedTerminalVote) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordVote(0, 29, ContextTransferVote::kCompleted); + reducer.recordVote(0, 29, ContextTransferVote::kCompleted); + EXPECT_ANY_THROW(reducer.recordVote(0, 29, ContextTransferVote::kFailed)); +} + +TEST(ContextTransferVoteReducerTest, AccumulatesInterleavedRequestsIndependently) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordVote(0, 31, ContextTransferVote::kCompleted); + reducer.recordVote(1, 37, ContextTransferVote::kFailed); + reducer.recordVote(1, 31, ContextTransferVote::kCompleted); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.completedRequestIds, std::unordered_set{31}); + EXPECT_TRUE(result.failedRequestIds.empty()); + + reducer.recordVote(0, 37, ContextTransferVote::kCompleted); + result = reducer.takeReady(); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{37}); +} + +TEST(ContextTransferVoteReducerTest, RejectsInvalidInput) +{ + EXPECT_ANY_THROW(ContextTransferVoteReducer(0)); + ContextTransferVoteReducer reducer(2); + EXPECT_ANY_THROW(reducer.recordVote(-1, 41, ContextTransferVote::kCompleted)); + EXPECT_ANY_THROW(reducer.recordVote(2, 41, ContextTransferVote::kCompleted)); + EXPECT_ANY_THROW(reducer.recordVote(0, 41, static_cast(99))); +} + +TEST(ContextTransferVoteReducerTest, TimeoutIsVisibleBeforeTerminalFailure) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordTimeout(43); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.timedOutRequestIds, std::unordered_set{43}); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_TRUE(result.failedRequestIds.empty()); + + result = reducer.takeReady(); + EXPECT_TRUE(result.timedOutRequestIds.empty()); + EXPECT_TRUE(result.failedRequestIds.empty()); + + reducer.recordVote(0, 43, ContextTransferVote::kCompleted); + reducer.recordVote(1, 43, ContextTransferVote::kCompleted); + result = reducer.takeReady(); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{43}); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_TRUE(result.timedOutRequestIds.empty()); +} + +TEST(ContextTransferVoteReducerTest, DuplicateTimeoutIsIdempotent) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordTimeout(47); + reducer.recordTimeout(47); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.timedOutRequestIds, std::unordered_set{47}); + + reducer.recordTimeout(47); + result = reducer.takeReady(); + EXPECT_TRUE(result.timedOutRequestIds.empty()); +} + +TEST(ContextTransferVoteReducerTest, TimeoutAfterPartialTerminalVotesStillWins) +{ + ContextTransferVoteReducer reducer(3); + reducer.recordVote(2, 53, ContextTransferVote::kCompleted); + reducer.recordTimeout(53); + reducer.recordVote(0, 53, ContextTransferVote::kCompleted); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.timedOutRequestIds, std::unordered_set{53}); + EXPECT_TRUE(result.failedRequestIds.empty()); + + reducer.recordVote(1, 53, ContextTransferVote::kCompleted); + result = reducer.takeReady(); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{53}); + EXPECT_TRUE(result.completedRequestIds.empty()); +} + +} // namespace +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index e294cdff9d49..5d1a958d37ce 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -31,6 +31,7 @@ #include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" @@ -43,6 +44,7 @@ #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/testing/kvCacheManagerTestUtil.h" +#include #include #include #include @@ -54,6 +56,7 @@ #include #include #include +#include #include "gtest/gtest.h" #include @@ -90,6 +93,95 @@ T serializeDeserialize(T const& val) } // namespace +TEST(ContextTransferCoordinatorTest, CommitsStaggeredSuccessAndFailureWithoutCollectivePolling) +{ + auto& world = tensorrt_llm::mpi::MpiComm::world(); + if (world.getSize() < 2) + { + GTEST_SKIP() << "mpirun with at least two processes is required to run this test."; + } + + auto comm = std::make_shared(std::addressof(world)); + { + ContextTransferCoordinator coordinator(comm); + auto waitForOutcome = [&](std::uint64_t const requestId, bool const expectFailure) + { + bool observed = false; + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (!observed && std::chrono::steady_clock::now() < deadline) + { + auto result = coordinator.poll(); + observed = expectFailure ? result.failedRequestIds.count(requestId) != 0 + : result.completedRequestIds.count(requestId) != 0; + if (!observed) + { + std::this_thread::yield(); + } + } + EXPECT_TRUE(observed); + }; + auto waitForTimeout = [&](std::uint64_t const requestId) + { + bool observed = false; + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (!observed && std::chrono::steady_clock::now() < deadline) + { + auto const result = coordinator.poll(); + observed = result.timedOutRequestIds.count(requestId) != 0; + if (!observed) + { + std::this_thread::yield(); + } + } + EXPECT_TRUE(observed); + }; + + constexpr std::uint64_t kCompletedRequestId = 644815201; + world.barrier(); + if (world.getRank() == world.getSize() - 1) + { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + coordinator.publishLocalOutcome(kCompletedRequestId, /*failed=*/false); + if (world.getRank() != world.getSize() - 1) + { + auto const earlyResult = coordinator.poll(); + EXPECT_TRUE(earlyResult.completedRequestIds.empty()); + EXPECT_TRUE(earlyResult.failedRequestIds.empty()); + } + waitForOutcome(kCompletedRequestId, /*expectFailure=*/false); + + constexpr std::uint64_t kFailedRequestId = 644815202; + world.barrier(); + coordinator.publishLocalOutcome(kFailedRequestId, /*failed=*/world.getRank() == 0); + waitForOutcome(kFailedRequestId, /*expectFailure=*/true); + + constexpr std::uint64_t kTimedOutRequestId = 644815203; + world.barrier(); + if (world.getRank() == 0) + { + coordinator.publishTimeout(kTimedOutRequestId); + coordinator.publishTimeout(kTimedOutRequestId); + } + waitForTimeout(kTimedOutRequestId); + coordinator.publishLocalOutcome(kTimedOutRequestId, /*failed=*/false); + waitForOutcome(kTimedOutRequestId, /*expectFailure=*/true); + + constexpr std::uint64_t kInvalidOrderingRequestId = 644815204; + world.barrier(); + coordinator.publishLocalOutcome(kInvalidOrderingRequestId, /*failed=*/false); + EXPECT_ANY_THROW(coordinator.publishTimeout(kInvalidOrderingRequestId)); + waitForOutcome(kInvalidOrderingRequestId, /*expectFailure=*/false); + + // Exercise asymmetric but orderly teardown after every decision has committed. + if (world.getRank() == 0) + { + std::this_thread::sleep_for(std::chrono::milliseconds(25)); + } + } + world.barrier(); +} + class RequestInfoTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) { public: diff --git a/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp b/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp index c8a04f658409..941d7cb53ed6 100644 --- a/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -24,6 +24,8 @@ #endif // ENABLE_MULTI_DEVICE #include +#include +#include namespace mpi = tensorrt_llm::mpi; namespace tr = tensorrt_llm::runtime; @@ -46,6 +48,29 @@ TEST(MPIUtils, WorldRankAndSize) EXPECT_LE(rank, size); } +#if ENABLE_MULTI_DEVICE +TEST(MPIUtils, AsyncSendCanBePolled) +{ + auto& comm = mpi::MpiComm::world(); + auto const rank = comm.getRank(); + auto const size = comm.getSize(); + auto const destination = (rank + 1) % size; + auto const source = (rank + size - 1) % size; + std::uint64_t const sentValue = static_cast(rank); + std::uint64_t receivedValue = 0; + + auto request + = comm.sendAsync(&sentValue, 1, mpi::MpiType::kUINT64, destination, mpi::MpiTag::kContextTransferEvent); + static_cast(comm.recv(&receivedValue, 1, mpi::MpiType::kUINT64, source, mpi::MpiTag::kContextTransferEvent)); + while (!request->isCompleted()) + { + std::this_thread::yield(); + } + + EXPECT_EQ(receivedValue, static_cast(source)); +} +#endif // ENABLE_MULTI_DEVICE + template void testBroadcast() { diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4b85232c482c..17993e315566 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -114,6 +114,15 @@ def _stats_buffer_is_unbounded(max_stats_len: int) -> bool: # Default: "0" (only rank 0 prints, matching existing behavior). PROFILE_LOG_RANKS_ENV_VAR_NAME = "TLLM_PROFILE_LOG_RANKS" +# Exact-workload, bounded phase tracing for NVBUG 6448152. This is enabled +# only by the diagnostic perf-sanity YAML and does not alter scheduler or +# transfer behavior. +NVBUG_6448152_PHASE_TRACE_ENV_VAR_NAME = \ + "TRTLLM_NVBUG_6448152_PHASE_TRACE" +NVBUG_6448152_PHASE_TRACE_MAX_REQUESTS = 1024 +NVBUG_6448152_PHASE_TRACE_MAX_TRANSITIONS = 65536 +NVBUG_6448152_PHASE_TRACE_MAX_CUDA_RECORDS = 64 + class PPCommTag(IntEnum): """ @@ -917,6 +926,23 @@ def on_detected(): self.gather_all_responses = False self.kv_cache_transceiver = kv_cache_transceiver + self._nvbug_6448152_phase_trace_enabled = os.environ.get( + NVBUG_6448152_PHASE_TRACE_ENV_VAR_NAME, "0") == "1" + self._nvbug_6448152_phase_trace_request_ids: set[int] = set() + self._nvbug_6448152_phase_trace_transitions: set[Tuple[str, int, + int]] = set() + self._nvbug_6448152_phase_trace_dropped_requests = 0 + self._nvbug_6448152_phase_trace_dropped_transitions = 0 + self._nvbug_6448152_phase_trace_counters: Dict[str, int] = {} + self._nvbug_6448152_phase_trace_cuda_enqueue_ns: Dict[Tuple[int, int], + int] = {} + self._nvbug_6448152_phase_trace_cuda_records: Dict[int, + Tuple[Tuple[int, + ...], + int, + int]] = {} + self._nvbug_6448152_phase_marker( + "python_mode", max_requests=NVBUG_6448152_PHASE_TRACE_MAX_REQUESTS) cache_transceiver_config = getattr(self.llm_args, "cache_transceiver_config", None) max_tokens_in_buffer = getattr(cache_transceiver_config, @@ -1075,7 +1101,191 @@ def _maybe_init_kv_connector_manager(self): self.kv_connector_manager.wait_for_initialization() + def _nvbug_6448152_phase_marker(self, + event: str, + request_id: Optional[int] = None, + per_iteration: bool = False, + transition_iteration: Optional[int] = None, + **fields) -> None: + """Emit one bounded, behavior-neutral phase marker.""" + if not getattr(self, "_nvbug_6448152_phase_trace_enabled", False): + return + + if request_id is not None: + request_id = int(request_id) + if request_id not in self._nvbug_6448152_phase_trace_request_ids: + if len(self._nvbug_6448152_phase_trace_request_ids + ) >= NVBUG_6448152_PHASE_TRACE_MAX_REQUESTS: + self._nvbug_6448152_phase_trace_dropped_requests += 1 + return + self._nvbug_6448152_phase_trace_request_ids.add(request_id) + transition = (event, request_id, + (self.iter_counter if transition_iteration is None + else transition_iteration) if per_iteration else -1) + if transition in self._nvbug_6448152_phase_trace_transitions: + return + if len(self._nvbug_6448152_phase_trace_transitions + ) >= NVBUG_6448152_PHASE_TRACE_MAX_TRANSITIONS: + self._nvbug_6448152_phase_trace_dropped_transitions += 1 + return + self._nvbug_6448152_phase_trace_transitions.add(transition) + + marker_fields = { + "event": event, + "world_rank": self.global_rank, + "pp_rank": self.dist.pp_rank, + "iteration": self.iter_counter, + "request_id": request_id if request_id is not None else -1, + "monotonic_ns": time.monotonic_ns(), + **fields, + } + payload = " ".join(f"{key}={value}" + for key, value in marker_fields.items()) + logger.info(f"NVBUG6448152_PHASE {payload}") + + def _nvbug_6448152_phase_context_requests(self, + event: str, + requests, + per_iteration: bool = False, + **fields) -> None: + if not getattr(self, "_nvbug_6448152_phase_trace_enabled", False): + return + for request in requests: + if request.is_context_only_request: + self._nvbug_6448152_phase_marker(event, + request.request_id, + per_iteration=per_iteration, + **fields) + + def _nvbug_6448152_phase_count(self, + counter: str, + value: int = 1, + **fields) -> None: + """Checkpoint aggregate counts only at powers of two.""" + if (not getattr(self, "_nvbug_6448152_phase_trace_enabled", False) + or value <= 0): + return + count = self._nvbug_6448152_phase_trace_counters.get(counter, 0) + value + self._nvbug_6448152_phase_trace_counters[counter] = count + if count & (count - 1) == 0: + self._nvbug_6448152_phase_marker(f"{counter}_checkpoint", + **{counter: count}, **fields) + + def _nvbug_6448152_phase_summary(self) -> None: + if not getattr(self, "_nvbug_6448152_phase_trace_enabled", False): + return + counter_fields = { + f"total_{counter}": value + for counter, value in + self._nvbug_6448152_phase_trace_counters.items() + } + self._nvbug_6448152_phase_marker( + "python_summary", + traced_requests=len(self._nvbug_6448152_phase_trace_request_ids), + transitions=len(self._nvbug_6448152_phase_trace_transitions), + dropped_requests=self._nvbug_6448152_phase_trace_dropped_requests, + dropped_transitions=( + self._nvbug_6448152_phase_trace_dropped_transitions), + stale_cuda_records=len( + self._nvbug_6448152_phase_trace_cuda_records), + stale_cuda_enqueues=len( + self._nvbug_6448152_phase_trace_cuda_enqueue_ns), + **counter_fields) + self._nvbug_6448152_phase_trace_cuda_records.clear() + self._nvbug_6448152_phase_trace_cuda_enqueue_ns.clear() + + def _nvbug_6448152_phase_record_cuda_enqueue( + self, scheduled_requests: ScheduledRequests) -> None: + if not getattr(self, "_nvbug_6448152_phase_trace_enabled", False): + return + originating_iteration = self.iter_counter + enqueue_monotonic_ns = time.monotonic_ns() + for request in scheduled_requests.context_requests: + if not request.is_context_only_request: + continue + request_id = int(request.request_id) + self._nvbug_6448152_phase_marker( + "cuda_forward_enqueue", + request_id, + per_iteration=True, + transition_iteration=originating_iteration, + originating_iteration=originating_iteration, + enqueue_monotonic_ns=enqueue_monotonic_ns) + if len(self._nvbug_6448152_phase_trace_cuda_enqueue_ns + ) < NVBUG_6448152_PHASE_TRACE_MAX_TRANSITIONS: + self._nvbug_6448152_phase_trace_cuda_enqueue_ns[( + request_id, originating_iteration)] = enqueue_monotonic_ns + else: + self._nvbug_6448152_phase_count( + "cuda_forward_enqueue_records_dropped") + + def _nvbug_6448152_phase_register_cuda_forward( + self, scheduled_requests: ScheduledRequests, + gpu_forward_end_event) -> None: + if not getattr(self, "_nvbug_6448152_phase_trace_enabled", False): + return + request_ids = tuple( + int(request.request_id) + for request in scheduled_requests.context_requests + if request.is_context_only_request) + if not request_ids: + return + + originating_iteration = self.iter_counter + enqueue_monotonic_ns = min( + (self._nvbug_6448152_phase_trace_cuda_enqueue_ns.pop( + (request_id, originating_iteration), time.monotonic_ns()) + for request_id in request_ids), + default=time.monotonic_ns()) + + if gpu_forward_end_event is None: + self._nvbug_6448152_phase_count("cuda_forward_event_unavailable", + requests=len(request_ids)) + return + if len(self._nvbug_6448152_phase_trace_cuda_records + ) >= NVBUG_6448152_PHASE_TRACE_MAX_CUDA_RECORDS: + self._nvbug_6448152_phase_count("cuda_forward_records_dropped") + return + self._nvbug_6448152_phase_trace_cuda_records[id( + gpu_forward_end_event)] = (request_ids, originating_iteration, + enqueue_monotonic_ns) + + def _nvbug_6448152_phase_observe_cuda_forward( + self, gpu_forward_end_event, + gpu_forward_time_ms: Optional[float]) -> None: + if (not getattr(self, "_nvbug_6448152_phase_trace_enabled", False) + or gpu_forward_end_event is None): + return + record = self._nvbug_6448152_phase_trace_cuda_records.pop( + id(gpu_forward_end_event), None) + if record is None: + return + request_ids, originating_iteration, enqueue_monotonic_ns = record + if gpu_forward_time_ms is None: + self._nvbug_6448152_phase_count( + "cuda_forward_elapsed_unavailable", + requests=len(request_ids), + originating_iteration=originating_iteration) + return + + observation_monotonic_ns = time.monotonic_ns() + for request_id in request_ids: + self._nvbug_6448152_phase_marker( + "cuda_forward_complete", + request_id, + per_iteration=True, + transition_iteration=originating_iteration, + originating_iteration=originating_iteration, + enqueue_monotonic_ns=enqueue_monotonic_ns, + observation_delay_ns=(observation_monotonic_ns - + enqueue_monotonic_ns), + gpu_elapsed_ns=int(gpu_forward_time_ms * 1_000_000)) + def _end_transfer_and_maybe_terminate(self, request: LlmRequest): + PyExecutor._nvbug_6448152_phase_context_requests( + self, + "python_global_commit_observed", [request], + state=request.state) transfer_failed = request.state == LlmRequestState.DISAGG_TRANS_ERROR if self.kv_cache_transceiver and request in self.active_requests: if transfer_failed: @@ -1691,6 +1901,15 @@ def profile_step(): host_step_time = (end_time - start_time) * 1000 # milliseconds self._latest_host_step_time_ms = host_step_time self._latest_prev_device_step_time_ms = prev_device_step_time + self._nvbug_6448152_phase_count( + "loop_timing_samples", + host_step_time_ms=host_step_time, + prev_device_step_time_ms=(prev_device_step_time + if prev_device_step_time + is not None else -1.0), + prev_device_step_time_valid=int( + prev_device_step_time is not None), + device_timing_lag_loops=1) if self.print_log and (log_all_ranks or self.dist.rank in log_ranks): @@ -2332,6 +2551,8 @@ def _process_iter_stats( # consumers which interpretation applies. iter_latency_ms = (iter_end_time - batch_state.iter_start_time) * 1e3 if batch_state.iter_stats is None: + self._nvbug_6448152_phase_observe_cuda_forward( + batch_state.gpu_forward_end_event, None) if batch_state.gpu_forward_events_from_perf_pool: self.perf_manager.release_forward_timing_events( batch_state.gpu_forward_start_event, @@ -2347,6 +2568,8 @@ def _process_iter_stats( gpu_forward_time_ms = self.perf_manager.try_compute_gpu_elapsed_time_ms( batch_state.gpu_forward_start_event, batch_state.gpu_forward_end_event) + self._nvbug_6448152_phase_observe_cuda_forward( + batch_state.gpu_forward_end_event, gpu_forward_time_ms) if batch_state.gpu_forward_events_from_perf_pool: self.perf_manager.release_forward_timing_events( batch_state.gpu_forward_start_event, @@ -2381,6 +2604,7 @@ def _process_iter_stats( gpu_forward_time_ms=gpu_forward_time_ms) def _executor_loop_cleanup(self): + PyExecutor._nvbug_6448152_phase_summary(self) # Wake any waiters in await_responses BEFORE potentially-blocking # work below. If wait_on_pp_send_handles hangs (e.g. after a # crash leaves PP send handles in a bad state), the await loop @@ -2424,6 +2648,15 @@ def _pp_schedule_and_propagate(self, microbatch_id: int): serializable_schedule = SerializableSchedulerOutput.from_scheduler_result( scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs, wait_for_disagg_gen_transfer_progress) + if getattr(self, "_nvbug_6448152_phase_trace_enabled", False): + for request_id in ( + serializable_schedule.context_requests_chunking + + serializable_schedule.context_requests_last_chunk): + self._nvbug_6448152_phase_marker( + "pp_schedule_produce", + request_id, + per_iteration=True, + microbatch_id=microbatch_id) # Broadcast within first tp+cp group before send/recv chain to other tp+cp groups if self.dist.is_first_pp_rank: @@ -2438,15 +2671,39 @@ def _pp_schedule_and_propagate(self, microbatch_id: int): # Other ranks receive the schedule result from the previous PP rank. if not self.dist.is_first_pp_rank: + phase_trace_enabled = getattr(self, + "_nvbug_6448152_phase_trace_enabled", + False) + recv_started_ns = time.monotonic_ns() if phase_trace_enabled else 0 with nvtx_range("recv_schedule_from_prev_pp"): serializable_schedule = self.dist.recv_object( self.dist.prev_pp_rank, PPCommTag.SCHEDULE_RESULT) + if phase_trace_enabled: + recv_duration_ns = time.monotonic_ns() - recv_started_ns + for request_id in ( + serializable_schedule.context_requests_chunking + + serializable_schedule.context_requests_last_chunk): + self._nvbug_6448152_phase_marker( + "pp_schedule_recv", + request_id, + per_iteration=True, + microbatch_id=microbatch_id, + recv_duration_ns=recv_duration_ns) # Propagate the schedule result to the next PP rank except the last PP rank. if not self.dist.is_last_pp_rank: self.wait_on_pp_send_handles(self.send_schedule_handles, microbatch_id) with nvtx_range("send_schedule_to_next_pp"): + if getattr(self, "_nvbug_6448152_phase_trace_enabled", False): + for request_id in ( + serializable_schedule.context_requests_chunking + + serializable_schedule.context_requests_last_chunk): + self._nvbug_6448152_phase_marker( + "pp_schedule_send", + request_id, + per_iteration=True, + microbatch_id=microbatch_id) self.send_schedule_handles[ microbatch_id] = self.dist.isend_object( serializable_schedule, self.dist.next_pp_rank, @@ -2457,6 +2714,11 @@ def _pp_schedule_and_propagate(self, microbatch_id: int): self.active_requests) wait_for_disagg_gen_transfer_progress = ( serializable_schedule.wait_for_disagg_gen_transfer_progress) + self._nvbug_6448152_phase_context_requests( + "pp_schedule_apply", + scheduled_batch.context_requests, + per_iteration=True, + microbatch_id=microbatch_id) return (scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs, wait_for_disagg_gen_transfer_progress) @@ -2490,6 +2752,9 @@ def _pp_retry_until_can_schedule(self, scheduled_batch): for retry_count in range(self.pp_scheduler_max_retry_count): if self.scheduler.can_schedule(scheduled_batch_requests): break + self._nvbug_6448152_phase_count( + "pp_schedule_retries", + scheduled_requests=len(scheduled_batch_requests)) logger.debug( f"Retrying to run first PP's schedule result ({retry_count + 1}/{self.pp_scheduler_max_retry_count})" ) @@ -2689,6 +2954,8 @@ def _executor_loop_pp(self): self._update_generation_requests_that_will_complete_next_iteration( scheduled_batch.generation_requests) + self._nvbug_6448152_phase_register_cuda_forward( + scheduled_batch, gpu_forward_end) batch_state = BatchStatePP( scheduled_requests=scheduled_batch, sample_state=sample_state, @@ -5090,6 +5357,10 @@ def _respond_if_invalid(request: LlmRequest) -> bool: new_requests_cur_rank = self._fetch_new_requests( self.waiting_queue, self.active_requests) + self._nvbug_6448152_phase_context_requests("request_fetch", + new_requests_cur_rank) + if not new_requests_cur_rank: + self._nvbug_6448152_phase_count("request_fetch_idle_iterations") validated_requests = [ request for request in new_requests_cur_rank @@ -5097,6 +5368,8 @@ def _respond_if_invalid(request: LlmRequest) -> bool: ] self.active_requests.extend(validated_requests) + self._nvbug_6448152_phase_context_requests("request_activate", + validated_requests) return validated_requests def _add_kv_cache_events(self): @@ -6035,7 +6308,11 @@ def kv_connector_request_finished(req: LlmRequest): # Order is important here: we need to start the transfer before responding # to make sure the blocks are stored for reuse before they are sent. self.async_transfer_manager.start_transfer(req) + PyExecutor._nvbug_6448152_phase_marker( + self, "python_send_async_submit", req.request_id) self.kv_cache_transceiver.respond_and_send_async(req) + PyExecutor._nvbug_6448152_phase_marker( + self, "python_send_async_return", req.request_id) if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: req.py_kv_transfer_start_time = time.monotonic() @@ -6221,6 +6498,11 @@ def _forward_step( num_accepted_tokens_device: Optional[torch.Tensor] = None): ExpertStatistic.set_iter(self.iter_counter) + self._nvbug_6448152_phase_context_requests( + "forward_entry", + scheduled_requests.context_requests, + per_iteration=True) + num_ctx_tokens = sum(req.context_chunk_size for req in scheduled_requests.context_requests) @@ -6253,6 +6535,8 @@ def forward(scheduled_requests, resource_manager, new_tensors_device, new_tensors_device, gather_context_logits, cache_indirection_buffer, num_accepted_tokens_device) + self._nvbug_6448152_phase_record_cuda_enqueue( + scheduled_requests) self._mark_cross_kv_projection_consumed(scheduled_requests) # Ensure the default stream waits for execution_stream to complete @@ -6524,6 +6808,8 @@ def _handle_errors(self, self.executor_request_queue.enqueue_shutdown_request() def _terminate_request(self, request: LlmRequest): + PyExecutor._nvbug_6448152_phase_context_requests( + self, "python_cleanup_dispatch", [request], state=request.state) # Dummy requests don't participate in disagg KV cache transfers, # so they must bypass the PP termination handler to avoid stale # sequences in the KV cache manager (the handler delays removal, @@ -6535,6 +6821,8 @@ def _terminate_request(self, request: LlmRequest): self._do_terminate_request(request) def _do_terminate_request(self, request: LlmRequest): + PyExecutor._nvbug_6448152_phase_context_requests( + self, "python_resource_free_enter", [request], state=request.state) self.resource_manager.free_resources(request) self._prefetched_request_ids.discard(request.py_request_id) self._disagg_timed_out_ctx_cancelled_ids.discard(request.py_request_id) @@ -6542,6 +6830,8 @@ def _do_terminate_request(self, request: LlmRequest): if self.gather_all_responses or self.dist.rank == 0: self.result_wait_queues.pop(request.py_request_id, None) + PyExecutor._nvbug_6448152_phase_context_requests( + self, "python_resource_free_exit", [request], state=request.state) def _is_request_in_transmission(self, request) -> bool: """Check if a request is currently in transmission state.""" diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml index 795a6192bf41..4952f52f78f9 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml @@ -36,7 +36,7 @@ environment: build_wheel: false work_dir: worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=1 - TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes + TRTLLM_ENABLE_PDL=1 TRTLLM_NVBUG_6448152_PHASE_TRACE=1 ENROOT_ALLOW_DEV=yes server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: nsys_on: false