From 7a865f3f0d80c804d2d2c4067b7afa52bfddbccf Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 17 Jul 2026 15:10:48 -0700 Subject: [PATCH 1/4] [NVBUG-6448152][test] establish post-#15139 merge treatment Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> From e83aa182f277b4a4a462baef43396efede4f5a27 Mon Sep 17 00:00:00 2001 From: Yanchao Lu Date: Mon, 29 Jun 2026 22:48:48 +0800 Subject: [PATCH 2/4] [None][ci] Disable home mount for sbatch L0 tests (#15713) Signed-off-by: Yanchao Lu (cherry picked from commit 833ddd2a59037a7dd961ded5a9ce4a6d010d1bcb) Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- jenkins/L0_Test.groovy | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 306b2bc3084f..f8296852fdd6 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -1,3 +1,19 @@ +/* + * Copyright (c) 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. + * 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. + */ + @Library(['bloom-jenkins-shared-lib@main', 'trtllm-jenkins-shared-lib@main']) _ import java.lang.InterruptedException @@ -980,6 +996,8 @@ def getMountListForSlurmTest(SlurmCluster cluster, boolean useSbatch = false) throw new Exception("Unsupported container runtime: ${cluster.containerRuntime}") } + // TODO: Add mounts for different cache directories like pip, triton, etc. + return mounts } @@ -1208,6 +1226,7 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG "--container-image=$containerImageArg", "--container-workdir=$jobWorkspace", "--container-mounts=$mounts", + "--no-container-mount-home", "--container-env=NVIDIA_IMEX_CHANNELS" ] envVarsToExport.each { varName, varValue -> From 4edb4d797f54192159c1fba93f8968cf90d3c6f4 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Sat, 18 Jul 2026 10:45:44 -0700 Subject: [PATCH 3/4] [NVBUG-6448152][test] backport sender wakeup fix Apply only the CacheSender readiness synchronization portion of #15737 to the exact #15139 boundary experiment. This keeps the historical control and treatment matched while removing the known lost-wakeup confound. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/dataTransceiver.cpp | 273 +++++++++--------- 1 file changed, 140 insertions(+), 133 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 853866687d53..b3eb1c0f9843 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -33,6 +33,8 @@ #include #include #include +#include +#include #include namespace tensorrt_llm::batch_manager @@ -286,7 +288,6 @@ class CacheSender::Impl TLLM_CHECK(mManager); TLLM_CHECK(mManager->getCommState().getSelfIdx() == selfIndex); TLLM_CUDA_CHECK(cudaGetDevice(&mDeviceId)); - mCurrentRequest = std::nullopt; mResponseFuture = std::async(std::launch::async, &Impl::response, this); int asyncSendThreadNum = common::getEnvKVCacheSendMaxConcurrenceNum(); for (int i = 0; i < asyncSendThreadNum; i++) @@ -303,12 +304,13 @@ class CacheSender::Impl auto future = promise.get_future(); llmRequest->setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); { - { - std::scoped_lock lkResp(mSenderMutex); - mReadyResponses.emplace(llmRequest->mRequestId, Response{llmRequest, std::move(promise)}); - } - std::unique_lock lkCond(mCondMutex); - mAnyReady = true; + 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)}); + TLLM_CHECK_WITH_INFO( + result.second, "Request %zu is already queued for KV cache transfer", llmRequest->mRequestId); } mSenderCv.notify_all(); return future; @@ -425,11 +427,11 @@ class CacheSender::Impl bool cancelRequest(LlmRequest const& llmRequest) { bool isCancelled = false; - std::scoped_lock lkResp(mSenderMutex); + std::scoped_lock lock(mSenderMutex); auto it = mReadyResponses.find(llmRequest.mRequestId); // If the request is not the current request and already in the ready queue, we can cancel it. if (it != mReadyResponses.end() - && (!mCurrentRequest.has_value() || getCurrentRequestId() != llmRequest.mRequestId)) + && (!mCurrentRequest.has_value() || mCurrentRequest.value() != llmRequest.mRequestId)) { mCancelledRequests.insert(llmRequest.mRequestId); isCancelled = true; @@ -546,191 +548,196 @@ class CacheSender::Impl void asyncSendAndRemoveResponse(RequestIdType id, Response resp) noexcept { - std::unique_lock lk(mAsyncSendResource.mMtxForQueue); - mAsyncSendResource.mSendQueue.emplace_back(std::move(resp)); - mAsyncSendResource.mCVforQueue.notify_one(); + try + { + std::unique_lock lock(mAsyncSendResource.mMtxForQueue); + mAsyncSendResource.mSendQueue.emplace_back(std::move(resp)); + mAsyncSendResource.mCVforQueue.notify_one(); + } + catch (std::exception const& err) + { + TLLM_LOG_ERROR("Failed to queue asynchronous KV cache send for request %zu: %s", id, err.what()); + failResponse(resp, std::current_exception()); + } } - void sendResponse(std::map::iterator it) + void sendResponse(RequestIdType reqId) { - auto reqId = mCurrentRequest.value(); - auto count = --mRemainSendCount[reqId]; - TLLM_CHECK(count >= 0); - if (count == 0) + bool isReady = true; { - mRemainSendCount.erase(reqId); - - // Check if the request is cancelled - bool isReady = true; + std::scoped_lock lock(mSenderMutex); + TLLM_CHECK(mCurrentRequest.has_value() && mCurrentRequest.value() == reqId); + TLLM_CHECK(mReadyResponses.find(reqId) != mReadyResponses.end()); + auto countIt = mRemainSendCount.find(reqId); + TLLM_CHECK(countIt != mRemainSendCount.end()); + auto const count = --countIt->second; + TLLM_CHECK(count >= 0); + if (count > 0) { - std::scoped_lock lk(mSenderMutex); - if (mCancelledRequests.find(reqId) != mCancelledRequests.end()) - { - isReady = false; - } + mCurrentRequest = std::nullopt; + return; } - sendReadySignal(reqId, isReady); + mRemainSendCount.erase(countIt); + isReady = mCancelledRequests.find(reqId) == mCancelledRequests.end(); + } - if (isReady) + // Keep mCurrentRequest set while notifying the peer so cancellation cannot change the decision after it has + // been made. The network operation must not run under mSenderMutex. + sendReadySignal(reqId, isReady); + + Response response; + { + std::scoped_lock lock(mSenderMutex); + auto it = mReadyResponses.find(reqId); + TLLM_CHECK(it != mReadyResponses.end()); + response = std::move(it->second); + mReadyResponses.erase(it); + mCancelledRequests.erase(reqId); + mCurrentRequest = std::nullopt; + } + + if (isReady) + { + if (dynamic_cast(mManager) != nullptr) { - if (dynamic_cast(mManager) != nullptr) - { - // our nixl impl seems only support recv and send in the same thread - // if we use zmq as control path, we may avoid this issue - sendAndRemoveResponse(it->first, std::move(it->second)); - } - else - { - // if we send data in another thread, multiple rank may send data for different requests at the same - // time with gen DP case. - asyncSendAndRemoveResponse(it->first, std::move(it->second)); - } - removeResponse(it); + // Our NIXL implementation only supports recv and send in the same thread. Using ZMQ for the control + // path may avoid this limitation. + sendAndRemoveResponse(reqId, std::move(response)); } else { - // TODO: if the generation does not require the kv cache, the request will - // not be removed from mCancelledRequests. This should be handled by timeout. - auto const cancelledReqId = mCurrentRequest.value(); - Response cancelledResponse; - { - std::scoped_lock lkResp(mSenderMutex); - auto it = mReadyResponses.find(cancelledReqId); - TLLM_CHECK(it != mReadyResponses.end()); - // Move out before erasing so the promise survives the - // map cleanup and can be resolved (vs. destroyed unfulfilled, - // which would surface as std::future_error: Broken promise). - cancelledResponse = std::move(it->second); - mReadyResponses.erase(it); - mCancelledRequests.erase(cancelledReqId); - mRemainSendCount.erase(cancelledReqId); - } - cancelledResponse.mPromise.set_exception(std::make_exception_ptr( - TLLM_REQUEST_EXCEPTION(cancelledReqId, common::RequestErrorCode::kNETWORK_ERROR, - "KV cache transfer for request %zu was cancelled", cancelledReqId))); - mCurrentRequest = std::nullopt; - - if (mReadyResponses.empty()) - { - std::unique_lock lk(mCondMutex); - mAnyReady = false; - } + // If we send data in another thread, multiple ranks may send data for different requests at the same + // time with generation attention DP. + asyncSendAndRemoveResponse(reqId, std::move(response)); } } - mCurrentRequest = std::nullopt; + else + { + response.mPromise.set_exception(std::make_exception_ptr(TLLM_REQUEST_EXCEPTION(reqId, + common::RequestErrorCode::kNETWORK_ERROR, "KV cache transfer for request %zu was cancelled", reqId))); + } } void response() noexcept { + std::exception_ptr responseException; try { tensorrt_llm::common::setThreadName("dataTransResp"); TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); - while (!mTerminate || !mAnyReady) + while (true) { - if (!mAnyReady) { - std::unique_lock lk(mCondMutex); - mSenderCv.wait(lk, [this]() { return (mAnyReady || mTerminate); }); + std::unique_lock lock(mSenderMutex); + mSenderCv.wait(lock, [this]() { return mTerminate || !mReadyResponses.empty(); }); + if (mTerminate) + { + break; + } } - if (mTerminate) + + auto const requestInfo = recvRequestInfo(); + if (mTerminate || !mManager->isRunning()) { break; } - if (!mReadyResponses.empty()) - { - auto const& requestInfo = recvRequestInfo(); - if (mTerminate || !mManager->isRunning()) - { - return; - } - auto reqId = requestInfo.getRequestId(); + auto const reqId = requestInfo.getRequestId(); - { - std::scoped_lock lk(mSenderMutex); - mCurrentRequest = reqId; - } - - if (mRemainSendCount.find(reqId) == mRemainSendCount.end()) - { - mRemainSendCount[reqId] = getCounterpartsCount(reqId); - } - } - auto it = getCurrentResponse(); - if (it != mReadyResponses.end()) + if (mRemainSendCount.find(reqId) == mRemainSendCount.end()) { - sendResponse(it); + mRemainSendCount[reqId] = getCounterpartsCount(reqId); } - else + { - auto it = getCurrentResponse(); - while (it == mReadyResponses.end()) + std::unique_lock lock(mSenderMutex); + mCurrentRequest = reqId; + mSenderCv.wait(lock, + [this, reqId]() { return mTerminate || mReadyResponses.find(reqId) != mReadyResponses.end(); }); + if (mTerminate) { - std::unique_lock lk(mCondMutex); - mSenderCv.wait(lk, [this]() { return (mAnyReady || mTerminate); }); - if (mTerminate) - { - break; - } - it = getCurrentResponse(); + mCurrentRequest = std::nullopt; + break; } - sendResponse(it); } + sendResponse(reqId); } } catch (std::exception const& err) { TLLM_LOG_ERROR("Exception in CacheSender response: %s", err.what()); - for (auto& it : mReadyResponses) - { - it.second.mPromise.set_exception(std::current_exception()); - } + responseException = std::current_exception(); } + + if (!responseException) + { + responseException + = std::make_exception_ptr(std::runtime_error("CacheSender terminated before response completed")); + } + { + std::scoped_lock lock(mSenderMutex); + mTerminate = true; + } + mSenderCv.notify_all(); + failPendingResponses(responseException); } void terminate() { { - std::unique_lock lk(mCondMutex); + std::scoped_lock lock(mSenderMutex); mTerminate = true; } - // We don't have to wait for the future. If another thread is sending data, it won't pay attention - // to the terminate flag. mSenderCv.notify_all(); - mAsyncSendResource.mTerminate = true; + if (mResponseFuture.valid()) + { + mResponseFuture.get(); + } + + std::deque pendingAsyncResponses; + { + std::scoped_lock lock(mAsyncSendResource.mMtxForQueue); + mAsyncSendResource.mTerminate = true; + pendingAsyncResponses.swap(mAsyncSendResource.mSendQueue); + } mAsyncSendResource.mCVforQueue.notify_all(); for (auto& future : mAsyncSendFutures) { future.get(); } - if (mResponseFuture.valid()) + auto const exception + = std::make_exception_ptr(std::runtime_error("CacheSender terminated before asynchronous send completed")); + for (auto& response : pendingAsyncResponses) { - mResponseFuture.get(); + failResponse(response, exception); } } - void removeResponse(std::map::iterator it) + void failResponse(Response& response, std::exception_ptr const& exception) noexcept { + try { - std::scoped_lock lkResp(mSenderMutex); - mReadyResponses.erase(it); + response.mPromise.set_exception(exception); } - if (mReadyResponses.empty()) + catch (std::future_error const& err) { - std::unique_lock lkCond(mCondMutex); - mAnyReady = false; + TLLM_LOG_ERROR("Failed to set CacheSender response exception: %s", err.what()); } } - [[nodiscard]] RequestIdType getCurrentRequestId() const + void failPendingResponses(std::exception_ptr const& exception) noexcept { - return mCurrentRequest.value(); - } - - [[nodiscard]] std::map::iterator getCurrentResponse() - { - std::scoped_lock lk(mSenderMutex); - return mReadyResponses.find(getCurrentRequestId()); + std::map pendingResponses; + { + std::scoped_lock lock(mSenderMutex); + pendingResponses.swap(mReadyResponses); + mCurrentRequest = std::nullopt; + mCancelledRequests.clear(); + mRemainSendCount.clear(); + } + for (auto& entry : pendingResponses) + { + failResponse(entry.second, exception); + } } public: @@ -746,9 +753,9 @@ class CacheSender::Impl std::optional mCurrentRequest; std::set mCancelledRequests; std::map mReadyResponses; - std::mutex mSenderMutex, mCondMutex; - std::atomic mAnyReady{false}, mTerminate{false}; - std::condition_variable mSenderCv, mResponderCv; + std::mutex mSenderMutex; + std::atomic mTerminate{false}; + std::condition_variable mSenderCv; std::future mResponseFuture; std::unordered_map mRemainSendCount; AsyncSendResource mAsyncSendResource; From dc06e82a182671a0d9626243d660a44a3b73bd1d Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Sun, 19 Jul 2026 15:23:07 -0700 Subject: [PATCH 4/4] [NVBUG-6448152][test] isolate PP rendezvous from global commit Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 22 + .../contextTransferCoordinator.h | 99 ++++ .../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 | 115 ++++- .../contextTransferCoordinator.cpp | 429 ++++++++++++++++++ .../unit_tests/batch_manager/CMakeLists.txt | 1 + .../contextTransferCoordinatorTest.cpp | 89 ++++ .../multi_gpu/cacheTransceiverTest.cpp | 60 +++ .../unit_tests/multi_gpu/mpiUtilsTest.cpp | 26 +- ...tx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml | 2 +- 12 files changed, 849 insertions(+), 18 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 80a06ea5ddf7..30fe338baa52 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()) @@ -298,6 +319,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..fc52e48c8b51 --- /dev/null +++ b/cpp/include/tensorrt_llm/batch_manager/contextTransferCoordinator.h @@ -0,0 +1,99 @@ +/* + * 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; +}; + +//! 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); + + [[nodiscard]] ContextTransferConsensusResult takeCompleted(); + + [[nodiscard]] std::size_t size() const noexcept; + + void clear() noexcept; + +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; +}; + +//! Diagnostic CTX protocol that asynchronously gathers PP votes at the last stage and broadcasts the decision. +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); + + //! 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. + 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..d8a9e111f575 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, + + // Diagnostic asynchronous context-transfer coordination. + kContextTransferVote = 1028, + kContextTransferCommit = 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 be12b743db1d..f5781ac71581 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 88e2484cc6f1..1389d0de068f 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 encoderBuffers.cpp diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 18c2b19e9eeb..0c65ddd94147 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" @@ -55,6 +56,7 @@ #include #include #include +#include #include #include @@ -68,6 +70,11 @@ namespace using RequestIdType = LlmRequest::RequestIdType; +bool isNvbug6448152ContextCoordinatorEnabled() +{ + return common::getBoolEnv("TRTLLM_NVBUG_6448152_CTX_COORDINATOR_CONSENSUS"); +} + enum class TransferConsensusState : std::uint64_t { kCompleted = 1, @@ -503,11 +510,52 @@ 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()); + // The perf-sanity worker environment is shared by CTX and GEN. Negotiate only on the eligible PP CTX service; + // the PP1 GEN service remains byte-for-byte on its existing status path. + bool const coordinatorRequested = isNvbug6448152ContextCoordinatorEnabled(); + bool const isPipelineService = worldConfig.getPipelineParallelism() > 1; + if (coordinatorRequested && isPipelineService) + { + TLLM_CHECK_WITH_INFO(useMPI(), "NVBUG6448152 context coordinator requires an MPI control plane."); + TLLM_CHECK_WITH_INFO(backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL, + "NVBUG6448152 context coordinator is qualified only for the NIXL backend."); + TLLM_CHECK_WITH_INFO(worldConfig.getTensorParallelism() == 1 && worldConfig.getContextParallelism() == 1, + "NVBUG6448152 context coordinator requires TP1/CP1/PP>1."); + TLLM_CHECK_WITH_INFO(!mCacheState->getParallelConfig().mEnableAttentionDP, + "NVBUG6448152 context coordinator does not support attention DP."); + } + + bool const coordinatorTopologyEligible = isPipelineService && useMPI() + && backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL + && worldConfig.getTensorParallelism() == 1 && worldConfig.getContextParallelism() == 1 + && !mCacheState->getParallelConfig().mEnableAttentionDP; + if (coordinatorTopologyEligible) + { + TLLM_CHECK(mGroupPipeParaComm != nullptr); + constexpr std::uint64_t kCoordinatorProtocolVersion = 1; + std::uint64_t const localVersion = coordinatorRequested ? kCoordinatorProtocolVersion : 0; + std::vector protocolVersions(static_cast(mGroupPipeParaComm->getSize())); + mGroupPipeParaComm->allgather(&localVersion, protocolVersions.data(), 1, mpi::MpiType::kUINT64); + TLLM_CHECK_WITH_INFO(std::all_of(protocolVersions.begin(), protocolVersions.end(), + [&](std::uint64_t const version) { return version == localVersion; }), + "NVBUG6448152 context coordinator mode/version differs across PP ranks."); + if (localVersion != 0) + { + mContextTransferCoordinator = std::make_unique(mGroupPipeParaComm); + TLLM_LOG_WARNING( + "NVBUG6448152_COORD event=mode_active world_rank=%d pp_rank=%d pp_size=%d coordinator_rank=%d " + "version=%lu decision=global resource_retention=global", + worldConfig.getRank(), mGroupPipeParaComm->getRank(), mGroupPipeParaComm->getSize(), + mGroupPipeParaComm->getSize() - 1, kCoordinatorProtocolVersion); + } + } + initializeCommState(); } CacheTransceiver::~CacheTransceiver() { + mContextTransferCoordinator.reset(); if (mWrapperLibHandle) { std::lock_guard lock(mDllMutex); @@ -764,6 +812,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. @@ -786,6 +845,8 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( if (blockAll || (toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end())) { auto const requestId = request->mRequestId; + bool terminal = false; + bool failed = false; try { // Wait for up to a specified timeout @@ -793,10 +854,8 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( if (status == std::future_status::ready || !senderFutureTimeoutMs.has_value()) { 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) { @@ -808,18 +867,23 @@ 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; + } + if (terminal) + { + auto terminalRequest = request; it = mSenderFutures.erase(it); + // Keep control-plane publication outside the transfer-future try/catch. A publication failure is a + // protocol error, not a different transfer outcome, and must never rewrite an immutable local vote. + recordOutcome(requestId, terminalRequest, failed); } } else @@ -829,8 +893,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..86ebb6de541f --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/contextTransferCoordinator.cpp @@ -0,0 +1,429 @@ +/* + * 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; +} + +ContextTransferConsensusResult ContextTransferVoteReducer::takeCompleted() +{ + ContextTransferConsensusResult 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; +} + +std::size_t ContextTransferVoteReducer::size() const noexcept +{ + return mRequestVotes.size(); +} + +void ContextTransferVoteReducer::clear() noexcept +{ + mRequestVotes.clear(); +} + +class ContextTransferCoordinator::Impl +{ +public: + explicit Impl(std::shared_ptr comm) + : mComm(std::move(comm)) + , mCoordinatorRank(mComm->getSize() - 1) + , mReducer(mComm->getSize()) + { + 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 pipeline-parallel participants."); + } + + 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::kContextTransferVote); + } + } + catch (...) + { + mPublishedLocalVotes.erase(voteIt); + throw; + } + logCheckpoint("local_vote", ++mPublishedVoteCount, requestId); + } + + [[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::kContextTransferVote); + } + + 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(); + } + } + TLLM_LOG_INFO( + "NVBUG6448152_COORD event=shutdown_summary rank=%d coordinator_rank=%d local_votes=%lu " + "peer_votes=%lu commits=%lu abandoned=%zu pending_sends=%zu", + mComm->getRank(), mCoordinatorRank, mPublishedVoteCount, mReceivedVoteCount, mCommittedCount, + mAbandonedRequestCount, mPendingSends.size()); + } + 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 kCloseMarker = 3; + 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 queueCommitForPeers(std::uint64_t const requestId, ContextTransferVote const outcome) + { + for (int peer = 0; peer < mComm->getSize(); ++peer) + { + if (peer != mCoordinatorRank) + { + queuePacket(requestId, static_cast(outcome), peer, mpi::MpiTag::kContextTransferCommit); + } + } + } + + 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::kContextTransferVote, &status)) + { + std::array packet{}; + mComm->recv( + packet.data(), packet.size(), mpi::MpiType::kUINT64, peer, mpi::MpiTag::kContextTransferVote); + 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."); + mReducer.recordVote(peer, packet.front(), static_cast(packet.back())); + ++mReceivedVoteCount; + } + } + } + + ContextTransferConsensusResult completeCoordinatorDecisions() + { + auto result = mReducer.takeCompleted(); + for (auto const requestId : result.failedRequestIds) + { + queueCommitForPeers(requestId, ContextTransferVote::kFailed); + mPublishedLocalVotes.erase(requestId); + logCheckpoint("global_commit", ++mCommittedCount, requestId); + } + for (auto const requestId : result.completedRequestIds) + { + queueCommitForPeers(requestId, ContextTransferVote::kCompleted); + mPublishedLocalVotes.erase(requestId); + logCheckpoint("global_commit", ++mCommittedCount, requestId); + } + return result; + } + + ContextTransferConsensusResult drainCommits() + { + ContextTransferConsensusResult result; + MPI_Status status{}; + while (mComm->iprobe(mCoordinatorRank, mpi::MpiTag::kContextTransferCommit, &status)) + { + std::array packet{}; + mComm->recv(packet.data(), packet.size(), mpi::MpiType::kUINT64, mCoordinatorRank, + mpi::MpiTag::kContextTransferCommit); + if (packet.back() == kCloseMarker) + { + TLLM_CHECK_WITH_INFO(!mCloseAcknowledged, "Received a duplicate coordinator close marker."); + mCloseAcknowledged = true; + mAbandonedRequestCount = mPublishedLocalVotes.size(); + mPublishedLocalVotes.clear(); + 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); + logCheckpoint("global_commit", ++mCommittedCount, packet.front()); + } + return result; + } + + ContextTransferConsensusResult progress() + { + reapCompletedSends(); + if (isCoordinator()) + { + drainVotes(); + auto result = completeCoordinatorDecisions(); + 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. + mAbandonedRequestCount = mReducer.size(); + mReducer.clear(); + mPublishedLocalVotes.clear(); + for (int peer = 0; peer < mCoordinatorRank; ++peer) + { + queuePacket(/*requestId=*/0, kCloseMarker, peer, mpi::MpiTag::kContextTransferCommit); + } + mCloseSent = true; + } + return result; + } + return drainCommits(); + } + + [[nodiscard]] bool shutdownComplete() const + { + if (isCoordinator()) + { + return mCloseSent && mPendingSends.empty(); + } + return mCloseAcknowledged && mPendingSends.empty(); + } + + void logCheckpoint(char const* event, std::uint64_t const count, std::uint64_t const requestId) const + { + if (count == 1) + { + TLLM_LOG_INFO("NVBUG6448152_COORD event=%s rank=%d coordinator_rank=%d count=%lu request_id=%lu", event, + mComm->getRank(), mCoordinatorRank, count, requestId); + } + } + + std::shared_ptr mComm; + int mCoordinatorRank; + ContextTransferVoteReducer mReducer; + std::unordered_map mPublishedLocalVotes; + std::unordered_set mClosedPeers; + std::list mPendingSends; + std::uint64_t mPublishedVoteCount{0}; + std::uint64_t mReceivedVoteCount{0}; + std::uint64_t mCommittedCount{0}; + std::size_t mAbandonedRequestCount{0}; + 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); +} + +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 2ee994f76226..55fdd5e6d6aa 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(kvCacheUtilsTest kvCacheUtilsTest.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..f2a4f71cf05b --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/contextTransferCoordinatorTest.cpp @@ -0,0 +1,89 @@ +/* + * 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.takeCompleted().completedRequestIds.empty()); + + reducer.recordVote(3, 17, ContextTransferVote::kCompleted); + auto const result = reducer.takeCompleted(); + EXPECT_EQ(result.completedRequestIds, std::unordered_set{17}); + EXPECT_TRUE(result.failedRequestIds.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.takeCompleted(); + 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.takeCompleted(); + EXPECT_EQ(result.completedRequestIds, std::unordered_set{31}); + EXPECT_TRUE(result.failedRequestIds.empty()); + + reducer.recordVote(0, 37, ContextTransferVote::kCompleted); + result = reducer.takeCompleted(); + 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))); +} + +} // 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 d1ca104cca1a..55ce4be35763 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,63 @@ 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); + }; + + 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); + + // 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..167351336c40 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,28 @@ 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::kContextTransferVote); + static_cast(comm.recv(&receivedValue, 1, mpi::MpiType::kUINT64, source, mpi::MpiTag::kContextTransferVote)); + while (!request->isCompleted()) + { + std::this_thread::yield(); + } + + EXPECT_EQ(receivedValue, static_cast(source)); +} +#endif // ENABLE_MULTI_DEVICE + template void testBroadcast() { 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..7c626a9fa8fa 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL.yaml @@ -36,7 +36,7 @@ environment: build_wheel: false work_dir: worker_env_var: TLLM_LOG_LEVEL=INFO TRTLLM_SERVER_DISABLE_GC=1 TRTLLM_WORKER_DISABLE_GC=1 TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS=1 - TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes + TRTLLM_ENABLE_PDL=1 ENROOT_ALLOW_DEV=yes TRTLLM_NVBUG_6448152_CTX_COORDINATOR_CONSENSUS=1 server_env_var: TRTLLM_SERVER_DISABLE_GC=1 profiling: nsys_on: false