From 8d5bf349388efa74a94b15d54bc94f3691e24743 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 30 Jun 2026 11:44:13 -0700 Subject: [PATCH 1/2] [TRTLLM-12721][fix] Squash disagg cancellation prerequisites Combine the current #15238, #15737, #15794, and #15795 snapshots into one parent commit for #15738. The source PRs remain independently reviewable and mergeable. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 22 +- cpp/include/tensorrt_llm/executor/executor.h | 4 +- .../tensorrt_llm/executor/transferAgent.h | 9 +- .../batch_manager/baseTransBuffer.cpp | 150 +++- .../batch_manager/baseTransBuffer.h | 33 +- .../batch_manager/cacheFormatter.cpp | 20 +- .../batch_manager/cacheTransceiver.cpp | 461 ++++++++++- .../batch_manager/cacheTransferLayer.cpp | 9 +- .../batch_manager/cacheTransferLayer.h | 5 +- .../batch_manager/cacheTransferState.h | 137 +++ .../batch_manager/dataTransceiver.cpp | 740 +++++++++++++---- .../batch_manager/dataTransceiver.h | 57 +- .../batch_manager/mlaCacheFormatter.cpp | 80 +- .../batch_manager/rnnCacheFormatter.cpp | 59 +- cpp/tensorrt_llm/common/envUtils.cpp | 6 + cpp/tensorrt_llm/common/envUtils.h | 3 + .../agent_utils/connection.cpp | 83 +- .../agent_utils/connection.h | 13 +- .../nixl_utils/transferAgent.cpp | 75 +- .../nixl_utils/transferAgent.h | 4 + .../batch_manager/cacheTransceiver.cpp | 8 +- .../unit_tests/batch_manager/CMakeLists.txt | 1 + .../batch_manager/bufferIndexHolderTest.cpp | 93 ++- .../batch_manager/cacheTransferStateTest.cpp | 102 +++ .../unit_tests/executor/agentCommTest.cpp | 15 +- .../_torch/auto_deploy/shim/ad_executor.py | 1 + tensorrt_llm/_torch/pyexecutor/_util.py | 24 +- .../pyexecutor/executor_request_queue.py | 20 +- .../_torch/pyexecutor/kv_cache_transceiver.py | 158 +++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 780 ++++++++++++++++-- .../accuracy/test_disaggregated_serving.py | 43 +- .../test_lists/test-db/l0_b200.yml | 3 + .../test_lists/test-db/l0_dgx_b200.yml | 1 + .../test_lists/test-db/l0_sanity_check.yml | 3 + tests/integration/test_lists/waives.txt | 11 - .../test_disagg_inflight_cancel_gate.py | 527 ++++++++++++ .../test_disagg_inflight_cancel_lifecycle.py | 727 ++++++++++++++++ .../executor/test_dual_pool_kv_cache.py | 16 +- .../_torch/executor/test_py_executor.py | 24 +- .../singlegpu/shim/test_create_ad_executor.py | 1 + .../others/test_kv_cache_transceiver.py | 337 +++++++- 41 files changed, 4422 insertions(+), 443 deletions(-) create mode 100644 cpp/tensorrt_llm/batch_manager/cacheTransferState.h create mode 100644 cpp/tests/unit_tests/batch_manager/cacheTransferStateTest.cpp create mode 100644 tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py create mode 100644 tests/unittest/_torch/executor/test_disagg_inflight_cancel_lifecycle.py diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 80a06ea5ddf7..efbb63eb67e6 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 @@ -226,6 +227,11 @@ class BaseCacheTransceiver [[nodiscard]] virtual bool checkGenTransferComplete() const = 0; virtual bool cancelRequest(std::shared_ptr llmRequest) = 0; + + [[nodiscard]] virtual bool hasPoisonedTransferBuffer() const + { + return false; + } }; class CacheTransceiver : public BaseCacheTransceiver @@ -238,7 +244,7 @@ class CacheTransceiver : public BaseCacheTransceiver = executor::kv_cache::CacheState::AttentionType::kDEFAULT, std::optional cacheTransceiverConfig = std::nullopt, rnn_state_manager::RnnStateManager* rnnStateManager = nullptr, - std::vector const& rnnLayerNumPerPP = {}); + std::vector const& rnnLayerNumPerPP = {}, bool enableInflightCancel = false); CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, std::vector numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, runtime::WorldConfig const& worldConfig, @@ -247,10 +253,11 @@ class CacheTransceiver : public BaseCacheTransceiver = executor::kv_cache::CacheState::AttentionType::kDEFAULT, std::optional cacheTransceiverConfig = std::nullopt, rnn_state_manager::RnnStateManager* rnnStateManager = nullptr, - std::vector const& rnnLayerNumPerPP = {}) + std::vector const& rnnLayerNumPerPP = {}, bool enableInflightCancel = false) : CacheTransceiver(cacheManager, executor::kv_cache::CacheState::ModelConfig{numKvHeadsPerLayer, sizePerHead, tokensPerBlock}, worldConfig, - attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnStateManager, rnnLayerNumPerPP) + attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnStateManager, rnnLayerNumPerPP, + enableInflightCancel) { } @@ -273,11 +280,18 @@ class CacheTransceiver : public BaseCacheTransceiver virtual bool cancelRequest(std::shared_ptr llmRequest) override; + [[nodiscard]] bool hasPoisonedTransferBuffer() const override; + private: void initializeCommState(); void setContextState(LlmRequest* llmRequest); + RequestStatuses checkContextTransferStatusWithInflightCancel( + std::optional const& atLeastRequestNum, bool markComplete); + + void checkGenTransferStatusWithInflightCancel(std::optional const& atLeastRequestNum); + std::unique_ptr mCacheSender; std::unique_ptr mCacheReceiver; // shared_ptr (not raw LlmRequest*) so the futures hold a strong reference for @@ -303,6 +317,8 @@ class CacheTransceiver : public BaseCacheTransceiver std::unique_ptr mCacheState; std::unique_ptr mManager; std::optional mCacheTransceiverConfig; + bool mInflightCancelEnabled{false}; + std::atomic mTopologyTransferBufferPoisoned{false}; std::vector> mCacheTransBufferManagers; std::vector mCacheTransBufferManagerPtrs; diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index 536849cdd87f..b5c0a96a305a 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -1528,8 +1528,8 @@ class CacheTransceiverConfig /// transfer may be degraded. std::optional mMaxTokensInBuffer; std::optional mKvTransferTimeoutMs; - // @brief Timeout in milliseconds to wait for the sender future to be ready when scheduled batch size is 0. This - // allows the request to be eventually cancelled by the user or because of kv_transfer_timeout_ms + // @brief Bounded wait in milliseconds for polling the sender future when the scheduled batch size is 0. This lets + // the executor observe cancellation and KV-transfer timeout state without blocking indefinitely. std::optional mKvTransferSenderFutureTimeoutMs; // @brief Bounded wait interval in milliseconds for polling KV transfer progress when active transfers block // disaggregated admission. diff --git a/cpp/include/tensorrt_llm/executor/transferAgent.h b/cpp/include/tensorrt_llm/executor/transferAgent.h index 532f0ae70c44..e1685c7c4ba5 100644 --- a/cpp/include/tensorrt_llm/executor/transferAgent.h +++ b/cpp/include/tensorrt_llm/executor/transferAgent.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025-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. @@ -356,6 +356,13 @@ class TransferStatus virtual ~TransferStatus() = default; [[nodiscard]] virtual bool isCompleted() const = 0; virtual TransferState wait(int64_t timeout_ms = -1) const = 0; + + /// Release the backend transfer request handle. A true return means the backend accepted the handle release; it + /// does not prove remote memory quiescence. + [[nodiscard]] virtual bool release() + { + return false; + } }; struct BaseAgentConfig diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp index fe47f6c531ad..8dd875c624e9 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.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"); @@ -21,11 +21,29 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/opUtils.h" +#include +#include #include namespace tensorrt_llm::batch_manager { +namespace +{ + +char const* bufferKindName(BufferKind kind) +{ + switch (kind) + { + case BufferKind::kKV: return "kv"; + case BufferKind::kKV_INDEXER: return "kv_indexer"; + case BufferKind::kRNN: return "rnn"; + } + return "unknown"; +} + +} // namespace + void BufferIndexHolder::release() noexcept { if (!mHeld || mMgr == nullptr) @@ -50,6 +68,30 @@ void BufferIndexHolder::release() noexcept mHeld = false; } +void BufferIndexHolder::poison() noexcept +{ + if (!mHeld || mMgr == nullptr) + { + return; + } + try + { + if (mIsRecv) + { + mMgr->poisonBufferIndexForRecv(mIndex); + } + else + { + mMgr->poisonBufferIndexForSend(mIndex); + } + } + catch (...) + { + // noexcept: poison is a fail-closed best effort from exception paths. + } + mHeld = false; +} + BaseTransBufferManager::BaseTransBufferManager( size_t transferBufferSize, nvinfer1::DataType dataType, std::optional maxNumTokens) : mDataType{dataType} @@ -78,9 +120,11 @@ BaseTransBufferManager::BaseTransBufferManager( allocateBuffer(); } -std::optional BaseTransBufferManager::assignBufferIndexForSend() +std::optional BaseTransBufferManager::assignBufferIndexForSend( + std::atomic const* perRequestCancel, int64_t waitSliceMs) { - return assignBufferIndex(mConcurrenceSendResource, mSendBufferCount, mOnlyUseDynamicBuffer); + return assignBufferIndex( + mConcurrenceSendResource, mSendBufferCount, mOnlyUseDynamicBuffer, perRequestCancel, waitSliceMs); } void BaseTransBufferManager::freeBufferIndexForSend(std::optional bufferId) @@ -88,9 +132,16 @@ void BaseTransBufferManager::freeBufferIndexForSend(std::optional bufferId) freeBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, mOnlyUseDynamicBuffer); } -std::optional BaseTransBufferManager::assignBufferIndexForRecv() +void BaseTransBufferManager::poisonBufferIndexForSend(std::optional bufferId) noexcept +{ + poisonBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, mOnlyUseDynamicBuffer, "send"); +} + +std::optional BaseTransBufferManager::assignBufferIndexForRecv( + std::atomic const* perRequestCancel, int64_t waitSliceMs) { - return assignBufferIndex(mConcurrenceRecvResource, mRecvBufferCount, mOnlyUseDynamicBuffer); + return assignBufferIndex( + mConcurrenceRecvResource, mRecvBufferCount, mOnlyUseDynamicBuffer, perRequestCancel, waitSliceMs); } void BaseTransBufferManager::freeBufferIndexForRecv(std::optional bufferId) @@ -98,6 +149,11 @@ void BaseTransBufferManager::freeBufferIndexForRecv(std::optional bufferId) freeBufferIndex(mConcurrenceRecvResource, bufferId, mRecvBufferCount, mOnlyUseDynamicBuffer); } +void BaseTransBufferManager::poisonBufferIndexForRecv(std::optional bufferId) noexcept +{ + poisonBufferIndex(mConcurrenceRecvResource, bufferId, mRecvBufferCount, mOnlyUseDynamicBuffer, "recv"); +} + std::tuple, size_t, bool> BaseTransBufferManager::getOrAllocateSendBuffers( std::optional bufferId, int targetNum, std::vector const& requestedNumberOfElements, runtime::BufferManager const& bufferManagerToUse) @@ -249,16 +305,44 @@ void BaseTransBufferManager::allocateBuffer() } } -std::optional BaseTransBufferManager::assignBufferIndex( - ConcurrenceResource& resource, size_t bufferCount, bool onlyUseDynamicBuffer) +std::optional BaseTransBufferManager::assignBufferIndex(ConcurrenceResource& resource, size_t bufferCount, + bool onlyUseDynamicBuffer, std::atomic const* perRequestCancel, int64_t waitSliceMs) { + auto const isCancelled = [perRequestCancel]() + { return perRequestCancel != nullptr && perRequestCancel->load(std::memory_order_relaxed); }; + if (isCancelled()) + { + TLLM_THROW("Cache transfer buffer acquisition cancelled"); + } if (onlyUseDynamicBuffer) { + TLLM_CHECK_WITH_INFO(!resource.mPoisoned.load(std::memory_order_relaxed), + "Cannot assign dynamic cache transfer buffer kind=%s because the transfer buffer pool is poisoned", + bufferKindName(getBufferKind())); return std::nullopt; } std::unique_lock lk(resource.mBuffersMutex); - resource.mBuffersCV.wait( - lk, [&resource, bufferCount]() { return static_cast(resource.mConcurrence) < bufferCount; }); + auto const predicate = [&resource, bufferCount]() + { + return resource.mPoisoned.load(std::memory_order_relaxed) + || static_cast(resource.mConcurrence) < bufferCount; + }; + auto const slice = std::chrono::milliseconds{waitSliceMs}; + while (!predicate()) + { + resource.mBuffersCV.wait_for(lk, slice); + if (isCancelled()) + { + TLLM_THROW("Cache transfer buffer acquisition cancelled"); + } + } + if (isCancelled()) + { + TLLM_THROW("Cache transfer buffer acquisition cancelled"); + } + TLLM_CHECK_WITH_INFO(!resource.mPoisoned.load(std::memory_order_relaxed), + "Cannot assign cache transfer buffer kind=%s because the transfer buffer pool is poisoned", + bufferKindName(getBufferKind())); int bufferId = -1; for (size_t i = 0; i < bufferCount; i++) { @@ -288,6 +372,12 @@ void BaseTransBufferManager::freeBufferIndex( TLLM_CHECK(static_cast(bufferId.value()) < bufferCount); { std::scoped_lock lk(resource.mBuffersMutex); + if (resource.mBufferIndexFlag[bufferId.value()] == 2) + { + TLLM_LOG_ERROR("Refusing to free poisoned cache transfer buffer kind=%s index=%d", + bufferKindName(getBufferKind()), bufferId.value()); + return; + } resource.mBufferIndexFlag[bufferId.value()] = 0; } resource.mConcurrence--; @@ -295,6 +385,48 @@ void BaseTransBufferManager::freeBufferIndex( } } +void BaseTransBufferManager::poisonBufferIndex(ConcurrenceResource& resource, std::optional bufferId, + size_t bufferCount, bool onlyUseDynamicBuffer, char const* direction) noexcept +{ + resource.mPoisoned.store(true, std::memory_order_relaxed); + + if (onlyUseDynamicBuffer) + { + TLLM_LOG_ERROR("Poisoned dynamic %s cache transfer buffer kind=%s; process restart is required", direction, + bufferKindName(getBufferKind())); + resource.mBuffersCV.notify_all(); + return; + } + + try + { + if (bufferId.has_value()) + { + TLLM_CHECK(static_cast(bufferId.value()) < bufferCount); + { + std::scoped_lock lk(resource.mBuffersMutex); + if (resource.mBufferIndexFlag[bufferId.value()] == 1) + { + resource.mBufferIndexFlag[bufferId.value()] = 2; + } + } + } + TLLM_LOG_ERROR("Poisoned %s cache transfer buffer kind=%s index=%d; process restart is required", direction, + bufferKindName(getBufferKind()), bufferId.value_or(-1)); + } + catch (std::exception const& e) + { + TLLM_LOG_ERROR("Exception while poisoning %s cache transfer buffer kind=%s index=%d: %s", direction, + bufferKindName(getBufferKind()), bufferId.value_or(-1), e.what()); + } + catch (...) + { + TLLM_LOG_ERROR("Unknown exception while poisoning %s cache transfer buffer kind=%s index=%d", direction, + bufferKindName(getBufferKind()), bufferId.value_or(-1)); + } + resource.mBuffersCV.notify_all(); +} + size_t BaseTransBufferManager::getRecvBufferCount() { return mRecvBufferCount; diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h index 2cbf9f514bd7..86a014d1ef96 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -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"); @@ -59,7 +59,7 @@ class BufferIndexHolder BufferIndexHolder(BaseTransBufferManager& mgr, std::optional index, bool isRecv) noexcept : mMgr(&mgr) , mIndex(index) - , mHeld(index.has_value()) + , mHeld(true) , mIsRecv(isRecv) { } @@ -120,6 +120,9 @@ class BufferIndexHolder /// @brief Release the slot now and disarm the destructor. Safe to call multiple times. void release() noexcept; + /// @brief Fail-closed release for an exit path where transfer-buffer quiescence is unknown. + void poison() noexcept; + private: BaseTransBufferManager* mMgr{nullptr}; std::optional mIndex{}; @@ -127,6 +130,8 @@ class BufferIndexHolder bool mIsRecv{true}; }; +inline constexpr int64_t kBufferAcquireSliceMs = 100; + /// @brief Base class for cache transfer buffer management. /// Handles buffer pool allocation, index assignment, and slicing. /// Derived classes provide cache-specific size calculations. @@ -139,20 +144,28 @@ class BaseTransBufferManager /// @brief Assign a buffer index for sending. /// @return Assigned buffer index, or nullopt if using dynamic buffers. - std::optional assignBufferIndexForSend(); + std::optional assignBufferIndexForSend( + std::atomic const* perRequestCancel = nullptr, int64_t waitSliceMs = kBufferAcquireSliceMs); /// @brief Free a buffer index used for sending. /// @param bufferId The buffer index to free. void freeBufferIndexForSend(std::optional bufferId); + /// @brief Poison a send buffer index after an unquiesced transfer exit. + void poisonBufferIndexForSend(std::optional bufferId) noexcept; + /// @brief Assign a buffer index for receiving. /// @return Assigned buffer index, or nullopt if using dynamic buffers. - std::optional assignBufferIndexForRecv(); + std::optional assignBufferIndexForRecv( + std::atomic const* perRequestCancel = nullptr, int64_t waitSliceMs = kBufferAcquireSliceMs); /// @brief Free a buffer index used for receiving. /// @param bufferId The buffer index to free. void freeBufferIndexForRecv(std::optional bufferId); + /// @brief Poison a receive buffer index after an unquiesced transfer exit. + void poisonBufferIndexForRecv(std::optional bufferId) noexcept; + /// @brief Get or allocate send buffers for cache transfer. /// @param bufferId The assigned buffer ID. /// @param targetNum Number of target sequences. @@ -191,6 +204,12 @@ class BaseTransBufferManager return mMaxNumTokens; } + [[nodiscard]] bool hasPoisonedBuffer() const noexcept + { + return mConcurrenceSendResource.mPoisoned.load(std::memory_order_relaxed) + || mConcurrenceRecvResource.mPoisoned.load(std::memory_order_relaxed); + } + protected: /// @brief Constructor - derived classes call this after computing buffer sizes. /// @param transferBufferSize Size of each transfer buffer in bytes. @@ -206,6 +225,7 @@ class BaseTransBufferManager std::mutex mBuffersMutex; std::condition_variable mBuffersCV; std::atomic mConcurrence{0}; + std::atomic mPoisoned{false}; }; std::tuple, size_t, bool> getOrAllocateBuffers(std::optional bufferId, @@ -213,9 +233,12 @@ class BaseTransBufferManager runtime::BufferManager const& bufferManagerToUse, ConcurrenceResource& concurrenceResource); void allocateBuffer(); - std::optional assignBufferIndex(ConcurrenceResource& resource, size_t bufferCount, bool onlyUseDynamicBuffer); + std::optional assignBufferIndex(ConcurrenceResource& resource, size_t bufferCount, bool onlyUseDynamicBuffer, + std::atomic const* perRequestCancel = nullptr, int64_t waitSliceMs = kBufferAcquireSliceMs); void freeBufferIndex( ConcurrenceResource& resource, std::optional bufferId, size_t bufferCount, bool onlyUseDynamicBuffer); + void poisonBufferIndex(ConcurrenceResource& resource, std::optional bufferId, size_t bufferCount, + bool onlyUseDynamicBuffer, char const* direction) noexcept; size_t mPreAllocBufferSize; size_t mRecvBufferCount; diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index 68a398f1b52d..881dc20119f0 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -522,7 +522,8 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio // cache blocks to the corresponding buffer. // 5. send the buffer to the corresponding target. Ideally, we send only once (one buffer) for each target. - auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend(); + auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); + auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend(sendCancelFlag); BufferIndexHolder sendHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/false); int peerDuplicateHeadFactor = targetInfo.mPeerDupHeadFactor; auto bufferTargetNum = targetNum / peerDuplicateHeadFactor; @@ -602,8 +603,19 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio == inputKvCacheBlocksPerWindow.begin()->second.front()->getDataType()); } - sendAllBuffers(session, deviceId, outputSplitCaches, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, - targetInfo, pickUpConnections); + try + { + sendAllBuffers(session, deviceId, outputSplitCaches, bufferCoverTargetNum, preAllocSendBuffer, + bufferManager, targetInfo, pickUpConnections); + } + catch (...) + { + if (agentConnection != nullptr && session.isInflightCancelEnabled()) + { + sendHolder.poison(); + } + throw; + } session.setTime(TransferSession::kTimeTransmissions); @@ -865,8 +877,8 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess else { cacheBufferId = mCacheTransBufferManager->assignBufferIndexForRecv(); + recvHolder = BufferIndexHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); } - recvHolder = BufferIndexHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); auto [recvSplitCachestmp, bufferCoverTargetNumtmp, onlyUseDynamicBuffer] = mCacheTransBufferManager->getOrAllocateRecvBuffers( cacheBufferId, static_cast(targetNum), bufferEleSizes, bufferManager); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 2c69af0059be..c2477f59a38d 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -35,6 +35,7 @@ #include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" +#include "tensorrt_llm/batch_manager/cacheTransferState.h" #include "tensorrt_llm/batch_manager/contextProgress.h" #include "tensorrt_llm/batch_manager/dataTransceiver.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" @@ -205,6 +206,137 @@ TransferConsensusOutcome reduceTransferStates(std::shared_ptr> gatherPackedTransferStatePayloads( + std::shared_ptr const& comm, std::vector const& localPayload) +{ + TLLM_CHECK(!localPayload.empty()); + if (comm == nullptr || comm->getSize() == 1) + { + return {localPayload}; + } + + constexpr std::uint64_t kPoisonedDescriptorBit = std::uint64_t{1} << 63; + auto const localStateSize = localPayload.size() - 1; + TLLM_CHECK(localStateSize < kPoisonedDescriptorBit); + std::uint64_t const localDescriptor + = localStateSize | ((localPayload.front() & detail::kTransferBufferPoisoned) != 0 ? kPoisonedDescriptorBit : 0); + std::vector descriptors(comm->getSize()); + if (useMPI()) + { + comm->allgather(&localDescriptor, descriptors.data(), 1, mpi::MpiType::kUINT64); + } + else + { + comm->allgather(&localDescriptor, std::ref(descriptors), {}); + } + + std::vector sizes(comm->getSize()); + std::vector displacements(comm->getSize()); + size_t totalSize = 0; + std::vector> rankPayloads(comm->getSize()); + for (int rank = 0; rank < comm->getSize(); ++rank) + { + auto const stateSize = descriptors[rank] & ~kPoisonedDescriptorBit; + TLLM_CHECK(stateSize <= static_cast(std::numeric_limits::max())); + sizes[rank] = static_cast(stateSize); + displacements[rank] = static_cast(totalSize); + totalSize += sizes[rank]; + rankPayloads[rank].push_back( + (descriptors[rank] & kPoisonedDescriptorBit) != 0 ? detail::kTransferBufferPoisoned : 0); + } + if (totalSize == 0) + { + return rankPayloads; + } + + std::vector localStates(localPayload.begin() + 1, localPayload.end()); + std::vector gatheredStates(totalSize); + if (useMPI()) + { + comm->allgatherv(localStates.data(), static_cast(localStates.size()), mpi::MpiType::kUINT64, + gatheredStates.data(), sizes, displacements, mpi::MpiType::kUINT64); + } + else + { + comm->allgatherv(std::ref(localStates), std::ref(gatheredStates), std::cref(sizes), {}); + } + for (size_t rank = 0; rank < sizes.size(); ++rank) + { + auto const begin = gatheredStates.begin() + displacements[rank]; + rankPayloads[rank].insert(rankPayloads[rank].end(), begin, begin + sizes[rank]); + } + return rankPayloads; +} + +detail::TransferTopologyOutcome exchangeTransferStates(std::shared_ptr const& comm, + std::unordered_set const& completedRequestIds, + std::unordered_set const& failedRequestIds, + std::unordered_set const& timedOutRequestIds, bool const localTransferBufferPoisoned) +{ + std::unordered_map statesByRequest; + for (auto const requestId : completedRequestIds) + { + statesByRequest[requestId] |= detail::kTransferCompleted; + } + for (auto const requestId : failedRequestIds) + { + statesByRequest[requestId] &= ~detail::kTransferCompleted; + statesByRequest[requestId] |= detail::kTransferFailed; + } + for (auto const requestId : timedOutRequestIds) + { + statesByRequest[requestId] |= detail::kTransferTimedOut; + } + + std::vector localPayload; + localPayload.reserve(1 + statesByRequest.size() * 2); + localPayload.push_back(localTransferBufferPoisoned ? detail::kTransferBufferPoisoned : 0); + for (auto const& [requestId, flags] : statesByRequest) + { + localPayload.push_back(requestId); + localPayload.push_back(flags); + } + return detail::summarizePackedTransferStates(gatherPackedTransferStatePayloads(comm, localPayload)); +} + +void validateInflightCancelModeAndCapability(std::shared_ptr const& comm, + bool const inflightCancelEnabled, bool const inflightCancelCapable) +{ + constexpr int kInflightCancelEnabled = 1 << 0; + constexpr int kInflightCancelCapable = 1 << 1; + int const localDescriptor + = (inflightCancelEnabled ? kInflightCancelEnabled : 0) | (inflightCancelCapable ? kInflightCancelCapable : 0); + + std::vector descriptors; + if (comm == nullptr || comm->getSize() == 1) + { + descriptors.push_back(localDescriptor); + } + else + { + descriptors.resize(comm->getSize()); + if (useMPI()) + { + comm->allgather(&localDescriptor, descriptors.data(), 1, mpi::MpiType::kINT32); + } + else + { + comm->allgather(&localDescriptor, std::ref(descriptors), {}); + } + } + + auto const isEnabled = [](int const descriptor) { return (descriptor & kInflightCancelEnabled) != 0; }; + auto const isCapable = [](int const descriptor) { return (descriptor & kInflightCancelCapable) != 0; }; + auto const anyEnabled = std::any_of(descriptors.begin(), descriptors.end(), isEnabled); + auto const allEnabled = std::all_of(descriptors.begin(), descriptors.end(), isEnabled); + TLLM_CHECK_WITH_INFO( + !anyEnabled || allEnabled, "enable_inflight_cancel must resolve identically on every local parallel rank."); + TLLM_CHECK_WITH_INFO(!allEnabled || std::all_of(descriptors.begin(), descriptors.end(), isCapable), + "enable_inflight_cancel is unsupported on one or more local parallel ranks; it requires asynchronous " + "NIXL/UCX cache transfer with preallocated buffers, a positive transfer timeout, overlap enabled, PP=1, " + "CP=1, attention DP disabled, layerwise transfer disabled, and zero-copy transfer disabled."); +} + void recordLocalTransferOutcome(RequestIdType requestId, std::shared_ptr request, bool failed, std::unordered_set& completedRequestIds, std::unordered_set& failedRequestIds, std::unordered_map>& requestsAwaitingConsensus) @@ -237,6 +369,9 @@ std::unique_ptr CacheTransceiverFactory::createCacheTransc runtime::WorldConfig const& worldConfig, executor::kv_cache::CacheState::AttentionType attentionType, std::optional cacheTransceiverConfig) { + TLLM_CHECK_WITH_INFO(!common::getEnvDisaggEnableInflightCancel(), + "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 is unsupported by the legacy TensorRT executor; unset it or set it " + "to 0 to continue without coordinated in-flight cancellation."); if (!cacheTransceiverConfig.has_value() || !cacheTransceiverConfig.value().getBackendType().has_value()) { TLLM_LOG_INFO("CacheTransceiver is disabled."); @@ -293,8 +428,10 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, executor::kv_cache::CacheState::AttentionType attentionType, std::optional cacheTransceiverConfig, - rnn_state_manager::RnnStateManager* rnnStateManager, std::vector const& rnnLayerNumPerPP) + rnn_state_manager::RnnStateManager* rnnStateManager, std::vector const& rnnLayerNumPerPP, + bool const enableInflightCancel) : mCacheTransceiverConfig{cacheTransceiverConfig} + , mInflightCancelEnabled{enableInflightCancel} , mRnnStateManager{rnnStateManager} { using tensorrt_llm::batch_manager::kv_cache_manager::CacheFormatter; @@ -307,6 +444,21 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa mGroupComm = std::make_shared(tensorrt_llm::pg_utils::get_world_pg()); } + using BackendType = executor::CacheTransceiverConfig::BackendType; + auto const backendType = mCacheTransceiverConfig.has_value() + ? mCacheTransceiverConfig->getBackendType().value_or(BackendType::DEFAULT) + : BackendType::DEFAULT; + bool const inflightCancelCapable = backendType == BackendType::NIXL && common::getEnvNixlBackend() == "UCX" + && mCacheTransceiverConfig->getMaxTokensInBuffer().value_or(0) > 0 + && mCacheTransceiverConfig->getKvTransferTimeoutMs().value_or(0) > 0 + && !common::getEnvDisableKVCacheTransferOverlap() && worldConfig.getPipelineParallelism() == 1 + && worldConfig.getContextParallelism() == 1 && !worldConfig.enableAttentionDP() + && !common::getEnvDisaggLayerwise() && !common::getEnvTryZCopyForKVCacheTransfer(); + validateInflightCancelModeAndCapability(mGroupComm, mInflightCancelEnabled, inflightCancelCapable); + + TLLM_CHECK_WITH_INFO(mCacheTransceiverConfig.has_value(), "CacheTransceiverConfig is not set."); + TLLM_CHECK_WITH_INFO(backendType != BackendType::DEFAULT, " CacheTransceiverConfig::BackendType is not set."); + if (worldConfig.isTensorParallel() || worldConfig.isContextParallel()) { mGroupTensorParaComm = std::make_shared( @@ -348,12 +500,6 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa } } bool isMLA = attentionType == executor::kv_cache::CacheState::AttentionType::kMLA; - TLLM_CHECK_WITH_INFO(mCacheTransceiverConfig.has_value(), "CacheTransceiverConfig is not set."); - auto backendType = mCacheTransceiverConfig.value().getBackendType(); - TLLM_CHECK_WITH_INFO( - backendType.has_value() && (backendType.value() != executor::CacheTransceiverConfig::BackendType::DEFAULT), - " CacheTransceiverConfig::BackendType is not set."); - std::optional maxNumTokens = mCacheTransceiverConfig.value().getMaxTokensInBuffer(); mCacheTransBufferManagers.push_back( @@ -443,7 +589,7 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa mCacheTransBufferManagerPtrs.push_back(mRnnCacheTransBufferManager.get()); } - if (backendType.value() == executor::CacheTransceiverConfig::BackendType::UCX) + if (backendType == BackendType::UCX) { std::lock_guard lock(mDllMutex); mWrapperLibHandle = dllOpen(UCX_WRAPPER_LIB_NAME); @@ -462,7 +608,7 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa mManager = makeUcxConnectionManager(); TLLM_LOG_INFO("UCX Connection Manager created"); } - else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL) + else if (backendType == BackendType::NIXL) { auto rnnState = mCacheState->hasRnnConfig() ? std::make_optional(mCacheState->getRnnCacheState()) : std::nullopt; @@ -470,7 +616,7 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa mCacheTransBufferManagerPtrs, *mCacheState, "nixl", rnnState); TLLM_LOG_INFO("NIXL Connection Manager created"); } - else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MOONCAKE) + else if (backendType == BackendType::MOONCAKE) { auto rnnState = mCacheState->hasRnnConfig() ? std::make_optional(mCacheState->getRnnCacheState()) : std::nullopt; @@ -478,7 +624,7 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa mCacheTransBufferManagerPtrs, *mCacheState, "mooncake", rnnState); TLLM_LOG_INFO("MOONCAKE Connection Manager created"); } - else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MPI) + else if (backendType == BackendType::MPI) { mMpiWorldComm = std::addressof(tensorrt_llm::mpi::MpiComm::world()); mManager = std::make_unique(mMpiWorldComm); @@ -515,8 +661,8 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa return nullptr; }; - auto makeCacheTransferLayer - = [&]() { return CacheTransferLayer(*mCacheState, makeFormatter(), makeRnnFormatter()); }; + auto makeCacheTransferLayer = [&]() + { return CacheTransferLayer(*mCacheState, makeFormatter(), makeRnnFormatter(), mInflightCancelEnabled); }; mCacheSender = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); mCacheReceiver = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); @@ -531,6 +677,10 @@ CacheTransceiver::~CacheTransceiver() mCacheSender.reset(); mCacheReceiver.reset(); + // Deregister transfer memory and destroy plugin-backed connection objects + // before their buffer owners or dynamic library are destroyed. + mManager.reset(); + if (mWrapperLibHandle) { std::lock_guard lock(mDllMutex); @@ -599,12 +749,21 @@ void CacheTransceiver::respondAndSendLayerWise( void CacheTransceiver::requestAndReceiveSync(std::shared_ptr llmRequest) { + TLLM_CHECK_WITH_INFO(!mInflightCancelEnabled, + "Synchronous generation transfer is not supported when in-flight cancellation is enabled; use the " + "asynchronous receive and finite status-poll path."); TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); { auto future = mCacheReceiver->receiveAsync(llmRequest); future.get(); } - llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); + // The default-off worker reports a cleanly rejected ready handshake by + // setting the request state and resolving its future. Preserve that error + // instead of converting the failed transfer into a false success. + if (llmRequest->getState() != LlmRequestState::kDISAGG_TRANS_ERROR) + { + llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); + } } void CacheTransceiver::requestAndReceiveAsync(std::shared_ptr llmRequest) @@ -721,9 +880,271 @@ void updateKVCacheTransferBW(std::shared_ptr const& mComm, } } +RequestStatuses CacheTransceiver::checkContextTransferStatusWithInflightCancel( + std::optional const& atLeastRequestNum, bool const markComplete) +{ + TLLM_CHECK_WITH_INFO(atLeastRequestNum.has_value(), + "In-flight cancellation requires a finite context-transfer status poll; pass 0 for a nonblocking poll."); + + auto const needsProgress = atLeastRequestNum.value() > 0; + auto const futureWaitInterval + = getTransferFutureWaitInterval(mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(), needsProgress); + auto const kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs().value(); + auto const syncComm = mGroupTensorParaComm; + + auto countReadyFutures = [this]() + { + return static_cast(std::count_if(mSenderFutures.begin(), mSenderFutures.end(), + [](auto& entry) + { return entry.second.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready; })); + }; + if (needsProgress) + { + auto const deadline = std::chrono::steady_clock::now() + futureWaitInterval; + while (countReadyFutures() < atLeastRequestNum.value() && std::chrono::steady_clock::now() < deadline) + { + auto const remaining + = std::chrono::duration_cast(deadline - std::chrono::steady_clock::now()); + std::this_thread::sleep_for(std::min(std::chrono::milliseconds(kTransferFuturePollIntervalMs), remaining)); + } + } + + for (auto const& [request, future] : mSenderFutures) + { + if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + { + continue; + } + auto const elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + if (elapsed.count() > kvTransferTimeoutMs && mTimedOutSenderIds.insert(request->mRequestId).second) + { + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld exceeded the local timeout; requesting topology-wide " + "cancellation.", + request->mRequestId); + } + } + + for (auto it = mSenderFutures.begin(); it != mSenderFutures.end();) + { + auto& [request, future] = *it; + if (future.wait_for(std::chrono::milliseconds(0)) != std::future_status::ready) + { + ++it; + continue; + } + auto const requestId = request->mRequestId; + try + { + future.get(); + auto const failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR; + recordLocalTransferOutcome(requestId, request, failed, mCompletedSenderRequestIds, mFailedSenderRequestIds, + mSenderRequestsAwaitingConsensus); + } + catch (std::exception const& error) + { + TLLM_LOG_ERROR("Error occurred during context transfer for request %ld: %s", requestId, error.what()); + recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, + mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); + } + it = mSenderFutures.erase(it); + } + + auto const localPoisoned = mTopologyTransferBufferPoisoned.load(std::memory_order_relaxed) + || std::any_of(mCacheTransBufferManagerPtrs.begin(), mCacheTransBufferManagerPtrs.end(), + [](BaseTransBufferManager const* manager) { return manager != nullptr && manager->hasPoisonedBuffer(); }); + auto const outcome = exchangeTransferStates( + syncComm, mCompletedSenderRequestIds, mFailedSenderRequestIds, mTimedOutSenderIds, localPoisoned); + if (outcome.transferBufferPoisoned) + { + mTopologyTransferBufferPoisoned.store(true, std::memory_order_relaxed); + } + + for (auto const requestId : outcome.failedRequestIds) + { + auto const futureIt = std::find_if(mSenderFutures.begin(), mSenderFutures.end(), + [requestId](auto const& entry) { return entry.first->mRequestId == requestId; }); + if (futureIt == mSenderFutures.end() + || futureIt->second.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + { + continue; + } + if (outcome.timedOutRequestIds.find(requestId) != outcome.timedOutRequestIds.end() + && mTimedOutSenderIds.insert(requestId).second) + { + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld exceeded the topology-wide timeout; cancelling it on " + "every participating rank.", + requestId); + } + (void) mCacheSender->cancelRequest(*futureIt->first); + } + + RequestStatuses requestsStatus; + for (auto const requestId : outcome.quiescedRequestIds) + { + auto const requestIt = mSenderRequestsAwaitingConsensus.find(requestId); + if (requestIt == mSenderRequestsAwaitingConsensus.end()) + { + continue; + } + if (outcome.failedRequestIds.find(requestId) != outcome.failedRequestIds.end()) + { + requestIt->second->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + requestsStatus.errorRequestIds.insert(requestId); + } + else + { + requestsStatus.completedRequestIds.insert(requestId); + if (markComplete) + { + requestIt->second->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); + } + } + mTimedOutSenderIds.erase(requestId); + eraseLocalTransferOutcome( + requestId, mCompletedSenderRequestIds, mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); + } + return requestsStatus; +} + +void CacheTransceiver::checkGenTransferStatusWithInflightCancel(std::optional const& atLeastRequestNum) +{ + TLLM_CHECK_WITH_INFO(atLeastRequestNum.has_value(), + "In-flight cancellation requires a finite generation-transfer status poll; pass 0 for a nonblocking poll."); + + auto const needsProgress = atLeastRequestNum.value() > 0; + auto const futureWaitInterval + = getTransferFutureWaitInterval(mCacheTransceiverConfig->getKvTransferPollIntervalMs(), needsProgress); + auto const kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs().value(); + auto const syncComm = mGroupComm; + + auto countReadyFutures = [this]() + { + return static_cast(std::count_if(mRequesterFutures.begin(), mRequesterFutures.end(), + [](auto& entry) + { return entry.second.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready; })); + }; + if (needsProgress) + { + auto const deadline = std::chrono::steady_clock::now() + futureWaitInterval; + while (countReadyFutures() < atLeastRequestNum.value() && std::chrono::steady_clock::now() < deadline) + { + auto const remaining + = std::chrono::duration_cast(deadline - std::chrono::steady_clock::now()); + std::this_thread::sleep_for(std::min(std::chrono::milliseconds(kTransferFuturePollIntervalMs), remaining)); + } + } + + for (auto const& [request, future] : mRequesterFutures) + { + if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + { + continue; + } + auto const elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + if (elapsed.count() > kvTransferTimeoutMs && mTimedOutRequesterIds.insert(request->mRequestId).second) + { + TLLM_LOG_WARNING( + "Generation KV cache transfer for request %ld exceeded the local timeout; requesting topology-wide " + "cancellation.", + request->mRequestId); + } + } + + for (auto it = mRequesterFutures.begin(); it != mRequesterFutures.end();) + { + auto& [request, future] = *it; + if (future.wait_for(std::chrono::milliseconds(0)) != std::future_status::ready) + { + ++it; + continue; + } + auto const requestId = request->mRequestId; + try + { + future.get(); + auto const failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR; + if (failed) + { + request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); + } + recordLocalTransferOutcome(requestId, request, failed, mCompletedRequesterRequestIds, + mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); + } + catch (std::exception const& error) + { + TLLM_LOG_ERROR("Error occurred during generation transfer for request %ld: %s", requestId, error.what()); + recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedRequesterRequestIds, + mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); + } + it = mRequesterFutures.erase(it); + } + + auto const localPoisoned = mTopologyTransferBufferPoisoned.load(std::memory_order_relaxed) + || std::any_of(mCacheTransBufferManagerPtrs.begin(), mCacheTransBufferManagerPtrs.end(), + [](BaseTransBufferManager const* manager) { return manager != nullptr && manager->hasPoisonedBuffer(); }); + auto const outcome = exchangeTransferStates( + syncComm, mCompletedRequesterRequestIds, mFailedRequesterRequestIds, mTimedOutRequesterIds, localPoisoned); + if (outcome.transferBufferPoisoned) + { + mTopologyTransferBufferPoisoned.store(true, std::memory_order_relaxed); + } + + for (auto const requestId : outcome.failedRequestIds) + { + auto const futureIt = std::find_if(mRequesterFutures.begin(), mRequesterFutures.end(), + [requestId](auto const& entry) { return entry.first->mRequestId == requestId; }); + if (futureIt == mRequesterFutures.end() + || futureIt->second.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + { + continue; + } + if (outcome.timedOutRequestIds.find(requestId) != outcome.timedOutRequestIds.end() + && mTimedOutRequesterIds.insert(requestId).second) + { + TLLM_LOG_WARNING( + "Generation KV cache transfer for request %ld exceeded the topology-wide timeout; cancelling it on " + "every participating rank.", + requestId); + } + (void) mCacheReceiver->cancelRequest(*futureIt->first); + } + + for (auto const requestId : outcome.quiescedRequestIds) + { + auto const requestIt = mRequesterRequestsAwaitingConsensus.find(requestId); + if (requestIt == mRequesterRequestsAwaitingConsensus.end()) + { + continue; + } + if (outcome.failedRequestIds.find(requestId) != outcome.failedRequestIds.end()) + { + requestIt->second->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + } + else + { + requestIt->second->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); + if (!common::getEnvKVCacheTimeOutputPath().empty()) + { + updateKVCacheTransferBW(syncComm, requestIt->second.get()); + } + } + mTimedOutRequesterIds.erase(requestId); + eraseLocalTransferOutcome( + requestId, mCompletedRequesterRequestIds, mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); + } +} + RequestStatuses CacheTransceiver::checkContextTransferStatus( std::optional const& atLeastRequestNum, bool markComplete) { + if (mInflightCancelEnabled) + { + return checkContextTransferStatusWithInflightCancel(atLeastRequestNum, markComplete); + } bool const blockAll = !atLeastRequestNum.has_value(); std::optional senderFutureTimeoutMs = std::nullopt; if (mCacheTransceiverConfig.has_value()) @@ -892,6 +1313,11 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastRequestNum) { + if (mInflightCancelEnabled) + { + checkGenTransferStatusWithInflightCancel(atLeastRequestNum); + return; + } bool const blockAll = !atLeastRequestNum.has_value(); bool const needsProgress = atLeastRequestNum.value_or(0) > 0; std::optional genTransferPollIntervalMs = std::nullopt; @@ -1104,6 +1530,13 @@ bool CacheTransceiver::checkGenTransferComplete() const return mRequesterFutures.empty() && mCompletedRequesterRequestIds.empty() && mFailedRequesterRequestIds.empty(); } +bool CacheTransceiver::hasPoisonedTransferBuffer() const +{ + return mTopologyTransferBufferPoisoned.load(std::memory_order_relaxed) + || std::any_of(mCacheTransBufferManagerPtrs.begin(), mCacheTransBufferManagerPtrs.end(), + [](BaseTransBufferManager const* manager) { return manager != nullptr && manager->hasPoisonedBuffer(); }); +} + bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) { if (llmRequest->isContextOnlyRequest()) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp index c013fd75c6e4..23a159b922ef 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp @@ -30,10 +30,12 @@ namespace tensorrt_llm::batch_manager { CacheTransferLayer::CacheTransferLayer(executor::kv_cache::CacheState cacheState, - std::unique_ptr kvFormatter, std::unique_ptr rnnFormatter) + std::unique_ptr kvFormatter, std::unique_ptr rnnFormatter, + bool const enableInflightCancel) : mCacheState{std::move(cacheState)} , mKvFormatter{std::move(kvFormatter)} , mRnnFormatter{std::move(rnnFormatter)} + , mEnableInflightCancel{enableInflightCancel} { TLLM_CHECK(mKvFormatter); } @@ -143,4 +145,9 @@ BaseCacheFormatter* CacheTransferLayer::getKvFormatter() const noexcept return mKvFormatter.get(); } +bool CacheTransferLayer::isInflightCancelEnabled() const noexcept +{ + return mEnableInflightCancel; +} + } // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h index 0506e98197e6..468b15bfdc4a 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h +++ b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h @@ -49,7 +49,7 @@ class CacheTransferLayer /// @param kvFormatter The KV cache formatter. /// @param rnnFormatter Optional RNN cache formatter (for separate RnnStateManager pool). CacheTransferLayer(executor::kv_cache::CacheState cacheState, std::unique_ptr kvFormatter, - std::unique_ptr rnnFormatter = nullptr); + std::unique_ptr rnnFormatter = nullptr, bool enableInflightCancel = false); ~CacheTransferLayer(); CacheTransferLayer(CacheTransferLayer&&) noexcept; @@ -86,10 +86,13 @@ class CacheTransferLayer [[nodiscard]] BaseCacheFormatter* getKvFormatter() const noexcept; + [[nodiscard]] bool isInflightCancelEnabled() const noexcept; + private: executor::kv_cache::CacheState mCacheState; std::unique_ptr mKvFormatter; std::unique_ptr mRnnFormatter; + bool mEnableInflightCancel{false}; }; } // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransferState.h b/cpp/tensorrt_llm/batch_manager/cacheTransferState.h new file mode 100644 index 000000000000..868c7122dc9f --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/cacheTransferState.h @@ -0,0 +1,137 @@ +/* + * 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 "tensorrt_llm/common/assert.h" + +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::detail +{ + +using TransferRequestId = std::uint64_t; + +enum TransferStateFlag : std::uint64_t +{ + kTransferCompleted = 1U << 0, + kTransferFailed = 1U << 1, + kTransferTimedOut = 1U << 2, +}; + +enum TransferGlobalFlag : std::uint64_t +{ + kTransferBufferPoisoned = 1U << 0, +}; + +struct TransferTopologyOutcome +{ + // Logical failure is a union: one failed or timed-out rank fails the request everywhere. + std::unordered_set failedRequestIds; + std::unordered_set timedOutRequestIds; + // Quiescence is an intersection: every rank has reached a local terminal state. + std::unordered_set quiescedRequestIds; + std::unordered_set completedRequestIds; + bool transferBufferPoisoned{false}; +}; + +enum class ReadySignalSummary +{ + kReady, + kNotReady, + kPartiallyReady, +}; + +inline ReadySignalSummary summarizeReadySignals(bool const allPeersReady, bool const anyPeerReady) +{ + TLLM_CHECK_WITH_INFO(!allPeersReady || anyPeerReady, "An all-ready response must contain at least one ready peer."); + if (allPeersReady) + { + return ReadySignalSummary::kReady; + } + return anyPeerReady ? ReadySignalSummary::kPartiallyReady : ReadySignalSummary::kNotReady; +} + +// Each rank payload starts with global flags, followed by request-id/state-flag pairs. +inline TransferTopologyOutcome summarizePackedTransferStates( + std::vector> const& rankPayloads) +{ + TLLM_CHECK_WITH_INFO(!rankPayloads.empty(), "Packed transfer state summary requires at least one rank payload."); + + struct Counts + { + int completed{0}; + int failed{0}; + int timedOut{0}; + }; + + std::unordered_map countsByRequest; + TransferTopologyOutcome outcome; + for (auto const& payload : rankPayloads) + { + TLLM_CHECK_WITH_INFO(!payload.empty() && ((payload.size() - 1) % 2 == 0), + "Packed transfer state payload must contain global flags followed by request/state pairs."); + outcome.transferBufferPoisoned + = outcome.transferBufferPoisoned || ((payload.front() & kTransferBufferPoisoned) != 0); + + std::unordered_set seenRequestIds; + for (size_t index = 1; index < payload.size(); index += 2) + { + auto const requestId = payload.at(index); + auto const flags = payload.at(index + 1); + TLLM_CHECK_WITH_INFO(seenRequestIds.insert(requestId).second, + "Packed transfer state payload contains duplicate request ID %lu.", requestId); + TLLM_CHECK_WITH_INFO(!((flags & kTransferCompleted) != 0 && (flags & kTransferFailed) != 0), + "Request %lu cannot be both locally completed and failed.", requestId); + + auto& counts = countsByRequest[requestId]; + counts.completed += (flags & kTransferCompleted) != 0; + counts.failed += (flags & kTransferFailed) != 0; + counts.timedOut += (flags & kTransferTimedOut) != 0; + } + } + + auto const rankCount = static_cast(rankPayloads.size()); + for (auto const& [requestId, counts] : countsByRequest) + { + auto const terminalCount = counts.completed + counts.failed; + bool const logicallyFailed = counts.failed > 0 || counts.timedOut > 0; + if (logicallyFailed) + { + outcome.failedRequestIds.insert(requestId); + } + if (counts.timedOut > 0) + { + outcome.timedOutRequestIds.insert(requestId); + } + if (terminalCount == rankCount) + { + outcome.quiescedRequestIds.insert(requestId); + if (!logicallyFailed && counts.completed == rankCount) + { + outcome.completedRequestIds.insert(requestId); + } + } + } + return outcome; +} + +} // namespace tensorrt_llm::batch_manager::detail diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 5aac98450be9..a1ea8f2f7f3f 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -18,6 +18,7 @@ #include "dataTransceiver.h" #include "tensorrt_llm/batch_manager/cacheFormatter.h" +#include "tensorrt_llm/batch_manager/cacheTransferState.h" #include "tensorrt_llm/batch_manager/common.h" #include "tensorrt_llm/batch_manager/kvCacheUtils.h" #include "tensorrt_llm/batch_manager/runtimeBuffers.h" @@ -34,6 +35,7 @@ #include #include #include +#include #include namespace tensorrt_llm::batch_manager @@ -212,6 +214,49 @@ struct ReceiveCacheResource } }; +namespace +{ + +class RecvBufferAdvertisementGuard +{ +public: + explicit RecvBufferAdvertisementGuard(std::vector& holders) noexcept + : mHolders{holders} + { + } + + RecvBufferAdvertisementGuard(RecvBufferAdvertisementGuard const&) = delete; + RecvBufferAdvertisementGuard& operator=(RecvBufferAdvertisementGuard const&) = delete; + + ~RecvBufferAdvertisementGuard() noexcept + { + if (!mDisarmed && mAdvertisementMayHaveOccurred) + { + for (auto& holder : mHolders) + { + holder.poison(); + } + } + } + + [[nodiscard]] bool* advertisementMayHaveOccurred() noexcept + { + return &mAdvertisementMayHaveOccurred; + } + + void disarm() noexcept + { + mDisarmed = true; + } + +private: + std::vector& mHolders; + bool mAdvertisementMayHaveOccurred{false}; + bool mDisarmed{false}; +}; + +} // namespace + RequestInfo::RequestInfo(LlmRequest::RequestIdType requestId, executor::DataTransceiverState transState) : mRequestId{requestId} , mTransState{std::move(transState)} @@ -281,13 +326,13 @@ class CacheSender::Impl Impl(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer) : mManager{manager} , mSelfState{cacheLayer.getCacheState(), executor::kv_cache::CommState{manager->getCommState()}} + , mInflightCancelEnabled{cacheLayer.isInflightCancelEnabled()} , mCacheTransferLayer{std::move(cacheLayer)} , mBufferManager{std::make_shared()} { 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,18 +348,33 @@ class CacheSender::Impl std::promise promise; auto future = promise.get_future(); llmRequest->setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); + (void) getOrCreateInFlightCancelFlag(llmRequest->mRequestId); { - { - 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; } + std::shared_ptr> getOrCreateInFlightCancelFlag(RequestIdType requestId) + { + std::lock_guard lg(mInFlightCancelMutex); + auto it = mInFlightCancelFlags.find(requestId); + if (it != mInFlightCancelFlags.end()) + { + return it->second; + } + auto flag = std::make_shared>(false); + mInFlightCancelFlags.emplace(requestId, flag); + return flag; + } + [[nodiscard]] executor::kv_cache::CommState const& getCommState() const { return mSelfState.getCommState().value(); @@ -335,11 +395,11 @@ class CacheSender::Impl void release(LlmRequest::RequestIdType requestId) { - std::unique_lock lk(mMtxForMap); - auto it = mRequestToSession.find(requestId); - TLLM_CHECK(it != mRequestToSession.end()); if (!common::getEnvKVCacheTimeOutputPath().empty()) { + std::unique_lock lk(mMtxForMap); + auto it = mRequestToSession.find(requestId); + TLLM_CHECK(it != mRequestToSession.end()); if (!mMeasuresFile.is_open()) { auto outputPath = getTransferOutputPath("send"); @@ -349,7 +409,38 @@ class CacheSender::Impl } it->second.exportMeasure(mMeasuresFile, true); } - mRequestToSession.erase(it); + { + std::unique_lock lk(mMtxForMap); + auto it = mRequestToSession.find(requestId); + TLLM_CHECK(it != mRequestToSession.end()); + mRequestToSession.erase(it); + } + { + std::lock_guard lg(mInFlightCancelMutex); + mInFlightCancelFlags.erase(requestId); + } + } + + void discardTransferState(LlmRequest::RequestIdType requestId) noexcept + { + try + { + std::unique_lock lk(mMtxForMap); + mRequestToSession.erase(requestId); + } + catch (...) + { + // noexcept cleanup after send failure. + } + try + { + std::lock_guard lg(mInFlightCancelMutex); + mInFlightCancelFlags.erase(requestId); + } + catch (...) + { + // noexcept cleanup after send failure. + } } [[nodiscard]] std::optional recvRequestInfo() @@ -392,15 +483,16 @@ class CacheSender::Impl TLLM_CHECK_WITH_INFO(peerIdx < static_cast(allCounterparts.size()), "Peer rank %d not found in expected counterparts", peerSelfIdx); + auto cancelFlag = getOrCreateInFlightCancelFlag(requestId); { std::unique_lock lk(mMtxForMap); auto it = mRequestToSession.find(requestId); if (it == mRequestToSession.end()) { auto session = TransferSession(std::vector(allCounterparts.size(), nullptr), - DataContext{tagFromRequestId(requestId), mTerminate}, allCounterparts, mSelfState, + DataContext{tagFromRequestId(requestId), *cancelFlag}, allCounterparts, mSelfState, info.getTransState(), mBufferManager, info.getIndexFromEnd(), info.getLastBlockKey(), nullptr, - !common::getEnvKVCacheTimeOutputPath().empty()); + !common::getEnvKVCacheTimeOutputPath().empty(), mInflightCancelEnabled); session.setTime(TransferSession::kTimeRequestInfo); it = mRequestToSession.emplace(requestId, std::move(session)).first; } @@ -426,16 +518,29 @@ class CacheSender::Impl bool cancelRequest(LlmRequest const& llmRequest) { bool isCancelled = false; - std::scoped_lock lkResp(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)) { - mCancelledRequests.insert(llmRequest.mRequestId); - isCancelled = true; + std::scoped_lock lock(mSenderMutex); + auto it = mReadyResponses.find(llmRequest.mRequestId); + // Until ready=true is committed to the peer, cancellation can still + // reject the request cleanly without starting a transfer. + if (it != mReadyResponses.end() + && (!mReadyCommittedRequest.has_value() || mReadyCommittedRequest.value() != llmRequest.mRequestId)) + { + mCancelledRequests.insert(llmRequest.mRequestId); + isCancelled = true; + } } - else + if (!isCancelled && mInflightCancelEnabled) + { + std::lock_guard lg(mInFlightCancelMutex); + auto flagIt = mInFlightCancelFlags.find(llmRequest.mRequestId); + if (flagIt != mInFlightCancelFlags.end()) + { + flagIt->second->store(true, std::memory_order_relaxed); + isCancelled = true; + } + } + if (!isCancelled) { TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); } @@ -535,207 +640,227 @@ class CacheSender::Impl catch (tensorrt_llm::common::RequestSpecificException const& e) { TLLM_LOG_ERROR("Exception in sendAndRemoveResponse: %s ", e.what()); + discardTransferState(id); auto new_exception = TLLM_REQUEST_EXCEPTION(id, e.getErrorCode(), "%s", e.what()); resp.mPromise.set_exception(std::make_exception_ptr(new_exception)); } catch (std::exception const& e) { TLLM_LOG_ERROR("Exception in sendAndRemoveResponse: %s request id: %ld", e.what(), id); + discardTransferState(id); resp.mPromise.set_exception(std::current_exception()); } } 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) { - 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); + mReadyCommittedRequest = reqId; + } + } + + // 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); + mReadyCommittedRequest = std::nullopt; + mCurrentRequest = std::nullopt; + } + + if (isReady) + { + if (dynamic_cast(mManager) != nullptr) + { + // 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))); + discardTransferState(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 requestInfo = recvRequestInfo(); + if (!requestInfo.has_value() || mTerminate || !mManager->isRunning()) { break; } - if (!mReadyResponses.empty()) - { - auto requestInfo = recvRequestInfo(); - if (!requestInfo.has_value() || mTerminate || !mManager->isRunning()) - { - return; - } - auto reqId = requestInfo->getRequestId(); - - { - std::scoped_lock lk(mSenderMutex); - mCurrentRequest = reqId; - } + auto const reqId = requestInfo->getRequestId(); - 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 lk(mCondMutex); - mSenderCv.wait(lk, [this]() { return (mAnyReady || mTerminate); }); - if (mTerminate) - { - break; - } - it = getCurrentResponse(); - } - if (mTerminate || it == mReadyResponses.end()) + std::unique_lock lock(mSenderMutex); + mCurrentRequest = reqId; + mSenderCv.wait(lock, + [this, reqId]() { return mTerminate || mReadyResponses.find(reqId) != mReadyResponses.end(); }); + if (mTerminate) { + mCurrentRequest = std::nullopt; + mReadyCommittedRequest = 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. + { + std::lock_guard lg(mInFlightCancelMutex); + for (auto& [id, flag] : mInFlightCancelFlags) + { + flag->store(true, std::memory_order_relaxed); + } + } + // Wake the sender loop and make in-flight agent transfers observe termination through + // their per-request cancellation flags. 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; + mReadyCommittedRequest = std::nullopt; + mCancelledRequests.clear(); + mRemainSendCount.clear(); + } + for (auto& entry : pendingResponses) + { + failResponse(entry.second, exception); + } } public: @@ -749,11 +874,12 @@ class CacheSender::Impl private: std::optional mCurrentRequest; + std::optional mReadyCommittedRequest; 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; @@ -763,10 +889,13 @@ class CacheSender::Impl executor::kv_cache::ConnectionManager* mManager; std::map mRequestToSession; executor::DataTransceiverState mSelfState; + bool mInflightCancelEnabled{false}; CacheTransferLayer mCacheTransferLayer; std::mutex mMtxForMap; runtime::BufferManager mBufferManager; std::ofstream mMeasuresFile; + std::mutex mInFlightCancelMutex; + std::unordered_map>> mInFlightCancelFlags; }; class CacheReceiver::Impl @@ -775,6 +904,7 @@ class CacheReceiver::Impl Impl(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer) : mManager{manager} , mSelfState{cacheLayer.getCacheState(), executor::kv_cache::CommState{manager->getCommState()}} + , mInflightCancelEnabled{cacheLayer.isInflightCancelEnabled()} , mCacheTransferLayer{std::move(cacheLayer)} , mBufferManager{std::make_shared()} { @@ -789,7 +919,15 @@ class CacheReceiver::Impl // TODO: Modify the implementation here to avoid frequent thread creation. // Capture by value so the async task owns a strong reference for its lifetime. auto llmRequestCopy = llmRequest; - return std::async(std::launch::async, [this, llmRequestCopy]() { requestSync(*llmRequestCopy); }); + return std::async(std::launch::async, + [this, llmRequestCopy]() + { + if (!requestSync(*llmRequestCopy) && mInflightCancelEnabled) + { + throw TLLM_REQUEST_EXCEPTION(llmRequestCopy->mRequestId, common::RequestErrorCode::kNETWORK_ERROR, + "Generation KV cache transfer failed for request %zu", llmRequestCopy->mRequestId); + } + }); } [[nodiscard]] std::future requestAndReceiveAsyncMultiThreads(std::shared_ptr const& llmRequest) @@ -814,9 +952,14 @@ class CacheReceiver::Impl mRequestFutures.emplace_back(std::move(requestFuture)); } auto& asyncResource = mInstanceToAsyncResource.at(processInfo); + auto cancelFlag = std::make_shared>(false); + { + std::lock_guard lg(mInFlightCancelMutex); + mInFlightCancelFlags[llmRequest->mRequestId] = cancelFlag; + } { std::unique_lock lck(asyncResource->mMtxForQueue); - asyncResource->mRequestsQueue.emplace_back(llmRequest, std::move(promise)); + asyncResource->mRequestsQueue.emplace_back(llmRequest, std::move(promise), cancelFlag); } asyncResource->mCVforQueue.notify_all(); return future; @@ -829,7 +972,16 @@ class CacheReceiver::Impl void receiveSync(TransferSession& session) { - mCacheTransferLayer.unformat(session); + try + { + mCacheTransferLayer.unformat(session); + session.releaseRecvBufferHolders(); + } + catch (...) + { + session.poisonRecvBufferHolders(); + throw; + } if (!common::getEnvKVCacheTimeOutputPath().empty()) { std::unique_lock lock(mMeasuresFileMutex); @@ -844,8 +996,12 @@ class CacheReceiver::Impl } } - TransferSession sendRequestInfo(LlmRequest const& llmRequest) + TransferSession sendRequestInfo(LlmRequest const& llmRequest, std::atomic const* perRequestCancel = nullptr, + std::vector> cacheBufferIds = {}, + RecvBufferAdvertisementGuard* externalRecvBufferGuard = nullptr) { + TLLM_CHECK_WITH_INFO(externalRecvBufferGuard == nullptr || !cacheBufferIds.empty(), + "An external receive-buffer advertisement guard requires preassigned buffer IDs."); uint64_t requestId = llmRequest.getContextPhaseParams().value().getReqId(); auto const& contextState = llmRequest.getDataTransceiverState(); auto const& commState = contextState.getCommState().value(); @@ -880,16 +1036,32 @@ class CacheReceiver::Impl } auto* agentConnectionManager = dynamic_cast(mManager); - std::vector> cacheBufferIds; + std::vector sessionRecvHolders; if (agentConnectionManager) { - for (auto& cacheTransBufferManager : agentConnectionManager->getCacheTransBufferManagers()) + if (cacheBufferIds.empty()) { - cacheBufferIds.push_back(cacheTransBufferManager->assignBufferIndexForRecv()); + auto const& managers = agentConnectionManager->getCacheTransBufferManagers(); + sessionRecvHolders.reserve(managers.size()); + for (auto* cacheTransBufferManager : managers) + { + auto rawIdx = cacheTransBufferManager->assignBufferIndexForRecv(perRequestCancel); + sessionRecvHolders.emplace_back(*cacheTransBufferManager, rawIdx, /*isRecv=*/true); + if (rawIdx.has_value()) + { + cacheBufferIds.push_back(static_cast(rawIdx.value())); + } + else + { + cacheBufferIds.push_back(std::nullopt); + } + } } TLLM_CHECK(!cacheBufferIds.empty()); } + RecvBufferAdvertisementGuard sessionRecvBufferGuard{sessionRecvHolders}; + auto& recvBufferGuard = externalRecvBufferGuard != nullptr ? *externalRecvBufferGuard : sessionRecvBufferGuard; auto allCounterparts = mCacheTransferLayer.computeCounterparts(mSelfState.getCommState().value().getSelfIdx(), contextState); @@ -914,6 +1086,7 @@ class CacheReceiver::Impl allConnections.emplace_back(connection); } + bool sentAgentRequestInfo = false; for (size_t ci = 0; ci < allCounterparts.size(); ci++) { auto rank = allCounterparts[ci]; @@ -965,8 +1138,14 @@ class CacheReceiver::Impl auto* agentConnection = dynamic_cast(connection); TLLM_CHECK(agentConnection != nullptr); + // Cancellation may abort before the first notification. Once one + // counterpart has received the request, finish the fanout so every + // sender can reach its expected terminal response count. + auto const* fanoutCancel = sentAgentRequestInfo ? nullptr : perRequestCancel; const_cast(agentConnection) - ->sendRequestAndBufferInfo(requestInfo, idsForRank, validConnectionIdx); + ->sendRequestAndBufferInfo(requestInfo, idsForRank, validConnectionIdx, fanoutCancel, + recvBufferGuard.advertisementMayHaveOccurred()); + sentAgentRequestInfo = true; } else { @@ -974,10 +1153,21 @@ class CacheReceiver::Impl } } auto const& resource = getReceiveCacheResource(llmRequest); - return TransferSession(std::move(allConnections), DataContext{tagFromRequestId(requestId), mTerminate}, - std::move(allCounterparts), mSelfState, contextState, resource->mBufferManager, - requestInfo.getIndexFromEnd(), requestInfo.getLastBlockKey(), &llmRequest, - !common::getEnvKVCacheTimeOutputPath().empty()); + TransferSession session = perRequestCancel != nullptr + ? TransferSession(std::move(allConnections), DataContext{tagFromRequestId(requestId), *perRequestCancel}, + std::move(allCounterparts), mSelfState, contextState, resource->mBufferManager, + requestInfo.getIndexFromEnd(), requestInfo.getLastBlockKey(), &llmRequest, + !common::getEnvKVCacheTimeOutputPath().empty(), mInflightCancelEnabled) + : TransferSession(std::move(allConnections), DataContext{tagFromRequestId(requestId), mTerminate}, + std::move(allCounterparts), mSelfState, contextState, resource->mBufferManager, + requestInfo.getIndexFromEnd(), requestInfo.getLastBlockKey(), &llmRequest, + !common::getEnvKVCacheTimeOutputPath().empty(), mInflightCancelEnabled); + session.adoptRecvBufferHolders(std::move(sessionRecvHolders)); + if (externalRecvBufferGuard == nullptr) + { + sessionRecvBufferGuard.disarm(); + } + return session; } std::unique_ptr const& getReceiveCacheResource(LlmRequest const& llmRequest) @@ -1021,6 +1211,7 @@ class CacheReceiver::Impl bool isCancelled = false; auto& asyncResource = mInstanceToAsyncResource.at(processInfo); + std::optional queuedCancelledReqId; { std::unique_lock lck(asyncResource->mMtxForQueue); auto it = std::find_if(asyncResource->mRequestsQueue.begin(), asyncResource->mRequestsQueue.end(), @@ -1047,45 +1238,112 @@ class CacheReceiver::Impl } asyncResource->mRequestsQueue.erase(it); isCancelled = true; + queuedCancelledReqId = llmRequest.mRequestId; } - else + } + if (queuedCancelledReqId.has_value()) + { + std::lock_guard lg(mInFlightCancelMutex); + mInFlightCancelFlags.erase(*queuedCancelledReqId); + } + if (!isCancelled && mInflightCancelEnabled) + { + std::lock_guard lg(mInFlightCancelMutex); + auto flagIt = mInFlightCancelFlags.find(llmRequest.mRequestId); + if (flagIt != mInFlightCancelFlags.end()) { - TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); + flagIt->second->store(true, std::memory_order_relaxed); + isCancelled = true; } } + if (!isCancelled) + { + TLLM_LOG_WARNING("Cannot cancel request %zu", llmRequest.mRequestId); + } return isCancelled; } - bool receiveReadySignal(TransferSession& session) + enum class ReadySignalResult + { + kReady, + kNotReady, + kPartiallyReady, + kCancelled, + }; + + ReadySignalResult receiveReadySignalDetailed(TransferSession& session, std::atomic const& perRequestCancel) { bool isReadyFinal = true; + bool anyReady = false; bool isReady = false; auto const& connections = session.getConnections(); for (size_t i = 0; i < connections.size(); i++) { + if (perRequestCancel.load(std::memory_order_relaxed)) + { + return ReadySignalResult::kCancelled; + } auto* agentConnectionManager = dynamic_cast(mManager); if (agentConnectionManager) { auto* agentConnection = dynamic_cast(connections.at(i)); TLLM_CHECK(agentConnection); - isReady = agentConnection->recvReadySignal( - executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG, mTerminate}); + auto ready = agentConnection->recvReadySignalWithStatus( + executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG, perRequestCancel}); + if (!ready.has_value()) + { + return ReadySignalResult::kCancelled; + } + isReady = ready.value(); } else { connections.at(i)->recv( executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG}, &isReady, sizeof(isReady)); + if (perRequestCancel.load(std::memory_order_relaxed)) + { + return ReadySignalResult::kCancelled; + } } isReadyFinal &= isReady; + anyReady |= isReady; + } + + switch (detail::summarizeReadySignals(isReadyFinal, anyReady)) + { + case detail::ReadySignalSummary::kReady: return ReadySignalResult::kReady; + case detail::ReadySignalSummary::kNotReady: return ReadySignalResult::kNotReady; + case detail::ReadySignalSummary::kPartiallyReady: return ReadySignalResult::kPartiallyReady; } + TLLM_THROW("Unknown ready-signal summary"); + } - return isReadyFinal; + bool receiveReadySignal(TransferSession& session) + { + auto const result = receiveReadySignalDetailed(session, mTerminate); + if (result == ReadySignalResult::kPartiallyReady || result == ReadySignalResult::kCancelled) + { + session.poisonRecvBufferHolders(); + } + else if (result == ReadySignalResult::kNotReady) + { + // No peer accepted the request, so no writer can have started. + session.releaseRecvBufferHolders(); + } + return result == ReadySignalResult::kReady; } ~Impl() { mTerminate.store(true); + { + std::lock_guard lg(mInFlightCancelMutex); + for (auto& [id, flag] : mInFlightCancelFlags) + { + flag->store(true, std::memory_order_relaxed); + } + } for (auto&& [processInfo, asyncResource] : mInstanceToAsyncResource) { asyncResource->mTerminate = true; @@ -1098,29 +1356,119 @@ class CacheReceiver::Impl } private: - void requestSync(LlmRequest& llmRequest) + [[nodiscard]] bool requestSync(LlmRequest& llmRequest, std::atomic const& perRequestCancel) { TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "Start calling requestSync for request ID: %zu, context request ID: %zu.", llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); - llmRequest.setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); + if (!mInflightCancelEnabled) + { + llmRequest.setKvCacheTransferStart(LlmRequest::getSteadyClockNow()); + } + if (perRequestCancel.load(std::memory_order_relaxed) || mTerminate.load(std::memory_order_relaxed)) + { + if (!mInflightCancelEnabled) + { + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + } + llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); + return false; + } TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); - auto session = sendRequestInfo(llmRequest); - session.setTime(TransferSession::kTimeRequestInfo); - bool isReady = receiveReadySignal(session); - if (!isReady) + + std::vector recvHolders; + std::vector> cacheBufferIds; + auto* agentConnectionManager = dynamic_cast(mManager); + if (agentConnectionManager) + { + auto const& managers = agentConnectionManager->getCacheTransBufferManagers(); + recvHolders.reserve(managers.size()); + cacheBufferIds.reserve(managers.size()); + for (auto* cacheTransBufferManager : managers) + { + auto rawIdx = cacheTransBufferManager->assignBufferIndexForRecv(&perRequestCancel); + recvHolders.emplace_back(*cacheTransBufferManager, rawIdx, /*isRecv=*/true); + if (rawIdx.has_value()) + { + cacheBufferIds.push_back(static_cast(rawIdx.value())); + } + else + { + cacheBufferIds.push_back(std::nullopt); + } + } + } + + RecvBufferAdvertisementGuard recvBufferGuard{recvHolders}; + + try + { + auto* externalRecvBufferGuard = agentConnectionManager != nullptr ? &recvBufferGuard : nullptr; + auto session + = sendRequestInfo(llmRequest, &perRequestCancel, std::move(cacheBufferIds), externalRecvBufferGuard); + session.setTime(TransferSession::kTimeRequestInfo); + auto readyResult = receiveReadySignalDetailed(session, perRequestCancel); + if (readyResult == ReadySignalResult::kCancelled) + { + if (!mInflightCancelEnabled) + { + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + } + llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); + return false; + } + if (readyResult == ReadySignalResult::kNotReady) + { + // Every peer declined the request, so no advertised buffer can have a writer. + recvBufferGuard.disarm(); + if (!mInflightCancelEnabled) + { + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + } + llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); + return false; + } + if (readyResult == ReadySignalResult::kPartiallyReady) + { + // A peer that answered ready may start writing immediately. Since another peer + // declined the request, no receive path will prove all advertised buffers + // quiescent; quarantine them rather than making them available for reuse. + if (!mInflightCancelEnabled) + { + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + } + llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); + return false; + } + + receiveSync(session); + llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); + + for (auto& holder : recvHolders) + { + holder.release(); + } + recvBufferGuard.disarm(); + } + catch (...) { - // Reuse the error state for the cancelled request. - llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + if (!mInflightCancelEnabled) + { + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + } llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); - return; + throw; } - receiveSync(session); - llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End calling requestSync for request ID: %zu, context request ID: %zu.", llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); + return true; + } + + [[nodiscard]] bool requestSync(LlmRequest& llmRequest) + { + return requestSync(llmRequest, mTerminate); } struct RequestAndPromise @@ -1129,16 +1477,20 @@ class CacheReceiver::Impl // protects worker-side dereferences and the promise itself from premature destruction. std::shared_ptr mRequest; std::unique_ptr> mPromise; + std::shared_ptr> mCancelFlag; RequestAndPromise() : mRequest(nullptr) , mPromise(nullptr) + , mCancelFlag(nullptr) { } - RequestAndPromise(std::shared_ptr request, std::unique_ptr>&& promise) + RequestAndPromise(std::shared_ptr request, std::unique_ptr>&& promise, + std::shared_ptr> cancelFlag) : mRequest(std::move(request)) , mPromise(std::move(promise)) + , mCancelFlag(std::move(cancelFlag)) { } @@ -1147,6 +1499,7 @@ class CacheReceiver::Impl RequestAndPromise(RequestAndPromise&& other) noexcept : mRequest(std::move(other.mRequest)) , mPromise(std::move(other.mPromise)) + , mCancelFlag(std::move(other.mCancelFlag)) { } @@ -1162,6 +1515,7 @@ class CacheReceiver::Impl mRequest = std::move(other.mRequest); mPromise = std::move(other.mPromise); + mCancelFlag = std::move(other.mCancelFlag); } return *this; } @@ -1205,8 +1559,20 @@ class CacheReceiver::Impl try { TLLM_CHECK_WITH_INFO(requestAndPromise.mRequest != nullptr, "requestAndPromise.mRequest is null"); - requestSync(*requestAndPromise.mRequest); - requestAndPromise.mPromise->set_value(); + TLLM_CHECK_WITH_INFO( + requestAndPromise.mCancelFlag != nullptr, "requestAndPromise.mCancelFlag is null"); + if (requestSync(*requestAndPromise.mRequest, *requestAndPromise.mCancelFlag) + || !mInflightCancelEnabled) + { + requestAndPromise.mPromise->set_value(); + } + else + { + requestAndPromise.mPromise->set_exception(std::make_exception_ptr(TLLM_REQUEST_EXCEPTION( + requestAndPromise.mRequest->mRequestId, common::RequestErrorCode::kNETWORK_ERROR, + "Generation KV cache transfer failed for request %zu", + requestAndPromise.mRequest->mRequestId))); + } } catch (tensorrt_llm::common::RequestSpecificException const& err) { @@ -1224,6 +1590,19 @@ class CacheReceiver::Impl requestAndPromise.mRequest->getContextPhaseParams().value().getReqId(), err.what()); requestAndPromise.mPromise->set_exception(std::current_exception()); } + catch (...) + { + TLLM_LOG_ERROR("Unknown exception in CacheReceiver request() loop"); + if (requestAndPromise.mPromise) + { + requestAndPromise.mPromise->set_exception(std::current_exception()); + } + } + if (requestAndPromise.mRequest != nullptr) + { + std::lock_guard lg(mInFlightCancelMutex); + mInFlightCancelFlags.erase(requestAndPromise.mRequest->mRequestId); + } } } } @@ -1244,6 +1623,7 @@ class CacheReceiver::Impl std::unordered_map> mInstanceToAsyncResource; executor::kv_cache::ConnectionManager* mManager; executor::DataTransceiverState mSelfState; + bool mInflightCancelEnabled{false}; CacheTransferLayer mCacheTransferLayer; std::unordered_map> mProcessToResources; std::mutex mProcessIoResouceMutex; @@ -1251,6 +1631,8 @@ class CacheReceiver::Impl std::ofstream mMeasuresFile; std::mutex mMeasuresFileMutex; std::atomic mTerminate{false}; + std::mutex mInFlightCancelMutex; + std::unordered_map>> mInFlightCancelFlags; }; void CacheSender::ImplDeleter::operator()(Impl* ptr) diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index 3362574da902..4099da5f9003 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -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"); @@ -22,6 +22,7 @@ #include #include +#include "tensorrt_llm/batch_manager/baseTransBuffer.h" #include "tensorrt_llm/batch_manager/cacheTransceiver.h" #include "tensorrt_llm/batch_manager/cacheTransferLayer.h" #include "tensorrt_llm/batch_manager/llmRequest.h" @@ -84,7 +85,8 @@ class TransferSession TransferSession(std::vector connections, DataContext dataContext, std::vector counterPartRanks, executor::DataTransceiverState const& selfState, executor::DataTransceiverState otherState, runtime::BufferManager const& bufferManager, int32_t indexFromEnd, - BlockKey const& lastBlockKey, LlmRequest const* llmRequest = nullptr, bool recordTiming = false) + BlockKey const& lastBlockKey, LlmRequest const* llmRequest = nullptr, bool recordTiming = false, + bool enableInflightCancel = false) : mConnections(std::move(connections)) , mCounterPartRanks(std::move(counterPartRanks)) , mDataContext(std::move(dataContext)) @@ -94,6 +96,7 @@ class TransferSession , mRequest(llmRequest) , mIndexFromEnd(indexFromEnd) , mLastBlockKey(lastBlockKey) + , mEnableInflightCancel(enableInflightCancel) { TLLM_CHECK(!mConnections.empty()); if (recordTiming) @@ -102,6 +105,19 @@ class TransferSession } } + TransferSession(TransferSession const&) = delete; + TransferSession& operator=(TransferSession const&) = delete; + TransferSession(TransferSession&&) noexcept = default; + TransferSession& operator=(TransferSession&&) = delete; + + ~TransferSession() + { + // An adopted receive slot has already been advertised to a peer. Unless + // a completed receive explicitly releases it, destruction cannot prove + // that the peer has stopped writing, so fail closed. + poisonRecvBufferHolders(); + } + [[nodiscard]] std::vector const& getConnections() const; // should be called only during the initialization of the TransferSession @@ -151,6 +167,41 @@ class TransferSession mCounterPartRanks = std::move(ranks); } + [[nodiscard]] bool isInflightCancelEnabled() const noexcept + { + return mEnableInflightCancel; + } + + /// @brief Adopt receive-buffer slots allocated while constructing this session. + /// + /// The asynchronous receive path keeps these holders outside the session so it can + /// quarantine them on cancellation. The public sendRequestInfo()/receiveSync() path + /// has no such outer owner, so the session owns and releases those slots instead. + void adoptRecvBufferHolders(std::vector holders) noexcept + { + mRecvBufferHolders = std::move(holders); + } + + /// @brief Release receive-buffer slots owned by this session after a successful receive. + void releaseRecvBufferHolders() noexcept + { + for (auto& holder : mRecvBufferHolders) + { + holder.release(); + } + mRecvBufferHolders.clear(); + } + + /// @brief Quarantine receive-buffer slots owned by this session when quiescence is unknown. + void poisonRecvBufferHolders() noexcept + { + for (auto& holder : mRecvBufferHolders) + { + holder.poison(); + } + mRecvBufferHolders.clear(); + } + private: std::vector mConnections; std::vector mCounterPartRanks; // Ranks corresponding to mConnections indices @@ -162,6 +213,8 @@ class TransferSession std::unique_ptr mTimes; int32_t mIndexFromEnd{0}; BlockKey mLastBlockKey{}; + bool mEnableInflightCancel{false}; + std::vector mRecvBufferHolders; }; using UniqueToken = tensorrt_llm::runtime::UniqueToken; diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index 812323e92b81..304db228425b 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.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"); @@ -253,7 +253,8 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses return bufferSizeForTarget; }; auto bufferEleSizes = getBufferSizeForTarget(); - auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend(); + auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); + auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend(sendCancelFlag); BufferIndexHolder sendHolder( *mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/false); auto result = mCacheTransBufferManagers[transferIndexerKCache]->getOrAllocateSendBuffers( @@ -339,48 +340,59 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses session.appendMeasure(startTime, endTime, outputSplitCaches.at(cacheIdx)->getSizeInBytes()); }; - if (pickUpConnections.size() > 1) + try { - if (!common::getEnvEnableReceiveKVCacheParallel()) + if (pickUpConnections.size() > 1) { - TLLM_LOG_DEBUG("Disable parallel receiving of the KV cache."); - for (size_t i = 0; i < pickUpConnections.size(); i++) + if (!common::getEnvEnableReceiveKVCacheParallel()) { - sendBufferFun(deviceId, pickUpConnections[i]); + TLLM_LOG_DEBUG("Disable parallel receiving of the KV cache."); + for (size_t i = 0; i < pickUpConnections.size(); i++) + { + sendBufferFun(deviceId, pickUpConnections[i]); + } } - } - else - { - // concurrency num - auto concurrencyNum - = std::min(std::max(static_cast(1), bufferCoverTargetNum), pPDomainSize * cPDomainSize); + else + { + // concurrency num + auto concurrencyNum + = std::min(std::max(static_cast(1), bufferCoverTargetNum), pPDomainSize * cPDomainSize); - auto remainSendNum = pickUpConnections.size(); + auto remainSendNum = pickUpConnections.size(); - while (remainSendNum > 0) - { - auto sendConcurrencyNum = std::min(remainSendNum, concurrencyNum); - std::vector> futures; - futures.reserve(sendConcurrencyNum); - for (size_t i = 0; i < sendConcurrencyNum; i++) - { - size_t idx = i + (pickUpConnections.size() - remainSendNum); - size_t connIdx = pickUpConnections[idx]; - TLLM_CHECK(idx < pickUpConnections.size()); - TLLM_CHECK(connIdx < session.getConnections().size()); - futures.push_back(std::async(std::launch::async, sendBufferFun, deviceId, connIdx)); - } - for (auto& future : futures) + while (remainSendNum > 0) { - future.get(); + auto sendConcurrencyNum = std::min(remainSendNum, concurrencyNum); + std::vector> futures; + futures.reserve(sendConcurrencyNum); + for (size_t i = 0; i < sendConcurrencyNum; i++) + { + size_t idx = i + (pickUpConnections.size() - remainSendNum); + size_t connIdx = pickUpConnections[idx]; + TLLM_CHECK(idx < pickUpConnections.size()); + TLLM_CHECK(connIdx < session.getConnections().size()); + futures.push_back(std::async(std::launch::async, sendBufferFun, deviceId, connIdx)); + } + for (auto& future : futures) + { + future.get(); + } + remainSendNum -= sendConcurrencyNum; } - remainSendNum -= sendConcurrencyNum; } } + else + { + sendBufferFun(deviceId, pickUpConnections[0]); + } } - else + catch (...) { - sendBufferFun(deviceId, pickUpConnections[0]); + if (agentConnection != nullptr && session.isInflightCancelEnabled()) + { + sendHolder.poison(); + } + throw; } sendHolder.release(); } @@ -497,9 +509,9 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s else { cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForRecv(); + recvHolder = BufferIndexHolder( + *mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/true); } - recvHolder - = BufferIndexHolder(*mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/true); auto targetNum = pickUpConnections.size(); diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp index 05742af17a28..af8943d8c055 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.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"); @@ -23,6 +23,7 @@ #include "tensorrt_llm/batch_manager/rnnStateManager.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/dataType.h" +#include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/nvtxUtils.h" #include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h" @@ -184,7 +185,9 @@ void RnnCacheFormatter::formatSlotMode(TransferSession& session) / targetInfo.mDomainTPSize; } - auto cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForSend(); + auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); + auto cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForSend(sendCancelFlag); + BufferIndexHolder sendHolder(*mRnnCacheTransBufferManager, cacheBufferId, /*isRecv=*/false); auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateSendBuffers( cacheBufferId, static_cast(bufferTargetNum), bufferSizesPerTarget, bufferManager); auto& outputBuffers = std::get<0>(allocationResult); @@ -222,12 +225,23 @@ void RnnCacheFormatter::formatSlotMode(TransferSession& session) auto preAllocSendBuffer = mRnnCacheTransBufferManager->getSendBuffer(cacheBufferId); - sendAllBuffers(session, deviceId, outputBuffers, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, - targetInfo, pickUpConnections); + try + { + sendAllBuffers(session, deviceId, outputBuffers, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, + targetInfo, pickUpConnections); + } + catch (...) + { + if (agentConnection != nullptr && session.isInflightCancelEnabled()) + { + sendHolder.poison(); + } + throw; + } session.setTime(TransferSession::kTimeTransmissions); - mRnnCacheTransBufferManager->freeBufferIndexForSend(cacheBufferId); + sendHolder.release(); session.setTime(TransferSession::kTimePostprocess); TLLM_LOG_DEBUG( @@ -348,6 +362,7 @@ void RnnCacheFormatter::unformatSlotMode(TransferSession& session) size_t remainNoCoverSourceNum = 0; size_t bufferCoverSourceNum = 0; std::optional cacheBufferId = std::nullopt; + BufferIndexHolder recvHolder; auto preAssignedRnnId = connections[pickUpConnections[0]]->getPreAssignedBufferId(static_cast(BufferKind::kRNN)); @@ -358,6 +373,7 @@ void RnnCacheFormatter::unformatSlotMode(TransferSession& session) else { cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForRecv(); + recvHolder = BufferIndexHolder(*mRnnCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); } auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateRecvBuffers( @@ -505,10 +521,7 @@ void RnnCacheFormatter::unformatSlotMode(TransferSession& session) bufferManager.getStream().synchronize(); - if (cacheBufferId.has_value()) - { - mRnnCacheTransBufferManager->freeBufferIndexForRecv(cacheBufferId); - } + recvHolder.release(); session.setTime(TransferSession::kTimePostprocess); TLLM_LOG_DEBUG( @@ -643,11 +656,16 @@ void RnnCacheFormatter::formatUnifiedPoolMode(TransferSession& session) bufferSizesPerTarget[t] = ssmBufBytes + convBufBytes; } - auto cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForSend(); + auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); + auto cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForSend(sendCancelFlag); + BufferIndexHolder sendHolder(*mRnnCacheTransBufferManager, cacheBufferId, /*isRecv=*/false); auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateSendBuffers( cacheBufferId, static_cast(bufferTargetNum), bufferSizesPerTarget, bufferManager); auto& outputBuffers = std::get<0>(allocationResult); auto& bufferCoverTargetNum = std::get<1>(allocationResult); + auto const* agentConnection = rnnSendConns.empty() + ? nullptr + : dynamic_cast(connections[rnnSendConns[0]]); // Split each outputBuffer into SSM and conv portions. std::vector ssmOutputBuffers(numTargets); @@ -677,11 +695,22 @@ void RnnCacheFormatter::formatUnifiedPoolMode(TransferSession& session) int deviceId; TLLM_CUDA_CHECK(cudaGetDevice(&deviceId)); auto preAllocSendBuffer = mRnnCacheTransBufferManager->getSendBuffer(cacheBufferId); - sendAllBuffers(session, deviceId, outputBuffers, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, - targetInfo, rnnSendConns); + try + { + sendAllBuffers(session, deviceId, outputBuffers, bufferCoverTargetNum, preAllocSendBuffer, + bufferManager, targetInfo, rnnSendConns); + } + catch (...) + { + if (agentConnection != nullptr && session.isInflightCancelEnabled()) + { + sendHolder.poison(); + } + throw; + } session.setTime(TransferSession::kTimeTransmissions); - mRnnCacheTransBufferManager->freeBufferIndexForSend(cacheBufferId); + sendHolder.release(); } } @@ -821,6 +850,7 @@ void RnnCacheFormatter::unformatUnifiedPoolMode(TransferSession& session) // Use pre-assigned buffer ID from NIXL connection if available (same as slot mode). std::optional cacheBufferId = std::nullopt; + BufferIndexHolder recvHolder; auto preAssignedRnnId = connections[rnnRecvConns[0]]->getPreAssignedBufferId(static_cast(BufferKind::kRNN)); if (preAssignedRnnId.has_value()) @@ -830,6 +860,7 @@ void RnnCacheFormatter::unformatUnifiedPoolMode(TransferSession& session) else { cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForRecv(); + recvHolder = BufferIndexHolder(*mRnnCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); } auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateRecvBuffers( @@ -876,7 +907,7 @@ void RnnCacheFormatter::unformatUnifiedPoolMode(TransferSession& session) bufferManager.getStream().synchronize(); - mRnnCacheTransBufferManager->freeBufferIndexForRecv(cacheBufferId); + recvHolder.release(); } } diff --git a/cpp/tensorrt_llm/common/envUtils.cpp b/cpp/tensorrt_llm/common/envUtils.cpp index 8987e5048c90..3a48bb602b48 100644 --- a/cpp/tensorrt_llm/common/envUtils.cpp +++ b/cpp/tensorrt_llm/common/envUtils.cpp @@ -395,6 +395,12 @@ bool getEnvTryZCopyForKVCacheTransfer() return zcopyForSysmmetricKVCache; } +bool getEnvDisaggEnableInflightCancel() +{ + static bool const enabled = getBoolEnv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL"); + return enabled; +} + bool getEnvForceDeterministic() { static bool const forceDeterministic = getBoolEnv("FORCE_DETERMINISTIC"); diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index d8e72b7b02cc..b1950194321c 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -110,6 +110,9 @@ std::string const& getEnvKVCacheTimeOutputPath(); bool getEnvTryZCopyForKVCacheTransfer(); +// Opt-in for disaggregated KV transfer in-flight cancellation and fail-closed transfer-buffer quarantine. +bool getEnvDisaggEnableInflightCancel(); + // Force deterministic behavior for all kernels. bool getEnvForceDeterministic(); diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp index a21d470b898a..487038f1ed68 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp @@ -189,8 +189,35 @@ void AgentConnection::send(DataContext const& ctx, void const* data, size_t size NotificationInfo notificationInfo{syncInfo}; std::stringstream ss; NotificationInfo::serialize(notificationInfo, ss); - TransferState transferState = status->wait(); + static constexpr int64_t kCancelPollTimeoutMs = 100; + TransferState transferState = TransferState::kIN_PROGRESS; + while (transferState == TransferState::kIN_PROGRESS) + { + transferState = status->wait(kCancelPollTimeoutMs); + if (transferState == TransferState::kIN_PROGRESS && ctx.getTransferTerminate().load(std::memory_order_relaxed)) + { + bool const released = status->release(); + TLLM_LOG_WARNING( + "AgentConnection::send cancelled while transfer was in progress (ctx tag=%d, remote=%s, " + "releaseAccepted=%d)", + ctx.getTag(), mRemoteAgentName.c_str(), released); + TLLM_CHECK_WITH_INFO( + released, "AgentConnection::send cancel could not release the backend transfer handle"); + TLLM_THROW("AgentConnection::send cancelled mid-transfer"); + } + } TLLM_CHECK_WITH_INFO(transferState == TransferState::kSUCCESS, "AgentConnection::send failed"); + if (ctx.getTransferTerminate().load(std::memory_order_relaxed)) + { + bool const released = status->release(); + TLLM_LOG_WARNING( + "AgentConnection::send cancelled after transfer completed but before notify (ctx tag=%d, remote=%s, " + "releaseAccepted=%d)", + ctx.getTag(), mRemoteAgentName.c_str(), released); + TLLM_CHECK_WITH_INFO( + released, "AgentConnection::send pre-notify cancel could not release the backend transfer handle"); + TLLM_THROW("AgentConnection::send cancelled pre-notify"); + } // TODO: there is a bug in request_with_notify https://github.com/ai-dynamo/nixl/pull/252 mAgentConnectionManager->getAgent()->notifySyncMessage(mRemoteAgentName, ss.str()); } @@ -199,11 +226,16 @@ void AgentConnection::recv(DataContext const& ctx, void* data, size_t size) cons { NotificationSyncInfo syncInfo{mAgentName, ctx}; - mAgentConnectionManager->waitForSyncInfo(mRemoteAgentName, syncInfo, ctx.getTransferTerminate()); + bool const received + = mAgentConnectionManager->waitForSyncInfo(mRemoteAgentName, syncInfo, ctx.getTransferTerminate()); + TLLM_CHECK_WITH_INFO(received, + "AgentConnection::recv ended before receiving sync notification (ctx tag=%d, remote=%s)", ctx.getTag(), + mRemoteAgentName.c_str()); } void AgentConnection::sendRequestAndBufferInfo(batch_manager::RequestInfo& requestInfo, - std::vector> const& cacheBufferIds, int connectionIdx) + std::vector> const& cacheBufferIds, int connectionIdx, + std::atomic const* perRequestCancel, bool* advertisementMayHaveOccurred) { TLLM_CHECK(!common::getEnvTryZCopyForKVCacheTransfer()); @@ -247,7 +279,6 @@ void AgentConnection::sendRequestAndBufferInfo(batch_manager::RequestInfo& reque { auto metadata = mAgentConnectionManager->getAgent()->getLocalAgentDesc().getBackendAgentDesc(); metadataOpt = metadata; - mNeedSendMetadata = false; } RequestAndBufferInfo requestAndBufferInfo{ @@ -255,7 +286,20 @@ void AgentConnection::sendRequestAndBufferInfo(batch_manager::RequestInfo& reque std::stringstream ss; NotificationInfo notificationInfo{requestAndBufferInfo}; NotificationInfo::serialize(notificationInfo, ss); + if (perRequestCancel != nullptr && perRequestCancel->load(std::memory_order_relaxed)) + { + TLLM_THROW("sendRequestAndBufferInfo cancelled before notify"); + } + if (advertisementMayHaveOccurred != nullptr) + { + // Once notification begins, an exception cannot prove that the peer did not receive the advertised buffers. + *advertisementMayHaveOccurred = true; + } mAgentConnectionManager->getAgent()->notifySyncMessage(mRemoteAgentName, ss.str()); + if (metadataOpt.has_value()) + { + mNeedSendMetadata = false; + } } void AgentConnection::setSenderState(std::vector cacheReceiverBufferDescs, int validSegmentIdx, @@ -291,9 +335,17 @@ void AgentConnection::sendReadySignal(DataContext const& ctx, bool isReady) cons } bool AgentConnection::recvReadySignal(DataContext const& ctx) const +{ + return recvReadySignalWithStatus(ctx).value_or(false); +} + +std::optional AgentConnection::recvReadySignalWithStatus(DataContext const& ctx) const { ReadySignalInfo readySignalInfo{mAgentName, ctx, false}; - mAgentConnectionManager->waitForReadySignal(mRemoteAgentName, readySignalInfo, ctx.getTransferTerminate()); + if (!mAgentConnectionManager->waitForReadySignal(mRemoteAgentName, readySignalInfo, ctx.getTransferTerminate())) + { + return std::nullopt; + } return readySignalInfo.mIsReady; } @@ -692,7 +744,7 @@ int AgentConnectionManager::getDeviceId() const } template -void AgentConnectionManager::waitForNotification( +bool AgentConnectionManager::waitForNotification( std::string const& remoteAgentName, NotificationType& expectedInfo, std::atomic const& terminateFlag) { while (!terminateFlag.load()) @@ -700,7 +752,7 @@ void AgentConnectionManager::waitForNotification( if (!mIsRunning) { - return; + return false; } updateUnhandledNotifications(); std::scoped_lock lock(mNotificationMutex); @@ -733,7 +785,7 @@ void AgentConnectionManager::waitForNotification( { it = mUnhandledNotifications.erase(it); } - return; + return true; } } } @@ -753,7 +805,7 @@ void AgentConnectionManager::waitForNotification( { it = mUnhandledNotifications.erase(it); } - return; + return true; } } } @@ -773,24 +825,25 @@ void AgentConnectionManager::waitForNotification( } } } + return false; } // Explicit template instantiations -template void AgentConnectionManager::waitForNotification( +template bool AgentConnectionManager::waitForNotification( std::string const& remoteAgentName, NotificationSyncInfo& expectedInfo, std::atomic const& terminateFlag); -template void AgentConnectionManager::waitForNotification( +template bool AgentConnectionManager::waitForNotification( std::string const& remoteAgentName, ReadySignalInfo& expectedInfo, std::atomic const& terminateFlag); -void AgentConnectionManager::waitForSyncInfo( +bool AgentConnectionManager::waitForSyncInfo( std::string const& remoteAgentName, NotificationSyncInfo& syncInfo, std::atomic const& terminateFlag) { - waitForNotification(remoteAgentName, syncInfo, terminateFlag); + return waitForNotification(remoteAgentName, syncInfo, terminateFlag); } -void AgentConnectionManager::waitForReadySignal( +bool AgentConnectionManager::waitForReadySignal( std::string const& remoteAgentName, ReadySignalInfo& readySignalInfo, std::atomic const& terminateFlag) { - waitForNotification(remoteAgentName, readySignalInfo, terminateFlag); + return waitForNotification(remoteAgentName, readySignalInfo, terminateFlag); } std::string const& AgentConnectionManager::getAgentName() const diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h index 25283d1341a1..06cf443ffd91 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h @@ -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"); @@ -259,14 +259,17 @@ class AgentConnection : public Connection std::string mAgentName, std::string mRemoteAgentName, AgentConnectionManager* mAgentConnectionManager); void send(DataContext const& ctx, void const* data, size_t size) const override; void recv(DataContext const& ctx, void* data, size_t size) const override; + //! Sets advertisementMayHaveOccurred immediately before notifying the peer, after all definite pre-notify exits. void sendRequestAndBufferInfo(batch_manager::RequestInfo& requestInfo, - std::vector> const& cacheBufferIds, int validConnectionIdx); + std::vector> const& cacheBufferIds, int validConnectionIdx, + std::atomic const* perRequestCancel = nullptr, bool* advertisementMayHaveOccurred = nullptr); void setSenderState(std::vector cacheReceiverBufferDescs, int valideSegmentIdx, std::vector> offsetRatios, std::vector bufferKinds); void setHasLoadRemoteAgent(bool hasLoadRemoteAgent); [[nodiscard]] bool hasLoadRemoteAgent() const; void sendReadySignal(DataContext const& ctx, bool isReady) const; bool recvReadySignal(DataContext const& ctx) const; + std::optional recvReadySignalWithStatus(DataContext const& ctx) const; void activateBuffer(uint8_t kind) const override; [[nodiscard]] std::optional getPreAssignedBufferId(uint8_t kind) const override; @@ -320,11 +323,11 @@ class AgentConnectionManager : public ConnectionManager [[nodiscard]] std::string const& getAgentName() const; template - void waitForNotification( + bool waitForNotification( std::string const& remoteAgentName, NotificationType& expectedInfo, std::atomic const& terminateFlag); - void waitForSyncInfo( + bool waitForSyncInfo( std::string const& remoteAgentName, NotificationSyncInfo& syncInfo, std::atomic const& terminateFlag); - void waitForReadySignal( + bool waitForReadySignal( std::string const& remoteAgentName, ReadySignalInfo& readySignalInfo, std::atomic const& terminateFlag); [[nodiscard]] bool isRunning() const override; diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp index 563b6e6f2a95..64f7a17e89cd 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.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"); @@ -328,19 +328,13 @@ NixlTransferStatus::NixlTransferStatus(std::weak_ptr agent, nixlXferR NixlTransferStatus::~NixlTransferStatus() noexcept { - if (mHandle == nullptr) - { - return; - } - // Skip release if the owning agent was reset; the underlying nixlXferReqH is already gone. - auto agent = mWeakAgent.lock(); - if (!agent) - { - return; - } try { - agent->releaseXferReq(mHandle); + if (!release()) + { + TLLM_LOG_WARNING( + "NIXL transfer handle release failed during destruction; backend handle may remain active"); + } } catch (std::exception const& e) { @@ -516,15 +510,24 @@ TransferState NixlTransferStatus::wait(int64_t timeout_ms) const while (true) { - auto agent = mWeakAgent.lock(); - if (!agent) + nixl_status_t status; { - // Owning agent was reset; report failure so callers don't deref a null status. - mLastStatus.store(static_cast(NIXL_ERR_INVALID_PARAM), std::memory_order_relaxed); - return TransferState::kFAILURE; + std::lock_guard lock(mHandleMutex); + if (mHandle == nullptr) + { + mLastStatus.store(static_cast(NIXL_ERR_INVALID_PARAM), std::memory_order_relaxed); + return TransferState::kFAILURE; + } + auto agent = mWeakAgent.lock(); + if (!agent) + { + // Owning agent was reset; report failure so callers don't deref a null status. + mLastStatus.store(static_cast(NIXL_ERR_INVALID_PARAM), std::memory_order_relaxed); + return TransferState::kFAILURE; + } + status = agent->getXferStatus(mHandle); + mLastStatus.store(static_cast(status), std::memory_order_relaxed); } - auto status = agent->getXferStatus(mHandle); - mLastStatus.store(static_cast(status), std::memory_order_relaxed); if (status == NIXL_SUCCESS) { return TransferState::kSUCCESS; @@ -566,6 +569,12 @@ std::string NixlTransferStatus::getLastStatusStr() const [[nodiscard]] bool NixlTransferStatus::isCompleted() const { + std::lock_guard lock(mHandleMutex); + if (mHandle == nullptr) + { + mLastStatus.store(static_cast(NIXL_ERR_INVALID_PARAM), std::memory_order_relaxed); + return false; + } auto agent = mWeakAgent.lock(); if (!agent) { @@ -577,6 +586,34 @@ std::string NixlTransferStatus::getLastStatusStr() const return status == NIXL_SUCCESS; } +[[nodiscard]] bool NixlTransferStatus::release() +{ + std::lock_guard lock(mHandleMutex); + if (mHandle == nullptr) + { + return true; + } + + auto agent = mWeakAgent.lock(); + if (!agent) + { + mHandle = nullptr; + mLastStatus.store(static_cast(NIXL_ERR_INVALID_PARAM), std::memory_order_relaxed); + return true; + } + + auto status = agent->releaseXferReq(mHandle); + mLastStatus.store(static_cast(status), std::memory_order_relaxed); + if (status == NIXL_SUCCESS) + { + mHandle = nullptr; + return true; + } + + TLLM_LOG_WARNING("NIXL releaseXferReq failed with status: %s", nixlEnumStrings::statusStr(status).c_str()); + return false; +} + NixlTransferAgent::NixlTransferAgent(BaseAgentConfig const& config) : mName{config.mName} { diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h index a8f3a0e71005..99ab21a32d7f 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h @@ -21,6 +21,7 @@ #include "tensorrt_llm/executor/transferAgent.h" #include #include +#include #include #include @@ -73,11 +74,14 @@ class NixlTransferStatus final : public TransferStatus [[nodiscard]] int getLastStatus() const noexcept; [[nodiscard]] std::string getLastStatusStr() const; + [[nodiscard]] bool release() override; + private: // weak_ptr so the status outliving the owning agent is safe (lock() returns null after reset). std::weak_ptr mWeakAgent; nixlXferReqH* mHandle{}; mutable std::atomic mLastStatus{0}; + mutable std::mutex mHandleMutex; }; class NixlTransferAgent final : public BaseTransferAgent diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index 5257c7adb1c9..828532a87684 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -110,7 +110,8 @@ void tb::CacheTransceiverBindings::initBindings(nb::module_& m) .def("check_gen_transfer_status", &BaseCacheTransceiver::checkGenTransferStatus, nb::call_guard()) .def("check_gen_transfer_complete", &BaseCacheTransceiver::checkGenTransferComplete) - .def("cancel_request", &BaseCacheTransceiver::cancelRequest); + .def("cancel_request", &BaseCacheTransceiver::cancelRequest) + .def("has_poisoned_transfer_buffer", &BaseCacheTransceiver::hasPoisonedTransferBuffer); nb::enum_(m, "AttentionType") .value("DEFAULT", executor::kv_cache::CacheState::AttentionType::kDEFAULT) @@ -120,11 +121,12 @@ void tb::CacheTransceiverBindings::initBindings(nb::module_& m) .def(nb::init, SizeType32, SizeType32, runtime::WorldConfig, std::vector, nvinfer1::DataType, executor::kv_cache::CacheState::AttentionType, std::optional, - tb::rnn_state_manager::RnnStateManager*, std::vector>(), + tb::rnn_state_manager::RnnStateManager*, std::vector, bool>(), nb::arg("cache_manager"), nb::arg("num_kv_heads_per_layer"), nb::arg("size_per_head"), nb::arg("tokens_per_block"), nb::arg("world_config"), nb::arg("attention_layer_num_per_pp"), nb::arg("dtype"), nb::arg("attention_type"), nb::arg("cache_transceiver_config") = std::nullopt, - nb::arg("rnn_state_manager") = nullptr, nb::arg("rnn_layer_num_per_pp") = std::vector{}); + nb::arg("rnn_state_manager") = nullptr, nb::arg("rnn_layer_num_per_pp") = std::vector{}, + nb::arg("enable_inflight_cancel") = false); nb::class_(m, "CacheTransceiverComm") .def( diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index f4e5e9fb2be0..b6b2672b2e0d 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -17,6 +17,7 @@ add_gtest(radixTreeTest radixTreeTest.cpp) add_gtest(blockKeyTest blockKeyTest.cpp) add_gtest(radixBlockTreeTest radixBlockTreeTest.cpp) add_gtest(cacheTransBufferTest cacheTransBufferTest.cpp) +add_gtest(cacheTransferStateTest cacheTransferStateTest.cpp) add_gtest(bufferIndexHolderTest bufferIndexHolderTest.cpp) add_gtest(capacitySchedulerTest capacitySchedulerTest.cpp) add_gtest(contextProgressTest contextProgressTest.cu) diff --git a/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp b/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp index d4d561f55d07..acb43b190f5c 100644 --- a/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp @@ -17,8 +17,10 @@ #include "tensorrt_llm/batch_manager/baseTransBuffer.h" #include "tensorrt_llm/batch_manager/cacheTransBuffer.h" +#include "tensorrt_llm/batch_manager/dataTransceiver.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include +#include #include #include #include @@ -80,6 +82,7 @@ class BufferIndexHolderLifecycleTest : public ::testing::TestWithParam(); + mBufferManager = std::make_unique(stream); auto const kvMaxNumTokens = tokensPerBlock * maxBlocksPerSeq; auto const totalNumBlocks = maxNumSequences * maxBlocksPerSeq; using BlocksPerWindow = std::map>; @@ -119,8 +122,19 @@ class BufferIndexHolderLifecycleTest : public ::testing::TestWithParam mKv; std::unique_ptr mTrans; + std::unique_ptr mBufferManager; + tensorrt_llm::executor::DataTransceiverState mSelfState; + std::atomic mCancel{false}; }; // Default-constructed holder owns nothing; destruction is a no-op. @@ -146,17 +160,31 @@ TEST_P(BufferIndexHolderLifecycleTest, DefaultConstructedExplicitReleaseIsNoOp) EXPECT_EQ(inUse(), before); } -// Holder built from a nullopt index is disarmed (mHeld == false). +// A nullopt index represents a held dynamic-buffer marker. Releasing it does +// not affect the fixed-slot concurrence counter. TEST_P(BufferIndexHolderLifecycleTest, NulloptIndexReleasesNothing) { int const before = inUse(); { BufferIndexHolder holder{mgr(), std::nullopt, isRecv()}; - EXPECT_FALSE(holder.held()); + EXPECT_TRUE(holder.held()); + EXPECT_FALSE(holder.index().has_value()); } EXPECT_EQ(inUse(), before); } +// Poisoning a nullopt dynamic-buffer marker records the fail-closed state +// without changing the fixed-slot concurrence counter. +TEST_P(BufferIndexHolderLifecycleTest, NulloptIndexPoisonMarksManager) +{ + int const before = inUse(); + BufferIndexHolder holder{mgr(), std::nullopt, isRecv()}; + holder.poison(); + EXPECT_FALSE(holder.held()); + EXPECT_TRUE(mgr().hasPoisonedBuffer()); + EXPECT_EQ(inUse(), before); +} + // RAII: a held slot is released when the holder goes out of scope. TEST_P(BufferIndexHolderLifecycleTest, ValidIndexReleasedOnDestruction) { @@ -286,6 +314,67 @@ TEST_P(BufferIndexHolderLifecycleTest, ExceptionUnwindStillReleases) EXPECT_EQ(inUse(), before); } +TEST_P(BufferIndexHolderLifecycleTest, TransferSessionExplicitReleaseFreesAdoptedRecvSlot) +{ + if (!isRecv()) + { + GTEST_SKIP() << "TransferSession owns receive slots only"; + } + + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + auto session = makeTransferSession(); + std::vector holders; + holders.emplace_back(mgr(), idx, /*isRecv=*/true); + session.adoptRecvBufferHolders(std::move(holders)); + + EXPECT_EQ(inUse(), before + 1); + session.releaseRecvBufferHolders(); + EXPECT_EQ(inUse(), before); + EXPECT_FALSE(mgr().hasPoisonedBuffer()); +} + +TEST_P(BufferIndexHolderLifecycleTest, TransferSessionExplicitPoisonQuarantinesAdoptedRecvSlot) +{ + if (!isRecv()) + { + GTEST_SKIP() << "TransferSession owns receive slots only"; + } + + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + auto session = makeTransferSession(); + std::vector holders; + holders.emplace_back(mgr(), idx, /*isRecv=*/true); + session.adoptRecvBufferHolders(std::move(holders)); + + session.poisonRecvBufferHolders(); + EXPECT_EQ(inUse(), before + 1); + EXPECT_TRUE(mgr().hasPoisonedBuffer()); +} + +TEST_P(BufferIndexHolderLifecycleTest, AbandonedTransferSessionQuarantinesAdvertisedRecvSlot) +{ + if (!isRecv()) + { + GTEST_SKIP() << "TransferSession owns receive slots only"; + } + + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + { + auto session = makeTransferSession(); + std::vector holders; + holders.emplace_back(mgr(), idx, /*isRecv=*/true); + session.adoptRecvBufferHolders(std::move(holders)); + } + EXPECT_EQ(inUse(), before + 1); + EXPECT_TRUE(mgr().hasPoisonedBuffer()); +} + INSTANTIATE_TEST_SUITE_P(SideVariants, BufferIndexHolderLifecycleTest, ::testing::Values(HolderCase{"send", Side::Send}, HolderCase{"recv", Side::Recv}), [](::testing::TestParamInfo const& info) { return info.param.name; }); diff --git a/cpp/tests/unit_tests/batch_manager/cacheTransferStateTest.cpp b/cpp/tests/unit_tests/batch_manager/cacheTransferStateTest.cpp new file mode 100644 index 000000000000..a782b0ffe8e5 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/cacheTransferStateTest.cpp @@ -0,0 +1,102 @@ +/* + * 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/cacheTransferState.h" + +#include + +#include +#include + +namespace tensorrt_llm::batch_manager::detail +{ +namespace +{ + +TEST(CacheTransferStateTest, TerminalMarkerRemainsUsefulAfterFutureIsConsumed) +{ + auto const firstPoll = summarizePackedTransferStates({{0, 17, kTransferCompleted}, {0}}); + EXPECT_TRUE(firstPoll.completedRequestIds.empty()); + EXPECT_TRUE(firstPoll.quiescedRequestIds.empty()); + + auto const secondPoll = summarizePackedTransferStates({{0, 17, kTransferCompleted}, {0, 17, kTransferCompleted}}); + EXPECT_EQ(secondPoll.completedRequestIds, std::unordered_set{17}); + EXPECT_EQ(secondPoll.quiescedRequestIds, std::unordered_set{17}); +} + +TEST(CacheTransferStateTest, FailureUnionPrecedesQuiescenceIntersection) +{ + auto const firstPoll = summarizePackedTransferStates({{0, 29, kTransferFailed}, {0}}); + EXPECT_EQ(firstPoll.failedRequestIds, std::unordered_set{29}); + EXPECT_TRUE(firstPoll.quiescedRequestIds.empty()); + + auto const secondPoll = summarizePackedTransferStates({{0, 29, kTransferFailed}, {0, 29, kTransferCompleted}}); + EXPECT_EQ(secondPoll.failedRequestIds, std::unordered_set{29}); + EXPECT_EQ(secondPoll.quiescedRequestIds, std::unordered_set{29}); + EXPECT_TRUE(secondPoll.completedRequestIds.empty()); +} + +TEST(CacheTransferStateTest, TimeoutVoteWinsOverLateSuccess) +{ + auto const outcome + = summarizePackedTransferStates({{0, 41, kTransferCompleted | kTransferTimedOut}, {0, 41, kTransferCompleted}}); + EXPECT_EQ(outcome.timedOutRequestIds, std::unordered_set{41}); + EXPECT_EQ(outcome.failedRequestIds, std::unordered_set{41}); + EXPECT_EQ(outcome.quiescedRequestIds, std::unordered_set{41}); + EXPECT_TRUE(outcome.completedRequestIds.empty()); +} + +TEST(CacheTransferStateTest, LargeRequestIdIsPreserved) +{ + auto const requestId = std::numeric_limits::max(); + auto const outcome + = summarizePackedTransferStates({{0, requestId, kTransferCompleted}, {0, requestId, kTransferCompleted}}); + EXPECT_EQ(outcome.completedRequestIds, std::unordered_set{requestId}); + EXPECT_EQ(outcome.quiescedRequestIds, std::unordered_set{requestId}); +} + +TEST(CacheTransferStateTest, MalformedPayloadIsRejected) +{ + EXPECT_THROW((void) summarizePackedTransferStates({{}}), std::exception); + EXPECT_THROW((void) summarizePackedTransferStates({{0, 17}}), std::exception); +} + +TEST(CacheTransferStateTest, DuplicateRequestIdOnOneRankIsRejected) +{ + EXPECT_THROW( + (void) summarizePackedTransferStates({{0, 17, kTransferCompleted, 17, kTransferFailed}}), std::exception); +} + +TEST(CacheTransferStateTest, PoisonIsTopologyWideIndependentOfRequestState) +{ + auto const outcome = summarizePackedTransferStates({{0}, {kTransferBufferPoisoned}}); + EXPECT_TRUE(outcome.transferBufferPoisoned); + EXPECT_TRUE(outcome.failedRequestIds.empty()); + EXPECT_TRUE(outcome.quiescedRequestIds.empty()); +} + +TEST(CacheTransferStateTest, ReadySignalsDistinguishNoneAllAndPartial) +{ + EXPECT_EQ(summarizeReadySignals(/*allPeersReady=*/true, /*anyPeerReady=*/true), ReadySignalSummary::kReady); + EXPECT_EQ(summarizeReadySignals(/*allPeersReady=*/false, /*anyPeerReady=*/false), ReadySignalSummary::kNotReady); + EXPECT_EQ( + summarizeReadySignals(/*allPeersReady=*/false, /*anyPeerReady=*/true), ReadySignalSummary::kPartiallyReady); + EXPECT_THROW((void) summarizeReadySignals(/*allPeersReady=*/true, /*anyPeerReady=*/false), std::exception); +} + +} // namespace +} // namespace tensorrt_llm::batch_manager::detail diff --git a/cpp/tests/unit_tests/executor/agentCommTest.cpp b/cpp/tests/unit_tests/executor/agentCommTest.cpp index 194d5267c6b3..252468447707 100644 --- a/cpp/tests/unit_tests/executor/agentCommTest.cpp +++ b/cpp/tests/unit_tests/executor/agentCommTest.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * 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"); @@ -198,7 +198,18 @@ TEST_P(AgentCommTest, AgentConnectionManagerConnect) // convert to AgentConnection auto agentConnection0 = const_cast( dynamic_cast(connection0)); - agentConnection0->sendRequestAndBufferInfo(sendRequestInfo, cacheBufferIds, validConnectionIdx); + // A pre-notify cancellation must not consume the one-time metadata. The + // retry still needs it so the sender can load the remote agent. + std::atomic perRequestCancel{true}; + bool advertisementMayHaveOccurred = false; + EXPECT_THROW(agentConnection0->sendRequestAndBufferInfo(sendRequestInfo, cacheBufferIds, validConnectionIdx, + &perRequestCancel, &advertisementMayHaveOccurred), + std::exception); + EXPECT_FALSE(advertisementMayHaveOccurred); + perRequestCancel.store(false, std::memory_order_relaxed); + agentConnection0->sendRequestAndBufferInfo( + sendRequestInfo, cacheBufferIds, validConnectionIdx, &perRequestCancel, &advertisementMayHaveOccurred); + EXPECT_TRUE(advertisementMayHaveOccurred); tensorrt_llm::batch_manager::RequestInfo recvRequestInfo; auto connection1 = connectionManager1->recvConnectionAndRequestInfo(recvRequestInfo, std::atomic(false)); diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index b272740deab4..561aa4169130 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -1282,6 +1282,7 @@ def create_autodeploy_executor( attention_type_cpp, cache_transceiver_config, mamba_cache_manager=None, + inflight_cancel_supported_by_executor=False, ) # Guided (structured) decoding. diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index ac0feab0347d..6a27dda1aac1 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1,3 +1,18 @@ +# 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. + import copy import dataclasses import os @@ -2193,8 +2208,13 @@ def create_py_executor_instance( mamba_cache_manager = kv_cache_manager kv_cache_transceiver = create_kv_cache_transceiver( - mapping, dist, kv_cache_manager, attention_type, - cache_transceiver_config, mamba_cache_manager) + mapping, + dist, + kv_cache_manager, + attention_type, + cache_transceiver_config, + mamba_cache_manager, + inflight_cancel_supported_by_executor=kv_connector_manager is None) waiting_queue_policy = (scheduler_config.waiting_queue_policy if scheduler_config is not None else diff --git a/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py b/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py index 5d1f19b2c6af..8d3ea4524253 100644 --- a/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py +++ b/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py @@ -1,3 +1,18 @@ +# 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. + import dataclasses import datetime import queue @@ -131,12 +146,15 @@ def enqueue_cancel_request(self, req_id: int): self.request_queue.put( RequestQueueItem(req_id, is_canceled_request=True)) - def enqueue_control_request(self, drain: bool = True): + def enqueue_control_request(self, drain: bool = True) -> bool: """Enqueue a control-action sentinel. See ``PyExecutor.control_action``.""" with self.enqueue_lock: + if not self.active: + return False self.request_queue.put( RequestQueueItem(id=CONTROL_REQUEST_ID, control_requires_drain=drain)) + return True def enqueue_shutdown_request(self): with self.enqueue_lock: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index 8e6b47c8bff8..7cda19640e17 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -1,3 +1,18 @@ +# 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. + from abc import ABC, abstractmethod from os import getenv from typing import Any, Dict, List, Optional @@ -19,6 +34,101 @@ CacheTransBufferManagerCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransBufferManager BackendTypeCpp = tensorrt_llm.bindings.executor.CacheTransceiverBackendType +_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV = "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL" +_NIXL_KVCACHE_BACKEND_ENV = "TRTLLM_NIXL_KVCACHE_BACKEND" +_DISABLE_KV_CACHE_TRANSFER_OVERLAP_ENV = "TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP" +_DISAGG_LAYERWISE_ENV = "TRTLLM_DISAGG_LAYERWISE" +_TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV = "TRTLLM_TRY_ZCOPY_FOR_KVCACHE_TRANSFER" +_SUPPORTED_INFLIGHT_CANCEL_NIXL_BACKEND = "UCX" +_CACHE_TRANSCEIVER_BACKEND_ENV_VARS = ( + ("TRTLLM_USE_NIXL_KVCACHE", "NIXL"), + ("TRTLLM_USE_UCX_KVCACHE", "UCX"), + ("TRTLLM_USE_MOONCAKE_KVCACHE", "MOONCAKE"), + ("TRTLLM_USE_MPI_KVCACHE", "MPI"), +) +_disagg_inflight_cancel_enabled_cache: Optional[bool] = None + + +def is_disagg_inflight_cancel_enabled() -> bool: + """Return whether disaggregated in-flight KV transfer cancellation is enabled.""" + global _disagg_inflight_cancel_enabled_cache + if _disagg_inflight_cancel_enabled_cache is None: + _disagg_inflight_cancel_enabled_cache = (getenv( + _DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "0") == "1") + if _disagg_inflight_cancel_enabled_cache: + logger.warning( + f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV}=1: disagg KV " + "transfer in-flight cancellation was requested. It is active " + "only for cache transceivers that advertise support.") + return _disagg_inflight_cancel_enabled_cache + + +def _is_disagg_inflight_cancel_config_supported( + cache_transceiver_config: CacheTransceiverConfig, + mapping: Mapping, + inflight_cancel_supported_by_executor: bool = False) -> bool: + runtime = cache_transceiver_config.transceiver_runtime or "CPP" + nixl_backend = getenv(_NIXL_KVCACHE_BACKEND_ENV, "UCX") + return (inflight_cancel_supported_by_executor and runtime == "CPP" + and cache_transceiver_config.backend == "NIXL" + and nixl_backend == _SUPPORTED_INFLIGHT_CANCEL_NIXL_BACKEND + and (cache_transceiver_config.max_tokens_in_buffer or 0) > 0 + and (cache_transceiver_config.kv_transfer_timeout_ms or 0) > 0 + and getenv(_DISABLE_KV_CACHE_TRANSFER_OVERLAP_ENV) != "1" + and mapping.pp_size == 1 and mapping.cp_size == 1 + and not mapping.enable_attention_dp + and getenv(_DISAGG_LAYERWISE_ENV) != "1" + and getenv(_TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV) != "1") + + +def _validate_disagg_inflight_cancel_config( + cache_transceiver_config: CacheTransceiverConfig, + mapping: Mapping, + inflight_cancel_supported_by_executor: bool = False) -> None: + if not is_disagg_inflight_cancel_enabled(): + return + + enabled_backend_env_vars = [ + env_name for env_name, _ in _CACHE_TRANSCEIVER_BACKEND_ENV_VARS + if getenv(env_name) == "1" + ] + if len(enabled_backend_env_vars) > 1: + raise ValueError( + f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV}=1 requires an " + "unambiguous cache transceiver backend, but multiple legacy " + f"backend selectors are enabled: {enabled_backend_env_vars}.") + + if _is_disagg_inflight_cancel_config_supported( + cache_transceiver_config, mapping, + inflight_cancel_supported_by_executor): + return + + runtime = cache_transceiver_config.transceiver_runtime or "CPP" + nixl_backend = getenv(_NIXL_KVCACHE_BACKEND_ENV, "UCX") + disable_overlap = getenv(_DISABLE_KV_CACHE_TRANSFER_OVERLAP_ENV) + layerwise = getenv(_DISAGG_LAYERWISE_ENV) + try_zcopy = getenv(_TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV) + raise ValueError( + f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV}=1 is experimental and " + "currently supported only by the standard PyTorch backend " + "PyExecutor with asynchronous transceiver_runtime='CPP', " + "backend='NIXL', and TRTLLM_NIXL_KVCACHE_BACKEND=UCX on ordinary " + "TP topology, using a positive max_tokens_in_buffer admission " + "budget, a transfer deadline, and preallocated staging buffers " + "without layerwise or zero-copy transfer; got " + f"transceiver_runtime={runtime!r}, " + f"backend={cache_transceiver_config.backend!r}, " + f"max_tokens_in_buffer={cache_transceiver_config.max_tokens_in_buffer!r}, " + f"kv_transfer_timeout_ms={cache_transceiver_config.kv_transfer_timeout_ms!r}, " + f"TRTLLM_NIXL_KVCACHE_BACKEND={nixl_backend!r}, " + f"pp_size={mapping.pp_size!r}, cp_size={mapping.cp_size!r}, " + f"enable_attention_dp={mapping.enable_attention_dp!r}, " + "inflight_cancel_supported_by_executor=" + f"{inflight_cancel_supported_by_executor!r}, " + f"{_DISABLE_KV_CACHE_TRANSFER_OVERLAP_ENV}={disable_overlap!r}, " + f"{_DISAGG_LAYERWISE_ENV}={layerwise!r}, " + f"{_TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV}={try_zcopy!r}.") + def mapping_to_world_config(mapping: Mapping) -> WorldConfig: @@ -37,7 +147,8 @@ def create_kv_cache_transceiver( kv_cache_manager: KVCacheManager, attention_type: AttentionTypeCpp, cache_transceiver_config: CacheTransceiverConfig, - mamba_cache_manager: Optional[BaseMambaCacheManager] = None): + mamba_cache_manager: Optional[BaseMambaCacheManager] = None, + inflight_cancel_supported_by_executor: bool = False): if cache_transceiver_config is None or cache_transceiver_config.backend is None: logger.info("cache_transceiver is disabled") return None @@ -47,13 +158,7 @@ def create_kv_cache_transceiver( # NIXL is the default backend for non hybrid models cache_transceiver_config.backend = "NIXL" # Ordered by priority - env_vars = [ - ("TRTLLM_USE_NIXL_KVCACHE", "NIXL"), - ("TRTLLM_USE_UCX_KVCACHE", "UCX"), - ("TRTLLM_USE_MOONCAKE_KVCACHE", "MOONCAKE"), - ("TRTLLM_USE_MPI_KVCACHE", "MPI"), - ] - for env_var, be_type in env_vars: + for env_var, be_type in _CACHE_TRANSCEIVER_BACKEND_ENV_VARS: if getenv(env_var) == "1": logger.warning( f"{env_var}=1 is set, but it's recommended to set cache_transceiver_config.backend in yaml config" @@ -61,6 +166,10 @@ def create_kv_cache_transceiver( cache_transceiver_config.backend = be_type break + _validate_disagg_inflight_cancel_config( + cache_transceiver_config, mapping, + inflight_cancel_supported_by_executor) + if cache_transceiver_config.backend == "MPI": logger.warning( "MPI CacheTransceiver is deprecated, UCX or NIXL is recommended") @@ -90,7 +199,8 @@ def create_kv_cache_transceiver( # Default: use C++ transceiver (transceiver_runtime is None or "CPP") return BindKvCacheTransceiver(mapping, dist, kv_cache_manager, attention_type, cache_transceiver_config, - mamba_cache_manager) + mamba_cache_manager, + inflight_cancel_supported_by_executor) class KvCacheTransceiver(ABC): @@ -123,6 +233,12 @@ def check_gen_transfer_complete(self): def cancel_request(self, req: LlmRequest): raise NotImplementedError + def supports_inflight_request_cancellation(self) -> bool: + return False + + def has_poisoned_transfer_buffer(self) -> bool: + return False + @abstractmethod def prepare_context_requests(self, requests: List[LlmRequest]): """ @@ -156,7 +272,11 @@ def __init__(self, kv_cache_manager: KVCacheManager, attention_type: AttentionTypeCpp, cache_transceiver_config: CacheTransceiverConfig, - mamba_cache_manager: Optional[BaseMambaCacheManager] = None): + mamba_cache_manager: Optional[BaseMambaCacheManager] = None, + inflight_cancel_supported_by_executor: bool = False): + _validate_disagg_inflight_cancel_config( + cache_transceiver_config, mapping, + inflight_cancel_supported_by_executor) world_config = mapping_to_world_config(mapping) # Filter out mamba/recurrent state layers (kv_heads == 0) so that # CacheState::ModelConfig::mNbKvHeadsPerLayer only contains attention @@ -180,6 +300,13 @@ def __init__(self, self.kv_transfer_timeout_ms = cache_transceiver_config.kv_transfer_timeout_ms self.kv_transfer_sender_future_timeout_ms = cache_transceiver_config.kv_transfer_sender_future_timeout_ms self.kv_transfer_poll_interval_ms = cache_transceiver_config.kv_transfer_poll_interval_ms + self._supports_inflight_request_cancellation = ( + _is_disagg_inflight_cancel_config_supported( + cache_transceiver_config, mapping, + inflight_cancel_supported_by_executor)) + self._inflight_request_cancellation_enabled = ( + is_disagg_inflight_cancel_enabled() + and self._supports_inflight_request_cancellation) # Get RNN state manager and layer distribution if mamba_cache_manager is provided. rnn_state_manager = None @@ -205,7 +332,8 @@ def __init__(self, tokens_per_block, world_config, pp_layer_num_per_pp_rank, dtype, attention_type, cache_transceiver_config._to_pybind(), rnn_state_manager, - rnn_layer_num_per_pp_rank) + rnn_layer_num_per_pp_rank, + self._inflight_request_cancellation_enabled) def respond_and_send_async(self, req: LlmRequest): return self.impl.respond_and_send_async(req) @@ -228,6 +356,14 @@ def check_gen_transfer_complete(self): def cancel_request(self, req: LlmRequest): return self.impl.cancel_request(req) + def supports_inflight_request_cancellation(self) -> bool: + return self._supports_inflight_request_cancellation + + def has_poisoned_transfer_buffer(self) -> bool: + if not self._inflight_request_cancellation_enabled: + return False + return self.impl.has_poisoned_transfer_buffer() + def prepare_context_requests(self, requests: List[LlmRequest]): # not implemented, an empty placeholder to allow being invoked unconditionally ... diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 188ab3302dda..ab9f5a04cc1d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1,3 +1,18 @@ +# 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. + import dataclasses import datetime import functools @@ -7,7 +22,7 @@ import traceback from contextlib import contextmanager from enum import IntEnum -from queue import Queue +from queue import Empty, Queue from typing import (TYPE_CHECKING, Callable, Dict, Iterable, List, Optional, Tuple, Union) @@ -65,7 +80,8 @@ from .hang_detector import HangDetector from .kv_cache_manager_v2 import KVCacheManagerV2 from .kv_cache_stats import append_kv_cache_iteration_stats -from .kv_cache_transceiver import KvCacheTransceiver +from .kv_cache_transceiver import (KvCacheTransceiver, + is_disagg_inflight_cancel_enabled) from .llm_request import (ATTENTION_DP_DUMMY_REQUEST_ID, MAX_SPEC_DECODE_POSITIONS, ExecutorRequest, LlmRequest, LlmRequestState, LlmResponse, @@ -411,6 +427,15 @@ class PyExecutor: # If the number of micro batches is too large, the executor will spend too much host memory (No additional GPU memory is required). # 1024 in-flight micro batches can avoid synchronization in most cases and keep host memory usage low. MIN_ASYNC_MICRO_BATCH_NUM = 1024 + # Poison cleanup is best effort. If a native transfer cannot quiesce in + # this window, preserve every transfer-owned resource and require the + # process supervisor to restart the worker. + FATAL_TRANSFER_CLEANUP_GRACE_PERIOD_S = 10.0 + # Unsafe executors intentionally remain strongly reachable until process + # teardown. Returning early from shutdown() alone is insufficient because + # a serving layer may drop its final reference during graceful shutdown. + _UNSAFE_TRANSFER_SHUTDOWN_QUARANTINE: List["PyExecutor"] = [] + _UNSAFE_TRANSFER_SHUTDOWN_QUARANTINE_LOCK = threading.Lock() def __init__( self, @@ -681,6 +706,21 @@ def __init__( self.is_shutdown = False self._fatal_error: Optional[BaseException] = None self._error_budget = ErrorBudget() + # Requests whose resources are still owned by a native KV transfer + # during topology-poison shutdown. They are kept outside the + # scheduler and reaped only after every native owner is terminal. + self._fatal_transfer_cleanup_requests: Dict[int, LlmRequest] = {} + # A detached context request may already have been terminated by the + # partial-reuse path while AsyncTransferManager still pins its blocks. + # Such requests need owner quiescence but must not be terminated twice. + self._fatal_transfer_cleanup_skip_termination_ids: set[int] = set() + self._fatal_transfer_shutdown = False + self._fatal_transfer_cleanup_complete = False + self._unsafe_transfer_shutdown = False + self._fatal_transfer_cleanup_deadline: Optional[float] = None + self._fatal_transfer_cleanup_timer: Optional[threading.Timer] = None + self._fatal_transfer_cleanup_lock = threading.Lock() + self._disagg_inflight_cancel_unsupported_logged = False self.max_batch_size = max_batch_size self.adp_ctx_waiting_iters_count = 0 self.adp_ctx_batching_wait_iters_count = 0 @@ -934,27 +974,41 @@ def _maybe_init_kv_connector_manager(self): self.kv_connector_manager.wait_for_initialization() - def _end_transfer_and_maybe_terminate(self, request: LlmRequest): - if self.kv_cache_transceiver and request in self.active_requests: + def _end_transfer_and_maybe_terminate(self, request: LlmRequest) -> None: + fatal_cleanup_pending = request.py_request_id in getattr( + self, "_fatal_transfer_cleanup_requests", {}) + if self.kv_cache_transceiver and (request in self.active_requests + or fatal_cleanup_pending): # Fast-transfer: KV transfer completed in the same iteration # before _handle_responses could run. Create the response now # while state is still TRANS_IN_PROGRESS (required by C++ # createResult). Then proceed with end_transfer + termination. - response = request.create_response(False, self.dist.rank) - if response: - response.result.cached_tokens = request.cached_tokens - self._maybe_attach_ctx_usage(request, response) - # Buffer the response instead of enqueueing immediately. - # With ADP, _enqueue_responses does a tp_gather collective. - # Calling it here would deadlock because only the owning DP - # rank reaches this point; the other DP rank never enters - # the matching collective. The buffer is flushed later at - # _flush_pending_transfer_responses where all ranks - # participate. - self._pending_transfer_responses.append( - (request.py_request_id, response)) + transfer_failed = request.state == LlmRequestState.DISAGG_TRANS_ERROR + user_cancel_pending = self._is_user_cancel_pending(request) + preserve_terminal_outcome = (transfer_failed or user_cancel_pending + or fatal_cleanup_pending) + if not preserve_terminal_outcome: + response = request.create_response(False, self.dist.rank) + if response: + response.result.cached_tokens = request.cached_tokens + self._maybe_attach_ctx_usage(request, response) + # Buffer the response instead of enqueueing immediately. + # With ADP, _enqueue_responses does a tp_gather collective. + # Calling it here would deadlock because only the owning DP + # rank reaches this point; the other DP rank never enters + # the matching collective. The buffer is flushed later at + # _flush_pending_transfer_responses where all ranks + # participate. + self._pending_transfer_responses.append( + (request.py_request_id, response)) if self.async_transfer_manager.end_transfer(request): - self.active_requests.remove(request) + # Error, cancellation, and fatal shutdown each have a + # centralized terminal-response path. Keep the request alive + # until that path observes that the last transfer owner is gone. + if preserve_terminal_outcome: + return + if request in self.active_requests: + self.active_requests.remove(request) self._terminate_request(request) return if self.async_transfer_manager.end_transfer(request): @@ -993,6 +1047,12 @@ def _event_loop_wrapper(self): except Exception as e: logger.error(f"Error in event loop: {e}") logger.error(traceback.format_exc()) + if getattr(self, "_fatal_transfer_shutdown", False): + # Any exception after topology poison means transfer ownership + # can no longer be proven quiescent. shutdown() must preserve + # registered buffers instead of running normal teardown. + self._mark_unsafe_transfer_shutdown( + "event loop failed during topology-poison shutdown", e) # Stash the original error so local consumers # (_await_single_response and BaseWorker.AwaitResponseHelper) # can surface it instead of letting callers hang. We do NOT @@ -1133,6 +1193,10 @@ def shutdown(self): self.executor_request_queue.enqueue_shutdown_request() self.shutdown_event.wait() if self.hang_detector.detected(): + if getattr(self, "_fatal_transfer_shutdown", False): + self._mark_unsafe_transfer_shutdown( + "hang detected during topology-poison cleanup", + complete_cleanup=False) # Early return here to avoid waiting for hanging threads. # Since `on_detected` has sent the error message as response, # this worker will be asked to shutdown immediately. @@ -1140,6 +1204,24 @@ def shutdown(self): # All threads and memory pools will be freed properly. logger.error("Hang detected, shutting down immediately.") return + transfer_shutdown_unsafe = ( + getattr(self, "_fatal_transfer_shutdown", False) and + (getattr(self, "_unsafe_transfer_shutdown", False) + or bool(getattr(self, "_fatal_transfer_cleanup_requests", {})))) + if transfer_shutdown_unsafe: + if not getattr(self, "_unsafe_transfer_shutdown", False): + self._mark_unsafe_transfer_shutdown( + "shutdown reached with native owners still pending", + complete_cleanup=False) + # Native code may still own registered KV buffers. Joining can + # block forever, while normal graph/resource teardown can free + # memory that the transfer worker is still accessing. The fatal + # response requires the serving process to be restarted, so leave + # those objects intentionally alive until process exit. + logger.error( + "Unsafe KV-transfer shutdown; skipping in-process resource " + "teardown and requiring process restart.") + return self.worker_thread.join() if self.dist.pp_size > 1: self.executed_batch_queue.put(None) @@ -2143,11 +2225,17 @@ def _executor_loop_pp(self): iter_start_time = time.time() iter_stats = None while True: - self.hang_detector.checkpoint() + if not getattr(self, "_fatal_transfer_shutdown", False): + self.hang_detector.checkpoint() profile_step() if self.enable_iter_perf_stats: iter_start_time = time.time() + if self._fatal_shutdown_complete(): + break + if self._progress_fatal_transfer_cleanup(): + continue + self._handle_disagg_cache_errors_synced() # Fetch new requests from request queue @@ -3120,11 +3208,17 @@ def _executor_loop(self): iter_stats = None can_forward = not self.is_benchmark_disagg while True: - self.hang_detector.checkpoint() + if not getattr(self, "_fatal_transfer_shutdown", False): + self.hang_detector.checkpoint() profile_step() if self.enable_iter_perf_stats: iter_start_time = time.time() + if self._fatal_shutdown_complete(): + break + if self._progress_fatal_transfer_cleanup(): + continue + if self._resource_governor_enabled: self._sync_and_process_resource_governor_queue() @@ -3523,10 +3617,33 @@ def control_action(self, *, drain: bool = True): are parked until the ``with`` block exits. """ + fatal_error = getattr(self, "_fatal_error", None) + if fatal_error is not None or getattr(self, "is_shutdown", False): + error = RuntimeError( + "Control action rejected because PyExecutor is shutting down") + if fatal_error is not None: + raise error from fatal_error + raise error + if self.dist.rank == 0: - self.executor_request_queue.enqueue_control_request(drain=drain) + if not self.executor_request_queue.enqueue_control_request( + drain=drain): + raise RuntimeError( + "Control action rejected because PyExecutor is shutting " + "down") self.control_request_barrier.wait() + fatal_error = getattr(self, "_fatal_error", None) + if fatal_error is not None or getattr(self, "is_shutdown", False): + # Fatal shutdown can release the barrier before the scheduler ever + # reaches this control action. Never yield exclusive access to a + # caller after the executor has become unusable. + self.control_request_barrier.clear() + error = RuntimeError( + "Control action aborted because PyExecutor is shutting down") + if fatal_error is not None: + raise error from fatal_error + raise error try: yield self @@ -3545,11 +3662,17 @@ def _executor_loop_overlap(self): previous_tensors_device = None can_forward = not self.is_benchmark_disagg while True: - self.hang_detector.checkpoint() + if not getattr(self, "_fatal_transfer_shutdown", False): + self.hang_detector.checkpoint() profile_step() if self.enable_iter_perf_stats: iter_start_time = time.time() + if self._fatal_shutdown_complete(): + break + if self._progress_fatal_transfer_cleanup(): + continue + if self._resource_governor_enabled: self._sync_and_process_resource_governor_queue() @@ -4501,19 +4624,66 @@ def _mark_cross_kv_projection_consumed( req.py_skip_cross_kv_projection = True @nvtx_range("_check_disagg_gen_transfer_status") - def _check_disagg_gen_transfer_status(self): + def _check_disagg_gen_transfer_status(self) -> None: + inflight_cancel_active = self._is_disagg_inflight_cancel_active() + canceled_without_forward = (self._handle_canceled_requests( + transfer_only=True) if inflight_cancel_active else []) + # Gen-transfer status performs cross-rank consensus internally. # Enter it symmetrically; ranks with no ready local future contribute # an empty ready set. self._check_disagg_gen_cache_transfer_status(0) + if getattr(self, "_fatal_error", None) is not None: + return + + if inflight_cancel_active: + # The first pass submits native user cancellation; the C++ poll + # commits topology-wide terminal state; this second pass finishes + # requests that no longer have a native or async owner. It must + # not submit a second cancellation in the same poll. + canceled_without_forward.extend( + self._handle_canceled_requests(transfer_only=True, + issue_native_cancel=False)) + self._finish_canceled_requests_without_forward( + canceled_without_forward) + + def _is_disagg_inflight_cancel_active(self) -> bool: + if not is_disagg_inflight_cancel_enabled(): + return False + + transceiver = getattr(self, "kv_cache_transceiver", None) + if transceiver is None: + return False + + supports = getattr(transceiver, + "supports_inflight_request_cancellation", None) + if callable(supports) and supports() is True: + return True + + if not getattr(self, "_disagg_inflight_cancel_unsupported_logged", + False): + logger.warning( + "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 was requested, but " + f"{type(transceiver).__name__} does not advertise in-flight " + "request cancellation support. Cancellation and transfer-buffer " + "quarantine are currently scoped to the C++ NIXL transceiver " + "with the UCX plugin; using the existing timeout and " + "cancellation behavior for this transceiver.") + self._disagg_inflight_cancel_unsupported_logged = True + return False @nvtx_range("_check_kv_transfer_timeout") - def _check_kv_transfer_timeout(self): + def _check_kv_transfer_timeout(self) -> None: if not self.kv_cache_transceiver: return timeout_ms = self.kv_cache_transceiver.kv_transfer_timeout_ms if timeout_ms is None: return + if self._is_disagg_inflight_cancel_active(): + # C++ owns deadline observation, topology voting, and cancellation + # while the feature is active. A Python rank-local timeout must + # never touch a transfer before that vote. + return def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: current_time = time.time() @@ -4522,8 +4692,8 @@ def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: elapsed_time = (current_time - req.py_kv_transfer_start_time) * 1000 if elapsed_time > timeout_ms and not req.py_kv_transfer_timed_out: logger.warning( - f"Terminating {type} request {req.py_request_id} due to KV cache transfer timeout" - ) + f"Observed timeout on {type} request {req.py_request_id} " + "due to KV cache transfer timeout") req.py_kv_transfer_timed_out = True for req in self.async_transfer_manager.requests_in_transfer().values(): @@ -4924,13 +5094,13 @@ def kv_connector_request_finished(req: LlmRequest): if self.kv_cache_transceiver: self._check_disagg_ctx_cache_transfer_status(0) - def _get_disagg_reqs_in_error_state(self): + def _get_disagg_reqs_in_error_state(self) -> List[LlmRequest]: return [ req for req in self.active_requests if req.state == LlmRequestState.DISAGG_TRANS_ERROR ] - def _check_cache_transfer_errors(self, error_msg_prefix: str): + def _check_cache_transfer_errors(self, error_msg_prefix: str) -> None: """Check and handle cache transfer errors. Under ADP this is a no-op: errors are handled by @@ -4938,17 +5108,45 @@ def _check_cache_transfer_errors(self, error_msg_prefix: str): """ if self.enable_attention_dp and self.dist.world_size != 1: return - error_requests = self._get_disagg_reqs_in_error_state() + canceled_req_ids = set(getattr(self, "canceled_req_ids", [])) + async_transfer_manager = getattr(self, "async_transfer_manager", None) + requests_in_transfer = (async_transfer_manager.requests_in_transfer() + if async_transfer_manager is not None else {}) + error_requests = [ + req for req in self._get_disagg_reqs_in_error_state() + if req.py_request_id not in requests_in_transfer + and not self._is_user_cancel_pending(req, canceled_req_ids) + ] if error_requests: - self._handle_errors( - f"Error in kv cache transfer for {error_msg_prefix}", - requests=error_requests, - charge_budget=False) + error_msg = f"Error in kv cache transfer for {error_msg_prefix}" + self._handle_errors(error_msg, + requests=error_requests, + charge_budget=False) + + def _handle_poisoned_transfer_buffer(self, request_role: str) -> bool: + """Enter fail-closed shutdown for a topology-wide poison latch.""" + if not self._is_disagg_inflight_cancel_active(): + return False + if not self.kv_cache_transceiver.has_poisoned_transfer_buffer(): + return False + + error_msg = (f"Error in kv cache transfer for {request_role}; poisoned " + "transfer buffer requires process restart") + self._handle_errors(error_msg, + charge_budget=False, + force_fatal=True, + defer_transfer_cleanup=True) + return True @nvtx_range("_check_disagg_ctx_cache_transfer_status") - def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): + def _check_disagg_ctx_cache_transfer_status(self, + atLeastNum: int = 0) -> None: finished_requests, error_requests = self.kv_cache_transceiver.check_context_transfer_status( atLeastNum) + # C++ updates a topology-wide sticky poison latch as part of the status + # exchange. Reading it is noncollective and must happen even when no + # request is quiescent yet. + poisoned = self._handle_poisoned_transfer_buffer("context requests") completed_req_ids = set(finished_requests + error_requests) @@ -4970,23 +5168,32 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): requests_in_transfer = self.async_transfer_manager.requests_in_transfer( ) - for request_id in list(requests_in_transfer.keys()): - request = requests_in_transfer[request_id] - if request.py_kv_transfer_timed_out and request_id not in completed_req_ids: + if not self._is_disagg_inflight_cancel_active(): + for request_id in list(requests_in_transfer.keys()): + request = requests_in_transfer[request_id] + if (not request.py_kv_transfer_timed_out + or request_id in completed_req_ids): + continue + is_cancelled = self.kv_cache_transceiver.cancel_request(request) - # If cancel is successful, mark as complete so it can be cleaned up - # Otherwise, try at next iteration - if is_cancelled: - request.py_kv_transfer_start_time = None - request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE + if not is_cancelled: + continue - self._end_transfer_and_maybe_terminate(request) + # Preserve the legacy timeout behavior when in-flight + # cancellation is disabled: a queued transfer that can be + # cancelled is immediately released from the async manager. + request.py_kv_transfer_start_time = None + request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE + self._end_transfer_and_maybe_terminate(request) - self._check_cache_transfer_errors("context requests") + if not poisoned: + self._check_cache_transfer_errors("context requests") @nvtx_range("_check_disagg_gen_cache_transfer_status") - def _check_disagg_gen_cache_transfer_status(self, atLeastNum: int = 0): + def _check_disagg_gen_cache_transfer_status(self, + atLeastNum: int = 0) -> None: result = self.kv_cache_transceiver.check_gen_transfer_status(atLeastNum) + poisoned = self._handle_poisoned_transfer_buffer("generation requests") if isinstance(result, tuple): _, _, cancelled_reqs = result user_canceled_set = set(self.canceled_req_ids) @@ -4994,7 +5201,8 @@ def _check_disagg_gen_cache_transfer_status(self, atLeastNum: int = 0): req_id = req.py_request_id if not req.is_child else req.parent_request_id if req_id not in user_canceled_set: req.state = LlmRequestState.DISAGG_TRANS_ERROR - self._check_cache_transfer_errors("generation requests") + if not poisoned: + self._check_cache_transfer_errors("generation requests") def _maybe_prefetch_next_iter_mm_encoders( self, scheduled_batch: ScheduledRequests) -> None: @@ -5238,7 +5446,9 @@ def _handle_errors(self, error_msg: Optional[str] = None, *, requests: Optional[List[LlmRequest]] = None, - charge_budget: bool = True) -> None: + charge_budget: bool = True, + force_fatal: bool = False, + defer_transfer_cleanup: bool = False) -> None: """Fail requests and optionally initiate shutdown on fatal errors. When ``charge_budget`` is True (the default), classifies the error @@ -5252,6 +5462,12 @@ def _handle_errors(self, Use this for request-scoped errors (validation, KV-transfer timeout, guided-decoder) that should not affect server health. + ``force_fatal`` is reserved for failures, such as a poisoned native + transfer buffer, whose severity is already known without consulting + the error budget. ``defer_transfer_cleanup`` keeps transfer-owned + requests out of the scheduler but does not free their resources until + the native owner becomes terminal. + .. note:: The ``charge_budget=False`` path reuses the full ``_handle_errors`` machinery (queue drain, response @@ -5270,24 +5486,64 @@ def _handle_errors(self, charge_budget: Whether to consume the error budget. Set to False for request-scoped errors that should not affect server health. + force_fatal: Enter fatal shutdown without charging the budget. + defer_transfer_cleanup: Retain resources owned by an in-flight + native or asynchronous transfer during fatal shutdown. """ - error_responses: Dict[int, LlmResponse] = {} error_msg = error_msg or "error" - is_fatal = (self._error_budget.consume(error_msg) - if charge_budget else False) - if is_fatal and self._error_budget.budget < 1e-9: + # A topology-poison latch is sticky and is observed on every status + # poll. Start fatal shutdown exactly once; later polls only make + # progress on deferred transfer cleanup. + if force_fatal and getattr(self, "_fatal_transfer_shutdown", False): + return + + budget_fatal = (self._error_budget.consume(error_msg) + if charge_budget else False) + is_fatal = force_fatal or budget_fatal + if budget_fatal and self._error_budget.budget < 1e-9: logger.error(f"Error budget exhausted " f"(budget={self._error_budget.budget:.3f}), " "treating as fatal") + responses: List[Tuple[int, LlmResponse]] = [] if is_fatal: + if not hasattr(self, "_fatal_transfer_cleanup_requests"): + self._fatal_transfer_cleanup_requests = {} + if not hasattr(self, + "_fatal_transfer_cleanup_skip_termination_ids"): + self._fatal_transfer_cleanup_skip_termination_ids = set() + if defer_transfer_cleanup: + self._fatal_transfer_shutdown = True + self._fatal_transfer_cleanup_complete = False + self._unsafe_transfer_shutdown = False self._fatal_error = RuntimeError(f"Fatal error: {error_msg}") self.is_shutdown = True logger.error( f"Fatal error detected, initiating shutdown: {error_msg}") requests = None + # Close the ingress queue under its lock before draining it. This + # removes the old race in which a producer could enqueue a request + # after the raw queue had been observed empty. + self.executor_request_queue.enqueue_shutdown_request() + # A drain-style control action may be waiting for the scheduler to + # reach its sentinel. Poison bypasses normal control handling, so + # release the waiter; control_action checks _fatal_error and raises + # before yielding access to the unusable executor. + control_requests = getattr(self, "control_requests", None) + if control_requests is not None: + control_requests.clear() + control_request_barrier = getattr(self, "control_request_barrier", + None) + if control_request_barrier is not None: + control_request_barrier.set() + + # Preserve responses for requests that completed before the fatal + # observation; they are no longer active and are safe to publish. + responses.extend(getattr(self, "_pending_transfer_responses", [])) + self._pending_transfer_responses = [] + # Drain waiting_queue so that queued-but-not-yet-activated # requests don't get picked up on the next iteration. # These are RequestQueueItems (not yet LlmRequests), so we @@ -5295,50 +5551,93 @@ def _handle_errors(self, # call _enqueue_responses once after the loop so every rank # enters the same number of collectives (attention-DP / # gather-all modes use collective gathers internally). - waiting_responses: List[Tuple[int, LlmResponse]] = [] while self.waiting_queue: item = self.waiting_queue.pop_request() if (self.gather_all_responses or self.dist.rank == 0) and item.request is not None: - waiting_responses.append( + responses.append( (item.id, LlmResponse(request_id=item.id, error_msg=error_msg, client_id=getattr(item.request, 'client_id', None)))) + # Requests fetched after a pending control sentinel are parked + # outside waiting_queue. They need the same terminal response and + # must not survive a fatal shutdown. + request_accumulated = getattr(self, "request_accumulated", []) + for item in request_accumulated: + if ((self.gather_all_responses or self.dist.rank == 0) + and item.request is not None): + responses.append( + (item.id, + LlmResponse(request_id=item.id, + error_msg=error_msg, + client_id=getattr(item.request, + 'client_id', None)))) + request_accumulated.clear() # Also drain executor_request_queue so items already queued # but not yet fetched by the main loop are not scheduled - # after the CUDA context is corrupted. Safe to use empty() - # here because is_shutdown is True and the queue's active - # flag is about to be set False, so no new items arrive. + # after the transfer buffer is poisoned. The queue is already + # inactive, so get_nowait() can drain it to a stable empty state. raw_queue = self.executor_request_queue.get_request_queue() - while not raw_queue.empty(): - item = raw_queue.get_nowait() + while True: + try: + item = raw_queue.get_nowait() + except Empty: + break if item.is_shutdown_request: continue if ((self.gather_all_responses or self.dist.rank == 0) and item.request is not None): - waiting_responses.append( + responses.append( (item.id, LlmResponse(request_id=item.id, error_msg=error_msg, client_id=getattr(item.request, 'client_id', None)))) - - if waiting_responses: - self._enqueue_responses(waiting_responses) - logger.info(f"Drained {len(waiting_responses)} queued requests " - "on fatal error") + if not defer_transfer_cleanup and self.dist.rank == 0: + # Generic fatal errors are not guaranteed to be observed on + # every rank. Preserve their normal broadcast shutdown path: + # rank 0 must fetch a sentinel on the next iteration instead + # of blocking forever on the now-empty inactive queue. + self.executor_request_queue.enqueue_shutdown_request() failed_requests = (list(self.active_requests) if requests is None else requests) + error_responses: Dict[int, LlmResponse] = {} + requests_to_terminate: List[LlmRequest] = [] for request in failed_requests: req_id = request.py_request_id - request.state = LlmRequestState.GENERATION_COMPLETE error_responses[req_id] = LlmResponse( request_id=req_id, error_msg=error_msg, client_id=request.py_client_id) + if is_fatal and defer_transfer_cleanup: + # Defer every active request until the next fatal-progress + # boundary. Besides native transfers, an overlap batch may + # still own request resources through sampler/D2H work that was + # launched just before a CTX status poll observed poison. + self._fatal_transfer_cleanup_requests[req_id] = request + else: + request.state = LlmRequestState.GENERATION_COMPLETE + requests_to_terminate.append(request) + if is_fatal and defer_transfer_cleanup: + # Finished context requests leave active_requests after publishing + # their response, but AsyncTransferManager can still own their KV + # blocks. Track those detached owners without emitting a second + # response. With partial reuse they were already terminated by + # _handle_responses, so only the final unpin remains. + async_transfer_manager = getattr(self, "async_transfer_manager", + None) + if async_transfer_manager is not None: + for req_id, request in list( + async_transfer_manager.requests_in_transfer().items()): + if req_id in self._fatal_transfer_cleanup_requests: + continue + self._fatal_transfer_cleanup_requests[req_id] = request + if async_transfer_manager.should_store_blocks: + self._fatal_transfer_cleanup_skip_termination_ids.add( + req_id) if requests is None: self.active_requests.clear() else: @@ -5346,12 +5645,221 @@ def _handle_errors(self, request for request in self.active_requests if request not in requests ] - self._enqueue_responses(list(error_responses.items())) - for request in failed_requests: + + responses.extend(error_responses.items()) + # One enqueue per error path keeps any response collective symmetric + # and prevents duplicate terminal delivery during sticky poison polls. + self._enqueue_responses(responses) + if is_fatal and defer_transfer_cleanup: + self._start_fatal_transfer_cleanup_watchdog() + for request in requests_to_terminate: self._terminate_request(request) - if self._fatal_error is not None: - self.executor_request_queue.enqueue_shutdown_request() + def _request_has_transfer_ownership(self, request: LlmRequest) -> bool: + async_transfer_manager = getattr(self, "async_transfer_manager", None) + if async_transfer_manager is not None: + requests_in_transfer = async_transfer_manager.requests_in_transfer() + if request.py_request_id in requests_in_transfer: + return True + return self._is_request_in_transmission(request) + + def _fatal_shutdown_complete(self) -> bool: + return (getattr(self, "_fatal_transfer_shutdown", False) + and getattr(self, "_fatal_transfer_cleanup_complete", False)) + + def _fatal_transfer_cleanup_expired(self) -> bool: + """Return whether this rank's poison-cleanup grace period expired.""" + deadline = getattr(self, "_fatal_transfer_cleanup_deadline", None) + return deadline is not None and time.monotonic() >= deadline + + def _start_fatal_transfer_cleanup_watchdog(self) -> None: + """Wake shutdown even if native poison cleanup blocks indefinitely.""" + lock = getattr(self, "_fatal_transfer_cleanup_lock", None) + if lock is None: + lock = threading.Lock() + self._fatal_transfer_cleanup_lock = lock + with lock: + if getattr(self, "_fatal_transfer_cleanup_timer", None) is not None: + return + grace_period = self.FATAL_TRANSFER_CLEANUP_GRACE_PERIOD_S + self._fatal_transfer_cleanup_deadline = (time.monotonic() + + grace_period) + + def on_timeout() -> None: + self._mark_unsafe_transfer_shutdown( + "native owners did not quiesce within " + f"{grace_period:g}s", + require_incomplete=True, + complete_cleanup=False) + + timer = threading.Timer(grace_period, on_timeout) + timer.daemon = True + self._fatal_transfer_cleanup_timer = timer + timer.start() + + def _claim_safe_fatal_transfer_cleanup(self) -> bool: + """Stop the watchdog after all ranks prove native quiescence.""" + lock = getattr(self, "_fatal_transfer_cleanup_lock", None) + if lock is None: + lock = threading.Lock() + self._fatal_transfer_cleanup_lock = lock + with lock: + if getattr(self, "_unsafe_transfer_shutdown", False): + self._fatal_transfer_cleanup_complete = True + return False + self._fatal_transfer_cleanup_complete = True + self._fatal_transfer_cleanup_deadline = None + timer = getattr(self, "_fatal_transfer_cleanup_timer", None) + self._fatal_transfer_cleanup_timer = None + if timer is not None: + timer.cancel() + return True + + def _mark_unsafe_transfer_shutdown(self, + reason: str, + error: Optional[BaseException] = None, + *, + require_incomplete: bool = False, + complete_cleanup: bool = True) -> None: + lock = getattr(self, "_fatal_transfer_cleanup_lock", None) + if lock is None: + lock = threading.Lock() + self._fatal_transfer_cleanup_lock = lock + with lock: + cleanup_complete = getattr(self, "_fatal_transfer_cleanup_complete", + False) + if require_incomplete and cleanup_complete: + return + if getattr(self, "_unsafe_transfer_shutdown", False): + if complete_cleanup: + self._fatal_transfer_cleanup_complete = True + return + self._unsafe_transfer_shutdown = True + if complete_cleanup: + self._fatal_transfer_cleanup_complete = True + timer = getattr(self, "_fatal_transfer_cleanup_timer", None) + self._fatal_transfer_cleanup_timer = None + if timer is not None: + timer.cancel() + detail = f": {error}" if error is not None else "" + logger.critical( + f"Fatal KV-transfer cleanup is unsafe ({reason}){detail}; " + "preserving transfer-owned resources until process restart.") + with PyExecutor._UNSAFE_TRANSFER_SHUTDOWN_QUARANTINE_LOCK: + if not any(executor is self for executor in + PyExecutor._UNSAFE_TRANSFER_SHUTDOWN_QUARANTINE): + PyExecutor._UNSAFE_TRANSFER_SHUTDOWN_QUARANTINE.append(self) + self.is_shutdown = True + if getattr(self, "_event_loop_error", None) is None: + self._event_loop_error = getattr( + self, "_fatal_error", + RuntimeError("Unsafe KV-transfer shutdown")) + shutdown_event = getattr(self, "shutdown_event", None) + if shutdown_event is not None: + shutdown_event.set() + + def _progress_fatal_transfer_cleanup(self) -> bool: + """Poll native owners and reap fatal requests only after quiescence.""" + if not getattr(self, "_fatal_transfer_shutdown", False): + return False + pending = getattr(self, "_fatal_transfer_cleanup_requests", {}) + cleanup_errors: List[str] = [] + + # A CTX poll can observe poison after overlap launched the next batch. + # That batch is installed as previous_batch before the loop reaches + # this fatal-progress boundary. Synchronize its sampler work and drop + # the normal response path before considering any resource releasable. + previous_batch = getattr(self, "previous_batch", None) + if previous_batch is not None: + try: + sampler_event = getattr(previous_batch.sample_state, + "sampler_event", None) + if sampler_event is not None: + sampler_event.synchronize() + self.previous_batch = None + except Exception as error: + cleanup_errors.append(f"overlap synchronization: {error}") + + # Keep the collective order fixed even after a local exception. A + # peer may still be inside either C++ status exchange and must not see + # this rank skip directly to the final Python state vote. + try: + self._cancel_fatal_transfer_owners() + except Exception as error: + cleanup_errors.append(f"native cancellation: {error}") + try: + self._check_disagg_ctx_cache_transfer_status(0) + except Exception as error: + cleanup_errors.append(f"context status: {error}") + try: + self._check_disagg_gen_cache_transfer_status(0) + except Exception as error: + cleanup_errors.append(f"generation status: {error}") + + try: + local_has_owner = any( + self._request_has_transfer_ownership(request) + for request in pending.values()) + except Exception as error: + cleanup_errors.append(f"ownership check: {error}") + local_has_owner = True + + local_state = (bool(cleanup_errors) + or getattr(self, "_unsafe_transfer_shutdown", False), + self._fatal_transfer_cleanup_expired(), local_has_owner, + tuple(sorted(pending))) + try: + states = ([local_state] if self._dist_size(self.dist, "world_size") + == 1 else self.dist.tp_allgather(local_state)) + except Exception as error: + self._mark_unsafe_transfer_shutdown( + "fatal-cleanup state consensus failed", error) + return True + + pending_id_sets = {state[3] for state in states} + any_cleanup_error = any(state[0] for state in states) + any_expired = any(state[1] for state in states) + pending_mismatch = len(pending_id_sets) != 1 + if any_cleanup_error or any_expired or pending_mismatch: + if any_cleanup_error: + reason = "native cancellation or status polling failed" + elif any_expired: + reason = ("native owners did not quiesce within " + f"{self.FATAL_TRANSFER_CLEANUP_GRACE_PERIOD_S:g}s") + else: + reason = "ranks disagreed on deferred request ownership" + detail = RuntimeError( + "; ".join(cleanup_errors)) if cleanup_errors else None + self._mark_unsafe_transfer_shutdown(reason, detail) + return True + + if any(state[2] for state in states): + # Fatal cleanup is not latency-sensitive; avoid a hot host spin + # while the cancellation worker reaches a terminal state. + time.sleep(0.001) + return True + + if not self._claim_safe_fatal_transfer_cleanup(): + return True + skip_termination_ids = getattr( + self, "_fatal_transfer_cleanup_skip_termination_ids", set()) + for request in list(pending.values()): + request_id = request.py_request_id + pending.pop(request_id, None) + if request_id in skip_termination_ids: + skip_termination_ids.discard(request_id) + continue + request.state = LlmRequestState.GENERATION_COMPLETE + self._terminate_request(request) + return True + + def _cancel_fatal_transfer_owners(self) -> None: + transceiver = getattr(self, "kv_cache_transceiver", None) + if transceiver is None: + return + for request in self._fatal_transfer_cleanup_requests.values(): + if self._is_request_in_transmission(request): + transceiver.cancel_request(request) def _terminate_request(self, request: LlmRequest): # Dummy requests don't participate in disagg KV cache transfers, @@ -5367,18 +5875,32 @@ def _terminate_request(self, request: LlmRequest): def _do_terminate_request(self, request: LlmRequest): self.resource_manager.free_resources(request) self._prefetched_request_ids.discard(request.py_request_id) - if self.gather_all_responses or self.dist.rank == 0: self.result_wait_queues.pop(request.py_request_id, None) - def _is_request_in_transmission(self, request) -> bool: + @staticmethod + def _request_cancel_id(request: LlmRequest) -> int: + return (request.parent_request_id if getattr(request, "is_child", False) + else request.py_request_id) + + def _is_user_cancel_pending( + self, + request: LlmRequest, + canceled_req_ids: Optional[set[int]] = None) -> bool: + ids = (set(self.canceled_req_ids) + if canceled_req_ids is None else canceled_req_ids) + return self._request_cancel_id(request) in ids + + def _is_request_in_transmission(self, request: LlmRequest) -> bool: """Check if a request is currently in transmission state.""" return (request.state == LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS or request.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS) - def _try_cancel_request(self, request) -> bool: + def _try_cancel_request(self, + request: LlmRequest, + issue_native_cancel: bool = True) -> bool: """Check if a request can be canceled and attempt cancellation if needed. Returns: @@ -5388,14 +5910,29 @@ def _try_cancel_request(self, request) -> bool: return True if not self._is_request_in_transmission(request): - return True + # A transceiver future can be terminal while a second async owner + # still pins the same context request. Do not turn the logical + # user cancellation into resource cleanup until the last owner is + # gone. + return request.py_request_id not in ( + self.async_transfer_manager.requests_in_transfer()) + + if self._is_disagg_inflight_cancel_active(): + if issue_native_cancel: + self.kv_cache_transceiver.cancel_request(request) + return False - return self.kv_cache_transceiver.cancel_request(request) + return (self.kv_cache_transceiver.cancel_request(request) + if issue_native_cancel else False) @nvtx_range("_handle_canceled_requests") - def _handle_canceled_requests(self): + def _handle_canceled_requests( + self, + *, + transfer_only: bool = False, + issue_native_cancel: bool = True) -> List[LlmRequest]: if len(self.canceled_req_ids) == 0: - return + return [] # Create set from list of canceled request ids to speed up canceled test canceled_req_ids_set = set(self.canceled_req_ids) @@ -5403,24 +5940,70 @@ def _handle_canceled_requests(self): # Remove canceled requests from the waiting queue self.waiting_queue.remove_by_ids(canceled_req_ids_set) - still_pending_canceled_ids = [] + previous_requests: List[LlmRequest] = [] + if transfer_only and getattr(self, "previous_batch", None) is not None: + previous_requests = self.previous_batch.scheduled_requests.all_requests( + ) + + finished_canceled_requests: List[LlmRequest] = [] for request in self.active_requests: - req_id = request.py_request_id if not request.is_child else request.parent_request_id + if transfer_only and request in previous_requests: + continue + req_id = self._request_cancel_id(request) if req_id not in canceled_req_ids_set: continue - is_cancelled = self._try_cancel_request(request) + is_cancelled = self._try_cancel_request(request, + issue_native_cancel) if is_cancelled: # Mark requests as finished, then, we reuse all existing code # to clean up the KV cache resources. request.finish_by_reason(FinishReason.CANCELLED) request.decoding_iter = request.py_decoding_iter - else: - still_pending_canceled_ids.append(req_id) + finished_canceled_requests.append(request) + + unfinished_cancel_ids = { + self._request_cancel_id(request) + for request in self.active_requests + if request not in finished_canceled_requests + } + # Preserve input order while retaining a parent cancellation until all + # of its child requests can finish. + self.canceled_req_ids = list( + dict.fromkeys(request_id for request_id in self.canceled_req_ids + if request_id in unfinished_cancel_ids)) + return finished_canceled_requests + + def _finish_canceled_requests_without_forward( + self, requests: List[LlmRequest]) -> None: + """Emit one terminal cancellation response after owners quiesce.""" + requests_to_terminate: List[LlmRequest] = [] + responses: List[Tuple[int, LlmResponse]] = [] + seen_request_ids: set[int] = set() + for request in requests: + request_id = request.py_request_id + if (request_id in seen_request_ids + or request not in self.active_requests): + continue + seen_request_ids.add(request_id) + request.draft_tokens = request.py_draft_tokens or [] + response = request.create_response(False, self.dist.rank) + if response: + response.result.cached_tokens = request.cached_tokens + self._maybe_attach_ctx_usage(request, response) + responses.append((request_id, response)) + requests_to_terminate.append(request) - # Clear list of requests marked for cancellation and add back those that failed to cancel. - self.canceled_req_ids.clear() - self.canceled_req_ids.extend(still_pending_canceled_ids) + if not requests_to_terminate: + return + + self._enqueue_responses(responses) + self.active_requests = [ + request for request in self.active_requests + if request not in requests_to_terminate + ] + for request in requests_to_terminate: + self._terminate_request(request) @nvtx_range("_enqueue_responses") def _enqueue_responses(self, responses: Iterable[Tuple[int, LlmResponse]]): @@ -5547,6 +6130,7 @@ def _handle_responses(self, emit_first_iter: bool = True): requests_finished_by_transfer = [] new_active_requests = [] timed_out_requests = [] + pending_canceled_req_ids = set(self.canceled_req_ids) logger.debug( f'------before _handle_responses, rank = {self.dist.rank}, output = {self.active_requests}' ) @@ -5560,11 +6144,31 @@ def _handle_responses(self, emit_first_iter: bool = True): requests_to_terminate.append(request) continue - # Check if generation request needs cleanup due to KV cache transfer timeout + if (self._request_cancel_id(request) in pending_canceled_req_ids + and self._is_request_in_transmission(request)): + # Native cancellation has been requested but C++ still owns + # the transfer. Do not race it with a successful context + # response; the centralized cancellation path emits exactly + # one terminal response after ownership is released. + new_active_requests.append(request) + continue + + # Check if a generation request needs cleanup due to KV cache transfer timeout. if request.py_kv_transfer_timed_out: - is_cancelled = self.kv_cache_transceiver.cancel_request(request) - if is_cancelled: - timed_out_requests.append(request) + defer_timeout_cleanup = ( + self._is_disagg_inflight_cancel_active() and + (request.is_disagg_generation_transmission_in_progress + or request.state == LlmRequestState.DISAGG_TRANS_ERROR)) + if not defer_timeout_cleanup: + is_cancelled = self.kv_cache_transceiver.cancel_request( + request) + if is_cancelled: + # _handle_errors enters response collectives under ADP. + # Defer it until the rank-uniform vote below. + timed_out_requests.append(request) + continue + + new_active_requests.append(request) continue if request.is_generation_only_request() and not request.is_finished: diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index c1d1e370ea2a..9963dca7a0bc 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -1,3 +1,18 @@ +# 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. + import concurrent import contextlib import functools @@ -54,6 +69,8 @@ def result(self): DEFAULT_SERVER_WAITING_TIMEOUT = 2100 # Timeout for the accuracy evaluation DEFAULT_ACC_EVALUATION_TIMEOUT = 1500 +# Timeout for each request sent to the disaggregated server +DEFAULT_REQUEST_TIMEOUT_S = 1800000 @functools.lru_cache(maxsize=1) @@ -152,6 +169,8 @@ def launch_disaggregated_llm( enable_perf=False, extra_env: Optional[Dict[str, str]] = None, gen_extra_env: Optional[Dict[str, str]] = None, + request_timeout_s: float = DEFAULT_REQUEST_TIMEOUT_S, + request_max_retries: Optional[int] = None, ): temp_dir = tempfile.TemporaryDirectory() disaggregated_serving_config_path = os.path.join( @@ -405,9 +424,14 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): f"Server is not ready after {server_waiting_timeout} seconds. Please check the logs for more details." ) - client = openai.OpenAI(api_key="1234567890", - base_url=f"http://localhost:{serve_port}/v1", - timeout=1800000) + client_kwargs: Dict[str, Any] = { + "api_key": "1234567890", + "base_url": f"http://localhost:{serve_port}/v1", + "timeout": request_timeout_s, + } + if request_max_retries is not None: + client_kwargs["max_retries"] = request_max_retries + client = openai.OpenAI(**client_kwargs) def send_request(prompt: str, sampling_params: SamplingParams, streaming: bool): @@ -1077,17 +1101,18 @@ def test_nixl_backend(self): @pytest.mark.skip_less_device_memory(60000) @skip_no_hopper def test_gen_only_sync(self): - """Test gen-only synchronous KV transfer path with NIXL Python transceiver. + """Test gen-only synchronous KV transfer path with NIXL C++ transceiver. Sets TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1 so the gen worker calls - request_and_receive_sync instead of the async path. Accuracy must be - identical to the standard async path. + the blocking request_and_receive_sync path. A bounded client timeout + makes a stuck transfer fail this test instead of waiting for its outer + one-hour timeout. """ ctx_server_config = { "disable_overlap_scheduler": True, "cache_transceiver_config": { "backend": "NIXL", - "transceiver_runtime": "PYTHON", + "transceiver_runtime": "CPP", "max_tokens_in_buffer": 4096, }, } @@ -1095,7 +1120,7 @@ def test_gen_only_sync(self): "disable_overlap_scheduler": True, "cache_transceiver_config": { "backend": "NIXL", - "transceiver_runtime": "PYTHON", + "transceiver_runtime": "CPP", "max_tokens_in_buffer": 4096, }, } @@ -1116,6 +1141,8 @@ def test_gen_only_sync(self): self.MODEL_PATH, # Apply to both servers: gen worker uses sync receive path. extra_env={"TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP": "1"}, + request_timeout_s=120, + request_max_retries=0, ) as llm: run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a958f8b24951..18db4fd5e224 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -15,6 +15,9 @@ l0_b200: backend: pytorch tests: # ------------- PyTorch tests --------------- + - unittest/others/test_kv_cache_transceiver.py::test_cpp_nixl_sync_transfer_stress + - unittest/others/test_kv_cache_transceiver.py::test_cpp_nixl_active_cancellation_reaches_terminal_error + - unittest/others/test_kv_cache_transceiver.py::test_cpp_nixl_sync_rejection_preserves_transfer_error - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_nvfp4 - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_nvfp4_streaming[stream_interval_4] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_nvfp4_streaming[stream_interval_64] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 2bd377836f93..7c06c07c042b 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -17,6 +17,7 @@ l0_dgx_b200: tests: - unittest/_torch/misc/test_autotuner.py::test_autotuner_distributed_strategy - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp2-TRTLLM] + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_only_sync # ------------- KV Cache V2 Scheduler IT (multi-GPU) --------------- - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2DSv3Lite::test_mtp_draft_tokens - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2DSv3Lite::test_mtp_chunked_draft_tokens diff --git a/tests/integration/test_lists/test-db/l0_sanity_check.yml b/tests/integration/test_lists/test-db/l0_sanity_check.yml index bc947305f276..aa0faa99cec0 100644 --- a/tests/integration/test_lists/test-db/l0_sanity_check.yml +++ b/tests/integration/test_lists/test-db/l0_sanity_check.yml @@ -32,6 +32,9 @@ l0_sanity_check: - examples/test_llm_api_with_mpi.py::test_llm_api_single_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] ISOLATION - unittest/others/test_kv_cache_transceiver.py::test_kv_cache_transceiver_single_process[NIXL-mha-ctx_fp16_gen_fp16] - unittest/others/test_kv_cache_transceiver.py::test_kv_cache_transceiver_single_process[UCX-mha-ctx_fp16_gen_fp16] + - unittest/others/test_kv_cache_transceiver.py::test_cpp_nixl_sync_transfer_stress + - unittest/others/test_kv_cache_transceiver.py::test_cpp_nixl_active_cancellation_reaches_terminal_error + - unittest/others/test_kv_cache_transceiver.py::test_cpp_nixl_sync_rejection_preserves_transfer_error - unittest/others/test_kv_cache_transceiver.py::test_cancel_request_in_transmission[mha] - unittest/others/test_kv_cache_transceiver.py::test_cancel_request_in_transmission[mla] - unittest/others/test_kv_cache_transceiver.py::test_async_transfer_keeps_llm_request_alive diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index e6931ca348d9..2c996c1f6bcb 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -378,10 +378,6 @@ perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_kimi-k25-thinking-fp4 perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6368078) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_128k8k_con64_ctx1_pp8_gen1_dep32_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6302903) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_1k1k_con1024_ctx1_dep4_gen1_dep32_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6302903) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_1k1k_con1024_ctx1_dep4_gen1_dep8_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6323889) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_1k1k_con2048_ctx2_dep4_gen1_dep16_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6323889) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_1k1k_con2048_ctx2_dep4_gen1_dep16_eplb288_mtp3_ccb-NIXL] SKIP (https://nvbugs/6323889) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_1k1k_con3072_ctx1_dep4_gen1_dep4_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6323889) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6302903) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_1k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6280721) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_1k1k_con2048_ctx1_dep4_gen1_dep4_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6374905) @@ -391,18 +387,11 @@ perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4 perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_32k4k_con256_ctx1_dep8_gen1_dep8_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6374893) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_8k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6280721) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_8k1k_con4096_ctx1_dep4_gen1_dep32_eplb256_mtp0_ccb-NIXL] SKIP (https://nvbugs/6302903) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_gpt-oss-120b-fp4_1k1k_con2048_ctx1_tp1_gen1_dep2_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6324123) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_gpt-oss-120b-fp4_1k1k_con512_ctx1_tp1_gen1_dep2_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6324123) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6324123) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_kimi-k25-thinking-fp4_1k1k_con4096_ctx1_dep4_gen1_dep8_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6323074) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6388238) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_deepseek-r1-fp4_1k1k_con3072_ctx1_dep4_gen1_dep4_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6323889) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_glm-5-fp4_1k1k_con4096_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL] SKIP (https://nvbugs/6324131) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL] SKIP (https://nvbugs/6324131) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_glm-5-fp4_8k1k_con1_ctx1_dep2_gen1_tep8_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6324131) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_kimi-k25-thinking-fp4_1k1k_con4096_ctx1_dep4_gen1_dep8_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6323074) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6323074) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL] SKIP (https://nvbugs/6323074) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6374893) perf/test_visual_gen_perf_sanity.py::test_visual_gen_e2e[vg_upload-ltx2_blackwell-ltx2_2stage_bf16_i2v_cfg2_ulysses4_compile_on] SKIP (https://nvbugs/6294413) perf/test_visual_gen_perf_sanity.py::test_visual_gen_e2e[vg_upload-ltx2_blackwell-ltx2_2stage_bf16_t2v_cfg2_ulysses4_compile_on] SKIP (https://nvbugs/6294413) diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py new file mode 100644 index 000000000000..0ebdce7015b8 --- /dev/null +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -0,0 +1,527 @@ +# 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. + +import sys +from types import SimpleNamespace +from unittest.mock import Mock, call + +import pytest + +from tensorrt_llm._torch.pyexecutor import kv_cache_transceiver as transceiver_module +from tensorrt_llm._torch.pyexecutor import py_executor as executor_module +from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import BindKvCacheTransceiver +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig + + +@pytest.fixture(autouse=True) +def _reset_inflight_cancel_env_cache(monkeypatch): + monkeypatch.delenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, raising=False) + monkeypatch.delenv(transceiver_module._NIXL_KVCACHE_BACKEND_ENV, raising=False) + monkeypatch.delenv(transceiver_module._DISABLE_KV_CACHE_TRANSFER_OVERLAP_ENV, raising=False) + monkeypatch.delenv(transceiver_module._DISAGG_LAYERWISE_ENV, raising=False) + monkeypatch.delenv(transceiver_module._TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV, raising=False) + for env_name, _ in transceiver_module._CACHE_TRANSCEIVER_BACKEND_ENV_VARS: + monkeypatch.delenv(env_name, raising=False) + monkeypatch.setattr(transceiver_module, "_disagg_inflight_cancel_enabled_cache", None) + + +def _supported_mapping(): + return SimpleNamespace(pp_size=1, cp_size=1, enable_attention_dp=False) + + +def _make_timeout_request(request_id=7, in_progress=False): + return SimpleNamespace( + is_attention_dp_dummy=False, + py_kv_transfer_timed_out=True, + py_request_id=request_id, + is_disagg_generation_transmission_in_progress=in_progress, + state=object(), + ) + + +def _make_response_handler_stub(active_requests, tp_allgather_result): + executor = object.__new__(PyExecutor) + executor.active_requests = list(active_requests) + executor.perf_manager = Mock() + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True + executor._disagg_inflight_cancel_unsupported_logged = False + executor.enable_attention_dp = True + executor.dist = SimpleNamespace( + rank=0, + world_size=2, + tp_allgather=Mock(return_value=tp_allgather_result), + ) + executor._enqueue_responses = Mock() + executor._terminate_request = Mock() + executor._handle_errors = Mock() + executor.canceled_req_ids = [] + executor._timeout_cleanup_order = Mock() + executor._timeout_cleanup_order.attach_mock(executor.dist.tp_allgather, "vote") + executor._timeout_cleanup_order.attach_mock(executor._handle_errors, "handle") + return executor + + +def test_flag_unset_short_circuits_before_capability_query(monkeypatch): + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor._disagg_inflight_cancel_unsupported_logged = False + monkeypatch.setattr(executor_module, "is_disagg_inflight_cancel_enabled", lambda: False) + + assert not PyExecutor._is_disagg_inflight_cancel_active(executor) + executor.kv_cache_transceiver.supports_inflight_request_cancellation.assert_not_called() + + +def test_unsupported_transceiver_warns_once(monkeypatch): + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = False + executor._disagg_inflight_cancel_unsupported_logged = False + monkeypatch.setattr(executor_module, "is_disagg_inflight_cancel_enabled", lambda: True) + warning = Mock() + monkeypatch.setattr(executor_module.logger, "warning", warning) + + assert not PyExecutor._is_disagg_inflight_cancel_active(executor) + assert not PyExecutor._is_disagg_inflight_cancel_active(executor) + + assert executor._disagg_inflight_cancel_unsupported_logged + warning.assert_called_once() + + +def test_flag_unset_generation_timeout_uses_rank_uniform_cleanup(): + request = _make_timeout_request() + executor = _make_response_handler_stub([request], [True, False]) + + PyExecutor._handle_responses(executor) + + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + assert executor.active_requests == [] + executor.dist.tp_allgather.assert_called_once_with(True) + executor._handle_errors.assert_called_once_with( + error_msg="Request timed out (KV transfer)", + requests=[request], + charge_budget=False, + ) + assert executor._timeout_cleanup_order.mock_calls == [ + call.vote(True), + call.handle( + error_msg="Request timed out (KV transfer)", + requests=[request], + charge_budget=False, + ), + ] + + +def test_flag_unset_generation_timeout_peer_enters_cleanup(): + executor = _make_response_handler_stub([], [False, True]) + + PyExecutor._handle_responses(executor) + + executor.kv_cache_transceiver.cancel_request.assert_not_called() + executor.dist.tp_allgather.assert_called_once_with(False) + executor._handle_errors.assert_called_once_with( + error_msg="Request timed out (KV transfer)", + requests=[], + charge_budget=False, + ) + assert executor._timeout_cleanup_order.mock_calls == [ + call.vote(False), + call.handle( + error_msg="Request timed out (KV transfer)", + requests=[], + charge_budget=False, + ), + ] + + +def test_flag_unset_context_timeout_preserves_legacy_cleanup(): + request = _make_timeout_request() + request.py_kv_transfer_start_time = 1.0 + request.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.check_context_transfer_status.return_value = ([], []) + executor.kv_cache_transceiver.cancel_request.return_value = True + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = { + request.py_request_id: request + } + executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True + executor._disagg_inflight_cancel_unsupported_logged = False + executor._end_transfer_and_maybe_terminate = Mock() + executor._check_cache_transfer_errors = Mock() + + PyExecutor._check_disagg_ctx_cache_transfer_status(executor, 0) + + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + assert request.py_kv_transfer_start_time is None + assert request.state == LlmRequestState.DISAGG_CONTEXT_COMPLETE + executor._end_transfer_and_maybe_terminate.assert_called_once_with(request) + + +def test_flag_unset_generation_driver_skips_cancel_pipeline(): + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True + executor._disagg_inflight_cancel_unsupported_logged = False + executor._check_disagg_gen_cache_transfer_status = Mock() + + PyExecutor._check_disagg_gen_transfer_status(executor) + + executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(0) + + +@pytest.mark.parametrize( + "backend,runtime,nixl_backend", + [ + ("UCX", None, "UCX"), + ("MPI", None, "UCX"), + ("MOONCAKE", None, "UCX"), + ("NIXL", "PYTHON", "UCX"), + ("NIXL", "CPP", "LIBFABRIC"), + ], +) +def test_feature_opt_in_rejects_unqualified_config(monkeypatch, backend, runtime, nixl_backend): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "1") + monkeypatch.setenv(transceiver_module._NIXL_KVCACHE_BACKEND_ENV, nixl_backend) + config = CacheTransceiverConfig(backend=backend, transceiver_runtime=runtime) + + with pytest.raises(ValueError, match="currently supported only"): + transceiver_module.create_kv_cache_transceiver( + _supported_mapping(), + Mock(), + Mock(), + Mock(), + config, + inflight_cancel_supported_by_executor=True, + ) + + +def test_feature_opt_in_accepts_cpp_nixl_ucx(monkeypatch): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "1") + monkeypatch.delenv(transceiver_module._NIXL_KVCACHE_BACKEND_ENV, raising=False) + config = CacheTransceiverConfig(backend="DEFAULT", max_tokens_in_buffer=512) + expected = object() + constructor = Mock(return_value=expected) + monkeypatch.setattr(transceiver_module, "BindKvCacheTransceiver", constructor) + + result = transceiver_module.create_kv_cache_transceiver( + _supported_mapping(), + Mock(), + Mock(), + Mock(), + config, + inflight_cancel_supported_by_executor=True, + ) + + assert result is expected + assert config.backend == "NIXL" + constructor.assert_called_once() + + +def test_feature_opt_in_rejects_ambiguous_legacy_backend_env(monkeypatch): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "1") + monkeypatch.setenv("TRTLLM_USE_NIXL_KVCACHE", "1") + monkeypatch.setenv("TRTLLM_USE_UCX_KVCACHE", "1") + config = CacheTransceiverConfig(backend="DEFAULT") + + with pytest.raises(ValueError, match="multiple legacy backend selectors"): + transceiver_module.create_kv_cache_transceiver( + _supported_mapping(), Mock(), Mock(), Mock(), config + ) + + +def test_direct_cpp_wrapper_rejects_python_runtime_opt_in(monkeypatch): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "1") + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") + + with pytest.raises(ValueError, match="currently supported only"): + BindKvCacheTransceiver( + _supported_mapping(), + Mock(), + Mock(), + Mock(), + config, + inflight_cancel_supported_by_executor=True, + ) + + +def test_flag_unset_preserves_existing_backend_selection(monkeypatch): + config = CacheTransceiverConfig(backend="UCX") + expected = object() + constructor = Mock(return_value=expected) + monkeypatch.setattr(transceiver_module, "BindKvCacheTransceiver", constructor) + + result = transceiver_module.create_kv_cache_transceiver(Mock(), Mock(), Mock(), Mock(), config) + + assert result is expected + assert config.backend == "UCX" + constructor.assert_called_once() + + +def test_flag_unset_preserves_python_transceiver(monkeypatch): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") + expected = object() + constructor = Mock(return_value=expected) + fake_module = SimpleNamespace(KvCacheTransceiverV2=constructor) + monkeypatch.setitem(sys.modules, "tensorrt_llm._torch.disaggregation.transceiver", fake_module) + + result = transceiver_module.create_kv_cache_transceiver(Mock(), Mock(), Mock(), Mock(), config) + + assert result is expected + constructor.assert_called_once() + + +def test_flag_unset_preserves_libfabric_selection(monkeypatch): + monkeypatch.setenv(transceiver_module._NIXL_KVCACHE_BACKEND_ENV, "LIBFABRIC") + config = CacheTransceiverConfig(backend="NIXL") + expected = object() + constructor = Mock(return_value=expected) + monkeypatch.setattr(transceiver_module, "BindKvCacheTransceiver", constructor) + + result = transceiver_module.create_kv_cache_transceiver(Mock(), Mock(), Mock(), Mock(), config) + + assert result is expected + constructor.assert_called_once() + + +@pytest.mark.parametrize( + "selector,expected_backend", + [ + ("TRTLLM_USE_NIXL_KVCACHE", "NIXL"), + ("TRTLLM_USE_UCX_KVCACHE", "UCX"), + ("TRTLLM_USE_MOONCAKE_KVCACHE", "MOONCAKE"), + ("TRTLLM_USE_MPI_KVCACHE", "MPI"), + ], +) +def test_flag_unset_preserves_legacy_backend_env(monkeypatch, selector, expected_backend): + monkeypatch.setenv(selector, "1") + config = CacheTransceiverConfig(backend="DEFAULT") + constructor = Mock(return_value=object()) + monkeypatch.setattr(transceiver_module, "BindKvCacheTransceiver", constructor) + + transceiver_module.create_kv_cache_transceiver(Mock(), Mock(), Mock(), Mock(), config) + + assert config.backend == expected_backend + constructor.assert_called_once() + + +def test_flag_unset_preserves_legacy_backend_env_precedence(monkeypatch): + for selector in ( + "TRTLLM_USE_MPI_KVCACHE", + "TRTLLM_USE_MOONCAKE_KVCACHE", + "TRTLLM_USE_UCX_KVCACHE", + "TRTLLM_USE_NIXL_KVCACHE", + ): + monkeypatch.setenv(selector, "1") + config = CacheTransceiverConfig(backend="DEFAULT") + monkeypatch.setattr(transceiver_module, "BindKvCacheTransceiver", Mock()) + + transceiver_module.create_kv_cache_transceiver(Mock(), Mock(), Mock(), Mock(), config) + + assert config.backend == "NIXL" + + +@pytest.mark.parametrize( + "backend,runtime,nixl_backend,disable_overlap,expected", + [ + ("NIXL", None, "UCX", None, True), + ("NIXL", "CPP", "UCX", None, True), + ("NIXL", "CPP", "UCX", "1", False), + ("NIXL", "PYTHON", "UCX", None, False), + ("NIXL", "CPP", "LIBFABRIC", None, False), + ("UCX", "CPP", "UCX", None, False), + ], +) +def test_cpp_capability_is_config_scoped( + monkeypatch, backend, runtime, nixl_backend, disable_overlap, expected +): + monkeypatch.setenv(transceiver_module._NIXL_KVCACHE_BACKEND_ENV, nixl_backend) + if disable_overlap is not None: + monkeypatch.setenv( + transceiver_module._DISABLE_KV_CACHE_TRANSFER_OVERLAP_ENV, + disable_overlap, + ) + config = CacheTransceiverConfig( + backend=backend, transceiver_runtime=runtime, max_tokens_in_buffer=512 + ) + mapping = _supported_mapping() + + assert ( + transceiver_module._is_disagg_inflight_cancel_config_supported( + config, mapping, inflight_cancel_supported_by_executor=True + ) + is expected + ) + + monkeypatch.setattr(transceiver_module, "mapping_to_world_config", lambda mapping: object()) + constructor = Mock(return_value=Mock()) + monkeypatch.setattr(transceiver_module, "CacheTransceiverCpp", constructor) + monkeypatch.setattr(CacheTransceiverConfig, "_to_pybind", lambda config: object()) + dist = Mock() + dist.pp_allgather.return_value = [1] + kv_cache_manager = SimpleNamespace( + total_num_kv_heads_per_layer=[1], + head_dim=64, + tokens_per_block=32, + dtype=object(), + num_kv_heads_per_layer=[1], + impl=object(), + ) + + transceiver = BindKvCacheTransceiver( + mapping, + dist, + kv_cache_manager, + Mock(), + config, + inflight_cancel_supported_by_executor=True, + ) + + assert transceiver.supports_inflight_request_cancellation() is expected + constructor.assert_called_once() + assert constructor.call_args.args[-1] is False + + +def test_cpp_constructor_receives_effective_opt_in(monkeypatch): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "1") + monkeypatch.setattr(transceiver_module, "mapping_to_world_config", lambda mapping: object()) + constructor = Mock(return_value=Mock()) + monkeypatch.setattr(transceiver_module, "CacheTransceiverCpp", constructor) + monkeypatch.setattr(CacheTransceiverConfig, "_to_pybind", lambda config: object()) + dist = Mock() + dist.pp_allgather.return_value = [1] + kv_cache_manager = SimpleNamespace( + total_num_kv_heads_per_layer=[1], + head_dim=64, + tokens_per_block=32, + dtype=object(), + num_kv_heads_per_layer=[1], + impl=object(), + ) + + BindKvCacheTransceiver( + _supported_mapping(), + dist, + kv_cache_manager, + Mock(), + CacheTransceiverConfig( + backend="NIXL", + transceiver_runtime="CPP", + max_tokens_in_buffer=512, + ), + inflight_cancel_supported_by_executor=True, + ) + + assert constructor.call_args.args[-1] is True + + +@pytest.mark.parametrize( + "qualifier", + [ + "max_tokens_in_buffer", + "kv_transfer_timeout_ms", + "pipeline_parallelism", + "context_parallelism", + "attention_dp", + "layerwise", + "zero_copy", + "executor", + ], +) +def test_feature_opt_in_rejects_unqualified_safety_path(monkeypatch, qualifier): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "1") + config_kwargs = { + "backend": "NIXL", + "transceiver_runtime": "CPP", + "max_tokens_in_buffer": 512, + } + mapping = _supported_mapping() + supported_by_executor = True + if qualifier == "max_tokens_in_buffer": + config_kwargs["max_tokens_in_buffer"] = None + elif qualifier == "kv_transfer_timeout_ms": + config_kwargs["kv_transfer_timeout_ms"] = None + elif qualifier == "pipeline_parallelism": + mapping.pp_size = 2 + elif qualifier == "context_parallelism": + mapping.cp_size = 2 + elif qualifier == "attention_dp": + mapping.enable_attention_dp = True + elif qualifier == "layerwise": + monkeypatch.setenv(transceiver_module._DISAGG_LAYERWISE_ENV, "1") + elif qualifier == "zero_copy": + monkeypatch.setenv(transceiver_module._TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV, "1") + elif qualifier == "executor": + supported_by_executor = False + + with pytest.raises(ValueError, match="currently supported only"): + transceiver_module.create_kv_cache_transceiver( + mapping, + Mock(), + Mock(), + Mock(), + CacheTransceiverConfig(**config_kwargs), + inflight_cancel_supported_by_executor=supported_by_executor, + ) + + +def test_flag_unset_keeps_unsupported_executor_inactive(monkeypatch): + config = CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="CPP", max_tokens_in_buffer=512 + ) + monkeypatch.setattr(transceiver_module, "mapping_to_world_config", lambda mapping: object()) + constructor = Mock(return_value=Mock()) + monkeypatch.setattr(transceiver_module, "CacheTransceiverCpp", constructor) + monkeypatch.setattr(CacheTransceiverConfig, "_to_pybind", lambda config: object()) + dist = Mock() + dist.pp_allgather.return_value = [1] + kv_cache_manager = SimpleNamespace( + total_num_kv_heads_per_layer=[1], + head_dim=64, + tokens_per_block=32, + dtype=object(), + num_kv_heads_per_layer=[1], + impl=object(), + ) + + transceiver = BindKvCacheTransceiver( + _supported_mapping(), dist, kv_cache_manager, Mock(), config + ) + + assert not transceiver.supports_inflight_request_cancellation() + assert constructor.call_args.args[-1] is False + + +def test_python_transceiver_capability_defaults_to_unsupported(): + from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 + + transceiver = object.__new__(KvCacheTransceiverV2) + + assert not transceiver.supports_inflight_request_cancellation() + assert not transceiver.has_poisoned_transfer_buffer() + + +def test_flag_unset_skips_cpp_poison_query(): + transceiver = SimpleNamespace( + impl=Mock(), + _inflight_request_cancellation_enabled=False, + ) + + assert not BindKvCacheTransceiver.has_poisoned_transfer_buffer(transceiver) + transceiver.impl.has_poisoned_transfer_buffer.assert_not_called() diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_lifecycle.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_lifecycle.py new file mode 100644 index 000000000000..d67e5f0387d2 --- /dev/null +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_lifecycle.py @@ -0,0 +1,727 @@ +# 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. + +import threading +from queue import Queue +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from tensorrt_llm._torch.pyexecutor import py_executor as executor_module +from tensorrt_llm._torch.pyexecutor.executor_request_queue import ( + ExecutorRequestQueue, + RequestQueueItem, +) +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm.bindings.executor import FinishReason + + +class _WaitingQueue: + def __init__(self, items=()): + self._items = list(items) + + def __bool__(self): + return bool(self._items) + + def pop_request(self): + return self._items.pop(0) + + def remove_by_ids(self, request_ids): + self._items = [item for item in self._items if item.id not in request_ids] + + +@pytest.fixture(autouse=True) +def _clear_unsafe_shutdown_quarantine(): + with PyExecutor._UNSAFE_TRANSFER_SHUTDOWN_QUARANTINE_LOCK: + PyExecutor._UNSAFE_TRANSFER_SHUTDOWN_QUARANTINE.clear() + yield + with PyExecutor._UNSAFE_TRANSFER_SHUTDOWN_QUARANTINE_LOCK: + PyExecutor._UNSAFE_TRANSFER_SHUTDOWN_QUARANTINE.clear() + + +def _request(request_id=7, state=LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS): + response = SimpleNamespace(result=SimpleNamespace(cached_tokens=None)) + request = SimpleNamespace( + py_request_id=request_id, + py_client_id=17, + is_child=False, + state=state, + py_decoding_iter=3, + py_draft_tokens=[], + cached_tokens=11, + create_response=Mock(return_value=response), + finish_by_reason=Mock(), + ) + request.finish_by_reason.side_effect = lambda _: setattr( + request, "state", LlmRequestState.GENERATION_COMPLETE + ) + return request, response + + +def _cancel_executor(request): + executor = object.__new__(PyExecutor) + executor.active_requests = [request] + executor.canceled_req_ids = [request.py_request_id] + executor.previous_batch = None + executor.waiting_queue = _WaitingQueue() + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = {} + executor.dist = SimpleNamespace(rank=0) + executor._maybe_attach_ctx_usage = Mock() + executor._enqueue_responses = Mock() + executor._terminate_request = Mock() + return executor + + +def _fatal_executor(active_requests=(), transfer_requests=(), *, should_store_blocks=False): + executor = object.__new__(PyExecutor) + executor.active_requests = list(active_requests) + executor.waiting_queue = _WaitingQueue() + raw_queue = Queue() + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_request_queue.return_value = raw_queue + executor.executor_request_queue.enqueue_shutdown_request.side_effect = lambda: raw_queue.put( + RequestQueueItem(-1) + ) + executor.previous_batch = None + executor.request_accumulated = [] + executor.control_requests = [] + executor.control_request_barrier = Mock() + executor.control_action_done = Mock() + executor._pending_transfer_responses = [] + executor._fatal_transfer_cleanup_requests = {} + executor._fatal_transfer_cleanup_skip_termination_ids = set() + executor._fatal_transfer_shutdown = False + executor._fatal_transfer_cleanup_complete = False + executor._unsafe_transfer_shutdown = False + executor._fatal_transfer_cleanup_deadline = None + executor._fatal_transfer_cleanup_timer = None + executor._fatal_transfer_cleanup_lock = threading.Lock() + executor._start_fatal_transfer_cleanup_watchdog = Mock( + side_effect=lambda: setattr(executor, "_fatal_transfer_cleanup_deadline", float("inf")) + ) + executor._fatal_error = None + executor._error_budget = Mock(budget=1.0) + executor.is_shutdown = False + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=0, world_size=1) + executor.kv_cache_transceiver = Mock() + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.should_store_blocks = should_store_blocks + executor.async_transfer_manager.requests_in_transfer.return_value = dict(transfer_requests) + executor._check_disagg_ctx_cache_transfer_status = Mock() + executor._check_disagg_gen_cache_transfer_status = Mock() + executor._enqueue_responses = Mock() + executor._terminate_request = Mock() + return executor + + +def test_context_success_waits_for_pending_user_cancel(): + request, response = _request(state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS) + executor = _cancel_executor(request) + executor._fatal_transfer_cleanup_requests = {} + executor._pending_transfer_responses = [] + executor._is_disagg_inflight_cancel_active = Mock(return_value=True) + + def finish_transfer(_): + request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE + return True + + executor.async_transfer_manager.end_transfer.side_effect = finish_transfer + + PyExecutor._end_transfer_and_maybe_terminate(executor, request) + + request.create_response.assert_not_called() + assert executor._pending_transfer_responses == [] + assert executor.active_requests == [request] + executor._terminate_request.assert_not_called() + + executor.async_transfer_manager.requests_in_transfer.return_value = {} + finished = PyExecutor._handle_canceled_requests(executor) + PyExecutor._finish_canceled_requests_without_forward(executor, finished) + + request.finish_by_reason.assert_called_once_with(FinishReason.CANCELLED) + request.create_response.assert_called_once_with(False, 0) + assert response.result.cached_tokens == request.cached_tokens + executor._enqueue_responses.assert_called_once_with([(request.py_request_id, response)]) + executor._terminate_request.assert_called_once_with(request) + assert executor.active_requests == [] + assert executor.canceled_req_ids == [] + + +def test_transfer_only_cancel_issues_once_then_finishes_after_poll(): + request, response = _request() + executor = _cancel_executor(request) + executor._fatal_error = None + executor._is_disagg_inflight_cancel_active = Mock(return_value=True) + executor._check_disagg_gen_cache_transfer_status = Mock( + side_effect=lambda _: setattr(request, "state", LlmRequestState.DISAGG_TRANS_ERROR) + ) + + PyExecutor._check_disagg_gen_transfer_status(executor) + + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + request.finish_by_reason.assert_called_once_with(FinishReason.CANCELLED) + executor._enqueue_responses.assert_called_once_with([(request.py_request_id, response)]) + executor._terminate_request.assert_called_once_with(request) + assert executor.canceled_req_ids == [] + + +def test_transfer_only_cancel_does_not_terminate_previous_forward(): + request, _ = _request() + executor = _cancel_executor(request) + executor.previous_batch = SimpleNamespace( + scheduled_requests=SimpleNamespace(all_requests=Mock(return_value=[request])) + ) + + finished = PyExecutor._handle_canceled_requests(executor, transfer_only=True) + + assert finished == [] + executor.kv_cache_transceiver.cancel_request.assert_not_called() + request.finish_by_reason.assert_not_called() + assert executor.canceled_req_ids == [request.py_request_id] + + +def test_cancel_waits_for_secondary_async_owner(): + request, _ = _request(state=LlmRequestState.DISAGG_TRANS_ERROR) + executor = _cancel_executor(request) + executor.async_transfer_manager.requests_in_transfer.return_value = { + request.py_request_id: request + } + + finished = PyExecutor._handle_canceled_requests(executor) + + assert finished == [] + request.finish_by_reason.assert_not_called() + assert executor.canceled_req_ids == [request.py_request_id] + + +def test_active_timeout_is_owned_by_cpp_not_python(): + request, _ = _request() + request.py_kv_transfer_start_time = 0.0 + request.py_kv_transfer_timed_out = False + request.is_disagg_generation_transmission_in_progress = True + executor = _cancel_executor(request) + executor.kv_cache_transceiver.kv_transfer_timeout_ms = 1 + executor._is_disagg_inflight_cancel_active = Mock(return_value=True) + + PyExecutor._check_kv_transfer_timeout(executor) + + assert not request.py_kv_transfer_timed_out + executor.async_transfer_manager.requests_in_transfer.assert_not_called() + executor.kv_cache_transceiver.cancel_request.assert_not_called() + + +def test_poison_fatal_closes_and_drains_queues_without_freeing_active_requests(): + transfer_request, _ = _request() + unrelated_transfer, _ = _request(request_id=8) + waiting_request = SimpleNamespace(client_id=31) + raw_request = SimpleNamespace(client_id=32) + accumulated_request = SimpleNamespace(client_id=33) + executor = object.__new__(PyExecutor) + executor.active_requests = [transfer_request, unrelated_transfer] + executor.waiting_queue = _WaitingQueue([RequestQueueItem(21, waiting_request)]) + raw_queue = Queue() + raw_queue.put(RequestQueueItem(22, raw_request)) + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_request_queue.return_value = raw_queue + executor.executor_request_queue.enqueue_shutdown_request.side_effect = lambda: raw_queue.put( + RequestQueueItem(-1) + ) + executor.request_accumulated = [RequestQueueItem(23, accumulated_request)] + executor.control_requests = [RequestQueueItem(-2)] + executor.control_request_barrier = Mock() + executor._start_fatal_transfer_cleanup_watchdog = Mock() + executor.previous_batch = None + executor._pending_transfer_responses = [] + executor._fatal_transfer_cleanup_requests = {} + executor._fatal_error = None + executor.is_shutdown = False + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=0) + executor.kv_cache_transceiver = Mock() + executor._enqueue_responses = Mock() + executor._terminate_request = Mock() + + PyExecutor._handle_errors( + executor, "poisoned", charge_budget=False, force_fatal=True, defer_transfer_cleanup=True + ) + PyExecutor._handle_errors( + executor, "poisoned", charge_budget=False, force_fatal=True, defer_transfer_cleanup=True + ) + + executor.executor_request_queue.enqueue_shutdown_request.assert_called_once_with() + assert not executor.waiting_queue + assert raw_queue.empty() + assert executor.request_accumulated == [] + assert executor.control_requests == [] + executor.control_request_barrier.set.assert_called_once_with() + assert executor.active_requests == [] + assert executor._fatal_transfer_cleanup_requests == { + transfer_request.py_request_id: transfer_request, + unrelated_transfer.py_request_id: unrelated_transfer, + } + executor._terminate_request.assert_not_called() + executor.kv_cache_transceiver.cancel_request.assert_not_called() + executor._enqueue_responses.assert_called_once() + response_ids = {request_id for request_id, _ in executor._enqueue_responses.call_args.args[0]} + assert response_ids == { + transfer_request.py_request_id, + unrelated_transfer.py_request_id, + 21, + 22, + 23, + } + + +@pytest.mark.parametrize("should_store_blocks", [False, True]) +def test_poison_tracks_detached_context_owner_until_quiescence( + should_store_blocks, +): + request, _ = _request(state=LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS) + executor = _fatal_executor( + transfer_requests={request.py_request_id: request}, + should_store_blocks=should_store_blocks, + ) + + PyExecutor._handle_errors( + executor, + "poisoned", + charge_budget=False, + force_fatal=True, + defer_transfer_cleanup=True, + ) + + assert executor._fatal_transfer_cleanup_requests == {request.py_request_id: request} + assert not PyExecutor._fatal_shutdown_complete(executor) + executor.kv_cache_transceiver.cancel_request.assert_not_called() + response_ids = { + request_id + for call_args in executor._enqueue_responses.call_args_list + for request_id, _ in call_args.args[0] + } + assert request.py_request_id not in response_ids + executor._terminate_request.assert_not_called() + + assert PyExecutor._progress_fatal_transfer_cleanup(executor) + assert request.py_request_id in executor._fatal_transfer_cleanup_requests + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + assert not PyExecutor._fatal_shutdown_complete(executor) + + request.state = LlmRequestState.DISAGG_TRANS_ERROR + executor.async_transfer_manager.requests_in_transfer.return_value = {} + + assert PyExecutor._progress_fatal_transfer_cleanup(executor) + assert executor._fatal_transfer_cleanup_requests == {} + assert PyExecutor._fatal_shutdown_complete(executor) + if should_store_blocks: + executor._terminate_request.assert_not_called() + else: + executor._terminate_request.assert_called_once_with(request) + + +def test_fatal_cleanup_reaps_only_after_native_owner_quiesces(monkeypatch): + request, _ = _request() + executor = _fatal_executor() + executor._fatal_transfer_cleanup_requests = {request.py_request_id: request} + executor._fatal_transfer_shutdown = True + executor._fatal_transfer_cleanup_deadline = float("inf") + executor._request_has_transfer_ownership = Mock(side_effect=[True, False]) + monkeypatch.setattr(executor_module.time, "sleep", Mock()) + + assert PyExecutor._progress_fatal_transfer_cleanup(executor) + assert request.py_request_id in executor._fatal_transfer_cleanup_requests + executor._terminate_request.assert_not_called() + + assert PyExecutor._progress_fatal_transfer_cleanup(executor) + assert executor._fatal_transfer_cleanup_requests == {} + executor._terminate_request.assert_called_once_with(request) + assert request.state == LlmRequestState.GENERATION_COMPLETE + assert executor.kv_cache_transceiver.cancel_request.call_count == 2 + + +def test_fatal_cleanup_deadline_marks_transfer_unsafe(monkeypatch): + request, _ = _request() + executor = _fatal_executor( + active_requests=[request], + transfer_requests={request.py_request_id: request}, + ) + monotonic = Mock(return_value=101.0) + monkeypatch.setattr(PyExecutor, "FATAL_TRANSFER_CLEANUP_GRACE_PERIOD_S", 0.5) + monkeypatch.setattr(executor_module.time, "monotonic", monotonic) + + PyExecutor._handle_errors( + executor, + "poisoned", + charge_budget=False, + force_fatal=True, + defer_transfer_cleanup=True, + ) + executor._fatal_transfer_cleanup_deadline = 100.0 + + assert PyExecutor._progress_fatal_transfer_cleanup(executor) + assert executor._unsafe_transfer_shutdown + assert executor._fatal_transfer_cleanup_complete + assert request.py_request_id in executor._fatal_transfer_cleanup_requests + executor._terminate_request.assert_not_called() + executor.kv_cache_transceiver.cancel_request.assert_called_once_with(request) + executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) + executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(0) + + +def test_fatal_cleanup_watchdog_wakes_shutdown_while_cleanup_is_blocked(monkeypatch): + request, _ = _request() + executor = _fatal_executor(transfer_requests={request.py_request_id: request}) + executor._fatal_transfer_cleanup_requests = {request.py_request_id: request} + executor._fatal_transfer_shutdown = True + executor.shutdown_event = threading.Event() + del executor._start_fatal_transfer_cleanup_watchdog + + timer = Mock() + timer_factory = Mock(return_value=timer) + monkeypatch.setattr(executor_module.threading, "Timer", timer_factory) + monkeypatch.setattr(executor_module.time, "monotonic", Mock(return_value=100.0)) + + entered = threading.Event() + release = threading.Event() + + def block_cleanup(): + entered.set() + release.wait() + + executor.previous_batch = SimpleNamespace( + sample_state=SimpleNamespace(sampler_event=Mock(side_effect=block_cleanup)) + ) + PyExecutor._start_fatal_transfer_cleanup_watchdog(executor) + watchdog_callback = timer_factory.call_args.args[1] + cleanup_thread = threading.Thread( + target=PyExecutor._progress_fatal_transfer_cleanup, + args=(executor,), + daemon=True, + ) + cleanup_thread.start() + assert entered.wait(timeout=1.0) + + watchdog_callback() + + assert executor.shutdown_event.wait(timeout=0.1) + assert executor._unsafe_transfer_shutdown + assert not executor._fatal_transfer_cleanup_complete + assert not PyExecutor._fatal_shutdown_complete(executor) + assert cleanup_thread.is_alive() + executor._terminate_request.assert_not_called() + assert any( + quarantined is executor for quarantined in PyExecutor._UNSAFE_TRANSFER_SHUTDOWN_QUARANTINE + ) + + release.set() + cleanup_thread.join(timeout=1.0) + assert not cleanup_thread.is_alive() + assert executor._fatal_transfer_cleanup_complete + assert PyExecutor._fatal_shutdown_complete(executor) + + +@pytest.mark.parametrize("failure_site", ["cancel", "context_status"]) +def test_fatal_cleanup_exception_marks_transfer_unsafe(failure_site): + request, _ = _request() + executor = _fatal_executor(transfer_requests={request.py_request_id: request}) + executor._fatal_transfer_cleanup_requests = {request.py_request_id: request} + executor._fatal_transfer_shutdown = True + executor._fatal_transfer_cleanup_complete = False + executor._fatal_transfer_cleanup_deadline = float("inf") + error = RuntimeError(f"{failure_site} failed") + if failure_site == "cancel": + executor.kv_cache_transceiver.cancel_request.side_effect = error + else: + executor._check_disagg_ctx_cache_transfer_status.side_effect = error + + assert PyExecutor._progress_fatal_transfer_cleanup(executor) + + assert executor._unsafe_transfer_shutdown + assert executor._fatal_transfer_cleanup_complete + assert request.py_request_id in executor._fatal_transfer_cleanup_requests + executor._terminate_request.assert_not_called() + executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) + executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(0) + + +def test_fatal_cleanup_empty_rank_participates_and_rejects_owner_mismatch(): + executor = _fatal_executor() + executor._fatal_transfer_shutdown = True + executor._fatal_transfer_cleanup_deadline = float("inf") + executor.dist = SimpleNamespace( + rank=0, + world_size=2, + tp_allgather=Mock( + return_value=[ + (False, False, False, ()), + (False, False, True, (7,)), + ] + ), + ) + + assert PyExecutor._progress_fatal_transfer_cleanup(executor) + + executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) + executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(0) + executor.dist.tp_allgather.assert_called_once() + assert executor._unsafe_transfer_shutdown + assert executor._fatal_transfer_cleanup_complete + + +@pytest.mark.parametrize( + ("unsafe", "has_pending"), + [(True, False), (False, True)], +) +def test_unsafe_transfer_shutdown_skips_resource_teardown(unsafe, has_pending): + executor = object.__new__(PyExecutor) + executor.executor_request_queue = Mock() + executor.shutdown_event = Mock() + executor.hang_detector = Mock() + executor.hang_detector.detected.return_value = False + executor._fatal_transfer_shutdown = True + executor._unsafe_transfer_shutdown = unsafe + executor._fatal_transfer_cleanup_requests = {7: Mock()} if has_pending else {} + executor.worker_thread = Mock() + manager = Mock() + executor.resource_manager = SimpleNamespace(resource_managers={"kv": manager}) + + PyExecutor.shutdown(executor) + + executor.executor_request_queue.enqueue_shutdown_request.assert_called_once_with() + executor.shutdown_event.wait.assert_called_once_with() + executor.worker_thread.join.assert_not_called() + manager.shutdown.assert_not_called() + + +def test_status_poll_uses_cpp_topology_poison_without_python_allreduce(): + executor = object.__new__(PyExecutor) + executor.active_requests = [] + executor.canceled_req_ids = [] + executor.enable_attention_dp = False + executor.dist = SimpleNamespace(world_size=2, tp_allreduce=Mock()) + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.requests_in_transfer.return_value = {} + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.check_gen_transfer_status.return_value = None + executor.kv_cache_transceiver.has_poisoned_transfer_buffer.return_value = False + executor._is_disagg_inflight_cancel_active = Mock(return_value=True) + + PyExecutor._check_disagg_gen_cache_transfer_status(executor, 0) + + executor.kv_cache_transceiver.has_poisoned_transfer_buffer.assert_called_once_with() + executor.dist.tp_allreduce.assert_not_called() + + +def test_poison_is_fatal_before_any_request_is_terminal(): + executor = object.__new__(PyExecutor) + executor.active_requests = [] + executor.canceled_req_ids = [] + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.check_gen_transfer_status.return_value = None + executor.kv_cache_transceiver.has_poisoned_transfer_buffer.return_value = True + executor._is_disagg_inflight_cancel_active = Mock(return_value=True) + executor._handle_errors = Mock() + + PyExecutor._check_disagg_gen_cache_transfer_status(executor, 0) + + executor._handle_errors.assert_called_once_with( + "Error in kv cache transfer for generation requests; poisoned " + "transfer buffer requires process restart", + charge_budget=False, + force_fatal=True, + defer_transfer_cleanup=True, + ) + + +def test_fatal_poison_is_started_only_once(): + executor = object.__new__(PyExecutor) + executor._fatal_error = RuntimeError("already fatal") + executor._fatal_transfer_shutdown = True + executor._error_budget = Mock() + + PyExecutor._handle_errors( + executor, "poisoned", charge_budget=False, force_fatal=True, defer_transfer_cleanup=True + ) + + executor._error_budget.consume.assert_not_called() + + +def test_generic_fatal_preserves_overlap_previous_batch(): + request, _ = _request() + previous_batch = SimpleNamespace(sample_state=SimpleNamespace(sampler_event=Mock())) + raw_queue = Queue() + executor = object.__new__(PyExecutor) + executor.active_requests = [request] + executor.waiting_queue = _WaitingQueue() + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_request_queue.return_value = raw_queue + executor.executor_request_queue.enqueue_shutdown_request.side_effect = lambda: raw_queue.put( + RequestQueueItem(-1) + ) + executor.previous_batch = previous_batch + executor._pending_transfer_responses = [] + executor._fatal_transfer_cleanup_requests = {} + executor._fatal_transfer_cleanup_skip_termination_ids = set() + executor._fatal_transfer_shutdown = False + executor._fatal_transfer_cleanup_complete = False + executor._unsafe_transfer_shutdown = False + executor._fatal_transfer_cleanup_deadline = None + executor._fatal_error = None + executor._error_budget = Mock(budget=0.0) + executor._error_budget.consume.return_value = True + executor.is_shutdown = False + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=0) + executor._enqueue_responses = Mock() + executor._terminate_request = Mock() + + PyExecutor._handle_errors(executor, "sampling failed") + + assert executor.previous_batch is previous_batch + previous_batch.sample_state.sampler_event.synchronize.assert_not_called() + assert not PyExecutor._fatal_shutdown_complete(executor) + executor.executor_request_queue.enqueue_shutdown_request.assert_called_with() + assert executor.executor_request_queue.enqueue_shutdown_request.call_count == 2 + assert raw_queue.get_nowait().is_shutdown_request + + +def test_control_action_aborts_after_fatal_shutdown(): + executor = object.__new__(PyExecutor) + executor.dist = SimpleNamespace(rank=0) + executor.executor_request_queue = Mock() + executor.control_request_barrier = Mock() + executor.control_action_done = Mock() + executor._fatal_error = RuntimeError("poisoned") + yielded = False + + with pytest.raises(RuntimeError): + with PyExecutor.control_action(executor): + yielded = True + + assert not yielded + executor.executor_request_queue.enqueue_control_request.assert_not_called() + executor.control_request_barrier.wait.assert_not_called() + executor.control_action_done.set.assert_not_called() + + +def test_control_action_rejects_ordinary_shutdown(): + executor = object.__new__(PyExecutor) + executor.dist = SimpleNamespace(rank=0) + executor.executor_request_queue = Mock() + executor.control_request_barrier = Mock() + executor.control_action_done = Mock() + executor._fatal_error = None + executor.is_shutdown = True + yielded = False + + with pytest.raises(RuntimeError): + with PyExecutor.control_action(executor): + yielded = True + + assert not yielded + executor.executor_request_queue.enqueue_control_request.assert_not_called() + executor.control_request_barrier.wait.assert_not_called() + executor.control_action_done.set.assert_not_called() + + +def test_control_action_rejects_queue_shutdown_race(): + executor = object.__new__(PyExecutor) + executor.dist = SimpleNamespace(rank=0) + executor.executor_request_queue = Mock() + executor.executor_request_queue.enqueue_control_request.return_value = False + executor.control_request_barrier = Mock() + executor.control_action_done = Mock() + executor._fatal_error = None + executor.is_shutdown = False + yielded = False + + with pytest.raises(RuntimeError): + with PyExecutor.control_action(executor): + yielded = True + + assert not yielded + executor.executor_request_queue.enqueue_control_request.assert_called_once_with(drain=True) + executor.control_request_barrier.wait.assert_not_called() + executor.control_action_done.set.assert_not_called() + + +def test_control_enqueue_and_shutdown_are_lock_ordered(): + for _ in range(20): + request_queue = ExecutorRequestQueue( + dist=SimpleNamespace(rank=0), + max_batch_size=8, + enable_iter_perf_stats=False, + batch_wait_timeout_ms=0, + ) + start = threading.Barrier(3) + control_accepted = [] + + def enqueue_control(): + start.wait() + control_accepted.append(request_queue.enqueue_control_request()) + + def enqueue_shutdown(): + start.wait() + request_queue.enqueue_shutdown_request() + + control_thread = threading.Thread(target=enqueue_control) + shutdown_thread = threading.Thread(target=enqueue_shutdown) + control_thread.start() + shutdown_thread.start() + start.wait() + control_thread.join(timeout=1.0) + shutdown_thread.join(timeout=1.0) + + assert not control_thread.is_alive() + assert not shutdown_thread.is_alive() + items = [] + raw_queue = request_queue.get_request_queue() + while not raw_queue.empty(): + items.append(raw_queue.get_nowait()) + if control_accepted == [True]: + assert [item.is_control_request for item in items] == [True, False] + assert items[-1].is_shutdown_request + else: + assert control_accepted == [False] + assert len(items) == 1 + assert items[0].is_shutdown_request + + +def test_waiting_control_action_aborts_when_fatal_shutdown_releases_barrier(): + executor = object.__new__(PyExecutor) + executor.dist = SimpleNamespace(rank=0) + executor.executor_request_queue = Mock() + executor.control_request_barrier = Mock() + executor.control_action_done = Mock() + executor._fatal_error = None + executor.control_request_barrier.wait.side_effect = lambda: setattr( + executor, "_fatal_error", RuntimeError("poisoned") + ) + yielded = False + + with pytest.raises(RuntimeError): + with PyExecutor.control_action(executor): + yielded = True + + assert not yielded + executor.executor_request_queue.enqueue_control_request.assert_called_once_with(drain=True) + executor.control_request_barrier.wait.assert_called_once_with() + executor.control_request_barrier.clear.assert_called_once_with() + executor.control_action_done.set.assert_not_called() diff --git a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py index 1d5c00617c6e..8dc33dc776f1 100644 --- a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py +++ b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py @@ -656,7 +656,13 @@ def test_cross_kv_cache_manager_is_stored(self): ) assert scheduler.cross_kv_cache_manager is cross_mgr - def test_factory_forwards_encoder_init_until_state_for_cross_pool(self): + @pytest.mark.parametrize( + "has_kv_connector,expected_inflight_cancel_support", + [(False, True), (True, False)], + ) + def test_factory_forwards_encoder_init_until_state_for_cross_pool( + self, has_kv_connector, expected_inflight_cancel_support + ): """The executor factory must widen V2 scheduling to ENCODER_INIT. Without this, V2 enc-dec requests are filtered by the default @@ -695,6 +701,7 @@ def test_factory_forwards_encoder_init_until_state_for_cross_pool(self): enable_early_first_token_response=False, kv_cache_config=SimpleNamespace(enable_kv_pool_rebalance=False), ) + kv_connector_manager = Mock() if has_kv_connector else None with ( patch( @@ -707,7 +714,7 @@ def test_factory_forwards_encoder_init_until_state_for_cross_pool(self): patch( "tensorrt_llm._torch.pyexecutor._util.create_kv_cache_transceiver", return_value=None, - ), + ) as create_transceiver, patch( "tensorrt_llm._torch.pyexecutor._util.PyExecutor", ), @@ -723,6 +730,7 @@ def test_factory_forwards_encoder_init_until_state_for_cross_pool(self): start_worker=False, sampler=Mock(), drafter=None, + kv_connector_manager=kv_connector_manager, max_seq_len=128, max_batch_size=8, max_beam_width=1, @@ -732,6 +740,10 @@ def test_factory_forwards_encoder_init_until_state_for_cross_pool(self): kwargs = scheduler_cls.call_args.kwargs assert kwargs["cross_kv_cache_manager"] is cross_mgr assert kwargs["no_schedule_until_state"] == LlmRequestState.ENCODER_INIT + assert ( + create_transceiver.call_args.kwargs["inflight_cancel_supported_by_executor"] + is expected_inflight_cancel_support + ) # --------------------------------------------------------------------------- diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index bf07bc9dbcc8..2ca9e10a744b 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1,3 +1,18 @@ +# 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. + """Tests for PyExecutor request handling functionality. This module tests the request handling logic that was moved from ExecutorRequestQueue @@ -1074,16 +1089,23 @@ def _err_req(): return _make_adp_request(LlmRequestState.DISAGG_TRANS_ERROR) +_DEFAULT_KV_CACHE_TRANSCEIVER = object() + + def _make_disagg_err_stub( *, enable_attention_dp=True, - kv_cache_transceiver=object(), + kv_cache_transceiver=_DEFAULT_KV_CACHE_TRANSCEIVER, world_size=2, active_requests=None, tp_allgather_result=None, ): stub = types.SimpleNamespace() stub.enable_attention_dp = enable_attention_dp + if kv_cache_transceiver is _DEFAULT_KV_CACHE_TRANSCEIVER: + kv_cache_transceiver = Mock() + kv_cache_transceiver.supports_inflight_request_cancellation.return_value = False + kv_cache_transceiver.has_poisoned_transfer_buffer.return_value = False stub.kv_cache_transceiver = kv_cache_transceiver stub.active_requests = active_requests if active_requests is not None else [] stub.dist = Mock() diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py b/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py index 4387d3d2730c..ec915999d80c 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py @@ -223,6 +223,7 @@ def test_create_executor_uses_cache_transceiver(cache_attention_type, expected_a _, _, passed_kv_cache_manager, attention_type, passed_config = create_transceiver.call_args.args assert passed_kv_cache_manager is kv_cache_manager assert attention_type == expected_attention_type + assert create_transceiver.call_args.kwargs["inflight_cancel_supported_by_executor"] is False assert passed_config is ad_config.cache_transceiver_config assert passed_config.max_tokens_in_buffer == mock_engine.cache_seq_interface.info.max_seq_len assert create_transceiver.call_args.kwargs["mamba_cache_manager"] is None diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index ccc059b1ab99..04f6b0d00274 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -1,4 +1,20 @@ +# 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. + import gc +import multiprocessing import sys import time import uuid @@ -11,6 +27,8 @@ import tensorrt_llm.bindings import tensorrt_llm.bindings.executor as trtllm from tensorrt_llm._torch.distributed import Distributed +from tensorrt_llm._torch.pyexecutor import \ + kv_cache_transceiver as transceiver_module from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import \ create_kv_cache_transceiver from tensorrt_llm._torch.pyexecutor.llm_request import (LlmRequest, @@ -32,10 +50,14 @@ KV_TRANSFER_COMPLETION_MARGIN_S = 10.0 -def create_kv_cache_manager(mapping, dtype): +def create_kv_cache_manager(mapping, + dtype, + max_tokens=256, + max_seq_len=256, + max_batch_size=1): return KVCacheManager( trtllm.KvCacheConfig( - max_tokens=256, + max_tokens=max_tokens, enable_block_reuse=False, ), tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, @@ -43,8 +65,8 @@ def create_kv_cache_manager(mapping, dtype): num_kv_heads=1, head_dim=1, tokens_per_block=8, - max_seq_len=256, - max_batch_size=1, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, mapping=mapping, dtype=dtype) @@ -224,6 +246,313 @@ def transfers_done(): kv_cache_transceiver_ctx) +def _run_cpp_nixl_sync_transfer_stress(): + """Exercise repeated empty-to-nonempty sender transitions in a child.""" + request_count = 64 + prompt_len = 16 + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + manager_kwargs = { + "max_tokens": request_count * prompt_len * 2, + "max_seq_len": prompt_len, + "max_batch_size": request_count, + } + kv_cache_manager_ctx = create_kv_cache_manager(mapping, DataType.HALF, + **manager_kwargs) + kv_cache_manager_gen = create_kv_cache_manager(mapping, DataType.HALF, + **manager_kwargs) + cache_transceiver_config = CacheTransceiverConfig(backend="NIXL", + transceiver_runtime="CPP", + max_tokens_in_buffer=512) + transceiver_ctx = create_kv_cache_transceiver(mapping, dist, + kv_cache_manager_ctx, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config) + transceiver_gen = create_kv_cache_transceiver(mapping, dist, + kv_cache_manager_gen, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config) + + try: + sampling_config = tensorrt_llm.bindings.SamplingConfig( + SamplingParams()._get_sampling_config()) + ctx_requests = [ + LlmRequest( + request_id=request_id, + max_new_tokens=1, + input_tokens=list(range(prompt_len)), + sampling_config=sampling_config, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + for request_id in range(request_count) + ] + kv_cache_manager_ctx.impl.add_sequence_batch( + [(request.py_request_id, request.prompt_len, 1) + for request in ctx_requests], ctx_requests) + fill_kv_cache_buffer(kv_cache_manager_ctx) + + transceiver_ctx.respond_and_send_async(ctx_requests[0]) + completed_ctx_ids = set() + + def poll_context_transfers(): + completed, failed = transceiver_ctx.check_context_transfer_status(1) + assert failed == [] + completed_ctx_ids.update(completed) + + for request_index, ctx_request in enumerate(ctx_requests): + gen_request = LlmRequest( + request_id=request_index, + max_new_tokens=1, + input_tokens=list(range(prompt_len)), + sampling_config=sampling_config, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + context_phase_params=ctx_request.context_phase_params) + kv_cache_manager_gen.impl.add_sequence_batch( + [(gen_request.py_request_id, gen_request.prompt_len, 1)], + [gen_request]) + + transceiver_gen.request_and_receive_sync(gen_request) + + # Queue the next response immediately so its insertion can overlap + # the previous response's sender-side cleanup. + if request_index + 1 < request_count: + transceiver_ctx.respond_and_send_async( + ctx_requests[request_index + 1]) + + assert (gen_request.state == + LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE) + + ctx_block_ids = kv_cache_manager_ctx.get_cache_indices(ctx_request) + gen_block_ids = kv_cache_manager_gen.get_cache_indices(gen_request) + assert torch.equal( + kv_cache_manager_ctx.get_unique_primary_pool()[ctx_block_ids], + kv_cache_manager_gen.get_unique_primary_pool()[gen_block_ids], + ), f"different KV-cache values for request {request_index}" + + expected_ctx_id = ctx_request.py_request_id + wait_for_transfer_completion(poll_context_transfers, + lambda request_id=expected_ctx_id: + request_id in completed_ctx_ids) + finally: + shutdown_transceivers(transceiver_gen, transceiver_ctx) + + +@pytest.mark.timeout(150) +def test_cpp_nixl_sync_transfer_stress(): + """C++ NIXL sync transfers must not lose sender wakeups between requests.""" + process = multiprocessing.get_context("spawn").Process( + target=_run_cpp_nixl_sync_transfer_stress, + name="cpp-nixl-sync-transfer-stress") + process.start() + process.join(timeout=120) + try: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if process.is_alive(): + process.kill() + process.join(timeout=5) + pytest.fail("C++ NIXL synchronous KV transfer stress test hung") + assert process.exitcode == 0, ( + f"C++ NIXL synchronous KV transfer child exited with " + f"code {process.exitcode}") + finally: + process.close() + + +@pytest.mark.timeout(120) +def test_cpp_nixl_active_cancellation_reaches_terminal_error(monkeypatch): + """Qualified C++ NIXL cancellation must quiesce both transfer sides.""" + _configure_cpp_nixl_cancel_env(monkeypatch, "1") + + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + kv_cache_manager_ctx = create_kv_cache_manager(mapping, DataType.HALF) + kv_cache_manager_gen = create_kv_cache_manager(mapping, DataType.HALF) + cache_transceiver_config = CacheTransceiverConfig( + backend="NIXL", + transceiver_runtime="CPP", + max_tokens_in_buffer=512, + kv_transfer_timeout_ms=5000, + kv_transfer_sender_future_timeout_ms=10, + kv_transfer_poll_interval_ms=10) + transceiver_ctx = create_kv_cache_transceiver( + mapping, + dist, + kv_cache_manager_ctx, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config, + inflight_cancel_supported_by_executor=True) + transceiver_gen = create_kv_cache_transceiver( + mapping, + dist, + kv_cache_manager_gen, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config, + inflight_cancel_supported_by_executor=True) + + try: + assert transceiver_ctx.supports_inflight_request_cancellation() + assert transceiver_gen.supports_inflight_request_cancellation() + assert transceiver_ctx._inflight_request_cancellation_enabled + assert transceiver_gen._inflight_request_cancellation_enabled + + fill_kv_cache_buffer(kv_cache_manager_ctx) + sampling_config = tensorrt_llm.bindings.SamplingConfig( + SamplingParams()._get_sampling_config()) + ctx_request = LlmRequest( + request_id=0, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=sampling_config, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + kv_cache_manager_ctx.impl.add_sequence_batch( + [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], + [ctx_request]) + transceiver_ctx.respond_and_send_async(ctx_request) + + # Cancel before starting GEN. This deterministically exercises the + # sender's queued-cancellation path instead of racing a tiny transfer. + assert ctx_request.state == LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + assert transceiver_ctx.cancel_request(ctx_request) + + gen_request = LlmRequest( + request_id=0, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=sampling_config, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + context_phase_params=ctx_request.context_phase_params) + kv_cache_manager_gen.impl.add_sequence_batch( + [(gen_request.py_request_id, gen_request.prompt_len, 1)], + [gen_request]) + transceiver_gen.request_and_receive_async(gen_request) + + failed_ctx_ids = set() + + def poll_transfers(): + completed, failed = transceiver_ctx.check_context_transfer_status(1) + assert completed == [] + failed_ctx_ids.update(failed) + transceiver_gen.check_gen_transfer_status(1) + + def cancellation_quiesced(): + return (ctx_request.py_request_id in failed_ctx_ids + and ctx_request.state == LlmRequestState.DISAGG_TRANS_ERROR + and gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR + and transceiver_gen.check_gen_transfer_complete()) + + wait_for_transfer_completion(poll_transfers, + cancellation_quiesced, + timeout_s=10.0) + + # A cancellation committed before ready=true never exposes a transfer + # buffer to a peer, so both sides should quiesce without poisoning it. + assert not transceiver_ctx.has_poisoned_transfer_buffer() + assert not transceiver_gen.has_poisoned_transfer_buffer() + finally: + shutdown_transceivers(transceiver_gen, transceiver_ctx) + + +def _configure_cpp_nixl_cancel_env(monkeypatch, mode): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, + mode) + monkeypatch.setenv(transceiver_module._NIXL_KVCACHE_BACKEND_ENV, "UCX") + monkeypatch.delenv( + transceiver_module._DISABLE_KV_CACHE_TRANSFER_OVERLAP_ENV, + raising=False) + monkeypatch.delenv(transceiver_module._DISAGG_LAYERWISE_ENV, raising=False) + monkeypatch.delenv(transceiver_module._TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV, + raising=False) + for env_name, _ in transceiver_module._CACHE_TRANSCEIVER_BACKEND_ENV_VARS: + monkeypatch.delenv(env_name, raising=False) + monkeypatch.setattr(transceiver_module, + "_disagg_inflight_cancel_mode_cache", None) + + +@pytest.mark.timeout(120) +def test_cpp_nixl_sync_rejection_preserves_transfer_error(monkeypatch): + """Default-off synchronous receive must not rewrite rejection as success.""" + _configure_cpp_nixl_cancel_env(monkeypatch, "0") + + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + kv_cache_manager_ctx = create_kv_cache_manager(mapping, DataType.HALF) + kv_cache_manager_gen = create_kv_cache_manager(mapping, DataType.HALF) + cache_transceiver_config = CacheTransceiverConfig( + backend="NIXL", + transceiver_runtime="CPP", + max_tokens_in_buffer=512, + kv_transfer_timeout_ms=5000, + kv_transfer_sender_future_timeout_ms=10, + kv_transfer_poll_interval_ms=10) + transceiver_ctx = create_kv_cache_transceiver( + mapping, + dist, + kv_cache_manager_ctx, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config, + inflight_cancel_supported_by_executor=True) + transceiver_gen = create_kv_cache_transceiver( + mapping, + dist, + kv_cache_manager_gen, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config, + inflight_cancel_supported_by_executor=True) + + try: + assert transceiver_ctx.supports_inflight_request_cancellation() + assert transceiver_gen.supports_inflight_request_cancellation() + assert not transceiver_ctx._inflight_request_cancellation_enabled + assert not transceiver_gen._inflight_request_cancellation_enabled + + fill_kv_cache_buffer(kv_cache_manager_ctx) + sampling_config = tensorrt_llm.bindings.SamplingConfig( + SamplingParams()._get_sampling_config()) + ctx_request = LlmRequest( + request_id=0, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=sampling_config, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + kv_cache_manager_ctx.impl.add_sequence_batch( + [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], + [ctx_request]) + transceiver_ctx.respond_and_send_async(ctx_request) + + # Cancel before starting GEN. This deterministically exercises the + # sender's queued-cancellation path instead of racing a tiny transfer. + assert ctx_request.state == LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + assert transceiver_ctx.cancel_request(ctx_request) + + gen_request = LlmRequest( + request_id=0, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=sampling_config, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + context_phase_params=ctx_request.context_phase_params) + kv_cache_manager_gen.impl.add_sequence_batch( + [(gen_request.py_request_id, gen_request.prompt_len, 1)], + [gen_request]) + transceiver_gen.request_and_receive_sync(gen_request) + + assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR + completed_ctx_ids, failed_ctx_ids = ( + transceiver_ctx.check_context_transfer_status(1)) + assert completed_ctx_ids == [] + assert set(failed_ctx_ids) == {ctx_request.py_request_id} + assert ctx_request.state == LlmRequestState.DISAGG_TRANS_ERROR + finally: + shutdown_transceivers(transceiver_gen, transceiver_ctx) + + @pytest.mark.timeout(120) @pytest.mark.parametrize("attention_type", [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], From 83a056568e65b4aada836b04b01b5df95e466a37 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 30 Jun 2026 11:28:38 -0700 Subject: [PATCH 2/2] [TRTLLM-12721][feat] Enable qualified disagg cancellation by default Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- cpp/tensorrt_llm/common/envUtils.cpp | 16 ++- docs/source/features/disagg-serving.md | 38 ++++++ .../_torch/auto_deploy/shim/ad_executor.py | 24 ++-- .../_torch/pyexecutor/kv_cache_transceiver.py | 75 +++++++++--- tensorrt_llm/_torch/pyexecutor/py_executor.py | 7 +- tensorrt_llm/llmapi/llm_args.py | 16 ++- .../test_disagg_inflight_cancel_gate.py | 111 ++++++++++++++---- .../singlegpu/shim/test_create_ad_executor.py | 35 ++++++ 8 files changed, 264 insertions(+), 58 deletions(-) diff --git a/cpp/tensorrt_llm/common/envUtils.cpp b/cpp/tensorrt_llm/common/envUtils.cpp index 3a48bb602b48..f7ebd84dabcc 100644 --- a/cpp/tensorrt_llm/common/envUtils.cpp +++ b/cpp/tensorrt_llm/common/envUtils.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include TRTLLM_NAMESPACE_BEGIN @@ -397,7 +398,20 @@ bool getEnvTryZCopyForKVCacheTransfer() bool getEnvDisaggEnableInflightCancel() { - static bool const enabled = getBoolEnv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL"); + static bool const enabled = []() + { + auto const value = getStrEnv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL"); + if (!value.has_value() || value.value() == "0") + { + return false; + } + if (value.value() == "1") + { + return true; + } + throw std::invalid_argument( + "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL must be unset, 0, or 1; got '" + value.value() + "'."); + }(); return enabled; } diff --git a/docs/source/features/disagg-serving.md b/docs/source/features/disagg-serving.md index f38e3c818f57..a50b6730e429 100644 --- a/docs/source/features/disagg-serving.md +++ b/docs/source/features/disagg-serving.md @@ -237,6 +237,44 @@ Example (two-node deployment): TRT-LLM uses some environment variables to control the behavior of disaggregated service. +* `TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL`: Controls cancellation of active KV cache transfers and quarantine of + transfer buffers whose completion cannot be proven. The value is strict and has the following behavior: + + | Value | Qualified configuration | Unqualified configuration | + | ----- | ----------------------- | ------------------------- | + | Unset | Enabled automatically | Coordinated in-flight cancellation remains disabled | + | `0` | Disabled | Disabled | + | `1` | Enabled and required | Startup fails with a configuration error | + + Any other value is invalid. A configuration is qualified only when all of the following are true: + + - The standard PyTorch `PyExecutor` is used without a simultaneous KV connector. + - The asynchronous C++ cache transceiver uses `backend: NIXL` with + `TRTLLM_NIXL_KVCACHE_BACKEND=UCX` (the default NIXL plugin). + - `max_tokens_in_buffer` and `kv_transfer_timeout_ms` are both positive. `PyExecutor` resolves an omitted + `max_tokens_in_buffer` to the model maximum sequence length before creating the transceiver; the transfer timeout + defaults to 60,000 ms. + - `PP=1`, `CP=1`, attention DP is disabled, and DWDP is disabled. Tensor parallelism is supported. + - Transfer overlap is enabled, and layerwise and zero-copy transfers are disabled. + + Consequently, AutoDeploy, the legacy TensorRT executor, the Python transceiver, direct UCX, MPI, Mooncake, + NIXL-libfabric, synchronous receive, PP, CP, attention DP, DWDP, layerwise transfer, zero-copy transfer, + no-deadline, unbounded-admission, and KV-connector paths do not enable coordinated in-flight cancellation when the + variable is unset; their ordinary successful-transfer path remains unchanged. Set the variable to `1` when + deployment validation should reject any such fallback, including a missing cache transceiver or use of the legacy + TensorRT executor. + + **Fail-closed behavior:** Cancelling a NIXL operation does not prove that its peer has stopped accessing the + advertised buffer. TRT-LLM therefore poisons the affected staging-buffer slot and terminates the executor instead + of reusing that memory. Deferred un-poison is not implemented, so an in-flight cancellation can require a process + or pod restart. + + **Rolling upgrades:** Before upgrading, set `TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=0` on every context and generation + worker. Upgrade all workers, verify that they use the same TensorRT-LLM version and transfer configuration, and only + then remove the override. Older workers interpret an unset variable as disabled, while workers with this policy can + enable it automatically. C++ validates the effective mode across local parallel ranks, but context and generation + services do not negotiate mixed-version policy. + * `TRTLLM_NIXL_KVCACHE_BACKEND`: When using NIXL as the cache transceiver backend, this variable specifies the underlying communication backend for NIXL. Valid options are: - `UCX` (default) - `LIBFABRIC` (available from v0.16.0) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index 561aa4169130..37824f8cefda 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -1249,7 +1249,7 @@ def create_autodeploy_executor( ) cache_transceiver_config = ad_config.cache_transceiver_config - kv_cache_transceiver = None + attention_type_cpp = AttentionTypeCpp.DEFAULT if cache_transceiver_config is not None and cache_transceiver_config.backend is not None: if isinstance(kv_cache_manager, BaseMambaCacheManager): # See https://github.com/NVIDIA/TensorRT-LLM/issues/14320. @@ -1275,15 +1275,19 @@ def create_autodeploy_executor( raise TypeError(f"attention_type must be AttentionType, got {cache_attention_type!r}") attention_type_cpp = _ATTENTION_TYPE_TO_CPP[cache_attention_type] - kv_cache_transceiver = create_kv_cache_transceiver( - dist_mapping, - dist, - kv_cache_manager, - attention_type_cpp, - cache_transceiver_config, - mamba_cache_manager=None, - inflight_cancel_supported_by_executor=False, - ) + # Call the shared factory even when no transceiver is configured so an + # explicit TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 cannot silently fall + # back to an unsupported AutoDeploy path. AUTO and disabled mode still + # return None and preserve the existing aggregate-serving behavior. + kv_cache_transceiver = create_kv_cache_transceiver( + dist_mapping, + dist, + kv_cache_manager, + attention_type_cpp, + cache_transceiver_config, + mamba_cache_manager=None, + inflight_cancel_supported_by_executor=False, + ) # Guided (structured) decoding. guided_decoder = None diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index 7cda19640e17..5fa166bf3631 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -14,6 +14,7 @@ # limitations under the License. from abc import ABC, abstractmethod +from enum import Enum from os import getenv from typing import Any, Dict, List, Optional @@ -46,21 +47,52 @@ ("TRTLLM_USE_MOONCAKE_KVCACHE", "MOONCAKE"), ("TRTLLM_USE_MPI_KVCACHE", "MPI"), ) -_disagg_inflight_cancel_enabled_cache: Optional[bool] = None + + +class _DisaggInflightCancelMode(Enum): + AUTO = "auto" + DISABLED = "0" + REQUIRED = "1" + + +_disagg_inflight_cancel_mode_cache: Optional[_DisaggInflightCancelMode] = None + + +def _get_disagg_inflight_cancel_mode() -> _DisaggInflightCancelMode: + global _disagg_inflight_cancel_mode_cache + if _disagg_inflight_cancel_mode_cache is not None: + return _disagg_inflight_cancel_mode_cache + + value = getenv(_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV) + if value is None: + mode = _DisaggInflightCancelMode.AUTO + elif value == "0": + mode = _DisaggInflightCancelMode.DISABLED + elif value == "1": + mode = _DisaggInflightCancelMode.REQUIRED + logger.warning( + f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV}=1: disagg KV " + "transfer in-flight cancellation was requested. It is active " + "only for cache transceivers that advertise support.") + else: + raise ValueError( + f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV} must be unset, 0, or 1; " + f"got {value!r}.") + + _disagg_inflight_cancel_mode_cache = mode + return mode def is_disagg_inflight_cancel_enabled() -> bool: - """Return whether disaggregated in-flight KV transfer cancellation is enabled.""" - global _disagg_inflight_cancel_enabled_cache - if _disagg_inflight_cancel_enabled_cache is None: - _disagg_inflight_cancel_enabled_cache = (getenv( - _DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "0") == "1") - if _disagg_inflight_cancel_enabled_cache: - logger.warning( - f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV}=1: disagg KV " - "transfer in-flight cancellation was requested. It is active " - "only for cache transceivers that advertise support.") - return _disagg_inflight_cancel_enabled_cache + """Return whether configuration-scoped in-flight cancellation is permitted.""" + return (_get_disagg_inflight_cancel_mode() + is not _DisaggInflightCancelMode.DISABLED) + + +def is_disagg_inflight_cancel_required() -> bool: + """Return whether in-flight cancellation was explicitly required.""" + return (_get_disagg_inflight_cancel_mode() + is _DisaggInflightCancelMode.REQUIRED) def _is_disagg_inflight_cancel_config_supported( @@ -77,6 +109,7 @@ def _is_disagg_inflight_cancel_config_supported( and getenv(_DISABLE_KV_CACHE_TRANSFER_OVERLAP_ENV) != "1" and mapping.pp_size == 1 and mapping.cp_size == 1 and not mapping.enable_attention_dp + and not getattr(mapping, "dwdp_enabled", False) and getenv(_DISAGG_LAYERWISE_ENV) != "1" and getenv(_TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV) != "1") @@ -85,7 +118,7 @@ def _validate_disagg_inflight_cancel_config( cache_transceiver_config: CacheTransceiverConfig, mapping: Mapping, inflight_cancel_supported_by_executor: bool = False) -> None: - if not is_disagg_inflight_cancel_enabled(): + if not is_disagg_inflight_cancel_required(): return enabled_backend_env_vars = [ @@ -109,8 +142,8 @@ def _validate_disagg_inflight_cancel_config( layerwise = getenv(_DISAGG_LAYERWISE_ENV) try_zcopy = getenv(_TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV) raise ValueError( - f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV}=1 is experimental and " - "currently supported only by the standard PyTorch backend " + f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV}=1 is currently supported " + "only by the standard PyTorch backend " "PyExecutor with asynchronous transceiver_runtime='CPP', " "backend='NIXL', and TRTLLM_NIXL_KVCACHE_BACKEND=UCX on ordinary " "TP topology, using a positive max_tokens_in_buffer admission " @@ -123,11 +156,14 @@ def _validate_disagg_inflight_cancel_config( f"TRTLLM_NIXL_KVCACHE_BACKEND={nixl_backend!r}, " f"pp_size={mapping.pp_size!r}, cp_size={mapping.cp_size!r}, " f"enable_attention_dp={mapping.enable_attention_dp!r}, " + f"dwdp_enabled={getattr(mapping, 'dwdp_enabled', False)!r}, " "inflight_cancel_supported_by_executor=" f"{inflight_cancel_supported_by_executor!r}, " f"{_DISABLE_KV_CACHE_TRANSFER_OVERLAP_ENV}={disable_overlap!r}, " f"{_DISAGG_LAYERWISE_ENV}={layerwise!r}, " - f"{_TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV}={try_zcopy!r}.") + f"{_TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV}={try_zcopy!r}. Unset " + f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV}, or set it to 0, to use " + "this configuration without coordinated in-flight cancellation.") def mapping_to_world_config(mapping: Mapping) -> WorldConfig: @@ -146,10 +182,15 @@ def create_kv_cache_transceiver( dist: Distributed, kv_cache_manager: KVCacheManager, attention_type: AttentionTypeCpp, - cache_transceiver_config: CacheTransceiverConfig, + cache_transceiver_config: Optional[CacheTransceiverConfig], mamba_cache_manager: Optional[BaseMambaCacheManager] = None, inflight_cancel_supported_by_executor: bool = False): if cache_transceiver_config is None or cache_transceiver_config.backend is None: + if is_disagg_inflight_cancel_required(): + raise ValueError( + f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV}=1 requires a " + "configured, supported cache transceiver; no cache transceiver " + "backend was configured.") logger.info("cache_transceiver is disabled") return None diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index ab9f5a04cc1d..7c323b6f0f03 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -81,7 +81,8 @@ from .kv_cache_manager_v2 import KVCacheManagerV2 from .kv_cache_stats import append_kv_cache_iteration_stats from .kv_cache_transceiver import (KvCacheTransceiver, - is_disagg_inflight_cancel_enabled) + is_disagg_inflight_cancel_enabled, + is_disagg_inflight_cancel_required) from .llm_request import (ATTENTION_DP_DUMMY_REQUEST_ID, MAX_SPEC_DECODE_POSITIONS, ExecutorRequest, LlmRequest, LlmRequestState, LlmResponse, @@ -4660,8 +4661,8 @@ def _is_disagg_inflight_cancel_active(self) -> bool: if callable(supports) and supports() is True: return True - if not getattr(self, "_disagg_inflight_cancel_unsupported_logged", - False): + if (is_disagg_inflight_cancel_required() and not getattr( + self, "_disagg_inflight_cancel_unsupported_logged", False)): logger.warning( "TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 was requested, but " f"{type(transceiver).__name__} does not advertise in-flight " diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 286205fb2eb6..2b26b30503e2 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3482,15 +3482,19 @@ class CacheTransceiverConfig(StrictBaseModel, PybindMirror): kv_transfer_timeout_ms: Optional[PositiveInt] = Field( default=60000, - description= - "Timeout in milliseconds for KV cache transfer. Requests exceeding this timeout will be cancelled." - ) + description=( + "Timeout in milliseconds for detecting stalled KV cache transfers. " + "Qualified C++ NIXL/UCX configurations use this deadline for " + "coordinated in-flight cancellation; other configurations handle " + "the timeout according to their transceiver and runtime.")) kv_transfer_sender_future_timeout_ms: Optional[PositiveInt] = Field( default=1000, - description= - "Timeout in milliseconds to wait for the sender future to be ready when scheduled batch size is 0. This allows the request to be eventually cancelled by the user or because of kv_transfer_timeout_ms" - ) + description=( + "Bounded wait in milliseconds for polling the sender future when " + "the scheduled batch size is 0. This lets the executor observe " + "cancellation and KV-transfer timeout state without blocking " + "indefinitely.")) kv_transfer_poll_interval_ms: Optional[PositiveInt] = Field( default=5000, diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py index 0ebdce7015b8..0ccb76badea0 100644 --- a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -36,11 +36,16 @@ def _reset_inflight_cancel_env_cache(monkeypatch): monkeypatch.delenv(transceiver_module._TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV, raising=False) for env_name, _ in transceiver_module._CACHE_TRANSCEIVER_BACKEND_ENV_VARS: monkeypatch.delenv(env_name, raising=False) - monkeypatch.setattr(transceiver_module, "_disagg_inflight_cancel_enabled_cache", None) + monkeypatch.setattr(transceiver_module, "_disagg_inflight_cancel_mode_cache", None) def _supported_mapping(): - return SimpleNamespace(pp_size=1, cp_size=1, enable_attention_dp=False) + return SimpleNamespace( + pp_size=1, + cp_size=1, + enable_attention_dp=False, + dwdp_enabled=False, + ) def _make_timeout_request(request_id=7, in_progress=False): @@ -77,7 +82,40 @@ def _make_response_handler_stub(active_requests, tp_allgather_result): return executor -def test_flag_unset_short_circuits_before_capability_query(monkeypatch): +@pytest.mark.parametrize( + "override,expected_enabled,expected_required,expected_warnings", + [ + (None, True, False, 0), + ("0", False, False, 0), + ("1", True, True, 1), + ], +) +def test_feature_policy_is_strict_tri_state( + monkeypatch, + override, + expected_enabled, + expected_required, + expected_warnings, +): + if override is not None: + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, override) + warning = Mock() + monkeypatch.setattr(transceiver_module.logger, "warning", warning) + + assert transceiver_module.is_disagg_inflight_cancel_enabled() is expected_enabled + assert transceiver_module.is_disagg_inflight_cancel_required() is expected_required + assert warning.call_count == expected_warnings + + +@pytest.mark.parametrize("override", ["", "2", "true"]) +def test_feature_policy_rejects_invalid_value(monkeypatch, override): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, override) + + with pytest.raises(ValueError, match="must be unset, 0, or 1"): + transceiver_module.is_disagg_inflight_cancel_enabled() + + +def test_feature_disabled_short_circuits_before_capability_query(monkeypatch): executor = object.__new__(PyExecutor) executor.kv_cache_transceiver = Mock() executor._disagg_inflight_cancel_unsupported_logged = False @@ -87,23 +125,35 @@ def test_flag_unset_short_circuits_before_capability_query(monkeypatch): executor.kv_cache_transceiver.supports_inflight_request_cancellation.assert_not_called() -def test_unsupported_transceiver_warns_once(monkeypatch): +@pytest.mark.parametrize("required,expected_warnings", [(False, 0), (True, 1)]) +def test_unsupported_transceiver_warns_only_when_required(monkeypatch, required, expected_warnings): executor = object.__new__(PyExecutor) executor.kv_cache_transceiver = Mock() executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = False executor._disagg_inflight_cancel_unsupported_logged = False monkeypatch.setattr(executor_module, "is_disagg_inflight_cancel_enabled", lambda: True) + monkeypatch.setattr(executor_module, "is_disagg_inflight_cancel_required", lambda: required) warning = Mock() monkeypatch.setattr(executor_module.logger, "warning", warning) assert not PyExecutor._is_disagg_inflight_cancel_active(executor) assert not PyExecutor._is_disagg_inflight_cancel_active(executor) - assert executor._disagg_inflight_cancel_unsupported_logged - warning.assert_called_once() + assert executor._disagg_inflight_cancel_unsupported_logged is required + assert warning.call_count == expected_warnings + +def test_auto_mode_supported_transceiver_is_active(): + executor = object.__new__(PyExecutor) + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True + executor._disagg_inflight_cancel_unsupported_logged = False -def test_flag_unset_generation_timeout_uses_rank_uniform_cleanup(): + assert PyExecutor._is_disagg_inflight_cancel_active(executor) + + +def test_feature_disabled_generation_timeout_uses_rank_uniform_cleanup(monkeypatch): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "0") request = _make_timeout_request() executor = _make_response_handler_stub([request], [True, False]) @@ -127,7 +177,8 @@ def test_flag_unset_generation_timeout_uses_rank_uniform_cleanup(): ] -def test_flag_unset_generation_timeout_peer_enters_cleanup(): +def test_feature_disabled_generation_timeout_peer_enters_cleanup(monkeypatch): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "0") executor = _make_response_handler_stub([], [False, True]) PyExecutor._handle_responses(executor) @@ -149,7 +200,8 @@ def test_flag_unset_generation_timeout_peer_enters_cleanup(): ] -def test_flag_unset_context_timeout_preserves_legacy_cleanup(): +def test_feature_disabled_context_timeout_preserves_legacy_cleanup(monkeypatch): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "0") request = _make_timeout_request() request.py_kv_transfer_start_time = 1.0 request.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS @@ -174,7 +226,8 @@ def test_flag_unset_context_timeout_preserves_legacy_cleanup(): executor._end_transfer_and_maybe_terminate.assert_called_once_with(request) -def test_flag_unset_generation_driver_skips_cancel_pipeline(): +def test_feature_disabled_generation_driver_skips_cancel_pipeline(monkeypatch): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "0") executor = object.__new__(PyExecutor) executor.kv_cache_transceiver = Mock() executor.kv_cache_transceiver.supports_inflight_request_cancellation.return_value = True @@ -234,6 +287,16 @@ def test_feature_opt_in_accepts_cpp_nixl_ucx(monkeypatch): constructor.assert_called_once() +@pytest.mark.parametrize("config", [None, CacheTransceiverConfig(backend=None)]) +def test_feature_opt_in_rejects_missing_transceiver(monkeypatch, config): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "1") + + with pytest.raises(ValueError, match="requires a configured, supported cache transceiver"): + transceiver_module.create_kv_cache_transceiver( + _supported_mapping(), Mock(), Mock(), Mock(), config + ) + + def test_feature_opt_in_rejects_ambiguous_legacy_backend_env(monkeypatch): monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "1") monkeypatch.setenv("TRTLLM_USE_NIXL_KVCACHE", "1") @@ -261,7 +324,8 @@ def test_direct_cpp_wrapper_rejects_python_runtime_opt_in(monkeypatch): ) -def test_flag_unset_preserves_existing_backend_selection(monkeypatch): +def test_feature_disabled_preserves_existing_backend_selection(monkeypatch): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "0") config = CacheTransceiverConfig(backend="UCX") expected = object() constructor = Mock(return_value=expected) @@ -274,7 +338,7 @@ def test_flag_unset_preserves_existing_backend_selection(monkeypatch): constructor.assert_called_once() -def test_flag_unset_preserves_python_transceiver(monkeypatch): +def test_auto_mode_preserves_python_transceiver(monkeypatch): config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") expected = object() constructor = Mock(return_value=expected) @@ -287,7 +351,7 @@ def test_flag_unset_preserves_python_transceiver(monkeypatch): constructor.assert_called_once() -def test_flag_unset_preserves_libfabric_selection(monkeypatch): +def test_auto_mode_preserves_libfabric_selection(monkeypatch): monkeypatch.setenv(transceiver_module._NIXL_KVCACHE_BACKEND_ENV, "LIBFABRIC") config = CacheTransceiverConfig(backend="NIXL") expected = object() @@ -309,7 +373,7 @@ def test_flag_unset_preserves_libfabric_selection(monkeypatch): ("TRTLLM_USE_MPI_KVCACHE", "MPI"), ], ) -def test_flag_unset_preserves_legacy_backend_env(monkeypatch, selector, expected_backend): +def test_auto_mode_preserves_legacy_backend_env(monkeypatch, selector, expected_backend): monkeypatch.setenv(selector, "1") config = CacheTransceiverConfig(backend="DEFAULT") constructor = Mock(return_value=object()) @@ -321,7 +385,7 @@ def test_flag_unset_preserves_legacy_backend_env(monkeypatch, selector, expected constructor.assert_called_once() -def test_flag_unset_preserves_legacy_backend_env_precedence(monkeypatch): +def test_auto_mode_preserves_legacy_backend_env_precedence(monkeypatch): for selector in ( "TRTLLM_USE_MPI_KVCACHE", "TRTLLM_USE_MOONCAKE_KVCACHE", @@ -395,11 +459,13 @@ def test_cpp_capability_is_config_scoped( assert transceiver.supports_inflight_request_cancellation() is expected constructor.assert_called_once() - assert constructor.call_args.args[-1] is False + assert constructor.call_args.args[-1] is expected -def test_cpp_constructor_receives_effective_opt_in(monkeypatch): - monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "1") +@pytest.mark.parametrize("override,expected_enabled", [(None, True), ("0", False), ("1", True)]) +def test_cpp_constructor_receives_effective_policy(monkeypatch, override, expected_enabled): + if override is not None: + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, override) monkeypatch.setattr(transceiver_module, "mapping_to_world_config", lambda mapping: object()) constructor = Mock(return_value=Mock()) monkeypatch.setattr(transceiver_module, "CacheTransceiverCpp", constructor) @@ -428,7 +494,7 @@ def test_cpp_constructor_receives_effective_opt_in(monkeypatch): inflight_cancel_supported_by_executor=True, ) - assert constructor.call_args.args[-1] is True + assert constructor.call_args.args[-1] is expected_enabled @pytest.mark.parametrize( @@ -439,6 +505,7 @@ def test_cpp_constructor_receives_effective_opt_in(monkeypatch): "pipeline_parallelism", "context_parallelism", "attention_dp", + "dwdp", "layerwise", "zero_copy", "executor", @@ -463,6 +530,8 @@ def test_feature_opt_in_rejects_unqualified_safety_path(monkeypatch, qualifier): mapping.cp_size = 2 elif qualifier == "attention_dp": mapping.enable_attention_dp = True + elif qualifier == "dwdp": + mapping.dwdp_enabled = True elif qualifier == "layerwise": monkeypatch.setenv(transceiver_module._DISAGG_LAYERWISE_ENV, "1") elif qualifier == "zero_copy": @@ -481,7 +550,7 @@ def test_feature_opt_in_rejects_unqualified_safety_path(monkeypatch, qualifier): ) -def test_flag_unset_keeps_unsupported_executor_inactive(monkeypatch): +def test_auto_mode_keeps_unsupported_executor_inactive(monkeypatch): config = CacheTransceiverConfig( backend="NIXL", transceiver_runtime="CPP", max_tokens_in_buffer=512 ) @@ -517,7 +586,7 @@ def test_python_transceiver_capability_defaults_to_unsupported(): assert not transceiver.has_poisoned_transfer_buffer() -def test_flag_unset_skips_cpp_poison_query(): +def test_feature_disabled_skips_cpp_poison_query(): transceiver = SimpleNamespace( impl=Mock(), _inflight_request_cancellation_enabled=False, diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py b/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py index ec915999d80c..f851fefe8ac1 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py @@ -23,6 +23,7 @@ from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import AttentionType from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs from tensorrt_llm._torch.auto_deploy.shim.ad_executor import create_autodeploy_executor +from tensorrt_llm._torch.pyexecutor import kv_cache_transceiver as transceiver_module from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import AttentionTypeCpp from tensorrt_llm.llmapi import CacheTransceiverConfig @@ -230,6 +231,40 @@ def test_create_executor_uses_cache_transceiver(cache_attention_type, expected_a assert result.kv_cache_transceiver is mock_transceiver +def test_create_executor_rejects_required_cancel_without_transceiver(monkeypatch): + monkeypatch.setenv( + transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, + "1", + ) + monkeypatch.setattr(transceiver_module, "_disagg_inflight_cancel_mode_cache", None) + ad_config = LlmArgs( + model="test-model", + max_batch_size=4, + max_seq_len=128, + max_input_len=64, + backend="_autodeploy", + cuda_graph_config={"max_batch_size": 4}, + ) + mock_engine, _ = make_mock_engine() + + with ( + patch("tensorrt_llm._torch.auto_deploy.shim.ad_executor.PyExecutor"), + patch( + "tensorrt_llm._torch.auto_deploy.shim.ad_executor.ADEngine.build_from_config", + return_value=mock_engine, + ), + patch( + "tensorrt_llm._torch.auto_deploy.llm_args.LlmArgs.create_factory", + return_value=MockFactory(vocab_size_padded=1000), + ), + ): + with pytest.raises( + ValueError, + match="requires a configured, supported cache transceiver", + ): + create_autodeploy_executor(ad_config, MockTokenizer()) + + @pytest.mark.parametrize( "cache_attention_type, expected_attention_type", [