From cbc7aa037a61997a3e227873fae1a2d9511a24da Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 24 Jul 2026 09:40:08 -0700 Subject: [PATCH 1/4] [NVBUG-6448152][test] apply audited current-runtime overlay Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- docker/common/install_nixl.sh | 2 +- jenkins/current_image_tags.properties | 10 +++++----- requirements-dev.txt | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docker/common/install_nixl.sh b/docker/common/install_nixl.sh index 6541ce50160d..51208e0ea89d 100644 --- a/docker/common/install_nixl.sh +++ b/docker/common/install_nixl.sh @@ -4,7 +4,7 @@ set -ex GITHUB_URL="https://github.com" UCX_INSTALL_PATH="/usr/local/ucx/" CUDA_PATH="/usr/local/cuda" -NIXL_VERSION="v1.3.0" +NIXL_VERSION="v1.3.1" NIXL_REPO="https://github.com/ai-dynamo/nixl.git" OLD_LD_LIBRARY_PATH=$LD_LIBRARY_PATH diff --git a/jenkins/current_image_tags.properties b/jenkins/current_image_tags.properties index 4732da2c556c..504fd8f234a8 100644 --- a/jenkins/current_image_tags.properties +++ b/jenkins/current_image_tags.properties @@ -13,8 +13,8 @@ # images are adopted from PostMerge pipelines, the abbreviated commit hash is used instead. IMAGE_NAME=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm -LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.05-py3-x86_64-ubuntu24.04-trt10.16.1.11-skip-tritondevel-202607111430-15087 -LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.05-py3-aarch64-ubuntu24.04-trt10.16.1.11-skip-tritondevel-202607111430-15087 -LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.2.1-devel-rocky8-x86_64-rocky8-py310-trt10.16.1.11-skip-tritondevel-202607111430-15087 -LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.2.1-devel-rocky8-x86_64-rocky8-py312-trt10.16.1.11-skip-tritondevel-202607111430-15087 -LLM_SBSA_WHEEL_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.2.1-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-trt10.16.1.11-skip-tritondevel-202607111430-15087 +LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.05-py3-x86_64-ubuntu24.04-trt10.16.1.11-skip-tritondevel-202607151440-16194 +LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.05-py3-sbsa-ubuntu24.04-trt10.16.1.11-skip-tritondevel-202607151440-16194 +LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.2.1-devel-rocky8-x86_64-rocky8-py310-trt10.16.1.11-skip-tritondevel-202607151440-16194 +LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.2.1-devel-rocky8-x86_64-rocky8-py312-trt10.16.1.11-skip-tritondevel-202607151440-16194 +LLM_SBSA_WHEEL_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.2.1-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-trt10.16.1.11-skip-tritondevel-202607151440-16194 diff --git a/requirements-dev.txt b/requirements-dev.txt index 6b16a4c29769..8c6766a590eb 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -42,7 +42,7 @@ opentelemetry-semantic-conventions-ai>=0.4.1 fuzzywuzzy==0.18.0 aiperf==0.8.0 nanobind>=2.9.0 -nixl-cu13==1.3.0 +nixl-cu13==1.3.1 cupti-python>=13.0,<13.2 nvidia-cuda-cupti>=13.0,<13.2 cxxfilt From 6fc7f330fcd621ddc632a56614f58f3208707cba Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 24 Jul 2026 09:40:31 -0700 Subject: [PATCH 2/4] [NVBUG-6448152][test] normalize CMake context only Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- cpp/tensorrt_llm/batch_manager/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index 88e2484cc6f1..128a24ee21e1 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -32,11 +32,11 @@ set(SRCS contextProgress.cpp dataTransceiver.cpp decoderBuffers.cpp + kvCacheManager.cpp encoderBuffers.cpp guidedDecoder.cpp handleContextLogits.cpp handleGenerationLogits.cpp - kvCacheManager.cpp kvCacheEventManager.cpp kvCacheTransferManager.cpp kvCacheManagerV2Utils.cpp From 9c8bf5952cc2d04d69efe309e42f010d74dc7053 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 24 Jul 2026 09:40:57 -0700 Subject: [PATCH 3/4] [NVBUG-6448152][test] apply semantically matched pre-cancellation consensus factor Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 22 + .../contextTransferCoordinator.h | 109 ++++ .../tensorrt_llm/runtime/utils/mpiTags.h | 8 +- .../tensorrt_llm/runtime/utils/mpiUtils.h | 13 +- cpp/tensorrt_llm/batch_manager/CMakeLists.txt | 3 +- .../batch_manager/cacheTransceiver.cpp | 103 +++- .../contextTransferCoordinator.cpp | 475 ++++++++++++++++++ .../unit_tests/batch_manager/CMakeLists.txt | 1 + .../contextTransferCoordinatorTest.cpp | 143 ++++++ .../multi_gpu/cacheTransceiverTest.cpp | 92 ++++ .../unit_tests/multi_gpu/mpiUtilsTest.cpp | 27 +- 11 files changed, 979 insertions(+), 17 deletions(-) create mode 100644 cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h create mode 100644 cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp create mode 100644 cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index a3b8cd348ce4..288af6219fc9 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 ContextTransferCoordinator; 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()) @@ -296,6 +317,7 @@ class CacheTransceiver : public BaseCacheTransceiver std::shared_ptr mGroupComm; std::shared_ptr mGroupTensorParaComm, mGroupPipeParaComm, mGroupDataComm, mGroupTPInDPComm; + std::unique_ptr mContextTransferCoordinator; executor::kv_cache::CommState const* mCommState; std::unique_ptr mCacheState; diff --git a/cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h b/cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h new file mode 100644 index 000000000000..9b156c8887a0 --- /dev/null +++ b/cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h @@ -0,0 +1,109 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager +{ + +class CacheTransceiverComm; + +enum class ContextTransferVote : std::uint64_t +{ + kCompleted = 1, + kFailed = 2, +}; + +struct ContextTransferConsensusResult +{ + std::unordered_set completedRequestIds; + std::unordered_set failedRequestIds; + std::unordered_set timedOutRequestIds; +}; + +//! Accumulates one immutable terminal vote per participant and request. +class ContextTransferVoteReducer +{ +public: + explicit ContextTransferVoteReducer(int participantCount); + + void recordVote(int participantRank, std::uint64_t requestId, ContextTransferVote vote); + + //! Record a sticky nonterminal timeout proposal. The final outcome still waits for every terminal vote. + void recordTimeout(std::uint64_t requestId); + + [[nodiscard]] ContextTransferConsensusResult takeReady(); + + void clear() noexcept; + +private: + struct RequestVotes + { + explicit RequestVotes(int participantCount) + : votes(static_cast(participantCount), 0) + { + } + + std::vector votes; + int terminalCount{0}; + bool failed{false}; + bool timedOut{false}; + bool timeoutPending{false}; + }; + + int mParticipantCount; + std::unordered_map mRequestVotes; +}; + +//! Coordinates asynchronous context-transfer outcomes across a rank group. +//! +//! Timeout and terminal events share one ordered stream toward the coordinator. Timeout updates and final commits +//! share another ordered stream toward followers, so a request that times out can never commit success first. +class ContextTransferCoordinator +{ +public: + explicit ContextTransferCoordinator(std::shared_ptr comm); + ~ContextTransferCoordinator(); + + ContextTransferCoordinator(ContextTransferCoordinator const&) = delete; + ContextTransferCoordinator& operator=(ContextTransferCoordinator const&) = delete; + + //! Publish this rank's immutable local terminal outcome without waiting for peers. + void publishLocalOutcome(std::uint64_t requestId, bool failed); + + //! Publish an idempotent, sticky, nonterminal timeout proposal without waiting for peers. + void publishTimeout(std::uint64_t requestId); + + //! Make nonblocking protocol progress and return newly committed global outcomes. + [[nodiscard]] ContextTransferConsensusResult poll(); + + //! Exchange ordered close markers so no active MPI request outlives its backing buffer. Failure aborts closed. + void shutdown() noexcept; + +private: + class Impl; + std::unique_ptr mImpl; +}; + +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h b/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h index 32c086c84ee9..49c50e49b2e2 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h +++ b/cpp/include/tensorrt_llm/runtime/utils/mpiTags.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -71,7 +71,11 @@ enum class MpiTag : int // KvCacheEventManager kKvCacheEventSize = 1026, - kKvCacheEvent = 1027 + kKvCacheEvent = 1027, + + // Asynchronous context-transfer coordination. + kContextTransferEvent = 1028, + kContextTransferUpdate = 1029 }; } // namespace tensorrt_llm::mpi diff --git a/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h b/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h index 75ec7a534815..b474b4b12cbf 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h +++ b/cpp/include/tensorrt_llm/runtime/utils/mpiUtils.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -235,6 +235,17 @@ class MpiRequest #endif } + [[nodiscard]] bool isCompleted() + { +#if ENABLE_MULTI_DEVICE + int completed = 0; + TLLM_MPI_CHECK(MPI_Test(&mRequest, &completed, MPI_STATUS_IGNORE)); + return completed != 0; +#else + TLLM_THROW("Multi device support is disabled."); +#endif + } + void cancel() { #if ENABLE_MULTI_DEVICE diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index 128a24ee21e1..154b8a4b8326 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & # AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); you may not @@ -30,6 +30,7 @@ set(SRCS capacityScheduler.cpp createNewDecoderRequests.cpp contextProgress.cpp + contextTransferCoordinator.cpp dataTransceiver.cpp decoderBuffers.cpp kvCacheManager.cpp diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 248ac55d7c90..c25e81632f71 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -36,6 +36,7 @@ #include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" #include "tensorrt_llm/batch_manager/contextProgress.h" +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" #include "tensorrt_llm/batch_manager/dataTransceiver.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/kvCacheType.h" @@ -56,6 +57,7 @@ #include #include #include +#include #include #include @@ -496,6 +498,34 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa mCacheSender = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); mCacheReceiver = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); + // Keep automatic enablement within the currently qualified C++ NIXL/UCX TP1/CP1 pipeline topology. + bool const coordinatorTopologyEligible = worldConfig.getPipelineParallelism() > 1 && useMPI() + && backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL + && common::getEnvNixlBackend() == "UCX" && worldConfig.getTensorParallelism() == 1 + && worldConfig.getContextParallelism() == 1 && !mCacheState->getParallelConfig().mEnableAttentionDP; + if (worldConfig.getPipelineParallelism() > 1 && useMPI()) + { + TLLM_CHECK(mGroupPipeParaComm != nullptr); + constexpr std::uint64_t kCoordinatorProtocolVersion = 1; + std::uint64_t const localVersion = coordinatorTopologyEligible ? kCoordinatorProtocolVersion : 0; + constexpr bool cancellationEnabled = false; + std::uint64_t const localProtocolMode = (localVersion << 1) | static_cast(cancellationEnabled); + std::vector protocolModes(static_cast(mGroupPipeParaComm->getSize())); + mGroupPipeParaComm->allgather(&localProtocolMode, protocolModes.data(), 1, mpi::MpiType::kUINT64); + TLLM_CHECK_WITH_INFO(std::all_of(protocolModes.begin(), protocolModes.end(), + [&](std::uint64_t const mode) { return mode == localProtocolMode; }), + "Context-transfer consensus protocol version or cancellation mode differs across PP ranks."); + if (localVersion != 0) + { + mContextTransferCoordinator = std::make_unique(mGroupPipeParaComm); + TLLM_LOG_INFO( + "Enable asynchronous context-transfer consensus version %llu for PP group of size %d; in-flight " + "cancellation=%s.", + static_cast(kCoordinatorProtocolVersion), mGroupPipeParaComm->getSize(), + cancellationEnabled ? "enabled" : "disabled"); + } + } + initializeCommState(); } @@ -505,6 +535,7 @@ CacheTransceiver::~CacheTransceiver() // plugin are still alive. The workers can access both during termination. mCacheSender.reset(); mCacheReceiver.reset(); + mContextTransferCoordinator.reset(); if (mWrapperLibHandle) { @@ -783,6 +814,17 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( toCompleteIdSet.insert(request->mRequestId); } + auto recordOutcome + = [&](RequestIdType const requestId, std::shared_ptr const& request, bool const failed) + { + recordLocalTransferOutcome(requestId, request, failed, mCompletedSenderRequestIds, mFailedSenderRequestIds, + mSenderRequestsAwaitingConsensus); + if (mContextTransferCoordinator) + { + mContextTransferCoordinator->publishLocalOutcome(requestId, failed); + } + }; + // Record local terminal outcomes for requests selected this round. The // request is reported only after all ranks in the sync group agree that the // request reached a terminal state. @@ -805,16 +847,16 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } if (blockAll || (toCompleteIdSet.find(requestId) != toCompleteIdSet.end())) { + bool terminal = false; + bool failed = false; try { auto const status = blockAll ? std::future_status::ready : future.wait_for(futureWaitInterval); if (status == std::future_status::ready) { future.get(); - bool const failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR; - recordLocalTransferOutcome(requestId, request, failed, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - it = mSenderFutures.erase(it); + failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR; + terminal = true; } else if (status == std::future_status::timeout) { @@ -828,18 +870,28 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { TLLM_LOG_ERROR( "Future returned unexpected status for request %ld. Recording as failed.", requestId); - - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); - it = mSenderFutures.erase(it); + failed = true; + terminal = true; } } catch (std::exception const& e) { TLLM_LOG_ERROR("Error occurred during context transfer for request %ld: %s", requestId, e.what()); - recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, - mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); + failed = true; + terminal = true; + } + catch (...) + { + TLLM_LOG_ERROR("Unknown error occurred during context transfer for request %ld", requestId); + failed = true; + terminal = true; + } + if (terminal) + { + auto terminalRequest = request; it = mSenderFutures.erase(it); + // Publish outside the transfer-future try/catch. A protocol error must not rewrite an immutable vote. + recordOutcome(requestId, terminalRequest, failed); } } else @@ -849,8 +901,35 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( } RequestStatuses requestsStatus{}; - auto const consensusOutcome - = reduceTransferStates(syncComm, mGroupPipeParaComm, mCompletedSenderRequestIds, mFailedSenderRequestIds); + TransferConsensusOutcome consensusOutcome; + if (mContextTransferCoordinator) + { + auto mergeCoordinatorOutcome = [&]() + { + auto coordinatorOutcome = mContextTransferCoordinator->poll(); + consensusOutcome.completedRequestIds.insert( + coordinatorOutcome.completedRequestIds.begin(), coordinatorOutcome.completedRequestIds.end()); + consensusOutcome.failedRequestIds.insert( + coordinatorOutcome.failedRequestIds.begin(), coordinatorOutcome.failedRequestIds.end()); + }; + do + { + mergeCoordinatorOutcome(); + if (blockAll + && consensusOutcome.completedRequestIds.size() + consensusOutcome.failedRequestIds.size() + < mSenderRequestsAwaitingConsensus.size()) + { + std::this_thread::yield(); + } + } while (blockAll + && consensusOutcome.completedRequestIds.size() + consensusOutcome.failedRequestIds.size() + < mSenderRequestsAwaitingConsensus.size()); + } + else + { + consensusOutcome + = reduceTransferStates(syncComm, mGroupPipeParaComm, mCompletedSenderRequestIds, mFailedSenderRequestIds); + } for (auto const requestId : consensusOutcome.failedRequestIds) { auto const requestIt = mSenderRequestsAwaitingConsensus.find(requestId); diff --git a/cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp b/cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp new file mode 100644 index 000000000000..909b7858c574 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp @@ -0,0 +1,475 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" + +#include "tensorrt_llm/batch_manager/cacheTransceiver.h" +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/logger.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager +{ + +ContextTransferVoteReducer::ContextTransferVoteReducer(int const participantCount) + : mParticipantCount(participantCount) +{ + TLLM_CHECK_WITH_INFO(participantCount > 0, "Context-transfer consensus requires at least one participant."); +} + +void ContextTransferVoteReducer::recordVote( + int const participantRank, std::uint64_t const requestId, ContextTransferVote const vote) +{ + TLLM_CHECK_WITH_INFO(participantRank >= 0 && participantRank < mParticipantCount, + "Context-transfer consensus participant rank is out of range."); + TLLM_CHECK_WITH_INFO(vote == ContextTransferVote::kCompleted || vote == ContextTransferVote::kFailed, + "Context-transfer consensus received an invalid vote."); + + auto [requestIt, inserted] = mRequestVotes.try_emplace(requestId, mParticipantCount); + static_cast(inserted); + auto& requestVotes = requestIt->second; + auto& recordedVote = requestVotes.votes.at(static_cast(participantRank)); + auto const packedVote = static_cast(vote); + if (recordedVote != 0) + { + TLLM_CHECK_WITH_INFO( + recordedVote == packedVote, "Context-transfer participant changed its terminal vote for a request."); + return; + } + + recordedVote = packedVote; + ++requestVotes.terminalCount; + requestVotes.failed = requestVotes.failed || vote == ContextTransferVote::kFailed; +} + +void ContextTransferVoteReducer::recordTimeout(std::uint64_t const requestId) +{ + auto [requestIt, inserted] = mRequestVotes.try_emplace(requestId, mParticipantCount); + static_cast(inserted); + auto& requestVotes = requestIt->second; + if (!requestVotes.timedOut) + { + requestVotes.timedOut = true; + requestVotes.timeoutPending = true; + } +} + +ContextTransferConsensusResult ContextTransferVoteReducer::takeReady() +{ + ContextTransferConsensusResult result; + for (auto requestIt = mRequestVotes.begin(); requestIt != mRequestVotes.end();) + { + auto& requestVotes = requestIt->second; + if (requestVotes.timeoutPending) + { + result.timedOutRequestIds.insert(requestIt->first); + requestVotes.timeoutPending = false; + } + if (requestVotes.terminalCount != mParticipantCount) + { + ++requestIt; + continue; + } + + auto& terminalRequestIds + = (requestVotes.failed || requestVotes.timedOut) ? result.failedRequestIds : result.completedRequestIds; + terminalRequestIds.insert(requestIt->first); + requestIt = mRequestVotes.erase(requestIt); + } + return result; +} + +void ContextTransferVoteReducer::clear() noexcept +{ + mRequestVotes.clear(); +} + +class ContextTransferCoordinator::Impl +{ +public: + explicit Impl(std::shared_ptr comm) + : mComm(std::move(comm)) + , mCoordinatorRank(0) + , mReducer(mComm ? mComm->getSize() : 1) + { + TLLM_CHECK_WITH_INFO(mComm != nullptr, "Context-transfer coordination requires a communicator."); + TLLM_CHECK_WITH_INFO(mComm->isMpi(), "Asynchronous context-transfer coordination requires MPI."); + TLLM_CHECK_WITH_INFO( + mComm->getSize() > 1, "Asynchronous context-transfer coordination requires multiple participants."); + mCoordinatorRank = mComm->getSize() - 1; + } + + void publishLocalOutcome(std::uint64_t const requestId, bool const failed) + { + TLLM_CHECK_WITH_INFO(!mShutdown, "Cannot publish a context-transfer vote after coordinator shutdown."); + auto const vote = failed ? ContextTransferVote::kFailed : ContextTransferVote::kCompleted; + auto const [voteIt, inserted] = mPublishedLocalVotes.emplace(requestId, vote); + TLLM_CHECK_WITH_INFO( + inserted || voteIt->second == vote, "This rank changed its terminal vote for a context transfer."); + if (!inserted) + { + return; + } + + try + { + if (isCoordinator()) + { + mReducer.recordVote(mComm->getRank(), requestId, vote); + } + else + { + queuePacket( + requestId, static_cast(vote), mCoordinatorRank, mpi::MpiTag::kContextTransferEvent); + } + } + catch (...) + { + mPublishedLocalVotes.erase(voteIt); + throw; + } + } + + void publishTimeout(std::uint64_t const requestId) + { + TLLM_CHECK_WITH_INFO(!mShutdown, "Cannot publish a context-transfer timeout after coordinator shutdown."); + TLLM_CHECK_WITH_INFO(mPublishedLocalVotes.find(requestId) == mPublishedLocalVotes.end(), + "A context-transfer timeout must be published before its immutable terminal vote."); + auto const [timeoutIt, inserted] = mPublishedTimeouts.insert(requestId); + if (!inserted) + { + return; + } + + try + { + if (isCoordinator()) + { + mReducer.recordTimeout(requestId); + } + else + { + queuePacket(requestId, kTimedOutSignal, mCoordinatorRank, mpi::MpiTag::kContextTransferEvent); + } + } + catch (...) + { + mPublishedTimeouts.erase(timeoutIt); + throw; + } + } + + [[nodiscard]] ContextTransferConsensusResult poll() + { + TLLM_CHECK_WITH_INFO(!mShutdown, "Cannot poll context-transfer coordination after shutdown."); + return progress(); + } + + void shutdown() noexcept + { + if (mShutdown) + { + return; + } + mShutdown = true; + + try + { + if (!isCoordinator()) + { + queuePacket(/*requestId=*/0, kCloseMarker, mCoordinatorRank, mpi::MpiTag::kContextTransferEvent); + } + + auto const deadline = std::chrono::steady_clock::now() + kShutdownTimeout; + while (!shutdownComplete()) + { + static_cast(progress()); + if (!shutdownComplete()) + { + if (std::chrono::steady_clock::now() >= deadline) + { + TLLM_LOG_ERROR( + "Timed out shutting down asynchronous context-transfer coordinator; rank=%d " + "peer_closes=%zu/%d ack=%d pending_sends=%zu. Aborting to avoid freeing active MPI " + "requests.", + mComm->getRank(), mClosedPeers.size(), mComm->getSize() - 1, mCloseAcknowledged, + mPendingSends.size()); + std::abort(); + } + std::this_thread::yield(); + } + } + } + catch (std::exception const& error) + { + TLLM_LOG_ERROR("Failed to shut down asynchronous context-transfer coordinator: %s", error.what()); + std::abort(); + } + catch (...) + { + TLLM_LOG_ERROR("Failed to shut down asynchronous context-transfer coordinator with an unknown error."); + std::abort(); + } + } + +private: + static constexpr std::size_t kPacketFieldCount = 2; + static constexpr std::uint64_t kTimedOutSignal = 3; + static constexpr std::uint64_t kCloseMarker = 4; + static constexpr auto kShutdownTimeout = std::chrono::seconds(30); + + struct PendingSend + { + std::array packet{}; + std::unique_ptr request; + }; + + [[nodiscard]] bool isCoordinator() const + { + return mComm->getRank() == mCoordinatorRank; + } + + void queuePacket(std::uint64_t const requestId, std::uint64_t const value, int const peer, mpi::MpiTag const tag) + { + mPendingSends.emplace_back(); + auto& pendingSend = mPendingSends.back(); + pendingSend.packet = {requestId, value}; + try + { + pendingSend.request = mComm->sendAsync( + pendingSend.packet.data(), pendingSend.packet.size(), mpi::MpiType::kUINT64, peer, tag); + } + catch (...) + { + mPendingSends.pop_back(); + throw; + } + } + + void queueUpdateForPeers(std::uint64_t const requestId, std::uint64_t const update) + { + for (int peer = 0; peer < mComm->getSize(); ++peer) + { + if (peer != mCoordinatorRank) + { + queuePacket(requestId, update, peer, mpi::MpiTag::kContextTransferUpdate); + } + } + } + + void reapCompletedSends() + { + for (auto sendIt = mPendingSends.begin(); sendIt != mPendingSends.end();) + { + TLLM_CHECK(sendIt->request); + if (sendIt->request->isCompleted()) + { + sendIt = mPendingSends.erase(sendIt); + } + else + { + ++sendIt; + } + } + } + + void drainVotes() + { + for (int peer = 0; peer < mComm->getSize(); ++peer) + { + if (peer == mCoordinatorRank) + { + continue; + } + + MPI_Status status{}; + while (mComm->iprobe(peer, mpi::MpiTag::kContextTransferEvent, &status)) + { + std::array packet{}; + mComm->recv( + packet.data(), packet.size(), mpi::MpiType::kUINT64, peer, mpi::MpiTag::kContextTransferEvent); + if (packet.back() == kCloseMarker) + { + TLLM_CHECK_WITH_INFO( + mClosedPeers.insert(peer).second, "Received a duplicate context-transfer close marker."); + continue; + } + TLLM_CHECK_WITH_INFO(mClosedPeers.find(peer) == mClosedPeers.end(), + "Received a context-transfer vote after its peer close marker."); + if (packet.back() == kTimedOutSignal) + { + mReducer.recordTimeout(packet.front()); + continue; + } + mReducer.recordVote(peer, packet.front(), static_cast(packet.back())); + } + } + } + + ContextTransferConsensusResult completeCoordinatorUpdates() + { + auto result = mReducer.takeReady(); + for (auto const requestId : result.timedOutRequestIds) + { + queueUpdateForPeers(requestId, kTimedOutSignal); + } + for (auto const requestId : result.failedRequestIds) + { + queueUpdateForPeers(requestId, static_cast(ContextTransferVote::kFailed)); + mPublishedLocalVotes.erase(requestId); + mPublishedTimeouts.erase(requestId); + } + for (auto const requestId : result.completedRequestIds) + { + queueUpdateForPeers(requestId, static_cast(ContextTransferVote::kCompleted)); + mPublishedLocalVotes.erase(requestId); + mPublishedTimeouts.erase(requestId); + } + return result; + } + + ContextTransferConsensusResult drainUpdates() + { + ContextTransferConsensusResult result; + MPI_Status status{}; + while (mComm->iprobe(mCoordinatorRank, mpi::MpiTag::kContextTransferUpdate, &status)) + { + std::array packet{}; + mComm->recv(packet.data(), packet.size(), mpi::MpiType::kUINT64, mCoordinatorRank, + mpi::MpiTag::kContextTransferUpdate); + if (packet.back() == kCloseMarker) + { + TLLM_CHECK_WITH_INFO(!mCloseAcknowledged, "Received a duplicate coordinator close marker."); + mCloseAcknowledged = true; + mPublishedLocalVotes.clear(); + mPublishedTimeouts.clear(); + continue; + } + + if (packet.back() == kTimedOutSignal) + { + result.timedOutRequestIds.insert(packet.front()); + continue; + } + + auto const outcome = static_cast(packet.back()); + TLLM_CHECK_WITH_INFO(outcome == ContextTransferVote::kCompleted || outcome == ContextTransferVote::kFailed, + "Received an invalid context-transfer commit outcome."); + auto const localVoteIt = mPublishedLocalVotes.find(packet.front()); + TLLM_CHECK_WITH_INFO(localVoteIt != mPublishedLocalVotes.end(), + "Received a context-transfer commit before publishing the local terminal vote."); + if (outcome == ContextTransferVote::kFailed) + { + result.failedRequestIds.insert(packet.front()); + } + else + { + result.completedRequestIds.insert(packet.front()); + } + mPublishedLocalVotes.erase(localVoteIt); + mPublishedTimeouts.erase(packet.front()); + } + return result; + } + + ContextTransferConsensusResult progress() + { + reapCompletedSends(); + if (isCoordinator()) + { + drainVotes(); + auto result = completeCoordinatorUpdates(); + if (mShutdown && !mCloseSent && mClosedPeers.size() == static_cast(mComm->getSize() - 1)) + { + // Shutdown is an explicit abort epoch. A peer close proves that no more votes will arrive from that + // peer, so incomplete requests cannot reach a global decision and are intentionally abandoned. + mReducer.clear(); + mPublishedLocalVotes.clear(); + mPublishedTimeouts.clear(); + for (int peer = 0; peer < mCoordinatorRank; ++peer) + { + queuePacket(/*requestId=*/0, kCloseMarker, peer, mpi::MpiTag::kContextTransferUpdate); + } + mCloseSent = true; + } + return result; + } + return drainUpdates(); + } + + [[nodiscard]] bool shutdownComplete() const + { + if (isCoordinator()) + { + return mCloseSent && mPendingSends.empty(); + } + return mCloseAcknowledged && mPendingSends.empty(); + } + + std::shared_ptr mComm; + int mCoordinatorRank; + ContextTransferVoteReducer mReducer; + std::unordered_map mPublishedLocalVotes; + std::unordered_set mPublishedTimeouts; + std::unordered_set mClosedPeers; + std::list mPendingSends; + bool mShutdown{false}; + bool mCloseSent{false}; + bool mCloseAcknowledged{false}; +}; + +ContextTransferCoordinator::ContextTransferCoordinator(std::shared_ptr comm) + : mImpl(std::make_unique(std::move(comm))) +{ +} + +ContextTransferCoordinator::~ContextTransferCoordinator() +{ + shutdown(); +} + +void ContextTransferCoordinator::publishLocalOutcome(std::uint64_t const requestId, bool const failed) +{ + mImpl->publishLocalOutcome(requestId, failed); +} + +void ContextTransferCoordinator::publishTimeout(std::uint64_t const requestId) +{ + mImpl->publishTimeout(requestId); +} + +ContextTransferConsensusResult ContextTransferCoordinator::poll() +{ + return mImpl->poll(); +} + +void ContextTransferCoordinator::shutdown() noexcept +{ + if (mImpl) + { + mImpl->shutdown(); + } +} + +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index f4e5e9fb2be0..9425d5d13c2a 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -20,6 +20,7 @@ add_gtest(cacheTransBufferTest cacheTransBufferTest.cpp) add_gtest(bufferIndexHolderTest bufferIndexHolderTest.cpp) add_gtest(capacitySchedulerTest capacitySchedulerTest.cpp) add_gtest(contextProgressTest contextProgressTest.cu) +add_gtest(contextTransferCoordinatorTest contextTransferCoordinatorTest.cpp) add_gtest(evictionPolicyTest evictionPolicyTest.cpp) add_gtest(kvCacheManagerTest kvCacheManagerTest.cpp) add_gtest(kvCacheManagerFabricMemoryTest kvCacheManagerFabricMemoryTest.cpp) diff --git a/cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp b/cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp new file mode 100644 index 000000000000..7b0f7aba3fc4 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp @@ -0,0 +1,143 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" + +#include + +namespace tensorrt_llm::batch_manager +{ +namespace +{ + +TEST(ContextTransferVoteReducerTest, WaitsForEveryParticipant) +{ + ContextTransferVoteReducer reducer(4); + reducer.recordVote(0, 17, ContextTransferVote::kCompleted); + reducer.recordVote(1, 17, ContextTransferVote::kCompleted); + reducer.recordVote(2, 17, ContextTransferVote::kCompleted); + EXPECT_TRUE(reducer.takeReady().completedRequestIds.empty()); + + reducer.recordVote(3, 17, ContextTransferVote::kCompleted); + auto const result = reducer.takeReady(); + EXPECT_EQ(result.completedRequestIds, std::unordered_set{17}); + EXPECT_TRUE(result.failedRequestIds.empty()); + EXPECT_TRUE(result.timedOutRequestIds.empty()); +} + +TEST(ContextTransferVoteReducerTest, FailureWinsAfterEveryParticipantVotes) +{ + ContextTransferVoteReducer reducer(4); + reducer.recordVote(0, 23, ContextTransferVote::kCompleted); + reducer.recordVote(1, 23, ContextTransferVote::kFailed); + reducer.recordVote(2, 23, ContextTransferVote::kCompleted); + reducer.recordVote(3, 23, ContextTransferVote::kCompleted); + + auto const result = reducer.takeReady(); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{23}); +} + +TEST(ContextTransferVoteReducerTest, RejectsChangedTerminalVote) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordVote(0, 29, ContextTransferVote::kCompleted); + reducer.recordVote(0, 29, ContextTransferVote::kCompleted); + EXPECT_ANY_THROW(reducer.recordVote(0, 29, ContextTransferVote::kFailed)); +} + +TEST(ContextTransferVoteReducerTest, AccumulatesInterleavedRequestsIndependently) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordVote(0, 31, ContextTransferVote::kCompleted); + reducer.recordVote(1, 37, ContextTransferVote::kFailed); + reducer.recordVote(1, 31, ContextTransferVote::kCompleted); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.completedRequestIds, std::unordered_set{31}); + EXPECT_TRUE(result.failedRequestIds.empty()); + + reducer.recordVote(0, 37, ContextTransferVote::kCompleted); + result = reducer.takeReady(); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{37}); +} + +TEST(ContextTransferVoteReducerTest, RejectsInvalidInput) +{ + EXPECT_ANY_THROW(ContextTransferVoteReducer(0)); + ContextTransferVoteReducer reducer(2); + EXPECT_ANY_THROW(reducer.recordVote(-1, 41, ContextTransferVote::kCompleted)); + EXPECT_ANY_THROW(reducer.recordVote(2, 41, ContextTransferVote::kCompleted)); + EXPECT_ANY_THROW(reducer.recordVote(0, 41, static_cast(99))); +} + +TEST(ContextTransferVoteReducerTest, TimeoutIsVisibleBeforeTerminalFailure) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordTimeout(43); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.timedOutRequestIds, std::unordered_set{43}); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_TRUE(result.failedRequestIds.empty()); + + result = reducer.takeReady(); + EXPECT_TRUE(result.timedOutRequestIds.empty()); + EXPECT_TRUE(result.failedRequestIds.empty()); + + reducer.recordVote(0, 43, ContextTransferVote::kCompleted); + reducer.recordVote(1, 43, ContextTransferVote::kCompleted); + result = reducer.takeReady(); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{43}); + EXPECT_TRUE(result.completedRequestIds.empty()); + EXPECT_TRUE(result.timedOutRequestIds.empty()); +} + +TEST(ContextTransferVoteReducerTest, DuplicateTimeoutIsIdempotent) +{ + ContextTransferVoteReducer reducer(2); + reducer.recordTimeout(47); + reducer.recordTimeout(47); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.timedOutRequestIds, std::unordered_set{47}); + + reducer.recordTimeout(47); + result = reducer.takeReady(); + EXPECT_TRUE(result.timedOutRequestIds.empty()); +} + +TEST(ContextTransferVoteReducerTest, TimeoutAfterPartialTerminalVotesStillWins) +{ + ContextTransferVoteReducer reducer(3); + reducer.recordVote(2, 53, ContextTransferVote::kCompleted); + reducer.recordTimeout(53); + reducer.recordVote(0, 53, ContextTransferVote::kCompleted); + + auto result = reducer.takeReady(); + EXPECT_EQ(result.timedOutRequestIds, std::unordered_set{53}); + EXPECT_TRUE(result.failedRequestIds.empty()); + + reducer.recordVote(1, 53, ContextTransferVote::kCompleted); + result = reducer.takeReady(); + EXPECT_EQ(result.failedRequestIds, std::unordered_set{53}); + EXPECT_TRUE(result.completedRequestIds.empty()); +} + +} // namespace +} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index 77f8ea03fb83..49afbf20dfae 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -31,6 +31,7 @@ #include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" +#include "tensorrt_llm/batch_manager/contextTransferCoordinator.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" @@ -42,6 +43,7 @@ #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/testing/kvCacheManagerTestUtil.h" +#include #include #include #include @@ -53,6 +55,7 @@ #include #include #include +#include #include "gtest/gtest.h" #include @@ -89,6 +92,95 @@ T serializeDeserialize(T const& val) } // namespace +TEST(ContextTransferCoordinatorTest, CommitsStaggeredSuccessAndFailureWithoutCollectivePolling) +{ + auto& world = tensorrt_llm::mpi::MpiComm::world(); + if (world.getSize() < 2) + { + GTEST_SKIP() << "mpirun with at least two processes is required to run this test."; + } + + auto comm = std::make_shared(std::addressof(world)); + { + ContextTransferCoordinator coordinator(comm); + auto waitForOutcome = [&](std::uint64_t const requestId, bool const expectFailure) + { + bool observed = false; + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (!observed && std::chrono::steady_clock::now() < deadline) + { + auto result = coordinator.poll(); + observed = expectFailure ? result.failedRequestIds.count(requestId) != 0 + : result.completedRequestIds.count(requestId) != 0; + if (!observed) + { + std::this_thread::yield(); + } + } + EXPECT_TRUE(observed); + }; + auto waitForTimeout = [&](std::uint64_t const requestId) + { + bool observed = false; + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (!observed && std::chrono::steady_clock::now() < deadline) + { + auto const result = coordinator.poll(); + observed = result.timedOutRequestIds.count(requestId) != 0; + if (!observed) + { + std::this_thread::yield(); + } + } + EXPECT_TRUE(observed); + }; + + constexpr std::uint64_t kCompletedRequestId = 644815201; + world.barrier(); + if (world.getRank() == world.getSize() - 1) + { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + coordinator.publishLocalOutcome(kCompletedRequestId, /*failed=*/false); + if (world.getRank() != world.getSize() - 1) + { + auto const earlyResult = coordinator.poll(); + EXPECT_TRUE(earlyResult.completedRequestIds.empty()); + EXPECT_TRUE(earlyResult.failedRequestIds.empty()); + } + waitForOutcome(kCompletedRequestId, /*expectFailure=*/false); + + constexpr std::uint64_t kFailedRequestId = 644815202; + world.barrier(); + coordinator.publishLocalOutcome(kFailedRequestId, /*failed=*/world.getRank() == 0); + waitForOutcome(kFailedRequestId, /*expectFailure=*/true); + + constexpr std::uint64_t kTimedOutRequestId = 644815203; + world.barrier(); + if (world.getRank() == 0) + { + coordinator.publishTimeout(kTimedOutRequestId); + coordinator.publishTimeout(kTimedOutRequestId); + } + waitForTimeout(kTimedOutRequestId); + coordinator.publishLocalOutcome(kTimedOutRequestId, /*failed=*/false); + waitForOutcome(kTimedOutRequestId, /*expectFailure=*/true); + + constexpr std::uint64_t kInvalidOrderingRequestId = 644815204; + world.barrier(); + coordinator.publishLocalOutcome(kInvalidOrderingRequestId, /*failed=*/false); + EXPECT_ANY_THROW(coordinator.publishTimeout(kInvalidOrderingRequestId)); + waitForOutcome(kInvalidOrderingRequestId, /*expectFailure=*/false); + + // Exercise asymmetric but orderly teardown after every decision has committed. + if (world.getRank() == 0) + { + std::this_thread::sleep_for(std::chrono::milliseconds(25)); + } + } + world.barrier(); +} + class RequestInfoTest : public ::testing::Test // NOLINT(cppcoreguidelines-pro-type-member-init) { public: diff --git a/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp b/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp index 221cd98b5f02..b8596d511fab 100644 --- a/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/mpiUtilsTest.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -25,6 +25,8 @@ #endif // ENABLE_MULTI_DEVICE #include +#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::kContextTransferEvent); + static_cast(comm.recv(&receivedValue, 1, mpi::MpiType::kUINT64, source, mpi::MpiTag::kContextTransferEvent)); + while (!request->isCompleted()) + { + std::this_thread::yield(); + } + + EXPECT_EQ(receivedValue, static_cast(source)); +} +#endif // ENABLE_MULTI_DEVICE + template void testBroadcast() { From 4b182f1118073faf43f7206808696c6fba02596e Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 24 Jul 2026 09:41:32 -0700 Subject: [PATCH 4/4] [NVBUG-6448152][test] restore native CMake source order Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- cpp/tensorrt_llm/batch_manager/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index 154b8a4b8326..1389d0de068f 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -33,11 +33,11 @@ set(SRCS contextTransferCoordinator.cpp dataTransceiver.cpp decoderBuffers.cpp - kvCacheManager.cpp encoderBuffers.cpp guidedDecoder.cpp handleContextLogits.cpp handleGenerationLogits.cpp + kvCacheManager.cpp kvCacheEventManager.cpp kvCacheTransferManager.cpp kvCacheManagerV2Utils.cpp