From fdc8ace6beca3cb5668f2913b6e685938957b03d Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 14 Jul 2026 17:20:36 -0700 Subject: [PATCH 1/4] [https://nvbugs/6448152][perf] publish PP transfer status from workers Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 33 +- .../batch_manager/transferStatusConsensus.h | 120 +++++ .../tensorrt_llm/runtime/utils/mpiTags.h | 7 +- .../tensorrt_llm/runtime/utils/mpiUtils.h | 16 +- cpp/tensorrt_llm/batch_manager/CMakeLists.txt | 3 +- .../batch_manager/cacheTransceiver.cpp | 486 +++++++++++++----- .../batch_manager/dataTransceiver.cpp | 40 +- .../batch_manager/dataTransceiver.h | 7 + .../batch_manager/transferStatusConsensus.cpp | 343 ++++++++++++ cpp/tensorrt_llm/runtime/utils/mpiUtils.cpp | 13 +- .../unit_tests/batch_manager/CMakeLists.txt | 1 + .../transferStatusConsensusTest.cpp | 164 ++++++ .../multi_gpu/cacheTransceiverTest.cpp | 165 +++++- .../unit_tests/multi_gpu/mpiUtilsTest.cpp | 27 +- .../defs/disaggregated/test_disaggregated.py | 28 +- .../others/test_kv_cache_transceiver.py | 6 +- 16 files changed, 1299 insertions(+), 160 deletions(-) create mode 100644 cpp/include/tensorrt_llm/batch_manager/transferStatusConsensus.h create mode 100644 cpp/tensorrt_llm/batch_manager/transferStatusConsensus.cpp create mode 100644 cpp/tests/unit_tests/batch_manager/transferStatusConsensusTest.cpp diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 0fcee1435005..b341dba3f826 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -26,6 +26,7 @@ #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/utils/pgUtils.h" +#include #include #include #include @@ -54,6 +55,7 @@ class BaseKVCacheManager; class CacheSender; class CacheReceiver; +class ContextTransferVoteMailbox; class CacheTransceiverComm { @@ -148,6 +150,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()) @@ -244,6 +265,7 @@ class CacheTransceiver : public BaseCacheTransceiver std::optional cacheTransceiverConfig = std::nullopt, std::vector const& rnnLayerNumPerPP = {}); + // Constructor used by the PyExecutor Nanobind path. The legacy C++ factory uses the ModelConfig overload above. CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, std::vector numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, runtime::WorldConfig const& worldConfig, std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, @@ -253,7 +275,8 @@ class CacheTransceiver : public BaseCacheTransceiver std::vector const& rnnLayerNumPerPP = {}) : CacheTransceiver(cacheManager, executor::kv_cache::CacheState::ModelConfig{numKvHeadsPerLayer, sizePerHead, tokensPerBlock}, worldConfig, - attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnLayerNumPerPP) + attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnLayerNumPerPP, + /*enableWorkerPublishedContextConsensus=*/true) { } @@ -279,6 +302,13 @@ class CacheTransceiver : public BaseCacheTransceiver [[nodiscard]] bool hasPoisonedTransferBuffer() const override; private: + CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, + executor::kv_cache::CacheState::ModelConfig const& cacheStateModelCfg, runtime::WorldConfig const& worldConfig, + std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + executor::kv_cache::CacheState::AttentionType attentionType, + std::optional cacheTransceiverConfig, + std::vector const& rnnLayerNumPerPP, bool enableWorkerPublishedContextConsensus); + void initializeCommState(); void setContextState(LlmRequest* llmRequest); @@ -306,6 +336,7 @@ class CacheTransceiver : public BaseCacheTransceiver std::shared_ptr mGroupComm; std::shared_ptr mGroupTensorParaComm, mGroupPipeParaComm, mGroupDataComm, mGroupTPInDPComm; + std::unique_ptr mContextTransferVoteMailbox; executor::kv_cache::CommState const* mCommState; std::unique_ptr mCacheState; diff --git a/cpp/include/tensorrt_llm/batch_manager/transferStatusConsensus.h b/cpp/include/tensorrt_llm/batch_manager/transferStatusConsensus.h new file mode 100644 index 000000000000..ed54c498666d --- /dev/null +++ b/cpp/include/tensorrt_llm/batch_manager/transferStatusConsensus.h @@ -0,0 +1,120 @@ +/* + * 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 + +namespace tensorrt_llm::batch_manager +{ + +class CacheTransceiverComm; + +struct WorkerPublishedConsensusConfig +{ + bool enabledForCppBinding{false}; + bool mpiControlPlane{false}; + bool nixlUcxBackend{false}; + int tensorParallelism{1}; + int contextParallelism{1}; + int pipelineParallelism{1}; + bool attentionDp{false}; + bool inflightCancellation{false}; + bool transferOverlap{true}; + bool layerwiseTransfer{false}; + bool mpiThreadMultiple{false}; +}; + +//! Return whether the worker-published protocol is supported by one rank's effective configuration. +[[nodiscard]] bool supportsWorkerPublishedConsensus(WorkerPublishedConsensusConfig const& config); + +//! Return whether every rank advertised the same nonzero protocol version. +[[nodiscard]] bool selectWorkerPublishedConsensus(std::vector const& protocolVersions); + +enum class TransferStatusVote : std::uint64_t +{ + kCompleted = 1, + kFailed = 2, +}; + +struct TransferStatusConsensusResult +{ + std::unordered_set completedRequestIds; + std::unordered_set failedRequestIds; +}; + +//! Reduces one immutable terminal vote per participant and request. +class TransferStatusConsensus +{ +public: + explicit TransferStatusConsensus(int participantCount); + + //! Record a participant's terminal vote. Identical duplicates are ignored; conflicting votes are rejected. + void recordVote(int participantRank, std::uint64_t requestId, TransferStatusVote vote); + + //! Return and remove every request for which all participants have cast a terminal vote. + [[nodiscard]] TransferStatusConsensusResult takeCompleted(); + +private: + struct RequestVotes + { + explicit RequestVotes(int participantCount) + : votes(static_cast(participantCount), 0) + { + } + + std::vector votes; + int terminalCount{0}; + bool failed{false}; + }; + + int mParticipantCount; + std::unordered_map mRequestVotes; +}; + +//! Exchanges immutable worker-published terminal votes without requiring peer scheduler participation. +class ContextTransferVoteMailbox +{ +public: + explicit ContextTransferVoteMailbox(std::shared_ptr comm); + ~ContextTransferVoteMailbox(); + + ContextTransferVoteMailbox(ContextTransferVoteMailbox const&) = delete; + ContextTransferVoteMailbox& operator=(ContextTransferVoteMailbox const&) = delete; + + //! Publish one terminal outcome to every peer. The first internal error is reported by a later scheduler poll. + void publishOutcomeToPeers(std::uint64_t requestId, bool failed) noexcept; + + //! Record the already-published local worker outcome in this rank's reducer. + [[nodiscard]] bool recordLocalOutcome(std::uint64_t requestId); + + //! Drain peer votes and return newly committed global outcomes. + [[nodiscard]] TransferStatusConsensusResult poll(); + + //! Stop accepting publications and exchange an ordered close marker with every peer. + 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..22b8f6e74fae 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,10 @@ enum class MpiTag : int // KvCacheEventManager kKvCacheEventSize = 1026, - kKvCacheEvent = 1027 + kKvCacheEvent = 1027, + + // Context-transfer terminal votes. + kContextTransferStatus = 1028 }; } // 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..0e2e84e11b16 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 @@ -467,6 +478,9 @@ int getNumNodes(); void initialize(MpiThreadSupport threadMode = MpiThreadSupport::THREAD_MULTIPLE, bool forwardAbortToParent = false); +//! \brief Return the MPI thread support level provided by the active runtime. +[[nodiscard]] MpiThreadSupport getThreadSupport(); + //! \brief Returns true iff the WideEP fault-tolerance MPI mode is enabled. //! //! Reads the `TLLM_FAULT_TOLERANCE_MODE` environment variable and returns diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index 88e2484cc6f1..c1ec106b3fd7 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 @@ -56,6 +56,7 @@ set(SRCS rnnCacheTransBuffer.cpp runtimeBuffers.cpp sequenceSlotManager.cpp + transferStatusConsensus.cpp transformerBuffers.cpp trtEncoderModel.cpp trtGptModelInflightBatching.cpp diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index aa05cd033892..9ab38925ff99 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -45,6 +45,7 @@ #include "tensorrt_llm/batch_manager/rnnCacheFormatter.h" #include "tensorrt_llm/batch_manager/rnnCacheTransBuffer.h" #include "tensorrt_llm/batch_manager/rnnStateManager.h" +#include "tensorrt_llm/batch_manager/transferStatusConsensus.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/executor/cache_transmission/mpi_utils/connection.h" @@ -71,6 +72,7 @@ namespace using RequestIdType = LlmRequest::RequestIdType; constexpr int kTransferFuturePollIntervalMs = 10; +constexpr int kContextConsensusPollIntervalMs = 1; // Finite status checks are scheduler polls, not terminal deadlines. Pure polls // use short slices; calls that ask for at least one completion keep bounded @@ -348,6 +350,17 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa executor::kv_cache::CacheState::AttentionType attentionType, std::optional cacheTransceiverConfig, std::vector const& rnnLayerNumPerPP) + : CacheTransceiver(cacheManager, cacheStateModelCfg, worldConfig, attentionLayerNumPerPP, dataType, attentionType, + cacheTransceiverConfig, rnnLayerNumPerPP, /*enableWorkerPublishedContextConsensus=*/false) +{ +} + +CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, + executor::kv_cache::CacheState::ModelConfig const& cacheStateModelCfg, runtime::WorldConfig const& worldConfig, + std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + executor::kv_cache::CacheState::AttentionType attentionType, + std::optional cacheTransceiverConfig, + std::vector const& rnnLayerNumPerPP, bool const enableWorkerPublishedContextConsensus) : mCacheTransceiverConfig{cacheTransceiverConfig} { using tensorrt_llm::batch_manager::kv_cache_manager::CacheFormatter; @@ -424,6 +437,58 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa mGroupComm->split(worldConfig.getPipelineParallelRank() * dpSize + dpRank, worldConfig.getRank())); } } + + if (useMPI() && worldConfig.getPipelineParallelism() > 1) + { + // This is a one-time startup agreement, not a scheduler-hot-path collective. Every MPI PP constructor + // participates so a mixed binding/configuration selects the synchronous fallback on every rank. + auto const threadSupport = mpi::getThreadSupport(); + auto const consensusConfig = [&]() + { + WorkerPublishedConsensusConfig config; + config.enabledForCppBinding = enableWorkerPublishedContextConsensus; + config.mpiControlPlane = true; + config.nixlUcxBackend = backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL + && common::getEnvNixlBackend() == "UCX"; + config.tensorParallelism = worldConfig.getTensorParallelism(); + config.contextParallelism = worldConfig.getContextParallelism(); + config.pipelineParallelism = worldConfig.getPipelineParallelism(); + config.attentionDp = mCacheState->getParallelConfig().mEnableAttentionDP; + config.inflightCancellation = common::getEnvDisaggEnableInflightCancel(); + config.transferOverlap = !common::getEnvDisableKVCacheTransferOverlap(); + config.layerwiseTransfer = common::getEnvDisaggLayerwise(); + config.mpiThreadMultiple = threadSupport == mpi::MpiThreadSupport::THREAD_MULTIPLE; + return config; + }(); + constexpr std::uint64_t kWorkerPublishedConsensusVersion = 1; + auto const localVersion + = supportsWorkerPublishedConsensus(consensusConfig) ? kWorkerPublishedConsensusVersion : 0; + TLLM_CHECK(mGroupPipeParaComm); + std::vector protocolVersions(static_cast(mGroupPipeParaComm->getSize())); + mGroupPipeParaComm->allgather(&localVersion, protocolVersions.data(), 1, mpi::MpiType::kUINT64); + auto const useWorkerPublishedContextConsensus = selectWorkerPublishedConsensus(protocolVersions); + if (useWorkerPublishedContextConsensus) + { + TLLM_LOG_INFO( + "Enabled worker-published context-transfer consensus v%lu for the C++ NIXL/UCX TP1/CP1 " + "pipeline-parallel path.", + kWorkerPublishedConsensusVersion); + } + else if (std::any_of(protocolVersions.begin(), protocolVersions.end(), + [](std::uint64_t const version) { return version != 0; })) + { + TLLM_LOG_WARNING( + "Pipeline ranks advertised different worker-published context-transfer consensus capabilities; " + "retaining synchronous context-transfer consensus on every rank."); + } + + // Construct the mailbox immediately after the all-rank agreement. If later transceiver initialization + // throws, member unwinding still sends this rank's ordered close marker to its peers. + if (useWorkerPublishedContextConsensus) + { + mContextTransferVoteMailbox = std::make_unique(mGroupPipeParaComm); + } + } bool isMLA = attentionType == executor::kv_cache::CacheState::AttentionType::kMLA; std::optional maxNumTokens = mCacheTransceiverConfig.value().getMaxTokensInBuffer(); @@ -576,6 +641,10 @@ CacheTransceiver::~CacheTransceiver() // Stop sender/receiver workers while the connection manager and transfer // plugin are still alive. The workers can access both during termination. mCacheSender.reset(); + if (mContextTransferVoteMailbox) + { + mContextTransferVoteMailbox->shutdown(); + } mCacheReceiver.reset(); if (mWrapperLibHandle) @@ -623,7 +692,16 @@ void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmReques return; } setContextState(llmRequest.get()); - auto future = mCacheSender->sendAsync(llmRequest); + CacheSender::CompletionCallback completionCallback; + if (mContextTransferVoteMailbox) + { + auto* mailbox = mContextTransferVoteMailbox.get(); + auto const requestId = llmRequest->mRequestId; + completionCallback + = [mailbox, requestId](bool const failed) { mailbox->publishOutcomeToPeers(requestId, failed); }; + } + auto future = completionCallback ? mCacheSender->sendAsync(llmRequest, std::move(completionCallback)) + : mCacheSender->sendAsync(llmRequest); mSenderFutures.emplace_back(std::move(llmRequest), std::move(future)); } @@ -731,6 +809,122 @@ std::vector gatherRequestIds( return retData; } +namespace +{ + +class ContextTransferConsensusBackend +{ +public: + ContextTransferConsensusBackend(ContextTransferVoteMailbox* mailbox, + std::shared_ptr const& intraStageComm, + std::shared_ptr const& pipeStageComm) + : mMailbox(mailbox) + , mIntraStageComm(intraStageComm) + , mPipeStageComm(pipeStageComm) + { + } + + [[nodiscard]] bool isWorkerPublished() const noexcept + { + return mMailbox != nullptr; + } + + template + [[nodiscard]] std::unordered_set selectReadyRequests(SenderFutures const& senderFutures) const + { + if (isWorkerPublished()) + { + std::unordered_set readyRequestIds; + for (auto const& [request, future] : senderFutures) + { + if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + { + readyRequestIds.insert(request->mRequestId); + } + } + return readyRequestIds; + } + + std::vector locallyReadyRequestIds; + for (auto const& [request, future] : senderFutures) + { + if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + { + locallyReadyRequestIds.push_back(request->mRequestId); + } + } + std::unordered_map frequencyMap; + if (mIntraStageComm && mIntraStageComm->getSize() > 1) + { + for (auto const requestId : gatherRequestIds(mIntraStageComm, locallyReadyRequestIds)) + { + ++frequencyMap[requestId]; + } + } + else + { + for (auto const requestId : locallyReadyRequestIds) + { + ++frequencyMap[requestId]; + } + } + + auto const requiredReadyCount = mIntraStageComm ? mIntraStageComm->getSize() : 1; + std::unordered_set readyRequestIds; + for (auto const& [requestId, readyCount] : frequencyMap) + { + if (readyCount == requiredReadyCount) + { + readyRequestIds.insert(requestId); + } + } + return readyRequestIds; + } + + [[nodiscard]] bool resolveLocalFailure(RequestIdType const requestId, bool const observedFailure) const + { + if (!isWorkerPublished()) + { + return observedFailure; + } + + auto const publishedFailure = mMailbox->recordLocalOutcome(requestId); + TLLM_CHECK_WITH_INFO(!observedFailure || publishedFailure, + "Context-transfer future failed after its worker published a successful terminal vote."); + return publishedFailure; + } + + [[nodiscard]] TransferConsensusOutcome poll(std::unordered_set const& completedRequestIds, + std::unordered_set const& failedRequestIds, + std::unordered_set const& timedOutRequestIds) const + { + if (!isWorkerPublished()) + { + return reduceTransferStates( + mIntraStageComm, mPipeStageComm, completedRequestIds, failedRequestIds, timedOutRequestIds); + } + + auto mailboxResult = mMailbox->poll(); + TransferConsensusOutcome outcome; + outcome.completedRequestIds = std::move(mailboxResult.completedRequestIds); + outcome.failedRequestIds = std::move(mailboxResult.failedRequestIds); + return outcome; + } + +private: + ContextTransferVoteMailbox* mMailbox; + std::shared_ptr const& mIntraStageComm; + std::shared_ptr const& mPipeStageComm; +}; + +std::unordered_set const& getEmptyRequestIdSet() +{ + static std::unordered_set const emptyRequestIds; + return emptyRequestIds; +} + +} // namespace + void updateKVCacheTransferBW(std::shared_ptr const& mComm, LlmRequest* request) { namespace su = executor::serialize_utils; @@ -796,80 +990,42 @@ void updateKVCacheTransferBW(std::shared_ptr const& mComm, } RequestStatuses CacheTransceiver::checkContextTransferStatus( - std::optional const& atLeastRequestNum, bool markComplete) + std::optional const& atLeastRequestNum, bool const markComplete) { bool const blockAll = !atLeastRequestNum.has_value(); bool const inflightCancelEnabled = common::getEnvDisaggEnableInflightCancel(); TLLM_CHECK_WITH_INFO(!inflightCancelEnabled || !blockAll, "In-flight cancellation requires a finite context-transfer status poll; pass 0 for a nonblocking poll."); - std::optional senderFutureTimeoutMs = std::nullopt; - if (mCacheTransceiverConfig.has_value()) - { - senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); - } + bool const needsProgress = atLeastRequestNum.value_or(0) > 0; - auto const futureWaitInterval = getTransferFutureWaitInterval(senderFutureTimeoutMs, needsProgress); - // Without the opt-in flag, deadline checks remain observe-only. With the - // flag, timeout IDs participate in the same topology consensus as terminal - // outcomes and request cancellation is requested on every nonterminal rank. - std::optional kvTransferTimeoutMs = std::nullopt; + std::optional senderFutureTimeoutMs; + std::optional kvTransferTimeoutMs; if (mCacheTransceiverConfig.has_value()) { + senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); } + auto const futureWaitInterval = getTransferFutureWaitInterval(senderFutureTimeoutMs, needsProgress); - auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupTPInDPComm : mGroupTensorParaComm; - std::vector contextCompleteRequestIds; - for (auto&& [request, future] : mSenderFutures) - { - if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) - { - contextCompleteRequestIds.push_back(request->mRequestId); - } - } - - std::unordered_map frequencyMap; - if ((syncComm) && syncComm->getSize() > 1) - { - auto gatherRequestIdVec = gatherRequestIds(syncComm, contextCompleteRequestIds); - for (auto&& requestId : gatherRequestIdVec) - { - frequencyMap[requestId]++; - } - } - else - { - for (auto&& requestId : contextCompleteRequestIds) - { - frequencyMap[requestId]++; - } - } - std::vector> freqVec(frequencyMap.begin(), frequencyMap.end()); - - std::sort(freqVec.begin(), freqVec.end(), - [](std::pair const& left, - std::pair const& right) { return left.second > right.second; }); - std::unordered_set toCompleteIdSet; - for (auto&& [requestId, freq] : freqVec) + auto const& intraStageComm + = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupTPInDPComm : mGroupTensorParaComm; + ContextTransferConsensusBackend consensusBackend{ + mContextTransferVoteMailbox.get(), intraStageComm, mGroupPipeParaComm}; + std::optional progressDeadline; + if (consensusBackend.isWorkerPublished() && needsProgress) { - if (freq == ((syncComm) ? syncComm->getSize() : 1)) - { - toCompleteIdSet.insert(requestId); - } + progressDeadline = std::chrono::steady_clock::now() + futureWaitInterval; } + auto toProcess = consensusBackend.selectReadyRequests(mSenderFutures); - // Make sure there are at least atLeastRequestNum requests in toCompleteIdSet. - // This will preserve the order of insertion for KVCache transfer requests. + // Preserve insertion order when a caller asks this poll to make progress on at least N requests. for (auto it = mSenderFutures.begin(); - atLeastRequestNum.value_or(0) > static_cast(toCompleteIdSet.size()) && it != mSenderFutures.end(); ++it) + atLeastRequestNum.value_or(0) > static_cast(toProcess.size()) && it != mSenderFutures.end(); ++it) { - auto& [request, future] = *it; - toCompleteIdSet.insert(request->mRequestId); + toProcess.insert(it->first->mRequestId); } - // 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. + // Future polling, timeout accounting, and local outcome retention are shared by every consensus backend. for (auto it = mSenderFutures.begin(); it != mSenderFutures.end();) { auto& [request, future] = *it; @@ -878,87 +1034,134 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( && future.wait_for(std::chrono::milliseconds(0)) != std::future_status::ready) { auto const elapsedMs = getTransferElapsedMs(request, LlmRequest::getSteadyClockNow()); - if (elapsedMs > kvTransferTimeoutMs.value()) + if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutSenderIds.insert(requestId).second) { - if (mTimedOutSenderIds.insert(requestId).second) - { - TLLM_LOG_WARNING( - "Context KV cache transfer for request %ld exceeded configured timeout: " - "elapsed %ld ms > limit %d ms (%s).", - requestId, elapsedMs, kvTransferTimeoutMs.value(), - inflightCancelEnabled ? "requesting cancellation" : "observe-only"); - } + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld exceeded configured timeout: " + "elapsed %ld ms > limit %d ms (%s).", + requestId, elapsedMs, kvTransferTimeoutMs.value(), + inflightCancelEnabled ? "requesting cancellation" : "observe-only"); } } - if (blockAll || (toCompleteIdSet.find(requestId) != toCompleteIdSet.end())) + + if (!blockAll && toProcess.find(requestId) == toProcess.end()) + { + ++it; + continue; + } + + auto waitInterval = futureWaitInterval; + if (progressDeadline.has_value()) + { + // Worker-published status polling has one total budget across local future waits and mailbox progress. + waitInterval = std::max(std::chrono::milliseconds(0), + std::chrono::duration_cast( + progressDeadline.value() - std::chrono::steady_clock::now())); + } + auto const status = blockAll ? std::future_status::ready : future.wait_for(waitInterval); + if (status == std::future_status::timeout) + { + TLLM_LOG_DEBUG( + "Context KV cache transfer for request %ld is not ready after %ld ms wait slice; keeping it " + "in progress.", + requestId, static_cast(waitInterval.count())); + ++it; + continue; + } + + bool observedFailure = status != std::future_status::ready; + if (!observedFailure) { try { - auto const status = blockAll ? std::future_status::ready : future.wait_for(futureWaitInterval); - if (status == std::future_status::ready) + future.get(); + if (kvTransferTimeoutMs.has_value()) { - future.get(); - if (kvTransferTimeoutMs.has_value()) + auto const elapsedMs = getTransferElapsedMs(request, request->getKvCacheTransferEnd()); + if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutSenderIds.insert(requestId).second) { - auto const elapsedMs = getTransferElapsedMs(request, request->getKvCacheTransferEnd()); - if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutSenderIds.insert(requestId).second) - { - TLLM_LOG_WARNING( - "Context KV cache transfer for request %ld completed after its deadline: " - "elapsed %ld ms > limit %d ms (%s).", - requestId, elapsedMs, kvTransferTimeoutMs.value(), - inflightCancelEnabled ? "failing request" : "observe-only"); - } + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld completed after its deadline: " + "elapsed %ld ms > limit %d ms (%s).", + requestId, elapsedMs, kvTransferTimeoutMs.value(), + inflightCancelEnabled ? "failing request" : "observe-only"); } - recordLocalTransferOutcome(requestId, request, /*failed=*/false, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - it = mSenderFutures.erase(it); - } - else if (status == std::future_status::timeout) - { - TLLM_LOG_DEBUG( - "Context KV cache transfer for request %ld is not ready after %ld ms wait slice; keeping it " - "in progress.", - requestId, static_cast(futureWaitInterval.count())); - ++it; - } - else - { - 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); } } - catch (std::exception const& e) + catch (std::exception const& error) { - 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); + TLLM_LOG_ERROR("Error occurred during context transfer for request %ld: %s", requestId, error.what()); + observedFailure = true; } catch (...) { TLLM_LOG_ERROR("Unknown error occurred during context transfer for request %ld", requestId); - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - it = mSenderFutures.erase(it); + observedFailure = true; } } else { - ++it; + TLLM_LOG_ERROR("Future returned unexpected status for request %ld. Recording as failed.", requestId); } + + auto completedRequest = request; + it = mSenderFutures.erase(it); + // Erase the consumed future before a backend invariant can throw, so no no-state future remains pollable. + auto const failed = consensusBackend.resolveLocalFailure(requestId, observedFailure); + recordLocalTransferOutcome(requestId, std::move(completedRequest), failed, mCompletedSenderRequestIds, + mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); } - RequestStatuses requestsStatus{}; - auto const consensusOutcome = reduceTransferStates(syncComm, mGroupPipeParaComm, mCompletedSenderRequestIds, - mFailedSenderRequestIds, inflightCancelEnabled ? mTimedOutSenderIds : std::unordered_set{}); - if (inflightCancelEnabled) + RequestStatuses requestsStatus; + auto commitConsensusOutcome = [&](TransferConsensusOutcome const& outcome) { - for (auto const requestId : consensusOutcome.timedOutRequestIds) + for (auto const requestId : sortedRequestIds(outcome.failedRequestIds)) + { + auto const requestIt = mSenderRequestsAwaitingConsensus.find(requestId); + if (requestIt == mSenderRequestsAwaitingConsensus.end()) + { + TLLM_CHECK_WITH_INFO(!consensusBackend.isWorkerPublished(), + "Missing retained context request %zu for a globally failed worker-published transfer.", + static_cast(requestId)); + continue; + } + requestIt->second->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + requestsStatus.errorRequestIds.insert(requestId); + mTimedOutSenderIds.erase(requestId); + mCancelRequestedSenderIds.erase(requestId); + eraseLocalTransferOutcome( + requestId, mCompletedSenderRequestIds, mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); + } + for (auto const requestId : sortedRequestIds(outcome.completedRequestIds)) + { + auto const requestIt = mSenderRequestsAwaitingConsensus.find(requestId); + if (requestIt == mSenderRequestsAwaitingConsensus.end()) + { + TLLM_CHECK_WITH_INFO(!consensusBackend.isWorkerPublished(), + "Missing retained context request %zu for a globally completed worker-published transfer.", + static_cast(requestId)); + continue; + } + requestsStatus.completedRequestIds.insert(requestId); + if (markComplete) + { + requestIt->second->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); + } + mTimedOutSenderIds.erase(requestId); + mCancelRequestedSenderIds.erase(requestId); + eraseLocalTransferOutcome( + requestId, mCompletedSenderRequestIds, mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); + } + }; + + auto requestTimedOutCancellations = [&](TransferConsensusOutcome const& outcome) + { + if (!inflightCancelEnabled) + { + return; + } + + for (auto const requestId : outcome.timedOutRequestIds) { auto const futureIt = std::find_if(mSenderFutures.begin(), mSenderFutures.end(), [requestId](auto const& entry) { return entry.first->mRequestId == requestId; }); @@ -979,37 +1182,42 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( TLLM_LOG_DEBUG("Context cancellation for request %ld was not accepted; will retry", requestId); } } - } - for (auto const requestId : sortedRequestIds(consensusOutcome.failedRequestIds)) + }; + + auto pollAndCommitConsensus = [&]() { - auto const requestIt = mSenderRequestsAwaitingConsensus.find(requestId); - if (requestIt == mSenderRequestsAwaitingConsensus.end()) - { - continue; - } - requestIt->second->setState(LlmRequestState::kDISAGG_TRANS_ERROR); - requestsStatus.errorRequestIds.insert(requestId); - mTimedOutSenderIds.erase(requestId); - mCancelRequestedSenderIds.erase(requestId); - eraseLocalTransferOutcome( - requestId, mCompletedSenderRequestIds, mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - } - for (auto const requestId : sortedRequestIds(consensusOutcome.completedRequestIds)) + auto const& timedOutRequestIds = inflightCancelEnabled ? mTimedOutSenderIds : getEmptyRequestIdSet(); + auto const outcome + = consensusBackend.poll(mCompletedSenderRequestIds, mFailedSenderRequestIds, timedOutRequestIds); + requestTimedOutCancellations(outcome); + commitConsensusOutcome(outcome); + }; + + // Collective consensus must be entered exactly once per status call. The mailbox can make independent progress + // and may be polled repeatedly without requiring a matching peer scheduler call. + pollAndCommitConsensus(); + if (consensusBackend.isWorkerPublished()) { - auto const requestIt = mSenderRequestsAwaitingConsensus.find(requestId); - if (requestIt == mSenderRequestsAwaitingConsensus.end()) + auto const committedCount + = [&]() { return requestsStatus.completedRequestIds.size() + requestsStatus.errorRequestIds.size(); }; + if (blockAll) { - continue; + while (!mSenderRequestsAwaitingConsensus.empty()) + { + std::this_thread::sleep_for(std::chrono::milliseconds(kContextConsensusPollIntervalMs)); + pollAndCommitConsensus(); + } } - requestsStatus.completedRequestIds.insert(requestId); - if (markComplete) + else if (needsProgress) { - requestIt->second->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); + while (committedCount() < static_cast(atLeastRequestNum.value()) + && !mSenderRequestsAwaitingConsensus.empty() + && std::chrono::steady_clock::now() < progressDeadline.value()) + { + std::this_thread::sleep_for(std::chrono::milliseconds(kContextConsensusPollIntervalMs)); + pollAndCommitConsensus(); + } } - mTimedOutSenderIds.erase(requestId); - mCancelRequestedSenderIds.erase(requestId); - eraseLocalTransferOutcome( - requestId, mCompletedSenderRequestIds, mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); } return requestsStatus; diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 109417965a75..2c1c82aa78a0 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -37,6 +37,7 @@ #include #include #include +#include namespace tensorrt_llm::batch_manager { @@ -341,7 +342,8 @@ class CacheSender::Impl } } - [[nodiscard]] std::future sendAsync(std::shared_ptr const& llmRequest) + [[nodiscard]] std::future sendAsync( + std::shared_ptr const& llmRequest, CacheSender::CompletionCallback completionCallback) { TLLM_CHECK(llmRequest != nullptr); std::promise promise; @@ -355,8 +357,8 @@ class CacheSender::Impl std::scoped_lock lock(mSenderMutex); TLLM_CHECK_WITH_INFO( !mTerminate, "Cannot enqueue request %zu after CacheSender termination", llmRequest->mRequestId); - auto const result - = mReadyResponses.emplace(llmRequest->mRequestId, Response{llmRequest, std::move(promise)}); + auto const result = mReadyResponses.emplace( + llmRequest->mRequestId, Response{llmRequest, std::move(promise), std::move(completionCallback)}); TLLM_CHECK_WITH_INFO( result.second, "Request %zu is already queued for KV cache transfer", llmRequest->mRequestId); } @@ -629,6 +631,7 @@ class CacheSender::Impl // protects worker-side dereferences and the promise itself from premature destruction. std::shared_ptr mRequest; std::promise mPromise; + CacheSender::CompletionCallback mCompletionCallback; }; struct AsyncSendResource @@ -677,6 +680,7 @@ class CacheSender::Impl TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); sendSync(*resp.mRequest); release(id); + notifyCompletion(resp, /*failed=*/false); resp.mPromise.set_value(); } catch (tensorrt_llm::common::RequestSpecificException const& e) @@ -922,6 +926,7 @@ class CacheSender::Impl void failResponse(Response& response, std::exception_ptr const& exception) noexcept { + notifyCompletion(response, /*failed=*/true); try { response.mPromise.set_exception(exception); @@ -932,6 +937,27 @@ class CacheSender::Impl } } + void notifyCompletion(Response& response, bool const failed) noexcept + { + auto completionCallback = std::exchange(response.mCompletionCallback, {}); + if (!completionCallback) + { + return; + } + try + { + completionCallback(failed); + } + catch (std::exception const& err) + { + TLLM_LOG_ERROR("CacheSender completion callback failed: %s", err.what()); + } + catch (...) + { + TLLM_LOG_ERROR("CacheSender completion callback failed with an unknown error"); + } + } + void failPendingResponses(std::exception_ptr const& exception) noexcept { std::map pendingResponses; @@ -1738,9 +1764,15 @@ CacheSender::CacheSender( { } +std::future CacheSender::sendAsync( + std::shared_ptr const& llmRequest, CompletionCallback completionCallback) const +{ + return mImpl->sendAsync(llmRequest, std::move(completionCallback)); +} + std::future CacheSender::sendAsync(std::shared_ptr const& llmRequest) const { - return mImpl->sendAsync(llmRequest); + return mImpl->sendAsync(llmRequest, {}); } executor::kv_cache::CommState const& CacheSender::getCommState() const diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index 8e84a71556af..89a6c5bbe3e5 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -17,6 +17,7 @@ #pragma once #include +#include #include #include #include @@ -263,6 +264,8 @@ class RequestInfo class CacheSender { public: + using CompletionCallback = std::function; + /// @brief Constructor. /// @param manager The connection manager. /// @param selfIndex The sequential index of the current executor process. @@ -277,6 +280,10 @@ class CacheSender /// @return Once the data is fully sent, the future object will become valid. [[nodiscard]] virtual std::future sendAsync(std::shared_ptr const& llmRequest) const; + /// @brief Asynchronously respond and report the terminal worker outcome before the returned future becomes ready. + [[nodiscard]] std::future sendAsync( + std::shared_ptr const& llmRequest, CompletionCallback completionCallback) const; + /// @brief Return the internal communicator status. /// @return The communicator status. [[nodiscard]] virtual executor::kv_cache::CommState const& getCommState() const; diff --git a/cpp/tensorrt_llm/batch_manager/transferStatusConsensus.cpp b/cpp/tensorrt_llm/batch_manager/transferStatusConsensus.cpp new file mode 100644 index 000000000000..769ab7334ae3 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/transferStatusConsensus.cpp @@ -0,0 +1,343 @@ +/* + * 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/transferStatusConsensus.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 +#include +#include + +namespace tensorrt_llm::batch_manager +{ + +bool supportsWorkerPublishedConsensus(WorkerPublishedConsensusConfig const& config) +{ + return config.enabledForCppBinding && config.mpiControlPlane && config.nixlUcxBackend + && config.tensorParallelism == 1 && config.contextParallelism == 1 && config.pipelineParallelism > 1 + && !config.attentionDp && !config.inflightCancellation && config.transferOverlap && !config.layerwiseTransfer + && config.mpiThreadMultiple; +} + +bool selectWorkerPublishedConsensus(std::vector const& protocolVersions) +{ + if (protocolVersions.empty() || protocolVersions.front() == 0) + { + return false; + } + return std::all_of(protocolVersions.begin(), protocolVersions.end(), + [&](std::uint64_t const version) { return version == protocolVersions.front(); }); +} + +TransferStatusConsensus::TransferStatusConsensus(int const participantCount) + : mParticipantCount(participantCount) +{ + TLLM_CHECK_WITH_INFO(participantCount > 0, "Transfer-status consensus requires at least one participant."); +} + +void TransferStatusConsensus::recordVote( + int const participantRank, std::uint64_t const requestId, TransferStatusVote const vote) +{ + TLLM_CHECK_WITH_INFO(participantRank >= 0 && participantRank < mParticipantCount, + "Transfer-status consensus participant rank is out of range."); + TLLM_CHECK_WITH_INFO(vote == TransferStatusVote::kCompleted || vote == TransferStatusVote::kFailed, + "Transfer-status 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, + "Transfer-status consensus participant changed its terminal vote for a request."); + return; + } + + recordedVote = packedVote; + ++requestVotes.terminalCount; + requestVotes.failed = requestVotes.failed || vote == TransferStatusVote::kFailed; +} + +TransferStatusConsensusResult TransferStatusConsensus::takeCompleted() +{ + TransferStatusConsensusResult result; + for (auto requestIt = mRequestVotes.begin(); requestIt != mRequestVotes.end();) + { + auto const& requestVotes = requestIt->second; + if (requestVotes.terminalCount != mParticipantCount) + { + ++requestIt; + continue; + } + + auto& terminalRequestIds = requestVotes.failed ? result.failedRequestIds : result.completedRequestIds; + terminalRequestIds.insert(requestIt->first); + requestIt = mRequestVotes.erase(requestIt); + } + return result; +} + +class ContextTransferVoteMailbox::Impl +{ +public: + explicit Impl(std::shared_ptr comm) + : mComm(std::move(comm)) + , mConsensus(mComm->getSize()) + { + TLLM_CHECK_WITH_INFO(mComm->isMpi(), "Worker-published context-transfer consensus requires MPI."); + TLLM_CHECK_WITH_INFO(mComm->getSize() > 1, + "Worker-published context-transfer consensus requires multiple pipeline-parallel participants."); + } + + void publishOutcomeToPeers(std::uint64_t const requestId, bool const failed) noexcept + { + std::lock_guard lock(mMutex); + if (mShutdown || mError) + { + return; + } + + try + { + auto const vote = failed ? TransferStatusVote::kFailed : TransferStatusVote::kCompleted; + auto const [voteIt, inserted] = mPublishedLocalVotes.emplace(requestId, vote); + TLLM_CHECK_WITH_INFO(inserted || voteIt->second == vote, + "Context-transfer worker changed its published terminal vote for a request."); + if (inserted) + { + queuePacketForPeers(requestId, static_cast(vote)); + } + } + catch (...) + { + mError = std::current_exception(); + } + } + + [[nodiscard]] bool recordLocalOutcome(std::uint64_t const requestId) + { + std::lock_guard lock(mMutex); + rethrowError(); + auto const voteIt = mPublishedLocalVotes.find(requestId); + TLLM_CHECK_WITH_INFO(voteIt != mPublishedLocalVotes.end(), + "Context-transfer future became ready before its worker published a terminal vote."); + auto const vote = voteIt->second; + mConsensus.recordVote(mComm->getRank(), requestId, vote); + mPublishedLocalVotes.erase(voteIt); + return vote == TransferStatusVote::kFailed; + } + + [[nodiscard]] TransferStatusConsensusResult poll() + { + std::lock_guard lock(mMutex); + rethrowError(); + reapCompletedSends(); + drainIncomingPackets(); + return mConsensus.takeCompleted(); + } + + void shutdown() noexcept + { + std::unique_lock lock(mMutex); + if (mShutdown) + { + return; + } + mShutdown = true; + + try + { + // Sender workers have stopped before shutdown. MPI preserves the order of messages from one source with + // the same tag, so receiving a close marker proves that every earlier vote from that peer was drained. + queuePacketForPeers(/*requestId=*/0, kCloseMarker); + auto const peerCount = static_cast(mComm->getSize() - 1); + auto const shutdownDeadline = std::chrono::steady_clock::now() + kShutdownTimeout; + while (!mPendingSends.empty() || mClosedPeers.size() != peerCount) + { + reapCompletedSends(); + drainIncomingPackets(); + if (!mPendingSends.empty() || mClosedPeers.size() != peerCount) + { + if (std::chrono::steady_clock::now() >= shutdownDeadline) + { + TLLM_LOG_ERROR( + "Timed out shutting down context-transfer vote mailbox; received %zu of %zu peer close " + "markers with %zu sends still pending. Aborting to avoid freeing active MPI requests.", + mClosedPeers.size(), peerCount, mPendingSends.size()); + std::abort(); + } + lock.unlock(); + std::this_thread::yield(); + lock.lock(); + } + } + } + catch (std::exception const& error) + { + TLLM_LOG_ERROR("Failed to shut down context-transfer vote mailbox: %s", error.what()); + std::abort(); + } + catch (...) + { + TLLM_LOG_ERROR("Failed to shut down context-transfer vote mailbox with an unknown error"); + std::abort(); + } + } + +private: + static constexpr std::size_t kPacketFieldCount = 2; + static constexpr std::uint64_t kCloseMarker = 3; + static constexpr auto kShutdownTimeout = std::chrono::seconds(30); + + struct PendingSend + { + std::array packet{}; + std::unique_ptr request; + }; + + void rethrowError() const + { + if (mError) + { + std::rethrow_exception(mError); + } + } + + void queuePacketForPeers(std::uint64_t const requestId, std::uint64_t const value) + { + for (int peer = 0; peer < mComm->getSize(); ++peer) + { + if (peer == mComm->getRank()) + { + continue; + } + + 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, mpi::MpiTag::kContextTransferStatus); + } + catch (...) + { + mPendingSends.pop_back(); + throw; + } + } + } + + 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 drainIncomingPackets() + { + for (int peer = 0; peer < mComm->getSize(); ++peer) + { + if (peer == mComm->getRank()) + { + continue; + } + + MPI_Status status{}; + while (mComm->iprobe(peer, mpi::MpiTag::kContextTransferStatus, &status)) + { + std::array packet{}; + mComm->recv( + packet.data(), packet.size(), mpi::MpiType::kUINT64, peer, mpi::MpiTag::kContextTransferStatus); + if (packet.back() == kCloseMarker) + { + TLLM_CHECK_WITH_INFO(mClosedPeers.insert(peer).second, + "Context-transfer vote mailbox received a duplicate close marker."); + continue; + } + TLLM_CHECK_WITH_INFO(mClosedPeers.find(peer) == mClosedPeers.end(), + "Context-transfer vote mailbox received a vote after its peer close marker."); + mConsensus.recordVote(peer, packet.front(), static_cast(packet.back())); + } + } + } + + std::shared_ptr mComm; + TransferStatusConsensus mConsensus; + std::unordered_map mPublishedLocalVotes; + std::unordered_set mClosedPeers; + std::list mPendingSends; + std::mutex mMutex; + std::exception_ptr mError; + bool mShutdown{false}; +}; + +ContextTransferVoteMailbox::ContextTransferVoteMailbox(std::shared_ptr comm) + : mImpl(std::make_unique(std::move(comm))) +{ +} + +ContextTransferVoteMailbox::~ContextTransferVoteMailbox() +{ + shutdown(); +} + +void ContextTransferVoteMailbox::publishOutcomeToPeers(std::uint64_t const requestId, bool const failed) noexcept +{ + mImpl->publishOutcomeToPeers(requestId, failed); +} + +bool ContextTransferVoteMailbox::recordLocalOutcome(std::uint64_t const requestId) +{ + return mImpl->recordLocalOutcome(requestId); +} + +TransferStatusConsensusResult ContextTransferVoteMailbox::poll() +{ + return mImpl->poll(); +} + +void ContextTransferVoteMailbox::shutdown() noexcept +{ + if (mImpl) + { + mImpl->shutdown(); + } +} + +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/runtime/utils/mpiUtils.cpp b/cpp/tensorrt_llm/runtime/utils/mpiUtils.cpp index 0f8f31082e96..1c51ddc16afd 100644 --- a/cpp/tensorrt_llm/runtime/utils/mpiUtils.cpp +++ b/cpp/tensorrt_llm/runtime/utils/mpiUtils.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. @@ -256,6 +256,17 @@ void initialize(MpiThreadSupport threadMode, bool forwardAbortToParent) mpiInitialized = true; } +MpiThreadSupport getThreadSupport() +{ +#if ENABLE_MULTI_DEVICE + int providedMode = MPI_THREAD_SINGLE; + TLLM_MPI_CHECK(MPI_Query_thread(&providedMode)); + return static_cast(providedMode); +#else + TLLM_THROW("Multi device support is disabled."); +#endif +} + void MpiComm::barrier() const { couldUseMPI(); diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index f4e5e9fb2be0..dc0ffd70b98b 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(transferStatusConsensusTest transferStatusConsensusTest.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/transferStatusConsensusTest.cpp b/cpp/tests/unit_tests/batch_manager/transferStatusConsensusTest.cpp new file mode 100644 index 000000000000..06226a4386ee --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/transferStatusConsensusTest.cpp @@ -0,0 +1,164 @@ +/* + * 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/transferStatusConsensus.h" + +#include + +#include +#include +#include + +namespace tensorrt_llm::batch_manager +{ +namespace +{ + +TEST(TransferStatusConsensusTest, CompletesOnlyAfterEveryParticipantVotes) +{ + TransferStatusConsensus consensus(3); + + consensus.recordVote(0, 11, TransferStatusVote::kCompleted); + consensus.recordVote(2, 11, TransferStatusVote::kCompleted); + auto result = consensus.takeCompleted(); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_TRUE(result.failedRequestIds.empty()); + + consensus.recordVote(1, 11, TransferStatusVote::kCompleted); + result = consensus.takeCompleted(); + EXPECT_EQ(result.completedRequestIds, std::unordered_set{11}); + EXPECT_TRUE(result.failedRequestIds.empty()); +} + +TEST(TransferStatusConsensusTest, FailureWinsAfterEveryParticipantIsTerminal) +{ + TransferStatusConsensus consensus(3); + + consensus.recordVote(2, 7, TransferStatusVote::kCompleted); + consensus.recordVote(0, 7, TransferStatusVote::kCompleted); + consensus.recordVote(1, 7, TransferStatusVote::kFailed); + auto const result = consensus.takeCompleted(); + + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{7}); +} + +TEST(TransferStatusConsensusTest, AccumulatesInterleavedRequestsIndependently) +{ + TransferStatusConsensus consensus(2); + + consensus.recordVote(0, 3, TransferStatusVote::kCompleted); + consensus.recordVote(1, 5, TransferStatusVote::kFailed); + consensus.recordVote(1, 3, TransferStatusVote::kCompleted); + auto result = consensus.takeCompleted(); + EXPECT_EQ(result.completedRequestIds, std::unordered_set{3}); + EXPECT_TRUE(result.failedRequestIds.empty()); + + consensus.recordVote(0, 5, TransferStatusVote::kCompleted); + result = consensus.takeCompleted(); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{5}); +} + +TEST(TransferStatusConsensusTest, AcceptsIdenticalDuplicateAndRejectsConflictingDuplicate) +{ + TransferStatusConsensus consensus(2); + + consensus.recordVote(0, 17, TransferStatusVote::kCompleted); + EXPECT_NO_THROW(consensus.recordVote(0, 17, TransferStatusVote::kCompleted)); + EXPECT_ANY_THROW(consensus.recordVote(0, 17, TransferStatusVote::kFailed)); + EXPECT_TRUE(consensus.takeCompleted().completedRequestIds.empty()); +} + +TEST(TransferStatusConsensusTest, RejectsInvalidParticipant) +{ + TransferStatusConsensus consensus(2); + + EXPECT_ANY_THROW(consensus.recordVote(-1, 9, TransferStatusVote::kCompleted)); + EXPECT_ANY_THROW(consensus.recordVote(2, 9, TransferStatusVote::kCompleted)); +} + +TEST(TransferStatusConsensusTest, RejectsInvalidVote) +{ + TransferStatusConsensus consensus(1); + + EXPECT_ANY_THROW(consensus.recordVote(0, 9, static_cast(99))); +} + +TEST(TransferStatusConsensusTest, RejectsEmptyParticipantSet) +{ + EXPECT_ANY_THROW(TransferStatusConsensus(0)); +} + +TEST(TransferStatusConsensusTest, DoesNotCommitFailureBeforeEveryParticipantIsTerminal) +{ + TransferStatusConsensus consensus(4); + + consensus.recordVote(0, 23, TransferStatusVote::kFailed); + consensus.recordVote(1, 23, TransferStatusVote::kCompleted); + consensus.recordVote(2, 23, TransferStatusVote::kCompleted); + auto result = consensus.takeCompleted(); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_TRUE(result.failedRequestIds.empty()); + + consensus.recordVote(3, 23, TransferStatusVote::kCompleted); + result = consensus.takeCompleted(); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{23}); +} + +TEST(WorkerPublishedConsensusConfigTest, RequiresEveryQualifiedCondition) +{ + WorkerPublishedConsensusConfig supported; + supported.enabledForCppBinding = true; + supported.mpiControlPlane = true; + supported.nixlUcxBackend = true; + supported.pipelineParallelism = 4; + supported.transferOverlap = true; + supported.mpiThreadMultiple = true; + ASSERT_TRUE(supportsWorkerPublishedConsensus(supported)); + + auto expectUnsupported = [&](std::function const& mutate) + { + auto config = supported; + mutate(config); + EXPECT_FALSE(supportsWorkerPublishedConsensus(config)); + }; + expectUnsupported([](auto& config) { config.enabledForCppBinding = false; }); + expectUnsupported([](auto& config) { config.mpiControlPlane = false; }); + expectUnsupported([](auto& config) { config.nixlUcxBackend = false; }); + expectUnsupported([](auto& config) { config.tensorParallelism = 2; }); + expectUnsupported([](auto& config) { config.contextParallelism = 2; }); + expectUnsupported([](auto& config) { config.pipelineParallelism = 1; }); + expectUnsupported([](auto& config) { config.attentionDp = true; }); + expectUnsupported([](auto& config) { config.inflightCancellation = true; }); + expectUnsupported([](auto& config) { config.transferOverlap = false; }); + expectUnsupported([](auto& config) { config.layerwiseTransfer = true; }); + expectUnsupported([](auto& config) { config.mpiThreadMultiple = false; }); +} + +TEST(WorkerPublishedConsensusConfigTest, SelectsOnlyOneNonzeroVersionOnEveryRank) +{ + EXPECT_FALSE(selectWorkerPublishedConsensus({})); + EXPECT_FALSE(selectWorkerPublishedConsensus({0, 0, 0, 0})); + EXPECT_FALSE(selectWorkerPublishedConsensus({1, 0, 1, 1})); + EXPECT_FALSE(selectWorkerPublishedConsensus({1, 2, 1, 1})); + EXPECT_TRUE(selectWorkerPublishedConsensus({1, 1, 1, 1})); +} + +} // 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 401cdbfd5d44..230f1bee6a4a 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -32,6 +32,7 @@ #include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" +#include "tensorrt_llm/batch_manager/transferStatusConsensus.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/envUtils.h" @@ -42,6 +43,8 @@ #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/testing/kvCacheManagerTestUtil.h" +#include +#include #include #include #include @@ -53,6 +56,7 @@ #include #include #include +#include #include "gtest/gtest.h" #include @@ -160,6 +164,106 @@ TEST_F(CacheConfigTest, EqualTo) EXPECT_EQ(state0, state1); } +#if ENABLE_MULTI_DEVICE +TEST(ContextTransferVoteMailboxTest, WorkerVotesProgressWithoutPeerSchedulerPolling) +{ + constexpr int kPipelineParallelism = 4; + constexpr std::uint64_t kBlockedRequestId = 41; + constexpr std::uint64_t kIndependentRequestId = 42; + constexpr int kFailureRank = 1; + constexpr auto kRankZeroSchedulerDelay = std::chrono::seconds(3); + constexpr auto kPressureTimeout = std::chrono::seconds(2); + constexpr auto kTeardownDelay = std::chrono::milliseconds(250); + constexpr auto kPollTimeout = std::chrono::seconds(5); + + tensorrt_llm::mpi::initialize(tensorrt_llm::mpi::MpiThreadSupport::THREAD_MULTIPLE); + if (tensorrt_llm::mpi::getThreadSupport() != tensorrt_llm::mpi::MpiThreadSupport::THREAD_MULTIPLE) + { + GTEST_SKIP() << "Test requires MPI_THREAD_MULTIPLE."; + } + auto& world = tensorrt_llm::mpi::MpiComm::world(); + if (world.getSize() < kPipelineParallelism || world.getSize() % kPipelineParallelism != 0) + { + GTEST_SKIP() << "Test requires a process count divisible by four."; + } + + auto ppMpiComm = std::make_shared( + world.split(world.getRank() / kPipelineParallelism, world.getRank())); + auto ppComm = std::make_shared(ppMpiComm); + ContextTransferVoteMailbox mailbox(ppComm); + auto const ppRank = ppComm->getRank(); + + std::thread worker( + [&]() + { + mailbox.publishOutcomeToPeers(kIndependentRequestId, ppRank == kFailureRank); + if (ppRank != kPipelineParallelism - 1) + { + mailbox.publishOutcomeToPeers(kBlockedRequestId, /*failed=*/false); + } + }); + worker.join(); + + if (ppRank == 0) + { + // Simulate rank 0 being inside model forward while its transfer worker still publishes to downstream ranks. + std::this_thread::sleep_for(kRankZeroSchedulerDelay); + } + EXPECT_EQ(mailbox.recordLocalOutcome(kIndependentRequestId), ppRank == kFailureRank); + if (ppRank != kPipelineParallelism - 1) + { + EXPECT_FALSE(mailbox.recordLocalOutcome(kBlockedRequestId)); + } + + bool independentFailed = false; + auto const independentDeadline = std::chrono::steady_clock::now() + (ppRank == 0 ? kPollTimeout : kPressureTimeout); + while (!independentFailed && std::chrono::steady_clock::now() < independentDeadline) + { + auto const result = mailbox.poll(); + EXPECT_EQ(result.completedRequestIds.count(kBlockedRequestId), 0); + EXPECT_EQ(result.failedRequestIds.count(kBlockedRequestId), 0); + independentFailed = result.failedRequestIds.count(kIndependentRequestId) != 0; + if (!independentFailed) + { + std::this_thread::yield(); + } + } + EXPECT_TRUE(independentFailed); + + // No rank may commit the blocked request until its last worker publishes, but the independent request above is + // not held behind it. This barrier is test coordination after the pressure property has already been verified. + ppMpiComm->barrier(); + if (ppRank == kPipelineParallelism - 1) + { + std::thread finalWorker([&]() { mailbox.publishOutcomeToPeers(kBlockedRequestId, /*failed=*/false); }); + finalWorker.join(); + EXPECT_FALSE(mailbox.recordLocalOutcome(kBlockedRequestId)); + } + + bool blockedCompleted = false; + auto const blockedDeadline = std::chrono::steady_clock::now() + kPollTimeout; + while (!blockedCompleted && std::chrono::steady_clock::now() < blockedDeadline) + { + auto const result = mailbox.poll(); + blockedCompleted = result.completedRequestIds.count(kBlockedRequestId) != 0; + if (!blockedCompleted) + { + std::this_thread::yield(); + } + } + EXPECT_TRUE(blockedCompleted); + + // Exercise skewed teardown: peers must keep draining until rank 0 sends its ordered close marker. + if (ppRank == 0) + { + std::this_thread::sleep_for(kTeardownDelay); + } + auto const shutdownStart = std::chrono::steady_clock::now(); + mailbox.shutdown(); + EXPECT_LT(std::chrono::steady_clock::now() - shutdownStart, kPollTimeout); +} +#endif // ENABLE_MULTI_DEVICE + // TODO: Restore multi-rank tests. // --------------------------------------- @@ -329,7 +433,8 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- return std::make_unique(mRequestId++, std::move(request)); } - void addRequestAndTransportCache(std::shared_ptr const& llmRequest) + void addRequestAndTransportCache( + std::shared_ptr const& llmRequest, CacheSender::CompletionCallback completionCallback = {}) { auto constexpr beamIdx{0}; auto constexpr beamWidth{1}; @@ -348,7 +453,9 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- TLLM_CUDA_CHECK(cudaMemset(it->data(), llmRequest->getPromptLen(), it->getSizeInBytes())); } } - mFutures.emplace_back(mSender->sendAsync(llmRequest)); + auto future = completionCallback ? mSender->sendAsync(llmRequest, std::move(completionCallback)) + : mSender->sendAsync(llmRequest); + mFutures.emplace_back(std::move(future)); } else { @@ -426,6 +533,60 @@ TEST_F(SymmetricalCacheTest, SimpleTest) } } +TEST_F(SymmetricalCacheTest, CompletionCallbackRunsExactlyOnceBeforeFutureReady) +{ + auto const worldSize = setUpCommunicator(); + if (worldSize != 2) + { + GTEST_SKIP() << "mpirun 2 processes is required to run this test."; + } + setUpCacheManager(); + setUpCacheTransceiver(); + auto request = std::shared_ptr(makeLlmRequest(10)); + + std::atomic callbackCount{0}; + std::atomic callbackReportedFailure{true}; + std::atomic callbackEntered{false}; + std::atomic allowCallbackReturn{false}; + CacheSender::CompletionCallback callback; + if (isSender) + { + callback = [&](bool const failed) + { + callbackReportedFailure.store(failed); + callbackCount.fetch_add(1); + callbackEntered.store(true); + while (!allowCallbackReturn.load()) + { + std::this_thread::yield(); + } + }; + } + addRequestAndTransportCache(request, std::move(callback)); + + if (isSender) + { + ASSERT_EQ(mFutures.size(), 1); + auto const callbackDeadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!callbackEntered.load() && std::chrono::steady_clock::now() < callbackDeadline) + { + std::this_thread::yield(); + } + EXPECT_TRUE(callbackEntered.load()); + if (callbackEntered.load()) + { + EXPECT_EQ(mFutures.front().wait_for(std::chrono::milliseconds(0)), std::future_status::timeout); + } + allowCallbackReturn.store(true); + mFutures.front().get(); + mFutures.clear(); + EXPECT_EQ(callbackCount.load(), 1); + EXPECT_FALSE(callbackReportedFailure.load()); + mSender.reset(); + EXPECT_EQ(callbackCount.load(), 1); + } +} + #if ENABLE_MULTI_DEVICE using AsymmetricTestParam = std::tuple +#include +#include namespace mpi = tensorrt_llm::mpi; namespace tr = tensorrt_llm::runtime; @@ -47,6 +49,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::kContextTransferStatus); + static_cast(comm.recv(&receivedValue, 1, mpi::MpiType::kUINT64, source, mpi::MpiTag::kContextTransferStatus)); + while (!request->isCompleted()) + { + std::this_thread::yield(); + } + + EXPECT_EQ(receivedValue, static_cast(source)); +} +#endif // ENABLE_MULTI_DEVICE + template void testBroadcast() { diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index cfe7f1153741..bdfb93ffb900 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -785,7 +785,8 @@ def run_disaggregated_test(example_dir, cwd=None, disagg_schedule_style=None, post_client_test=None, - assert_gen_log_contains=None): + assert_gen_log_contains=None, + expected_context_worker_log=None): """Run disaggregated test using service discovery instead of MPI. If assert_gen_log_contains is set, the generation-worker logs are captured and, after the @@ -805,7 +806,8 @@ def run_disaggregated_test(example_dir, config, ctx_workers, gen_workers, disagg_server, server_port, work_dir = \ setup_disagg_cluster(config_file, model_name=model_path, env=run_env, cwd=cwd, schedule_style=disagg_schedule_style, - save_log=assert_gen_log_contains is not None) + save_log=assert_gen_log_contains is not None + or expected_context_worker_log is not None) server_host = config.get("hostname", "localhost") @@ -853,6 +855,15 @@ def run_disaggregated_test(example_dir, f"expected marker {assert_gen_log_contains!r} in a generation-worker log, " f"but none of {len(logs)} log(s) contained it (bounce did not engage)" ) + if expected_context_worker_log is not None: + context_worker_logs = [] + for worker in ctx_workers: + assert worker.log_path is not None + with open(worker.log_path, 'r', errors='replace') as log_file: + context_worker_logs.append(log_file.read()) + assert any(expected_context_worker_log in log + for log in context_worker_logs), \ + f"Expected context-worker log not found: {expected_context_worker_log!r}" finally: terminate(*ctx_workers, *gen_workers, disagg_server) shutil.rmtree(work_dir, ignore_errors=True) @@ -1575,11 +1586,14 @@ def test_disaggregated_ctxpp4_gentp4(disaggregated_test_root, llm_venv, llama_model_root): setup_model_symlink(llm_venv, llama_model_root, "TinyLlama/TinyLlama-1.1B-Chat-v1.0") - run_disaggregated_test(disaggregated_example_root, - "ctxpp4_gentp4", - env=llm_venv._new_env, - model_path=llama_model_root, - cwd=llm_venv.get_working_directory()) + run_disaggregated_test( + disaggregated_example_root, + "ctxpp4_gentp4", + env=llm_venv._new_env, + model_path=llama_model_root, + cwd=llm_venv.get_working_directory(), + expected_context_worker_log= + "Enabled worker-published context-transfer consensus") @skip_no_hopper diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index cd84403c6104..fcdfa5dc90cd 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -517,7 +517,8 @@ def test_async_transfer_keeps_llm_request_alive(): completed_ctx_ids = set() def poll_transfers(): - completed, failed = transceiver_ctx.check_context_transfer_status(1) + completed, failed = transceiver_ctx.impl.check_context_transfer_status( + 1, True) assert failed == [] completed_ctx_ids.update(completed) transceiver_gen.check_gen_transfer_status(1) @@ -528,6 +529,9 @@ def transfers_done(): == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE) wait_for_transfer_completion(poll_transfers, transfers_done) + assert ctx_request.state == LlmRequestState.DISAGG_CONTEXT_COMPLETE + assert transceiver_ctx.impl.check_context_transfer_status(0, + True) == ([], []) def _build_ctx_request_for_timeout_test(request_id: int) -> LlmRequest: From 656873e38d1eb1a446c649eda7da4cb56ef844f1 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 15 Jul 2026 12:29:17 -0700 Subject: [PATCH 2/4] [https://nvbugs/6448152][perf] diagnose local context completion Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 4 ++ .../batch_manager/cacheTransceiver.cpp | 67 +++++++++++++++++-- .../defs/disaggregated/test_disaggregated.py | 25 ++++++- ...tx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml | 2 +- 4 files changed, 90 insertions(+), 8 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index b341dba3f826..062f0325d23d 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -329,6 +329,8 @@ class CacheTransceiver : public BaseCacheTransceiver std::unordered_set mCompletedSenderRequestIds; std::unordered_set mFailedSenderRequestIds; std::unordered_map> mSenderRequestsAwaitingConsensus; + // Diagnostic-only tombstones for requests reported locally before PP consensus commits. + std::unordered_set mDiagnosticLocallyReportedSenderRequestIds; std::unordered_set mCompletedRequesterRequestIds; std::unordered_set mFailedRequesterRequestIds; std::unordered_map> mRequesterRequestsAwaitingConsensus; @@ -337,6 +339,8 @@ class CacheTransceiver : public BaseCacheTransceiver std::shared_ptr mGroupComm; std::shared_ptr mGroupTensorParaComm, mGroupPipeParaComm, mGroupDataComm, mGroupTPInDPComm; std::unique_ptr mContextTransferVoteMailbox; + bool mDiagnosticEarlyLocalContextCompletion{false}; + bool mDiagnosticObservedConsensusGap{false}; executor::kv_cache::CommState const* mCommState; std::unique_ptr mCacheState; diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 9ab38925ff99..518bdda6d1fa 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -73,6 +73,7 @@ using RequestIdType = LlmRequest::RequestIdType; constexpr int kTransferFuturePollIntervalMs = 10; constexpr int kContextConsensusPollIntervalMs = 1; +constexpr char const* kDiagnosticEarlyLocalContextCompletionEnv = "TRTLLM_DIAGNOSTIC_EARLY_LOCAL_CONTEXT_COMPLETION"; // Finite status checks are scheduler polls, not terminal deadlines. Pure polls // use short slices; calls that ask for at least one completion keep bounded @@ -486,9 +487,24 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa // throws, member unwinding still sends this rank's ordered close marker to its peers. if (useWorkerPublishedContextConsensus) { + auto const localDiagnosticMode + = static_cast(common::getBoolEnv(kDiagnosticEarlyLocalContextCompletionEnv)); + std::vector diagnosticModes(static_cast(mGroupPipeParaComm->getSize())); + mGroupPipeParaComm->allgather(&localDiagnosticMode, diagnosticModes.data(), 1, mpi::MpiType::kUINT64); + TLLM_CHECK_WITH_INFO(std::all_of(diagnosticModes.begin(), diagnosticModes.end(), + [&](std::uint64_t const mode) { return mode == localDiagnosticMode; }), + "%s must be set consistently on every context pipeline rank.", + kDiagnosticEarlyLocalContextCompletionEnv); + mDiagnosticEarlyLocalContextCompletion = localDiagnosticMode != 0; mContextTransferVoteMailbox = std::make_unique(mGroupPipeParaComm); } } + if (mDiagnosticEarlyLocalContextCompletion) + { + TLLM_LOG_WARNING( + "DIAGNOSTIC ONLY: reporting successful local context-transfer completion before PP consensus. " + "Any local or late global failure is fatal; do not use this mode in production."); + } bool isMLA = attentionType == executor::kv_cache::CacheState::AttentionType::kMLA; std::optional maxNumTokens = mCacheTransceiverConfig.value().getMaxTokensInBuffer(); @@ -680,6 +696,15 @@ void CacheTransceiver::setContextState(LlmRequest* llmRequest) void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isContextOnlyRequest()); + if (mDiagnosticEarlyLocalContextCompletion) + { + auto const requestId = llmRequest->mRequestId; + TLLM_CHECK_WITH_INFO(mSenderRequestsAwaitingConsensus.find(requestId) == mSenderRequestsAwaitingConsensus.end() + && mDiagnosticLocallyReportedSenderRequestIds.find(requestId) + == mDiagnosticLocallyReportedSenderRequestIds.end(), + "DIAGNOSTIC ONLY: context request ID %zu was reused before its prior PP consensus retired.", + static_cast(requestId)); + } llmRequest->setState(LlmRequestState::kDISAGG_CONTEXT_TRANS_IN_PROGRESS); // If context phase params is already set, it means that the KV cache // transfer is already in progress. @@ -1017,6 +1042,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( progressDeadline = std::chrono::steady_clock::now() + futureWaitInterval; } auto toProcess = consensusBackend.selectReadyRequests(mSenderFutures); + RequestStatuses requestsStatus; // Preserve insertion order when a caller asks this poll to make progress on at least N requests. for (auto it = mSenderFutures.begin(); @@ -1110,9 +1136,23 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( auto const failed = consensusBackend.resolveLocalFailure(requestId, observedFailure); recordLocalTransferOutcome(requestId, std::move(completedRequest), failed, mCompletedSenderRequestIds, mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); + if (mDiagnosticEarlyLocalContextCompletion) + { + TLLM_CHECK_WITH_INFO(!failed, + "DIAGNOSTIC ONLY: context request %zu failed locally; the early-release experiment supports " + "successful transfers only.", + static_cast(requestId)); + TLLM_CHECK_WITH_INFO(mDiagnosticLocallyReportedSenderRequestIds.insert(requestId).second, + "Context request %zu was reported locally more than once in diagnostic mode.", + static_cast(requestId)); + requestsStatus.completedRequestIds.insert(requestId); + if (markComplete) + { + mSenderRequestsAwaitingConsensus.at(requestId)->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); + } + } } - RequestStatuses requestsStatus; auto commitConsensusOutcome = [&](TransferConsensusOutcome const& outcome) { for (auto const requestId : sortedRequestIds(outcome.failedRequestIds)) @@ -1125,6 +1165,10 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( static_cast(requestId)); continue; } + TLLM_CHECK_WITH_INFO(mDiagnosticLocallyReportedSenderRequestIds.erase(requestId) == 0, + "DIAGNOSTIC ONLY: context request %zu failed PP consensus after local completion was already " + "reported; the early-release experiment is invalid.", + static_cast(requestId)); requestIt->second->setState(LlmRequestState::kDISAGG_TRANS_ERROR); requestsStatus.errorRequestIds.insert(requestId); mTimedOutSenderIds.erase(requestId); @@ -1142,10 +1186,14 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( static_cast(requestId)); continue; } - requestsStatus.completedRequestIds.insert(requestId); - if (markComplete) + auto const wasReportedLocally = mDiagnosticLocallyReportedSenderRequestIds.erase(requestId) > 0; + if (!wasReportedLocally) { - requestIt->second->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); + requestsStatus.completedRequestIds.insert(requestId); + if (markComplete) + { + requestIt->second->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); + } } mTimedOutSenderIds.erase(requestId); mCancelRequestedSenderIds.erase(requestId); @@ -1196,7 +1244,16 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( // Collective consensus must be entered exactly once per status call. The mailbox can make independent progress // and may be polled repeatedly without requiring a matching peer scheduler call. pollAndCommitConsensus(); - if (consensusBackend.isWorkerPublished()) + if (mDiagnosticEarlyLocalContextCompletion && !mDiagnosticObservedConsensusGap + && !mDiagnosticLocallyReportedSenderRequestIds.empty()) + { + TLLM_LOG_WARNING( + "DIAGNOSTIC OBSERVED: %zu locally completed context transfer(s) will be released while still awaiting " + "PP consensus.", + mDiagnosticLocallyReportedSenderRequestIds.size()); + mDiagnosticObservedConsensusGap = true; + } + if (consensusBackend.isWorkerPublished() && !mDiagnosticEarlyLocalContextCompletion) { auto const committedCount = [&]() { return requestsStatus.completedRequestIds.size() + requestsStatus.errorRequestIds.size(); }; diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index bdfb93ffb900..3d3146abd1af 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -861,9 +861,9 @@ def run_disaggregated_test(example_dir, assert worker.log_path is not None with open(worker.log_path, 'r', errors='replace') as log_file: context_worker_logs.append(log_file.read()) - assert any(expected_context_worker_log in log + assert all(expected_context_worker_log in log for log in context_worker_logs), \ - f"Expected context-worker log not found: {expected_context_worker_log!r}" + f"Expected context-worker log not found on every worker: {expected_context_worker_log!r}" finally: terminate(*ctx_workers, *gen_workers, disagg_server) shutil.rmtree(work_dir, ignore_errors=True) @@ -1596,6 +1596,27 @@ def test_disaggregated_ctxpp4_gentp4(disaggregated_test_root, llm_venv, "Enabled worker-published context-transfer consensus") +@pytest.mark.skip_less_device(8) +@pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], + indirect=True) +def test_disaggregated_ctxpp4_gentp4_diagnostic_local_completion( + disaggregated_test_root, llm_venv, disaggregated_example_root, + llama_model_root): + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") + env = llm_venv._new_env.copy() + env["TRTLLM_DIAGNOSTIC_EARLY_LOCAL_CONTEXT_COMPLETION"] = "1" + run_disaggregated_test( + disaggregated_example_root, + "ctxpp4_gentp4", + env=env, + model_path=llama_model_root, + cwd=llm_venv.get_working_directory(), + expected_context_worker_log= + "DIAGNOSTIC ONLY: reporting successful local context-transfer completion" + ) + + @skip_no_hopper @pytest.mark.skip_less_device(4) @pytest.mark.skip( 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..44f08822f075 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 @@ -35,7 +35,7 @@ environment: trtllm_repo: '' 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 + worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_DIAGNOSTIC_EARLY_LOCAL_CONTEXT_COMPLETION=1 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 server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: From 1522ea72308bcd8103de6fc877456515ebeb6f9f Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 16 Jul 2026 01:36:35 -0700 Subject: [PATCH 3/4] [https://nvbugs/6448152][perf] diagnose PP consensus traffic Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 2 + .../batch_manager/cacheTransceiver.cpp | 127 ++++++++++++++---- .../defs/disaggregated/test_disaggregated.py | 6 +- ...tx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml | 2 +- 4 files changed, 108 insertions(+), 29 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 062f0325d23d..370451368a66 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -340,7 +340,9 @@ class CacheTransceiver : public BaseCacheTransceiver std::shared_ptr mGroupTensorParaComm, mGroupPipeParaComm, mGroupDataComm, mGroupTPInDPComm; std::unique_ptr mContextTransferVoteMailbox; bool mDiagnosticEarlyLocalContextCompletion{false}; + bool mDiagnosticNoContextConsensusTraffic{false}; bool mDiagnosticObservedConsensusGap{false}; + bool mDiagnosticObservedNoConsensusCompletion{false}; executor::kv_cache::CommState const* mCommState; std::unique_ptr mCacheState; diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 518bdda6d1fa..3b3b338136de 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -74,6 +74,15 @@ using RequestIdType = LlmRequest::RequestIdType; constexpr int kTransferFuturePollIntervalMs = 10; constexpr int kContextConsensusPollIntervalMs = 1; constexpr char const* kDiagnosticEarlyLocalContextCompletionEnv = "TRTLLM_DIAGNOSTIC_EARLY_LOCAL_CONTEXT_COMPLETION"; +constexpr char const* kDiagnosticNoContextConsensusTrafficEnv + = "TRTLLM_DIAGNOSTIC_DISABLE_PP_CONTEXT_CONSENSUS_TRAFFIC"; + +enum class ContextConsensusDiagnosticMode : std::uint64_t +{ + kNone = 0, + kEarlyLocalCompletion = 1, + kNoConsensusTraffic = 2, +}; // Finite status checks are scheduler polls, not terminal deadlines. Pure polls // use short slices; calls that ask for at least one completion keep bounded @@ -468,12 +477,43 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa std::vector protocolVersions(static_cast(mGroupPipeParaComm->getSize())); mGroupPipeParaComm->allgather(&localVersion, protocolVersions.data(), 1, mpi::MpiType::kUINT64); auto const useWorkerPublishedContextConsensus = selectWorkerPublishedConsensus(protocolVersions); + + // Agree on diagnostic mode before constructing the mailbox. The no-traffic experiment deliberately retains + // this one-time startup agreement while removing all per-request consensus publications and polls. + auto const diagnosticEarlyLocalCompletion = common::getBoolEnv(kDiagnosticEarlyLocalContextCompletionEnv); + auto const diagnosticNoConsensusTraffic = common::getBoolEnv(kDiagnosticNoContextConsensusTrafficEnv); + TLLM_CHECK_WITH_INFO(!diagnosticEarlyLocalCompletion || !diagnosticNoConsensusTraffic, + "%s and %s cannot both be enabled.", kDiagnosticEarlyLocalContextCompletionEnv, + kDiagnosticNoContextConsensusTrafficEnv); + auto const diagnosticMode = diagnosticNoConsensusTraffic + ? ContextConsensusDiagnosticMode::kNoConsensusTraffic + : (diagnosticEarlyLocalCompletion ? ContextConsensusDiagnosticMode::kEarlyLocalCompletion + : ContextConsensusDiagnosticMode::kNone); + auto const localDiagnosticMode = static_cast(diagnosticMode); + std::vector diagnosticModes(static_cast(mGroupPipeParaComm->getSize())); + mGroupPipeParaComm->allgather(&localDiagnosticMode, diagnosticModes.data(), 1, mpi::MpiType::kUINT64); + TLLM_CHECK_WITH_INFO(std::all_of(diagnosticModes.begin(), diagnosticModes.end(), + [&](std::uint64_t const mode) { return mode == localDiagnosticMode; }), + "Context-consensus diagnostic mode must be set consistently on every context pipeline rank."); + TLLM_CHECK_WITH_INFO( + diagnosticMode == ContextConsensusDiagnosticMode::kNone || useWorkerPublishedContextConsensus, + "Context-consensus diagnostics require the qualified C++ NIXL/UCX TP1/CP1 pipeline-parallel path."); + mDiagnosticEarlyLocalContextCompletion + = diagnosticMode == ContextConsensusDiagnosticMode::kEarlyLocalCompletion; + mDiagnosticNoContextConsensusTraffic = diagnosticMode == ContextConsensusDiagnosticMode::kNoConsensusTraffic; + if (useWorkerPublishedContextConsensus) { - TLLM_LOG_INFO( - "Enabled worker-published context-transfer consensus v%lu for the C++ NIXL/UCX TP1/CP1 " - "pipeline-parallel path.", - kWorkerPublishedConsensusVersion); + if (!mDiagnosticNoContextConsensusTraffic) + { + // If later transceiver initialization throws, member unwinding still sends this rank's ordered close + // marker to its peers. + mContextTransferVoteMailbox = std::make_unique(mGroupPipeParaComm); + TLLM_LOG_INFO( + "Enabled worker-published context-transfer consensus v%lu for the C++ NIXL/UCX TP1/CP1 " + "pipeline-parallel path.", + kWorkerPublishedConsensusVersion); + } } else if (std::any_of(protocolVersions.begin(), protocolVersions.end(), [](std::uint64_t const version) { return version != 0; })) @@ -482,22 +522,6 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa "Pipeline ranks advertised different worker-published context-transfer consensus capabilities; " "retaining synchronous context-transfer consensus on every rank."); } - - // Construct the mailbox immediately after the all-rank agreement. If later transceiver initialization - // throws, member unwinding still sends this rank's ordered close marker to its peers. - if (useWorkerPublishedContextConsensus) - { - auto const localDiagnosticMode - = static_cast(common::getBoolEnv(kDiagnosticEarlyLocalContextCompletionEnv)); - std::vector diagnosticModes(static_cast(mGroupPipeParaComm->getSize())); - mGroupPipeParaComm->allgather(&localDiagnosticMode, diagnosticModes.data(), 1, mpi::MpiType::kUINT64); - TLLM_CHECK_WITH_INFO(std::all_of(diagnosticModes.begin(), diagnosticModes.end(), - [&](std::uint64_t const mode) { return mode == localDiagnosticMode; }), - "%s must be set consistently on every context pipeline rank.", - kDiagnosticEarlyLocalContextCompletionEnv); - mDiagnosticEarlyLocalContextCompletion = localDiagnosticMode != 0; - mContextTransferVoteMailbox = std::make_unique(mGroupPipeParaComm); - } } if (mDiagnosticEarlyLocalContextCompletion) { @@ -505,6 +529,12 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa "DIAGNOSTIC ONLY: reporting successful local context-transfer completion before PP consensus. " "Any local or late global failure is fatal; do not use this mode in production."); } + if (mDiagnosticNoContextConsensusTraffic) + { + TLLM_LOG_WARNING( + "DIAGNOSTIC ONLY: committing successful local context-transfer completion without publishing or " + "polling PP consensus. Any local failure is fatal; do not use this mode in production."); + } bool isMLA = attentionType == executor::kv_cache::CacheState::AttentionType::kMLA; std::optional maxNumTokens = mCacheTransceiverConfig.value().getMaxTokensInBuffer(); @@ -842,11 +872,14 @@ class ContextTransferConsensusBackend public: ContextTransferConsensusBackend(ContextTransferVoteMailbox* mailbox, std::shared_ptr const& intraStageComm, - std::shared_ptr const& pipeStageComm) + std::shared_ptr const& pipeStageComm, bool const diagnosticLocalOnly) : mMailbox(mailbox) , mIntraStageComm(intraStageComm) , mPipeStageComm(pipeStageComm) + , mDiagnosticLocalOnly(diagnosticLocalOnly) { + TLLM_CHECK_WITH_INFO(!mDiagnosticLocalOnly || mMailbox == nullptr, + "Diagnostic local-only context completion cannot use a worker-published consensus mailbox."); } [[nodiscard]] bool isWorkerPublished() const noexcept @@ -854,10 +887,20 @@ class ContextTransferConsensusBackend return mMailbox != nullptr; } + [[nodiscard]] bool isDiagnosticLocalOnly() const noexcept + { + return mDiagnosticLocalOnly; + } + + [[nodiscard]] bool makesIndependentProgress() const noexcept + { + return isWorkerPublished() || isDiagnosticLocalOnly(); + } + template [[nodiscard]] std::unordered_set selectReadyRequests(SenderFutures const& senderFutures) const { - if (isWorkerPublished()) + if (makesIndependentProgress()) { std::unordered_set readyRequestIds; for (auto const& [request, future] : senderFutures) @@ -923,6 +966,14 @@ class ContextTransferConsensusBackend std::unordered_set const& failedRequestIds, std::unordered_set const& timedOutRequestIds) const { + if (isDiagnosticLocalOnly()) + { + TransferConsensusOutcome outcome; + outcome.completedRequestIds = completedRequestIds; + outcome.failedRequestIds = failedRequestIds; + outcome.timedOutRequestIds = timedOutRequestIds; + return outcome; + } if (!isWorkerPublished()) { return reduceTransferStates( @@ -940,6 +991,7 @@ class ContextTransferConsensusBackend ContextTransferVoteMailbox* mMailbox; std::shared_ptr const& mIntraStageComm; std::shared_ptr const& mPipeStageComm; + bool mDiagnosticLocalOnly; }; std::unordered_set const& getEmptyRequestIdSet() @@ -1035,9 +1087,9 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( auto const& intraStageComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupTPInDPComm : mGroupTensorParaComm; ContextTransferConsensusBackend consensusBackend{ - mContextTransferVoteMailbox.get(), intraStageComm, mGroupPipeParaComm}; + mContextTransferVoteMailbox.get(), intraStageComm, mGroupPipeParaComm, mDiagnosticNoContextConsensusTraffic}; std::optional progressDeadline; - if (consensusBackend.isWorkerPublished() && needsProgress) + if (consensusBackend.makesIndependentProgress() && needsProgress) { progressDeadline = std::chrono::steady_clock::now() + futureWaitInterval; } @@ -1079,7 +1131,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( auto waitInterval = futureWaitInterval; if (progressDeadline.has_value()) { - // Worker-published status polling has one total budget across local future waits and mailbox progress. + // Independently progressing backends have one total budget across local future waits and status progress. waitInterval = std::max(std::chrono::milliseconds(0), std::chrono::duration_cast( progressDeadline.value() - std::chrono::steady_clock::now())); @@ -1241,6 +1293,31 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( commitConsensusOutcome(outcome); }; + if (mDiagnosticNoContextConsensusTraffic) + { + TLLM_CHECK(consensusBackend.isDiagnosticLocalOnly()); + TLLM_CHECK(!mContextTransferVoteMailbox); + TLLM_CHECK_WITH_INFO(!inflightCancelEnabled, + "DIAGNOSTIC ONLY: context consensus traffic cannot be disabled with in-flight cancellation enabled."); + TLLM_CHECK_WITH_INFO(mFailedSenderRequestIds.empty(), + "DIAGNOSTIC ONLY: a context transfer failed locally while PP consensus traffic was disabled."); + TLLM_CHECK_WITH_INFO(mTimedOutSenderIds.empty(), + "DIAGNOSTIC ONLY: a context transfer exceeded its deadline while PP consensus traffic was disabled."); + auto const localOutcome + = consensusBackend.poll(mCompletedSenderRequestIds, mFailedSenderRequestIds, mTimedOutSenderIds); + auto const localCompletedCount = localOutcome.completedRequestIds.size(); + commitConsensusOutcome(localOutcome); + if (!mDiagnosticObservedNoConsensusCompletion && localCompletedCount > 0) + { + TLLM_LOG_WARNING( + "DIAGNOSTIC OBSERVED: context transfer completed with no runtime PP consensus traffic " + "(%zu successful local completion(s)).", + localCompletedCount); + mDiagnosticObservedNoConsensusCompletion = true; + } + return requestsStatus; + } + // Collective consensus must be entered exactly once per status call. The mailbox can make independent progress // and may be polled repeatedly without requiring a matching peer scheduler call. pollAndCommitConsensus(); diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 3d3146abd1af..46131760aad3 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -1599,13 +1599,13 @@ def test_disaggregated_ctxpp4_gentp4(disaggregated_test_root, llm_venv, @pytest.mark.skip_less_device(8) @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], indirect=True) -def test_disaggregated_ctxpp4_gentp4_diagnostic_local_completion( +def test_disaggregated_ctxpp4_gentp4_diagnostic_no_consensus_traffic( disaggregated_test_root, llm_venv, disaggregated_example_root, llama_model_root): setup_model_symlink(llm_venv, llama_model_root, "TinyLlama/TinyLlama-1.1B-Chat-v1.0") env = llm_venv._new_env.copy() - env["TRTLLM_DIAGNOSTIC_EARLY_LOCAL_CONTEXT_COMPLETION"] = "1" + env["TRTLLM_DIAGNOSTIC_DISABLE_PP_CONTEXT_CONSENSUS_TRAFFIC"] = "1" run_disaggregated_test( disaggregated_example_root, "ctxpp4_gentp4", @@ -1613,7 +1613,7 @@ def test_disaggregated_ctxpp4_gentp4_diagnostic_local_completion( model_path=llama_model_root, cwd=llm_venv.get_working_directory(), expected_context_worker_log= - "DIAGNOSTIC ONLY: reporting successful local context-transfer completion" + "DIAGNOSTIC OBSERVED: context transfer completed with no runtime PP consensus traffic" ) 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 44f08822f075..b803a927b5ec 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 @@ -35,7 +35,7 @@ environment: trtllm_repo: '' build_wheel: false work_dir: - worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_DIAGNOSTIC_EARLY_LOCAL_CONTEXT_COMPLETION=1 TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=1 + worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_DIAGNOSTIC_DISABLE_PP_CONTEXT_CONSENSUS_TRAFFIC=1 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 server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: From f645cb7d21c332825aad587c61a2335e37066c55 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 16 Jul 2026 18:03:00 -0700 Subject: [PATCH 4/4] [https://nvbugs/6448152][perf] diagnose 2x admission budget Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 79 ++++++++++++++++++- ...tx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml | 2 +- .../_torch/executor/test_py_executor.py | 70 +++++++++++++++- 3 files changed, 145 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a792e8f9a634..0920e4971807 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -247,6 +247,29 @@ def _strip_py_multimodal_data_post_prefill(request: LlmRequest) -> None: strip_mm_data_for_generation(mm_data) +_DIAGNOSTIC_DISAGG_ADMISSION_BUDGET_MULTIPLIER_ENV = ( + "TRTLLM_DIAGNOSTIC_DISAGG_ADMISSION_BUDGET_MULTIPLIER") + + +def _get_diagnostic_disagg_admission_budget_multiplier() -> int: + raw_multiplier = os.getenv( + _DIAGNOSTIC_DISAGG_ADMISSION_BUDGET_MULTIPLIER_ENV) + if raw_multiplier is None: + return 1 + + try: + multiplier = int(raw_multiplier) + except ValueError as error: + raise ValueError( + f"{_DIAGNOSTIC_DISAGG_ADMISSION_BUDGET_MULTIPLIER_ENV} must be " + "a positive integer") from error + if multiplier <= 0: + raise ValueError( + f"{_DIAGNOSTIC_DISAGG_ADMISSION_BUDGET_MULTIPLIER_ENV} must be " + "a positive integer") + return multiplier + + @dataclasses.dataclass class DisaggTransferAdmissionResult: admitted_requests: List[LlmRequest] @@ -263,11 +286,31 @@ def is_blocked_by_active_transfers(self) -> bool: class DisaggTransferAdmissionController: """FCFS admission gate for disaggregated generation KV transfers.""" - def __init__(self, max_tokens_in_buffer: Optional[int], - tokens_per_block: Optional[int]) -> None: - self.max_transfer_blocks = self._to_block_budget( + def __init__(self, + max_tokens_in_buffer: Optional[int], + tokens_per_block: Optional[int], + admission_budget_multiplier: int = 1) -> None: + if (isinstance(admission_budget_multiplier, bool) + or not isinstance(admission_budget_multiplier, int) + or admission_budget_multiplier <= 0): + raise ValueError("admission_budget_multiplier must be a positive " + "integer") + + self.base_max_transfer_blocks = self._to_block_budget( max_tokens_in_buffer, tokens_per_block) + self.admission_budget_multiplier = admission_budget_multiplier + self.max_transfer_blocks = (None if self.base_max_transfer_blocks + is None else self.base_max_transfer_blocks * + admission_budget_multiplier) self.tokens_per_block = tokens_per_block or 0 + self.diagnostic_budget_multiplier_observed = False + if (self.admission_budget_multiplier != 1 + and self.base_max_transfer_blocks is not None): + logger.info( + "DIAGNOSTIC ONLY: disaggregated admission base window=" + f"{self.base_max_transfer_blocks} transfer blocks, effective " + f"window={self.max_transfer_blocks} transfer blocks, " + f"multiplier={self.admission_budget_multiplier}") def enabled(self) -> bool: return self.max_transfer_blocks is not None @@ -346,8 +389,34 @@ def select(self, active_requests: Iterable[LlmRequest], result.deferred_request_count = len(candidates) - len( result.admitted_requests) + self._observe_diagnostic_budget_multiplier(result) return result + def _observe_diagnostic_budget_multiplier( + self, result: DisaggTransferAdmissionResult) -> None: + if (self.diagnostic_budget_multiplier_observed + or self.admission_budget_multiplier == 1 + or self.base_max_transfer_blocks is None): + return + + used_blocks = (result.active_transfer_blocks + + result.admitted_transfer_blocks) + oversized_idle_head_only = (result.active_transfer_blocks == 0 + and len(result.admitted_requests) == 1 + and result.admitted_transfer_blocks + > self.base_max_transfer_blocks) + if (result.admitted_transfer_blocks == 0 + or used_blocks <= self.base_max_transfer_blocks + or oversized_idle_head_only): + return + + self.diagnostic_budget_multiplier_observed = True + logger.info( + "DIAGNOSTIC OBSERVED: disaggregated admission used " + f"{used_blocks} transfer blocks, exceeding the base window of " + f"{self.base_max_transfer_blocks} blocks with admission budget " + f"multiplier={self.admission_budget_multiplier}") + @dataclasses.dataclass class ScheduledBatchStats: @@ -929,7 +998,8 @@ def on_detected(): tokens_per_block = getattr(self.kv_cache_manager, "tokens_per_block", None) self._disagg_transfer_admission_controller = DisaggTransferAdmissionController( - max_tokens_in_buffer, tokens_per_block) + max_tokens_in_buffer, tokens_per_block, + _get_diagnostic_disagg_admission_budget_multiplier()) self.is_benchmark_disagg = (self.benchmark_req_queues_size > 0 and self.kv_cache_transceiver is not None) # True while the benchmark disagg fill phase is in progress (waiting @@ -3361,6 +3431,7 @@ def _get_disagg_transfer_admission_controller( return DisaggTransferAdmissionController( getattr(cache_transceiver_config, "max_tokens_in_buffer", None), getattr(kv_cache_manager, "tokens_per_block", None), + _get_diagnostic_disagg_admission_budget_multiplier(), ) @staticmethod 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 b803a927b5ec..1e62116e4b1c 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 @@ -35,7 +35,7 @@ environment: trtllm_repo: '' build_wheel: false work_dir: - worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_DIAGNOSTIC_DISABLE_PP_CONTEXT_CONSENSUS_TRAFFIC=1 TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=1 + worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_DIAGNOSTIC_DISABLE_PP_CONTEXT_CONSENSUS_TRAFFIC=1 TRTLLM_DIAGNOSTIC_DISAGG_ADMISSION_BUDGET_MULTIPLIER=2 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 server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index a0d77293b616..1ae5d391d93f 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -26,7 +26,11 @@ RequestQueueItem, ) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig -from tensorrt_llm._torch.pyexecutor.py_executor import DisaggTransferAdmissionController, PyExecutor +from tensorrt_llm._torch.pyexecutor.py_executor import ( + DisaggTransferAdmissionController, + PyExecutor, + _get_diagnostic_disagg_admission_budget_multiplier, +) from tensorrt_llm._torch.pyexecutor.resource_manager import NoFreeSlotsError, ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ( FCFSWaitingQueue, @@ -441,10 +445,74 @@ def _make_disagg_transfer_request( def _clear_disagg_transfer_mode_env(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", raising=False) monkeypatch.delenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", raising=False) + monkeypatch.delenv("TRTLLM_DIAGNOSTIC_DISAGG_ADMISSION_BUDGET_MULTIPLIER", raising=False) @pytest.mark.usefixtures("_clear_disagg_transfer_mode_env") class TestDisaggTransferAdmissionController: + def test_diagnostic_multiplier_doubles_logical_budget(self, monkeypatch): + log_info = Mock() + monkeypatch.setattr("tensorrt_llm._torch.pyexecutor.py_executor.logger.info", log_info) + controller = DisaggTransferAdmissionController( + max_tokens_in_buffer=32, + tokens_per_block=32, + admission_budget_multiplier=2, + ) + active = _make_disagg_transfer_request(1, 32, in_progress=True) + candidate = _make_disagg_transfer_request(2, 32) + + result = controller.select(active_requests=[active], candidates=[candidate]) + + assert controller.base_max_transfer_blocks == 1 + assert controller.max_transfer_blocks == 2 + assert result.admitted_requests == [candidate] + assert result.active_transfer_blocks + result.admitted_transfer_blocks == 2 + assert controller.diagnostic_budget_multiplier_observed + assert ["DIAGNOSTIC ONLY", "DIAGNOSTIC OBSERVED"] == [ + call.args[0].split(":", maxsplit=1)[0] for call in log_info.call_args_list + ] + + controller.select(active_requests=[active], candidates=[candidate]) + assert log_info.call_count == 2 + + @pytest.mark.parametrize("multiplier", [0, -1, True, 1.5, "2"]) + def test_rejects_invalid_admission_budget_multiplier(self, multiplier): + with pytest.raises(ValueError, match="positive integer"): + DisaggTransferAdmissionController( + max_tokens_in_buffer=32, + tokens_per_block=32, + admission_budget_multiplier=multiplier, + ) + + @pytest.mark.parametrize("raw_multiplier", ["0", "-1", "", "1.5", "invalid"]) + def test_rejects_invalid_diagnostic_multiplier_env(self, monkeypatch, raw_multiplier): + monkeypatch.setenv("TRTLLM_DIAGNOSTIC_DISAGG_ADMISSION_BUDGET_MULTIPLIER", raw_multiplier) + + with pytest.raises(ValueError, match="positive integer"): + _get_diagnostic_disagg_admission_budget_multiplier() + + def test_reads_diagnostic_multiplier_env(self, monkeypatch): + monkeypatch.setenv("TRTLLM_DIAGNOSTIC_DISAGG_ADMISSION_BUDGET_MULTIPLIER", "2") + + assert _get_diagnostic_disagg_admission_budget_multiplier() == 2 + + def test_idle_oversized_head_does_not_report_multiplier_exercised(self, monkeypatch): + log_info = Mock() + monkeypatch.setattr("tensorrt_llm._torch.pyexecutor.py_executor.logger.info", log_info) + controller = DisaggTransferAdmissionController( + max_tokens_in_buffer=32, + tokens_per_block=32, + admission_budget_multiplier=2, + ) + log_info.reset_mock() + oversized = _make_disagg_transfer_request(1, 96) + + result = controller.select(active_requests=[], candidates=[oversized]) + + assert result.admitted_requests == [oversized] + assert not controller.diagnostic_budget_multiplier_observed + log_info.assert_not_called() + def test_disabled_preserves_candidates(self): controller = DisaggTransferAdmissionController( max_tokens_in_buffer=None, tokens_per_block=32