From edf970576b44214ef7a45f4a27b73fb2bbdde256 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 30 Jun 2026 15:14:16 -0700 Subject: [PATCH 1/2] [TRTLLM-12721][fix] Harden disagg cancellation consensus and buffer ownership Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.h | 37 + .../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 | 608 +++++++++++++-- .../batch_manager/cacheTransferLayer.cpp | 9 +- .../batch_manager/cacheTransferLayer.h | 5 +- .../batch_manager/cacheTransferState.h | 138 ++++ .../batch_manager/dataTransceiver.cpp | 736 +++++++++++++----- .../batch_manager/dataTransceiver.h | 55 +- .../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 | 80 +- .../agent_utils/connection.h | 12 +- .../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 | 87 ++- .../batch_manager/cacheTransferStateTest.cpp | 101 +++ .../unit_tests/executor/agentCommTest.cpp | 36 +- .../unit_tests/executor/transferAgentTest.cpp | 38 +- .../_torch/auto_deploy/shim/ad_executor.py | 1 + tensorrt_llm/_torch/pyexecutor/_util.py | 27 +- .../_torch/pyexecutor/kv_cache_transceiver.py | 214 ++++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 238 +++++- .../accuracy/test_disaggregated_serving.py | 47 +- .../test_lists/qa/llm_function_core.txt | 3 +- .../test_lists/test-db/l0_b200.yml | 1 + .../test_lists/test-db/l0_dgx_b200.yml | 2 + .../test_lists/test-db/l0_dgx_h100.yml | 3 +- .../test_lists/test-db/l0_sanity_check.yml | 1 + .../test_disagg_inflight_cancel_gate.py | 590 ++++++++++++++ .../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 | 253 +++++- 40 files changed, 3374 insertions(+), 437 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 diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 80a06ea5ddf7..1457806e323e 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 @@ -240,6 +246,14 @@ class CacheTransceiver : public BaseCacheTransceiver rnn_state_manager::RnnStateManager* rnnStateManager = nullptr, std::vector const& rnnLayerNumPerPP = {}); + CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, + executor::kv_cache::CacheState::ModelConfig const& cacheStateModelCfg, runtime::WorldConfig const& worldConfig, + std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + executor::kv_cache::CacheState::AttentionType attentionType, + std::optional cacheTransceiverConfig, + rnn_state_manager::RnnStateManager* rnnStateManager, std::vector const& rnnLayerNumPerPP, + bool enableInflightCancel); + CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, std::vector numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, runtime::WorldConfig const& worldConfig, std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, @@ -254,6 +268,20 @@ class CacheTransceiver : public BaseCacheTransceiver { } + CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, std::vector numKvHeadsPerLayer, + SizeType32 sizePerHead, SizeType32 tokensPerBlock, runtime::WorldConfig const& worldConfig, + 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, + bool enableInflightCancel) + : CacheTransceiver(cacheManager, + executor::kv_cache::CacheState::ModelConfig{numKvHeadsPerLayer, sizePerHead, tokensPerBlock}, worldConfig, + attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnStateManager, rnnLayerNumPerPP, + enableInflightCancel) + { + } + virtual ~CacheTransceiver(); void respondAndSendAsync(std::shared_ptr llmRequest) override; @@ -273,11 +301,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 +338,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/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 e9daca6f5447..59abe09e057c 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -529,7 +529,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; @@ -609,8 +610,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); @@ -879,8 +891,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..b45ff6d2a9d8 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,138 @@ 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())); + TLLM_CHECK(totalSize <= static_cast(std::numeric_limits::max()) - static_cast(stateSize)); + 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 +370,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 not supported by the legacy C++ executor; use the standard " + "PyExecutor path after its cancellation lifecycle support is enabled."); if (!cacheTransceiverConfig.has_value() || !cacheTransceiverConfig.value().getBackendType().has_value()) { TLLM_LOG_INFO("CacheTransceiver is disabled."); @@ -294,7 +430,20 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa executor::kv_cache::CacheState::AttentionType attentionType, std::optional cacheTransceiverConfig, rnn_state_manager::RnnStateManager* rnnStateManager, std::vector const& rnnLayerNumPerPP) + : CacheTransceiver(cacheManager, cacheStateModelCfg, worldConfig, attentionLayerNumPerPP, dataType, attentionType, + std::move(cacheTransceiverConfig), rnnStateManager, rnnLayerNumPerPP, /*enableInflightCancel=*/false) +{ +} + +CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, + executor::kv_cache::CacheState::ModelConfig const& cacheStateModelCfg, runtime::WorldConfig const& worldConfig, + std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + executor::kv_cache::CacheState::AttentionType attentionType, + std::optional cacheTransceiverConfig, + 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 +456,22 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa mGroupComm = std::make_shared(tensorrt_llm::pg_utils::get_world_pg()); } + auto const backendType + = mCacheTransceiverConfig.has_value() ? mCacheTransceiverConfig->getBackendType() : std::nullopt; + bool const inflightCancelCapable = backendType.has_value() + && backendType.value() == executor::CacheTransceiverConfig::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.has_value() && (backendType.value() != executor::CacheTransceiverConfig::BackendType::DEFAULT), + " CacheTransceiverConfig::BackendType is not set."); + if (worldConfig.isTensorParallel() || worldConfig.isContextParallel()) { mGroupTensorParaComm = std::make_shared( @@ -348,12 +513,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,85 +602,104 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa mCacheTransBufferManagerPtrs.push_back(mRnnCacheTransBufferManager.get()); } - if (backendType.value() == executor::CacheTransceiverConfig::BackendType::UCX) - { - std::lock_guard lock(mDllMutex); - mWrapperLibHandle = dllOpen(UCX_WRAPPER_LIB_NAME); - TLLM_CHECK_WITH_INFO( - mWrapperLibHandle != nullptr, "UCX wrapper library is not open correctly. error : %s", dlerror()); - auto load_sym = [](void* handle, char const* name) - { - void* ret = dllGetSym(handle, name); - TLLM_CHECK_WITH_INFO(ret != nullptr, - "Unable to load UCX wrapper library symbol, possible cause is that TensorRT LLM library is not " - "built with UCX support, please rebuild in UCX-enabled environment."); - return ret; - }; - std::unique_ptr (*makeUcxConnectionManager)(); - *(void**) (&makeUcxConnectionManager) = load_sym(mWrapperLibHandle, "makeUcxConnectionManager"); - mManager = makeUcxConnectionManager(); - TLLM_LOG_INFO("UCX Connection Manager created"); - } - else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL) - { - auto rnnState - = mCacheState->hasRnnConfig() ? std::make_optional(mCacheState->getRnnCacheState()) : std::nullopt; - mManager = std::make_unique( - mCacheTransBufferManagerPtrs, *mCacheState, "nixl", rnnState); - TLLM_LOG_INFO("NIXL Connection Manager created"); - } - else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MOONCAKE) - { - auto rnnState - = mCacheState->hasRnnConfig() ? std::make_optional(mCacheState->getRnnCacheState()) : std::nullopt; - mManager = std::make_unique( - mCacheTransBufferManagerPtrs, *mCacheState, "mooncake", rnnState); - TLLM_LOG_INFO("MOONCAKE Connection Manager created"); - } - else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MPI) - { - mMpiWorldComm = std::addressof(tensorrt_llm::mpi::MpiComm::world()); - mManager = std::make_unique(mMpiWorldComm); - TLLM_LOG_INFO("MPI Connection Manager created"); - } - else - { - TLLM_THROW("Unsupported cache transceiver backend type "); - } - - auto makeFormatter = [cacheManager, isMLA, this]() + try { - std::vector kvBufferPtrs; - kvBufferPtrs.reserve(mCacheTransBufferManagers.size()); - for (auto& mgr : mCacheTransBufferManagers) + if (backendType.value() == executor::CacheTransceiverConfig::BackendType::UCX) { - kvBufferPtrs.push_back(mgr.get()); + std::lock_guard lock(mDllMutex); + mWrapperLibHandle = dllOpen(UCX_WRAPPER_LIB_NAME); + TLLM_CHECK_WITH_INFO( + mWrapperLibHandle != nullptr, "UCX wrapper library is not open correctly. error : %s", dlerror()); + auto load_sym = [](void* handle, char const* name) + { + void* ret = dllGetSym(handle, name); + TLLM_CHECK_WITH_INFO(ret != nullptr, + "Unable to load UCX wrapper library symbol, possible cause is that TensorRT LLM library is not " + "built with UCX support, please rebuild in UCX-enabled environment."); + return ret; + }; + std::unique_ptr (*makeUcxConnectionManager)(); + *(void**) (&makeUcxConnectionManager) = load_sym(mWrapperLibHandle, "makeUcxConnectionManager"); + mManager = makeUcxConnectionManager(); + TLLM_LOG_INFO("UCX Connection Manager created"); } - return createCacheFormatter(cacheManager, kvBufferPtrs, isMLA); - }; - - auto makeRnnFormatter = [this, cacheManager]() -> std::unique_ptr - { - if (mRnnStateManager != nullptr && mRnnCacheTransBufferManager != nullptr) + else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL) { - // Slot-based path (CppMambaCacheManager) - return std::make_unique(mRnnStateManager, mRnnCacheTransBufferManager.get()); + auto rnnState + = mCacheState->hasRnnConfig() ? std::make_optional(mCacheState->getRnnCacheState()) : std::nullopt; + mManager = std::make_unique( + mCacheTransBufferManagerPtrs, *mCacheState, "nixl", rnnState); + TLLM_LOG_INFO("NIXL Connection Manager created"); } - // Unified pool path (CppMambaHybridCacheManager) - if (mCacheState->hasRnnConfig() && mRnnCacheTransBufferManager != nullptr) + else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MOONCAKE) { - return std::make_unique(cacheManager, mRnnCacheTransBufferManager.get()); + auto rnnState + = mCacheState->hasRnnConfig() ? std::make_optional(mCacheState->getRnnCacheState()) : std::nullopt; + mManager = std::make_unique( + mCacheTransBufferManagerPtrs, *mCacheState, "mooncake", rnnState); + TLLM_LOG_INFO("MOONCAKE Connection Manager created"); + } + else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MPI) + { + mMpiWorldComm = std::addressof(tensorrt_llm::mpi::MpiComm::world()); + mManager = std::make_unique(mMpiWorldComm); + TLLM_LOG_INFO("MPI Connection Manager created"); + } + else + { + TLLM_THROW("Unsupported cache transceiver backend type "); } - return nullptr; - }; - auto makeCacheTransferLayer - = [&]() { return CacheTransferLayer(*mCacheState, makeFormatter(), makeRnnFormatter()); }; + auto makeFormatter = [cacheManager, isMLA, this]() + { + std::vector kvBufferPtrs; + kvBufferPtrs.reserve(mCacheTransBufferManagers.size()); + for (auto& mgr : mCacheTransBufferManagers) + { + kvBufferPtrs.push_back(mgr.get()); + } + return createCacheFormatter(cacheManager, kvBufferPtrs, isMLA); + }; - mCacheSender = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); - mCacheReceiver = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); + auto makeRnnFormatter = [this, cacheManager]() -> std::unique_ptr + { + if (mRnnStateManager != nullptr && mRnnCacheTransBufferManager != nullptr) + { + // Slot-based path (CppMambaCacheManager) + return std::make_unique(mRnnStateManager, mRnnCacheTransBufferManager.get()); + } + // Unified pool path (CppMambaHybridCacheManager) + if (mCacheState->hasRnnConfig() && mRnnCacheTransBufferManager != nullptr) + { + return std::make_unique(cacheManager, mRnnCacheTransBufferManager.get()); + } + return nullptr; + }; + + auto makeCacheTransferLayer = [&]() + { return CacheTransferLayer(*mCacheState, makeFormatter(), makeRnnFormatter(), mInflightCancelEnabled); }; - initializeCommState(); + mCacheSender = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); + mCacheReceiver + = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); + + initializeCommState(); + } + catch (...) + { + // A constructor failure skips ~CacheTransceiver(), so tear dependencies + // down explicitly while registered buffers and plugin code are alive. + mCacheSender.reset(); + mCacheReceiver.reset(); + mManager.reset(); + if (mWrapperLibHandle) + { + std::lock_guard lock(mDllMutex); + dllClose(mWrapperLibHandle); + mWrapperLibHandle = nullptr; + } + throw; + } } CacheTransceiver::~CacheTransceiver() @@ -530,6 +708,9 @@ CacheTransceiver::~CacheTransceiver() // plugin are still alive. The workers can access both during termination. mCacheSender.reset(); mCacheReceiver.reset(); + // Deregister transport memory while the registered buffer managers and the + // dynamically loaded UCX implementation are both still alive. + mManager.reset(); if (mWrapperLibHandle) { @@ -599,6 +780,9 @@ 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); @@ -721,9 +905,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.allRanksTerminalRequestIds) + { + 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.allRanksTerminalRequestIds) + { + 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 +1338,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 +1555,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..ffb2dc6c62f4 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/cacheTransferState.h @@ -0,0 +1,138 @@ +/* + * 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, +}; + +enum class ReadySignalSummary +{ + kAllReady, + kAllNotReady, + kMixed, +}; + +inline ReadySignalSummary summarizeReadySignals(bool const anyReady, bool const anyNotReady) +{ + TLLM_CHECK_WITH_INFO(anyReady || anyNotReady, "Ready-signal summary requires at least one peer response."); + if (anyReady && anyNotReady) + { + return ReadySignalSummary::kMixed; + } + return anyReady ? ReadySignalSummary::kAllReady : ReadySignalSummary::kAllNotReady; +} + +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; + // Local ownership settlement is an intersection: every rank has reached a + // local terminal state. This does not prove remote-memory quiescence. + std::unordered_set allRanksTerminalRequestIds; + std::unordered_set completedRequestIds; + bool transferBufferPoisoned{false}; +}; + +// 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.allRanksTerminalRequestIds.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 93812056abf6..2c14b3400aa5 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 @@ -281,13 +283,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 +305,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 +352,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 +366,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 +440,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 +475,43 @@ 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); + // Active mode must settle a queued request immediately so every local rank can reach a terminal state + // even when the peer never sends RequestInfo. Keep a tombstone so a late peer still receives ready=false. + // Legacy mode preserves its existing promise lifetime and settles only after the peer handshake. + if (it != mReadyResponses.end() + && (!mCurrentRequest.has_value() || mCurrentRequest.value() != llmRequest.mRequestId)) + { + mCancelledRequests.insert(llmRequest.mRequestId); + if (mInflightCancelEnabled) + { + failResponse(it->second, + std::make_exception_ptr( + TLLM_REQUEST_EXCEPTION(llmRequest.mRequestId, common::RequestErrorCode::kNETWORK_ERROR, + "KV cache transfer for request %zu was cancelled before the peer became ready", + llmRequest.mRequestId))); + mReadyResponses.erase(it); + } + isCancelled = true; + } } - else + if (isCancelled) + { + mSenderCv.notify_all(); + } + 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 +611,234 @@ 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); + std::scoped_lock lock(mSenderMutex); + TLLM_CHECK(mCurrentRequest.has_value() && mCurrentRequest.value() == reqId); + auto const responseIt = mReadyResponses.find(reqId); + auto const isCancelled = mCancelledRequests.find(reqId) != mCancelledRequests.end(); + TLLM_CHECK(responseIt != mReadyResponses.end() || isCancelled); + auto countIt = mRemainSendCount.find(reqId); + TLLM_CHECK(countIt != mRemainSendCount.end()); + auto const count = --countIt->second; + TLLM_CHECK(count >= 0); + if (count > 0) + { + mCurrentRequest = std::nullopt; + return; + } + mRemainSendCount.erase(countIt); + isReady = responseIt != mReadyResponses.end() && !isCancelled; + } - // Check if the request is cancelled - bool isReady = true; + // 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); + + std::optional response; + { + std::scoped_lock lock(mSenderMutex); + auto it = mReadyResponses.find(reqId); + if (it != mReadyResponses.end()) { - std::scoped_lock lk(mSenderMutex); - if (mCancelledRequests.find(reqId) != mCancelledRequests.end()) - { - isReady = false; - } + response.emplace(std::move(it->second)); + mReadyResponses.erase(it); } - sendReadySignal(reqId, isReady); + mCancelledRequests.erase(reqId); + mCurrentRequest = std::nullopt; + } - if (isReady) + if (isReady) + { + TLLM_CHECK(response.has_value()); + if (dynamic_cast(mManager) != nullptr) { - if (dynamic_cast(mManager) != nullptr) - { - // our nixl impl seems only support recv and send in the same thread - // if we use zmq as control path, we may avoid this issue - sendAndRemoveResponse(it->first, std::move(it->second)); - } - else - { - // if we send data in another thread, multiple rank may send data for different requests at the same - // time with gen DP case. - asyncSendAndRemoveResponse(it->first, std::move(it->second)); - } - removeResponse(it); + // Our NIXL implementation only supports recv and send in the same thread. Using ZMQ for the control + // path may avoid this limitation. + sendAndRemoveResponse(reqId, std::move(*response)); } else { - // TODO: if the generation does not require the kv cache, the request will - // not be removed from mCancelledRequests. This should be handled by timeout. - auto const cancelledReqId = mCurrentRequest.value(); - Response cancelledResponse; - { - std::scoped_lock lkResp(mSenderMutex); - auto it = mReadyResponses.find(cancelledReqId); - TLLM_CHECK(it != mReadyResponses.end()); - // Move out before erasing so the promise survives the - // map cleanup and can be resolved (vs. destroyed unfulfilled, - // which would surface as std::future_error: Broken promise). - cancelledResponse = std::move(it->second); - mReadyResponses.erase(it); - mCancelledRequests.erase(cancelledReqId); - mRemainSendCount.erase(cancelledReqId); - } - cancelledResponse.mPromise.set_exception(std::make_exception_ptr( - TLLM_REQUEST_EXCEPTION(cancelledReqId, common::RequestErrorCode::kNETWORK_ERROR, - "KV cache transfer for request %zu was cancelled", cancelledReqId))); - mCurrentRequest = std::nullopt; - - if (mReadyResponses.empty()) - { - std::unique_lock lk(mCondMutex); - mAnyReady = false; - } + // If we send data in another thread, multiple ranks may send data for different requests at the same + // time with generation attention DP. + asyncSendAndRemoveResponse(reqId, std::move(*response)); + } + } + else + { + if (response.has_value()) + { + failResponse(*response, + std::make_exception_ptr(TLLM_REQUEST_EXCEPTION(reqId, common::RequestErrorCode::kNETWORK_ERROR, + "KV cache transfer for request %zu was cancelled", reqId))); } + discardTransferState(reqId); } - mCurrentRequest = std::nullopt; } 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() || !mCancelledRequests.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(); + auto const reqId = requestInfo->getRequestId(); - { - std::scoped_lock lk(mSenderMutex); - mCurrentRequest = reqId; - } - - if (mRemainSendCount.find(reqId) == mRemainSendCount.end()) - { - mRemainSendCount[reqId] = getCounterpartsCount(reqId); - } - } - auto it = getCurrentResponse(); - if (it != mReadyResponses.end()) + if (mRemainSendCount.find(reqId) == mRemainSendCount.end()) { - sendResponse(it); + mRemainSendCount[reqId] = getCounterpartsCount(reqId); } - else + { - auto it = getCurrentResponse(); - while (it == mReadyResponses.end()) - { - std::unique_lock lk(mCondMutex); - mSenderCv.wait(lk, [this]() { return (mAnyReady || mTerminate); }); - if (mTerminate) + std::unique_lock lock(mSenderMutex); + mCurrentRequest = reqId; + mSenderCv.wait(lock, + [this, reqId]() { - break; - } - it = getCurrentResponse(); - } - if (mTerminate || it == mReadyResponses.end()) + return mTerminate || mReadyResponses.find(reqId) != mReadyResponses.end() + || mCancelledRequests.find(reqId) != mCancelledRequests.end(); + }); + if (mTerminate) { + mCurrentRequest = std::nullopt; break; } - sendResponse(it); } + sendResponse(reqId); } } catch (std::exception const& err) { TLLM_LOG_ERROR("Exception in CacheSender response: %s", err.what()); - for (auto& it : mReadyResponses) - { - it.second.mPromise.set_exception(std::current_exception()); - } + responseException = std::current_exception(); + } + + if (!responseException) + { + responseException + = std::make_exception_ptr(std::runtime_error("CacheSender terminated before response completed")); + } + { + std::scoped_lock lock(mSenderMutex); + mTerminate = true; } + mSenderCv.notify_all(); + failPendingResponses(responseException); } void terminate() { { - std::unique_lock lk(mCondMutex); + std::scoped_lock lock(mSenderMutex); mTerminate = true; } - // We don't have to wait for the future. If another thread is sending data, it won't pay attention - // to the terminate flag. + { + 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 - { - return mCurrentRequest.value(); - } - - [[nodiscard]] std::map::iterator getCurrentResponse() + void failPendingResponses(std::exception_ptr const& exception) noexcept { - std::scoped_lock lk(mSenderMutex); - return mReadyResponses.find(getCurrentRequestId()); + std::map pendingResponses; + { + std::scoped_lock lock(mSenderMutex); + pendingResponses.swap(mReadyResponses); + mCurrentRequest = std::nullopt; + mCancelledRequests.clear(); + mRemainSendCount.clear(); + } + for (auto& entry : pendingResponses) + { + failResponse(entry.second, exception); + } } public: @@ -749,11 +852,13 @@ class CacheSender::Impl private: std::optional mCurrentRequest; + // A canceled request remains here until a late peer is explicitly rejected. Retiring no-peer tombstones safely + // requires the cross-service cancellation acknowledgement/version handshake that gates production activation. 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 +868,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 +883,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 +898,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 +931,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,22 +951,41 @@ class CacheReceiver::Impl void receiveSync(TransferSession& session) { - mCacheTransferLayer.unformat(session); - if (!common::getEnvKVCacheTimeOutputPath().empty()) + try { - std::unique_lock lock(mMeasuresFileMutex); - if (!mMeasuresFile.is_open()) + mCacheTransferLayer.unformat(session); + // unformat synchronizes transport/postprocessing before returning. + // Telemetry failures after this point cannot make the slot unsafe. + session.releaseRecvBufferHolders(); + if (!common::getEnvKVCacheTimeOutputPath().empty()) { - auto outputPath = getTransferOutputPath("recv"); - mMeasuresFile.open(outputPath); - TLLM_CHECK_WITH_INFO( - mMeasuresFile.is_open(), "Failed to open transfer output file: %s", outputPath.string().c_str()); + std::unique_lock lock(mMeasuresFileMutex); + if (!mMeasuresFile.is_open()) + { + auto outputPath = getTransferOutputPath("recv"); + mMeasuresFile.open(outputPath); + TLLM_CHECK_WITH_INFO(mMeasuresFile.is_open(), "Failed to open transfer output file: %s", + outputPath.string().c_str()); + } + session.exportMeasure(mMeasuresFile, false); + } + } + catch (...) + { + if (session.isInflightCancelEnabled()) + { + session.poisonRecvBufferHolders(); + } + else + { + session.releaseRecvBufferHolders(); } - session.exportMeasure(mMeasuresFile, false); + throw; } } - TransferSession sendRequestInfo(LlmRequest const& llmRequest) + TransferSession sendRequestInfo(LlmRequest const& llmRequest, std::atomic const* perRequestCancel = nullptr, + std::vector> cacheBufferIds = {}) { uint64_t requestId = llmRequest.getContextPhaseParams().value().getReqId(); auto const& contextState = llmRequest.getDataTransceiverState(); @@ -886,12 +1027,27 @@ class CacheReceiver::Impl } auto* agentConnectionManager = dynamic_cast(mManager); - std::vector> cacheBufferIds; + std::vector allocatedRecvHolders; if (agentConnectionManager) { - for (auto& cacheTransBufferManager : agentConnectionManager->getCacheTransBufferManagers()) + if (cacheBufferIds.empty()) { - cacheBufferIds.push_back(cacheTransBufferManager->assignBufferIndexForRecv()); + auto const& managers = agentConnectionManager->getCacheTransBufferManagers(); + allocatedRecvHolders.reserve(managers.size()); + cacheBufferIds.reserve(managers.size()); + for (auto* cacheTransBufferManager : managers) + { + auto rawIdx = cacheTransBufferManager->assignBufferIndexForRecv(perRequestCancel); + allocatedRecvHolders.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()); } @@ -920,6 +1076,14 @@ class CacheReceiver::Impl allConnections.emplace_back(connection); } + auto const& resource = getReceiveCacheResource(llmRequest); + auto const& cancelFlag = perRequestCancel != nullptr ? *perRequestCancel : mTerminate; + TransferSession session(allConnections, DataContext{tagFromRequestId(requestId), cancelFlag}, allCounterparts, + mSelfState, contextState, resource->mBufferManager, requestInfo.getIndexFromEnd(), + requestInfo.getLastBlockKey(), &llmRequest, !common::getEnvKVCacheTimeOutputPath().empty(), + mInflightCancelEnabled); + session.setRecvBufferHolders(std::move(allocatedRecvHolders)); + for (size_t ci = 0; ci < allCounterparts.size(); ci++) { auto rank = allCounterparts[ci]; @@ -972,18 +1136,14 @@ class CacheReceiver::Impl TLLM_CHECK(agentConnection != nullptr); const_cast(agentConnection) - ->sendRequestAndBufferInfo(requestInfo, idsForRank, validConnectionIdx); + ->sendRequestAndBufferInfo(requestInfo, idsForRank, validConnectionIdx, perRequestCancel); } else { sendRequestInfo(connection, requestInfo); } } - 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()); + return session; } std::unique_ptr const& getReceiveCacheResource(LlmRequest const& llmRequest) @@ -1027,6 +1187,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(), @@ -1053,45 +1214,114 @@ 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, + kMixed, + kCancelled, + }; + + ReadySignalResult receiveReadySignalDetailed(TransferSession& session, std::atomic const& perRequestCancel) { - bool isReadyFinal = true; bool isReady = false; + bool anyReady = false; + bool anyNotReady = 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; + anyNotReady |= !isReady; + } + + switch (detail::summarizeReadySignals(anyReady, anyNotReady)) + { + case detail::ReadySignalSummary::kAllReady: return ReadySignalResult::kReady; + case detail::ReadySignalSummary::kAllNotReady: return ReadySignalResult::kNotReady; + case detail::ReadySignalSummary::kMixed: return ReadySignalResult::kMixed; } + TLLM_THROW("Unexpected ready-signal summary"); + } - return isReadyFinal; + bool receiveReadySignal(TransferSession& session) + { + auto const result = receiveReadySignalDetailed(session, mTerminate); + if (result == ReadySignalResult::kNotReady) + { + // An explicit peer rejection proves that no transfer started. + session.releaseRecvBufferHolders(); + } + else if (result == ReadySignalResult::kMixed && session.isInflightCancelEnabled()) + { + // At least one peer may already be writing while another rejected + // the request. Remote quiescence is unknown, so fail closed. + session.poisonRecvBufferHolders(); + } + 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; @@ -1104,29 +1334,125 @@ 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); + } + } + } + + auto poisonRecvHolders = [&recvHolders]() + { + for (auto& holder : recvHolders) + { + holder.poison(); + } + }; + + try + { + auto session = sendRequestInfo(llmRequest, &perRequestCancel, std::move(cacheBufferIds)); + session.setRecvBufferHolders(std::move(recvHolders)); + 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) + { + // ready=false is a remote pre-transfer rejection, so the + // preassigned slots are safe to reuse even in active mode. + session.releaseRecvBufferHolders(); + if (!mInflightCancelEnabled) + { + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + } + llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); + return false; + } + if (readyResult == ReadySignalResult::kMixed) + { + if (mInflightCancelEnabled) + { + session.poisonRecvBufferHolders(); + } + else + { + session.releaseRecvBufferHolders(); + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + } + llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); + return false; + } + + receiveSync(session); + llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); + } + catch (...) { - // Reuse the error state for the cancelled request. - llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + if (mInflightCancelEnabled) + { + poisonRecvHolders(); + } + 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 @@ -1135,16 +1461,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)) { } @@ -1153,6 +1483,7 @@ class CacheReceiver::Impl RequestAndPromise(RequestAndPromise&& other) noexcept : mRequest(std::move(other.mRequest)) , mPromise(std::move(other.mPromise)) + , mCancelFlag(std::move(other.mCancelFlag)) { } @@ -1168,6 +1499,7 @@ class CacheReceiver::Impl mRequest = std::move(other.mRequest); mPromise = std::move(other.mPromise); + mCancelFlag = std::move(other.mCancelFlag); } return *this; } @@ -1211,8 +1543,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) { @@ -1230,6 +1574,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); + } } } } @@ -1250,6 +1607,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; @@ -1257,6 +1615,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..028403b224e8 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&&) noexcept = delete; + + ~TransferSession() + { + if (mEnableInflightCancel) + { + poisonRecvBufferHolders(); + } + } + [[nodiscard]] std::vector const& getConnections() const; // should be called only during the initialization of the TransferSession @@ -151,6 +167,39 @@ class TransferSession mCounterPartRanks = std::move(ranks); } + [[nodiscard]] bool isInflightCancelEnabled() const noexcept + { + return mEnableInflightCancel; + } + + /// Transfer ownership of preassigned receive-buffer slots to this session. + /// Formatters borrow these slots and must not release them independently. + void setRecvBufferHolders(std::vector holders) + { + TLLM_CHECK(mRecvBufferHolders.empty()); + mRecvBufferHolders = std::move(holders); + } + + /// Release receive-buffer slots after the transfer and postprocessing have completed. + void releaseRecvBufferHolders() noexcept + { + for (auto& holder : mRecvBufferHolders) + { + holder.release(); + } + mRecvBufferHolders.clear(); + } + + /// Quarantine receive-buffer slots when transfer quiescence is not known. + 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 +211,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 ad55455ad3c2..27ecacc2e2a2 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( @@ -346,48 +347,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(); } @@ -511,9 +523,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..c713eb5ec61e 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) { 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,17 @@ 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"); + } mAgentConnectionManager->getAgent()->notifySyncMessage(mRemoteAgentName, ss.str()); + // A cancelled or failed first notification never reached the peer, so the + // next request must retry the metadata needed to load this agent. + if (metadataOpt.has_value()) + { + mNeedSendMetadata = false; + } } void AgentConnection::setSenderState(std::vector cacheReceiverBufferDescs, int validSegmentIdx, @@ -291,9 +332,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 +741,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 +749,7 @@ void AgentConnectionManager::waitForNotification( if (!mIsRunning) { - return; + return false; } updateUnhandledNotifications(); std::scoped_lock lock(mNotificationMutex); @@ -733,7 +782,7 @@ void AgentConnectionManager::waitForNotification( { it = mUnhandledNotifications.erase(it); } - return; + return true; } } } @@ -753,7 +802,7 @@ void AgentConnectionManager::waitForNotification( { it = mUnhandledNotifications.erase(it); } - return; + return true; } } } @@ -773,24 +822,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..4b167556cc8d 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"); @@ -260,13 +260,15 @@ class AgentConnection : public Connection 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; void sendRequestAndBufferInfo(batch_manager::RequestInfo& requestInfo, - std::vector> const& cacheBufferIds, int validConnectionIdx); + std::vector> const& cacheBufferIds, int validConnectionIdx, + std::atomic const* perRequestCancel = 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 +322,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..ff956caa8c73 100644 --- a/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp @@ -17,6 +17,7 @@ #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 @@ -50,6 +51,20 @@ class ObservableTransBufferManager : public CacheTransBufferManager } }; +class DynamicTransBufferManager : public BaseTransBufferManager +{ +public: + DynamicTransBufferManager() + : BaseTransBufferManager(/*transferBufferSize=*/0, nvinfer1::DataType::kFLOAT) + { + } + + [[nodiscard]] BufferKind getBufferKind() const override + { + return BufferKind::kKV; + } +}; + } // namespace enum class Side @@ -146,17 +161,37 @@ TEST_P(BufferIndexHolderLifecycleTest, DefaultConstructedExplicitReleaseIsNoOp) EXPECT_EQ(inUse(), before); } -// Holder built from a nullopt index is disarmed (mHeld == false). +// A nullopt index still represents ownership: dynamic-buffer users must be able +// to poison their pool through the same holder type. TEST_P(BufferIndexHolderLifecycleTest, NulloptIndexReleasesNothing) { int const before = inUse(); { BufferIndexHolder holder{mgr(), std::nullopt, isRecv()}; - EXPECT_FALSE(holder.held()); + EXPECT_TRUE(holder.held()); } EXPECT_EQ(inUse(), before); } +TEST_P(BufferIndexHolderLifecycleTest, NulloptIndexPoisonsDynamicPool) +{ + DynamicTransBufferManager manager; + auto index = isRecv() ? manager.assignBufferIndexForRecv() : manager.assignBufferIndexForSend(); + ASSERT_FALSE(index.has_value()); + + BufferIndexHolder holder{manager, index, isRecv()}; + holder.poison(); + EXPECT_TRUE(manager.hasPoisonedBuffer()); + if (isRecv()) + { + EXPECT_THROW((void) manager.assignBufferIndexForRecv(), std::exception); + } + else + { + EXPECT_THROW((void) manager.assignBufferIndexForSend(), std::exception); + } +} + // RAII: a held slot is released when the holder goes out of scope. TEST_P(BufferIndexHolderLifecycleTest, ValidIndexReleasedOnDestruction) { @@ -286,6 +321,54 @@ TEST_P(BufferIndexHolderLifecycleTest, ExceptionUnwindStillReleases) EXPECT_EQ(inUse(), before); } +TEST_P(BufferIndexHolderLifecycleTest, TransferSessionReleasesRecvSlotOnSuccess) +{ + if (!isRecv()) + { + GTEST_SKIP() << "TransferSession owns receive slots only"; + } + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + + executor::DataTransceiverState selfState; + executor::DataTransceiverState otherState; + runtime::BufferManager bufferManager{std::make_shared()}; + std::vector connections{nullptr}; + TransferSession session(std::move(connections), DataContext{0}, std::vector{0}, selfState, + std::move(otherState), bufferManager, 0, BlockKey{}, nullptr, false, false); + std::vector holders; + holders.emplace_back(mgr(), idx, /*isRecv=*/true); + session.setRecvBufferHolders(std::move(holders)); + + EXPECT_EQ(inUse(), before + 1); + session.releaseRecvBufferHolders(); + EXPECT_EQ(inUse(), before); +} + +TEST_P(BufferIndexHolderLifecycleTest, ActiveTransferSessionPoisonsUnreleasedRecvSlot) +{ + if (!isRecv()) + { + GTEST_SKIP() << "TransferSession owns receive slots only"; + } + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + { + executor::DataTransceiverState selfState; + executor::DataTransceiverState otherState; + runtime::BufferManager bufferManager{std::make_shared()}; + std::vector connections{nullptr}; + TransferSession session(std::move(connections), DataContext{0}, std::vector{0}, selfState, + std::move(otherState), bufferManager, 0, BlockKey{}, nullptr, false, true); + std::vector holders; + holders.emplace_back(mgr(), idx, /*isRecv=*/true); + session.setRecvBufferHolders(std::move(holders)); + } + + 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..d244b7134b87 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/cacheTransferStateTest.cpp @@ -0,0 +1,101 @@ +/* + * 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.allRanksTerminalRequestIds.empty()); + + auto const secondPoll = summarizePackedTransferStates({{0, 17, kTransferCompleted}, {0, 17, kTransferCompleted}}); + EXPECT_EQ(secondPoll.completedRequestIds, std::unordered_set{17}); + EXPECT_EQ(secondPoll.allRanksTerminalRequestIds, std::unordered_set{17}); +} + +TEST(CacheTransferStateTest, FailureUnionPrecedesAllRanksTerminalIntersection) +{ + auto const firstPoll = summarizePackedTransferStates({{0, 29, kTransferFailed}, {0}}); + EXPECT_EQ(firstPoll.failedRequestIds, std::unordered_set{29}); + EXPECT_TRUE(firstPoll.allRanksTerminalRequestIds.empty()); + + auto const secondPoll = summarizePackedTransferStates({{0, 29, kTransferFailed}, {0, 29, kTransferCompleted}}); + EXPECT_EQ(secondPoll.failedRequestIds, std::unordered_set{29}); + EXPECT_EQ(secondPoll.allRanksTerminalRequestIds, 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.allRanksTerminalRequestIds, 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.allRanksTerminalRequestIds, 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.allRanksTerminalRequestIds.empty()); +} + +TEST(CacheTransferStateTest, ReadySignalSummaryDistinguishesMixedPeers) +{ + EXPECT_EQ(summarizeReadySignals(/*anyReady=*/true, /*anyNotReady=*/false), ReadySignalSummary::kAllReady); + EXPECT_EQ(summarizeReadySignals(/*anyReady=*/false, /*anyNotReady=*/true), ReadySignalSummary::kAllNotReady); + EXPECT_EQ(summarizeReadySignals(/*anyReady=*/true, /*anyNotReady=*/true), ReadySignalSummary::kMixed); + EXPECT_THROW((void) summarizeReadySignals(/*anyReady=*/false, /*anyNotReady=*/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..45acb130ea54 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"); @@ -18,6 +18,8 @@ #include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h" #include +#include + using namespace tensorrt_llm::batch_manager::kv_cache_manager; using namespace tensorrt_llm::runtime; using namespace tensorrt_llm::executor::kv_cache; @@ -231,5 +233,37 @@ TEST_P(AgentCommTest, AgentConnectionManagerConnect) TLLM_LOG_INFO("after finish"); } +TEST_P(AgentCommTest, CancelledFirstNotificationRetriesAgentMetadata) +{ + std::vector bufferManagers{mTransBufferManager.get()}; + auto connectionManager0 = std::make_unique(bufferManagers, *mCacheState, backend); + auto connectionManager1 = std::make_unique(bufferManagers, *mCacheState, backend); + auto const commState0 = connectionManager0->getCommState(); + auto const commState1 = connectionManager1->getCommState(); + auto const* connection0 = connectionManager0->getConnections(commState1).at(0); + auto* agentConnection0 = const_cast(dynamic_cast(connection0)); + ASSERT_NE(agentConnection0, nullptr); + + std::vector> cacheBufferIds{std::optional{0}}; + int constexpr validConnectionIdx = 0; + std::atomic cancelled{true}; + tensorrt_llm::executor::DataTransceiverState dataTransceiverState0{*mCacheState, commState0}; + tensorrt_llm::batch_manager::RequestInfo cancelledRequestInfo{2, dataTransceiverState0}; + EXPECT_THROW(agentConnection0->sendRequestAndBufferInfo( + cancelledRequestInfo, cacheBufferIds, validConnectionIdx, &cancelled), + std::exception); + + cancelled.store(false, std::memory_order_relaxed); + uint64_t constexpr retryRequestId = 3; + tensorrt_llm::batch_manager::RequestInfo retryRequestInfo{retryRequestId, dataTransceiverState0}; + ASSERT_NO_THROW( + agentConnection0->sendRequestAndBufferInfo(retryRequestInfo, cacheBufferIds, validConnectionIdx, &cancelled)); + + tensorrt_llm::batch_manager::RequestInfo receivedRequestInfo; + std::atomic terminate{false}; + ASSERT_NE(connectionManager1->recvConnectionAndRequestInfo(receivedRequestInfo, terminate), nullptr); + EXPECT_EQ(receivedRequestInfo.getRequestId(), retryRequestId); +} + INSTANTIATE_TEST_SUITE_P(AvailableBackends, AgentCommTest, ::testing::ValuesIn(getAvailableBackends()), [](::testing::TestParamInfo const& info) { return info.param; }); diff --git a/cpp/tests/unit_tests/executor/transferAgentTest.cpp b/cpp/tests/unit_tests/executor/transferAgentTest.cpp index b915ec3bb9c4..fba92841d8f2 100644 --- a/cpp/tests/unit_tests/executor/transferAgentTest.cpp +++ b/cpp/tests/unit_tests/executor/transferAgentTest.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"); @@ -160,6 +160,36 @@ TEST_P(TransferAgentTest, Basic2) xferAgent0->invalidateRemoteAgent(agent1); } +TEST_P(TransferAgentTest, NixlTransferStatusReleaseIsIdempotent) +{ + if (backend != "nixl") + { + GTEST_SKIP() << "Explicit transfer-handle release is implemented only by NIXL"; + } + + std::string const agent0{"agent0"}, agent1{"agent1"}; + BaseAgentConfig config0{agent0, true, false, true}, config1{agent1, true, false, true}; + auto xferAgent0 = makeTransferAgent(config0); + auto xferAgent1 = makeTransferAgent(config1); + std::vector memory0(100, 10); + std::vector memory1(100, 1); + RegisteredHostMemory regMem0(MemoryDescs{MemoryType::kDRAM, {MemoryDesc{memory0}}}, xferAgent0.get()); + RegisteredHostMemory regMem1(MemoryDescs{MemoryType::kDRAM, {MemoryDesc{memory1}}}, xferAgent1.get()); + + xferAgent0->loadRemoteAgent(agent1, xferAgent1->getLocalConnectionInfo()); + while (!xferAgent0->checkRemoteDescs(agent1, regMem1.getDescs())) + { + } + TransferRequest writeReq{TransferOp::kWRITE, regMem0.getDescs(), regMem1.getDescs(), agent1}; + auto status = xferAgent0->submitTransferRequests(writeReq); + ASSERT_EQ(status->wait(), TransferState::kSUCCESS); + + EXPECT_TRUE(status->release()); + EXPECT_TRUE(status->release()); + EXPECT_FALSE(status->isCompleted()); + EXPECT_EQ(status->wait(0), TransferState::kFAILURE); +} + TEST_P(TransferAgentTest, DeviceMemory) { @@ -421,8 +451,10 @@ TEST_P(TransferAgentTest, StatusOutlivesAgent) // kFAILURE/false instead of dereferencing the freed agent. EXPECT_FALSE(status->isCompleted()); EXPECT_EQ(status->wait(0), TransferState::kFAILURE); - // `status` destructor runs at scope exit: weak_ptr.lock() == nullptr -> - // early return (no releaseXferReq on a dangling agent). + // Releasing an orphaned status is safe and idempotent: no backend handle is + // dereferenced after the owning agent has gone away. + EXPECT_TRUE(status->release()); + EXPECT_TRUE(status->release()); } // Concurrent submitTransferRequests (#14137): submit holds a std::shared_lock and 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 151091593fc5..4997e659ffa6 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 @@ -2205,8 +2220,16 @@ 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, + # This protocol-only prerequisite remains dormant until the dependent + # PyExecutor lifecycle change can guarantee safe timeout, cancellation, + # poison, and terminal-response ownership. + inflight_cancel_supported_by_executor=False) waiting_queue_policy = (scheduler_config.waiting_queue_policy if scheduler_config is not None else diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index 8e6b47c8bff8..05723c7ca863 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,132 @@ 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 _get_disagg_inflight_cancel_config_error( + cache_transceiver_config: Optional[CacheTransceiverConfig], + mapping: Mapping, + inflight_cancel_supported_by_executor: bool = False) -> Optional[str]: + if cache_transceiver_config is None or cache_transceiver_config.backend is None: + return (f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV}=1 requires a " + "configured cache transceiver.") + + 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: + return (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 None + + 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) + return ( + 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 _validate_disagg_inflight_cancel_config( + cache_transceiver_config: Optional[CacheTransceiverConfig], + mapping: Mapping, + dist: Distributed, + inflight_cancel_supported_by_executor: bool = False) -> None: + requested = is_disagg_inflight_cancel_enabled() + local_error = (_get_disagg_inflight_cancel_config_error( + cache_transceiver_config, mapping, + inflight_cancel_supported_by_executor) if requested else None) + local_descriptor = (requested, local_error) + + world_size = getattr(mapping, "world_size", 1) + descriptors = (dist.allgather(local_descriptor) + if isinstance(world_size, int) and world_size > 1 else + [local_descriptor]) + requested_by_rank = [descriptor[0] for descriptor in descriptors] + if any(requested_by_rank) and not all(requested_by_rank): + raise ValueError( + f"{_DISAGG_INFLIGHT_CANCEL_ENABLED_ENV} must resolve identically " + f"on every worker rank; got requested={requested_by_rank}.") + + errors = [(rank, descriptor[1]) + for rank, descriptor in enumerate(descriptors) + if descriptor[1] is not None] + if errors: + details = "; ".join(f"rank {rank}: {error}" for rank, error in errors) + raise ValueError( + "Disaggregated in-flight cancellation is unsupported on one or " + f"more worker ranks: {details}") + def mapping_to_world_config(mapping: Mapping) -> WorldConfig: @@ -37,23 +178,15 @@ def create_kv_cache_transceiver( kv_cache_manager: KVCacheManager, attention_type: AttentionTypeCpp, cache_transceiver_config: CacheTransceiverConfig, - mamba_cache_manager: Optional[BaseMambaCacheManager] = None): - if cache_transceiver_config is None or cache_transceiver_config.backend is None: - logger.info("cache_transceiver is disabled") - return None - - if cache_transceiver_config.backend == "DEFAULT": + mamba_cache_manager: Optional[BaseMambaCacheManager] = None, + inflight_cancel_supported_by_executor: bool = False): + if (cache_transceiver_config is not None + and cache_transceiver_config.backend == "DEFAULT"): # When cache_transceiver_config.backend is not set, fallback to env_vars settings # 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 +194,17 @@ def create_kv_cache_transceiver( cache_transceiver_config.backend = be_type break + # Every rank enters the same preflight before one rank can reject a local + # environment/configuration mismatch and strand peers in a later native + # startup collective. + _validate_disagg_inflight_cancel_config( + cache_transceiver_config, mapping, dist, + inflight_cancel_supported_by_executor) + + if cache_transceiver_config is None or cache_transceiver_config.backend is None: + logger.info("cache_transceiver is disabled") + return None + if cache_transceiver_config.backend == "MPI": logger.warning( "MPI CacheTransceiver is deprecated, UCX or NIXL is recommended") @@ -88,9 +232,14 @@ def create_kv_cache_transceiver( cache_transceiver_config) # 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) + return BindKvCacheTransceiver(mapping, + dist, + kv_cache_manager, + attention_type, + cache_transceiver_config, + mamba_cache_manager, + inflight_cancel_supported_by_executor, + _inflight_cancel_config_prevalidated=True) class KvCacheTransceiver(ABC): @@ -123,6 +272,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 +311,14 @@ 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, + *, + _inflight_cancel_config_prevalidated: bool = False): + if not _inflight_cancel_config_prevalidated: + _validate_disagg_inflight_cancel_config( + cache_transceiver_config, mapping, dist, + 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 +342,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 +374,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 +398,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..a691b0894d29 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 @@ -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, @@ -681,6 +697,9 @@ def __init__( self.is_shutdown = False self._fatal_error: Optional[BaseException] = None self._error_budget = ErrorBudget() + self._disagg_timed_out_ctx_cancelled_ids: set[int] = set() + self._disagg_timed_out_gen_cancelled_ids: set[int] = set() + 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 @@ -4507,6 +4526,149 @@ def _check_disagg_gen_transfer_status(self): # an empty ready set. self._check_disagg_gen_cache_transfer_status(0) + if self._is_disagg_inflight_cancel_active(): + self._cancel_timed_out_gen_transfers() + self._check_gen_cache_transfer_errors_consensus() + + return + + 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("_cancel_timed_out_gen_transfers") + def _cancel_timed_out_gen_transfers(self) -> None: + """Request cancellation for timed-out generation transfers. + + This path is deliberately rank-synchronized. Under attention-DP, each + TP rank can own a different request subset, but later error responses + still pass through a TP collective. The timeout/cancel decision must + therefore be based on the TP-wide request-id union rather than a + rank-local timeout observation. + """ + timeout_ms = self.kv_cache_transceiver.kv_transfer_timeout_ms + if timeout_ms is None: + return + + requests_in_transfer = { + req.py_request_id: req + for req in self.active_requests + if req.is_disagg_generation_transmission_in_progress + } + current_time = time.time() + for request in requests_in_transfer.values(): + if request.py_kv_transfer_start_time is None: + continue + elapsed_time = ((current_time - request.py_kv_transfer_start_time) * + 1000) + if (elapsed_time > timeout_ms + and not request.py_kv_transfer_timed_out): + logger.warning( + f"Requesting cancellation for generation request " + f"{request.py_request_id} due to KV cache transfer timeout") + request.py_kv_transfer_timed_out = True + + user_canceled_ids = set(self.canceled_req_ids) + local_timed_out_ids = sorted( + request_id for request_id, request in requests_in_transfer.items() + if request.py_kv_transfer_timed_out + and request_id not in user_canceled_ids + and request_id not in self._disagg_timed_out_gen_cancelled_ids) + + if self.dist.tp_size > 1: + any_timed_out = self.dist.tp_allreduce(int( + bool(local_timed_out_ids)), + op=ReduceOp.MAX) + else: + any_timed_out = int(bool(local_timed_out_ids)) + if not any_timed_out: + return + + if self.dist.tp_size > 1: + gathered_timed_out_ids = self.dist.tp_allgather(local_timed_out_ids) + timed_out_ids = sorted(set().union(*gathered_timed_out_ids)) + else: + timed_out_ids = local_timed_out_ids + + for request_id in timed_out_ids: + request = requests_in_transfer.get(request_id) + if request is None: + continue + + # A peer rank may have crossed the timeout first. Mirror the + # TP-wide decision locally so a failed cancel attempt keeps + # retrying even if this rank's wall clock had not yet expired. + request.py_kv_transfer_timed_out = True + + if request_id in self._disagg_timed_out_gen_cancelled_ids: + continue + + is_cancelled = self.kv_cache_transceiver.cancel_request(request) + if is_cancelled: + self._disagg_timed_out_gen_cancelled_ids.add(request_id) + logger.warning( + f"Cancelled timed-out generation KV transfer for request " + f"{request.py_request_id}; waiting for C++ transfer " + "status to report final cleanup") + + @nvtx_range("_check_gen_cache_transfer_errors_consensus") + def _check_gen_cache_transfer_errors_consensus(self) -> None: + """Flush generation transfer errors through a TP-uniform path.""" + error_requests = [ + req for req in self._get_disagg_reqs_in_error_state() + if req.is_generation_only_request() + ] + poisoned_transfer_buffer = ( + self.kv_cache_transceiver is not None + and self.kv_cache_transceiver.has_poisoned_transfer_buffer()) + local_needs_flush = bool(error_requests) or poisoned_transfer_buffer + + if self.dist.tp_size > 1: + any_needs_flush = self.dist.tp_allreduce(int(local_needs_flush), + op=ReduceOp.MAX) + else: + any_needs_flush = int(local_needs_flush) + if not any_needs_flush: + return + + if self.dist.tp_size > 1: + any_poisoned_transfer_buffer = self.dist.tp_allreduce( + int(poisoned_transfer_buffer), op=ReduceOp.MAX) + else: + any_poisoned_transfer_buffer = int(poisoned_transfer_buffer) + + error_msg = "Error in kv cache transfer for generation requests" + if any_poisoned_transfer_buffer: + error_msg += "; poisoned transfer buffer requires process restart" + self._fatal_error = RuntimeError(f"Fatal error: {error_msg}") + self.is_shutdown = True + + self._handle_errors( + error_msg, + requests=None if any_poisoned_transfer_buffer else error_requests, + charge_budget=False) + @nvtx_range("_check_kv_transfer_timeout") def _check_kv_transfer_timeout(self): if not self.kv_cache_transceiver: @@ -4521,8 +4683,11 @@ def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: return elapsed_time = (current_time - req.py_kv_transfer_start_time) * 1000 if elapsed_time > timeout_ms and not req.py_kv_transfer_timed_out: + verb = ("Requesting cancellation for" + if self._is_disagg_inflight_cancel_active() else + "Observed timeout on") logger.warning( - f"Terminating {type} request {req.py_request_id} due to KV cache transfer timeout" + f"{verb} {type} request {req.py_request_id} due to KV cache transfer timeout" ) req.py_kv_transfer_timed_out = True @@ -4940,9 +5105,18 @@ def _check_cache_transfer_errors(self, error_msg_prefix: str): return error_requests = self._get_disagg_reqs_in_error_state() if error_requests: + poisoned_transfer_buffer = ( + self.kv_cache_transceiver is not None + and self.kv_cache_transceiver.has_poisoned_transfer_buffer()) + error_msg = f"Error in kv cache transfer for {error_msg_prefix}" + if poisoned_transfer_buffer: + error_msg += ( + "; poisoned transfer buffer requires process restart") + self._fatal_error = RuntimeError(f"Fatal error: {error_msg}") + self.is_shutdown = True self._handle_errors( - f"Error in kv cache transfer for {error_msg_prefix}", - requests=error_requests, + error_msg, + requests=None if poisoned_transfer_buffer else error_requests, charge_budget=False) @nvtx_range("_check_disagg_ctx_cache_transfer_status") @@ -4972,15 +5146,27 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): 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: - 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 request.py_kv_transfer_timed_out + or request_id in completed_req_ids + or request_id in self._disagg_timed_out_ctx_cancelled_ids): + continue + + is_cancelled = self.kv_cache_transceiver.cancel_request(request) + if not is_cancelled: + continue - self._end_transfer_and_maybe_terminate(request) + if self._is_disagg_inflight_cancel_active(): + self._disagg_timed_out_ctx_cancelled_ids.add(request_id) + logger.warning(f"Cancelled timed-out context KV transfer for " + f"request {request.py_request_id}; waiting for " + "C++ transfer status to report final cleanup") + else: + # 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") @@ -4994,7 +5180,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 self._is_disagg_inflight_cancel_active(): + self._check_cache_transfer_errors("generation requests") def _maybe_prefetch_next_iter_mm_encoders( self, scheduled_batch: ScheduledRequests) -> None: @@ -5367,6 +5554,8 @@ 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) + self._disagg_timed_out_ctx_cancelled_ids.discard(request.py_request_id) + self._disagg_timed_out_gen_cancelled_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) @@ -5390,6 +5579,10 @@ def _try_cancel_request(self, request) -> bool: if not self._is_request_in_transmission(request): return True + if self._is_disagg_inflight_cancel_active(): + self.kv_cache_transceiver.cancel_request(request) + return False + return self.kv_cache_transceiver.cancel_request(request) @nvtx_range("_handle_canceled_requests") @@ -5560,11 +5753,22 @@ 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 + # 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..64d37dfee4dc 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): @@ -1076,18 +1100,21 @@ def test_nixl_backend(self): @pytest.mark.skip_less_device(2) @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. + @pytest.mark.parametrize("transceiver_runtime", ["PYTHON", "CPP"], + ids=["python", "cpp"]) + def test_gen_only_sync(self, transceiver_runtime): + """Test gen-only synchronous KV transfer with both NIXL transceivers. 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": transceiver_runtime, "max_tokens_in_buffer": 4096, }, } @@ -1095,7 +1122,7 @@ def test_gen_only_sync(self): "disable_overlap_scheduler": True, "cache_transceiver_config": { "backend": "NIXL", - "transceiver_runtime": "PYTHON", + "transceiver_runtime": transceiver_runtime, "max_tokens_in_buffer": 4096, }, } @@ -1116,6 +1143,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/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 680f11f27afe..d95916a792dd 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -19,7 +19,8 @@ accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp1cp4] accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp2cp2] accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp2tp1cp2] -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_only_sync +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_only_sync[python] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_only_sync[cpp] accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_guided_decoding[llguidance-mtp_nextn=0] accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_guided_decoding[llguidance-mtp_nextn=2] accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_guided_decoding[xgrammar-mtp_nextn=0] diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a958f8b24951..635c0a6be839 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -15,6 +15,7 @@ l0_b200: backend: pytorch tests: # ------------- PyTorch tests --------------- + - unittest/others/test_kv_cache_transceiver.py::test_cpp_nixl_sync_transfer_stress - 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 1d4cee037e69..634ca4773e19 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,8 @@ 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[python] + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_only_sync[cpp] # ------------- 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_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index d65e8ba0cd70..8b8d5bd2e95e 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -29,7 +29,8 @@ l0_dgx_h100: - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_chunked_prefill - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_nixl_backend - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_nixl_backend - - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_only_sync + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_only_sync[python] + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_only_sync[cpp] - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_kv_cache_v2_nixl_python - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ngram - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_kv_cache_v2_nixl_python 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..bd85859bb165 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,7 @@ 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_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/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..3b9aff69a552 --- /dev/null +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -0,0 +1,590 @@ +# 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._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._disagg_timed_out_ctx_cancelled_ids = set() + 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) + assert request.py_request_id not in executor._disagg_timed_out_ctx_cancelled_ids + + +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() + executor._cancel_timed_out_gen_transfers = Mock() + executor._check_gen_cache_transfer_errors_consensus = Mock() + + PyExecutor._check_disagg_gen_transfer_status(executor) + + executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(0) + executor._cancel_timed_out_gen_transfers.assert_not_called() + executor._check_gen_cache_transfer_errors_consensus.assert_not_called() + + +@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() + assert constructor.call_args.kwargs["_inflight_cancel_config_prevalidated"] is True + + +def test_feature_preflight_rejects_asymmetric_rank_capability(monkeypatch): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "1") + mapping = _supported_mapping() + mapping.world_size = 2 + config = CacheTransceiverConfig(backend="NIXL", max_tokens_in_buffer=512) + dist = Mock() + dist.allgather.return_value = [ + (True, None), + (True, "rank-local capability is unsupported"), + ] + + with pytest.raises(ValueError, match="unsupported on one or more worker ranks"): + transceiver_module._validate_disagg_inflight_cancel_config( + config, + mapping, + dist, + inflight_cancel_supported_by_executor=True, + ) + + dist.allgather.assert_called_once_with((True, None)) + + +def test_feature_preflight_rejects_asymmetric_rank_mode(monkeypatch): + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, "1") + mapping = _supported_mapping() + mapping.world_size = 2 + config = CacheTransceiverConfig(backend="NIXL", max_tokens_in_buffer=512) + dist = Mock() + dist.allgather.return_value = [(True, None), (False, None)] + + with pytest.raises(ValueError, match="must resolve identically"): + transceiver_module._validate_disagg_inflight_cancel_config( + config, + mapping, + dist, + inflight_cancel_supported_by_executor=True, + ) + + +def test_feature_preflight_disabled_rank_rejects_enabled_peer(): + mapping = _supported_mapping() + mapping.world_size = 2 + config = CacheTransceiverConfig(backend="NIXL", max_tokens_in_buffer=512) + dist = Mock() + dist.allgather.return_value = [(False, None), (True, None)] + + with pytest.raises(ValueError, match="must resolve identically"): + transceiver_module._validate_disagg_inflight_cancel_config( + config, + mapping, + dist, + inflight_cancel_supported_by_executor=True, + ) + + dist.allgather.assert_called_once_with((False, None)) + + +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_dual_pool_kv_cache.py b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py index 1d5c00617c6e..15ff99917b9b 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, False), (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..ce3712398db5 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,121 @@ 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) @pytest.mark.parametrize("attention_type", [AttentionTypeCpp.DEFAULT, AttentionTypeCpp.MLA], @@ -293,6 +430,114 @@ def test_cancel_request_in_transmission(attention_type): assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR +@pytest.mark.timeout(120) +def test_queued_context_cancel_settles_and_rejects_late_peer(monkeypatch): + """A queued CTX cancellation must not wait for GEN RequestInfo. + + The local sender future becomes terminal immediately, while a tombstone + remains long enough to reject a generation peer that arrives later. + """ + monkeypatch.setenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, + "1") + monkeypatch.setattr(transceiver_module, + "_disagg_inflight_cancel_enabled_cache", None) + 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) + config = CacheTransceiverConfig(backend="NIXL", max_tokens_in_buffer=512) + transceiver_ctx = create_kv_cache_transceiver( + mapping, + dist, + kv_cache_manager_ctx, + AttentionTypeCpp.DEFAULT, + config, + inflight_cancel_supported_by_executor=True) + transceiver_gen = create_kv_cache_transceiver( + mapping, + dist, + kv_cache_manager_gen, + AttentionTypeCpp.DEFAULT, + config, + inflight_cancel_supported_by_executor=True) + + try: + fill_kv_cache_buffer(kv_cache_manager_ctx) + ctx_request = _build_ctx_request_for_timeout_test(request_id=43) + 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) + + assert transceiver_ctx.cancel_request(ctx_request) + failed_ctx_ids = set() + + def poll_context(): + _, failed = transceiver_ctx.check_context_transfer_status(0) + failed_ctx_ids.update(failed) + + wait_for_transfer_completion( + poll_context, + lambda: ctx_request.py_request_id in failed_ctx_ids, + timeout_s=5, + ) + assert ctx_request.state == LlmRequestState.DISAGG_TRANS_ERROR + + sampling_config = tensorrt_llm.bindings.SamplingConfig( + SamplingParams()._get_sampling_config()) + gen_request = LlmRequest( + request_id=43, + 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) + + wait_for_transfer_completion( + lambda: transceiver_gen.check_gen_transfer_status(0), + lambda: gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR, + timeout_s=5, + ) + assert not transceiver_gen.has_poisoned_transfer_buffer() + finally: + shutdown_transceivers(transceiver_gen, transceiver_ctx) + + +@pytest.mark.timeout(60) +def test_default_off_queued_context_cancel_preserves_pending_future( + monkeypatch): + """Legacy queued cancellation remains pending until the peer handshake.""" + monkeypatch.delenv(transceiver_module._DISAGG_INFLIGHT_CANCEL_ENABLED_ENV, + raising=False) + monkeypatch.setattr(transceiver_module, + "_disagg_inflight_cancel_enabled_cache", None) + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + kv_cache_manager = create_kv_cache_manager(mapping, DataType.HALF) + config = CacheTransceiverConfig(backend="NIXL", max_tokens_in_buffer=512) + transceiver = create_kv_cache_transceiver(mapping, dist, kv_cache_manager, + AttentionTypeCpp.DEFAULT, config) + try: + fill_kv_cache_buffer(kv_cache_manager) + request = _build_ctx_request_for_timeout_test(request_id=44) + kv_cache_manager.impl.add_sequence_batch( + [(request.py_request_id, request.prompt_len, 1)], [request]) + transceiver.respond_and_send_async(request) + + assert transceiver.cancel_request(request) + completed, failed = transceiver.check_context_transfer_status(0) + assert completed == [] + assert failed == [] + finally: + shutdown_transceivers(transceiver) + + @pytest.mark.timeout(120) def test_async_transfer_keeps_llm_request_alive(): """Async entry points must hold a strong shared_ptr to the LlmRequest. From 25aac382f44ac30e8d8580bf49fa28635cb79c29 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 30 Jun 2026 15:22:58 -0700 Subject: [PATCH 2/2] [TRTLLM-12721][fix] Negotiate disaggregated cancellation protocol and mode Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/cacheTransceiver.cpp | 6 +- .../batch_manager/cacheTransferLayer.cpp | 24 +- .../batch_manager/cacheTransferLayer.h | 9 +- .../batch_manager/dataTransceiver.cpp | 157 +++++-- .../batch_manager/dataTransceiver.h | 26 +- .../agent_utils/connection.cpp | 74 +++- .../agent_utils/connection.h | 9 +- .../agent_utils/peerProtocol.h | 388 ++++++++++++++++++ .../unit_tests/batch_manager/CMakeLists.txt | 1 + .../batch_manager/bufferIndexHolderTest.cpp | 58 +++ .../batch_manager/peerProtocolTest.cpp | 267 ++++++++++++ .../unit_tests/executor/agentCommTest.cpp | 38 +- .../executor/serializeUtilsTest.cpp | 64 +++ .../multi_gpu/cacheTransceiverTest.cpp | 136 +++++- 14 files changed, 1196 insertions(+), 61 deletions(-) create mode 100644 cpp/tensorrt_llm/executor/cache_transmission/agent_utils/peerProtocol.h create mode 100644 cpp/tests/unit_tests/batch_manager/peerProtocolTest.cpp diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index b45ff6d2a9d8..81849cd774c7 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -627,8 +627,11 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa { auto rnnState = mCacheState->hasRnnConfig() ? std::make_optional(mCacheState->getRnnCacheState()) : std::nullopt; + auto const peerCancellationMode = mInflightCancelEnabled + ? executor::kv_cache::PeerCancellationMode::kEnabled + : executor::kv_cache::PeerCancellationMode::kBaseline; mManager = std::make_unique( - mCacheTransBufferManagerPtrs, *mCacheState, "nixl", rnnState); + mCacheTransBufferManagerPtrs, *mCacheState, "nixl", rnnState, peerCancellationMode); TLLM_LOG_INFO("NIXL Connection Manager created"); } else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MOONCAKE) @@ -660,7 +663,6 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa } return createCacheFormatter(cacheManager, kvBufferPtrs, isMLA); }; - auto makeRnnFormatter = [this, cacheManager]() -> std::unique_ptr { if (mRnnStateManager != nullptr && mRnnCacheTransBufferManager != nullptr) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp index 23a159b922ef..a61a8651439a 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp @@ -25,6 +25,7 @@ #include "tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h" #include +#include namespace tensorrt_llm::batch_manager { @@ -45,7 +46,7 @@ CacheTransferLayer::~CacheTransferLayer() = default; CacheTransferLayer::CacheTransferLayer(CacheTransferLayer&&) noexcept = default; CacheTransferLayer& CacheTransferLayer::operator=(CacheTransferLayer&&) noexcept = default; -void CacheTransferLayer::validateSupport(executor::DataTransceiverState const& peerState) const +void CacheTransferLayer::validateCacheSupport(executor::DataTransceiverState const& peerState) const { TLLM_CHECK_WITH_INFO(mKvFormatter->inquireSupport(mCacheState, peerState.getCacheState().value()), "Disagg server does not currently support these cacheState, please check the cacheState of the context and " @@ -75,6 +76,27 @@ void CacheTransferLayer::validateSupport(executor::DataTransceiverState const& p } } +executor::kv_cache::PeerProtocolCompatibility CacheTransferLayer::getPeerProtocolCompatibility( + executor::DataTransceiverState const& peerState) const +{ + if (!peerState.getCommState().has_value() || !peerState.getCommState()->isAgentState()) + { + return { + true, false, false, false, std::nullopt, "peer protocol negotiation is not applicable to this transport"}; + } + + std::vector peerAgentNames; + auto const& peerAgentStates = peerState.getCommState()->getAgentState(); + peerAgentNames.reserve(peerAgentStates.size()); + for (auto const& peerAgentState : peerAgentStates) + { + peerAgentNames.emplace_back(peerAgentState.mAgentName); + } + auto const localMode = mEnableInflightCancel ? executor::kv_cache::PeerCancellationMode::kEnabled + : executor::kv_cache::PeerCancellationMode::kBaseline; + return executor::kv_cache::validatePeerProtocol(localMode, peerAgentNames); +} + std::vector CacheTransferLayer::computeCounterparts( SizeType32 selfIdx, executor::DataTransceiverState const& peerState) const { diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h index 468b15bfdc4a..a4241366b21c 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h +++ b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h @@ -18,6 +18,7 @@ #pragma once #include "tensorrt_llm/batch_manager/rnnCacheFormatter.h" +#include "tensorrt_llm/executor/cache_transmission/agent_utils/peerProtocol.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/runtime/common.h" @@ -55,9 +56,13 @@ class CacheTransferLayer CacheTransferLayer(CacheTransferLayer&&) noexcept; CacheTransferLayer& operator=(CacheTransferLayer&&) noexcept; - /// @brief Validates all cache types against the peer state. Throws on incompatibility. + /// @brief Validates cache layout support independently of peer protocol negotiation. /// @param peerState The peer's DataTransceiverState. - void validateSupport(executor::DataTransceiverState const& peerState) const; + void validateCacheSupport(executor::DataTransceiverState const& peerState) const; + + /// @brief Check the peer's advertised wire protocol and effective cancellation mode. + [[nodiscard]] executor::kv_cache::PeerProtocolCompatibility getPeerProtocolCompatibility( + executor::DataTransceiverState const& peerState) const; /// @brief Computes counterparts for all cache types and returns the union (KV + RNN if present). /// @param selfIdx The sequential index of the current executor process. diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 2c14b3400aa5..669dbedaf2d9 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -371,6 +371,7 @@ class CacheSender::Impl auto it = mRequestToSession.find(requestId); TLLM_CHECK(it != mRequestToSession.end()); mRequestToSession.erase(it); + mPeerProtocolErrors.erase(requestId); } { std::lock_guard lg(mInFlightCancelMutex); @@ -384,6 +385,7 @@ class CacheSender::Impl { std::unique_lock lk(mMtxForMap); mRequestToSession.erase(requestId); + mPeerProtocolErrors.erase(requestId); } catch (...) { @@ -429,7 +431,8 @@ class CacheSender::Impl } auto requestId = info.getRequestId(); - mCacheTransferLayer.validateSupport(info.getTransState()); + mCacheTransferLayer.validateCacheSupport(info.getTransState()); + auto const peerProtocolCompatibility = mCacheTransferLayer.getPeerProtocolCompatibility(info.getTransState()); auto allCounterparts = mCacheTransferLayer.computeCounterparts( mSelfState.getCommState().value().getSelfIdx(), info.getTransState()); @@ -450,10 +453,22 @@ class CacheSender::Impl DataContext{tagFromRequestId(requestId), *cancelFlag}, allCounterparts, mSelfState, info.getTransState(), mBufferManager, info.getIndexFromEnd(), info.getLastBlockKey(), nullptr, !common::getEnvKVCacheTimeOutputPath().empty(), mInflightCancelEnabled); + // Version 1 has no version-dependent dispatch. Future versions must branch on this session-bound + // value instead of the local binary's maximum supported version. + session.setPeerProtocolVersion(peerProtocolCompatibility.selectedVersion); session.setTime(TransferSession::kTimeRequestInfo); it = mRequestToSession.emplace(requestId, std::move(session)).first; } it->second.setConnection(peerIdx, connection); + if (!peerProtocolCompatibility.compatible) + { + mPeerProtocolErrors.emplace(requestId, peerProtocolCompatibility.reason); + } + } + if (!peerProtocolCompatibility.compatible) + { + TLLM_LOG_WARNING("Rejecting KV cache transfer for request %zu before DMA: %s", requestId, + peerProtocolCompatibility.reason.c_str()); } return info; } @@ -640,6 +655,15 @@ class CacheSender::Impl void sendResponse(RequestIdType reqId) { + std::optional peerProtocolError; + { + std::scoped_lock lock(mMtxForMap); + auto const errorIt = mPeerProtocolErrors.find(reqId); + if (errorIt != mPeerProtocolErrors.end()) + { + peerProtocolError = errorIt->second; + } + } bool isReady = true; { std::scoped_lock lock(mSenderMutex); @@ -657,7 +681,7 @@ class CacheSender::Impl return; } mRemainSendCount.erase(countIt); - isReady = responseIt != mReadyResponses.end() && !isCancelled; + isReady = responseIt != mReadyResponses.end() && !isCancelled && !peerProtocolError.has_value(); } // Keep mCurrentRequest set while notifying the peer so cancellation cannot change the decision after it has @@ -697,9 +721,12 @@ class CacheSender::Impl { if (response.has_value()) { + auto const errorMessage = peerProtocolError.has_value() + ? "KV cache transfer rejected due to peer protocol mismatch: " + peerProtocolError.value() + : "KV cache transfer was cancelled"; failResponse(*response, std::make_exception_ptr(TLLM_REQUEST_EXCEPTION(reqId, common::RequestErrorCode::kNETWORK_ERROR, - "KV cache transfer for request %zu was cancelled", reqId))); + "%s for request %zu", errorMessage.c_str(), reqId))); } discardTransferState(reqId); } @@ -867,6 +894,7 @@ class CacheSender::Impl executor::kv_cache::ConnectionManager* mManager; std::map mRequestToSession; + std::unordered_map mPeerProtocolErrors; executor::DataTransceiverState mSelfState; bool mInflightCancelEnabled{false}; CacheTransferLayer mCacheTransferLayer; @@ -984,14 +1012,41 @@ class CacheReceiver::Impl } } + executor::kv_cache::PeerProtocolCompatibility validateRequestSupport(LlmRequest const& llmRequest) const + { + auto const& contextState = llmRequest.getDataTransceiverState(); + mCacheTransferLayer.validateCacheSupport(contextState); + auto const compatibility = mCacheTransferLayer.getPeerProtocolCompatibility(contextState); + if (!compatibility.compatible && !compatibility.allPeersCanRejectBeforeDma) + { + TLLM_THROW("Cannot safely reject incompatible legacy or malformed CTX peer before buffer advertisement: %s", + compatibility.reason.c_str()); + } + return compatibility; + } + TransferSession sendRequestInfo(LlmRequest const& llmRequest, std::atomic const* perRequestCancel = nullptr, - std::vector> cacheBufferIds = {}) + std::vector> cacheBufferIds = {}, + executor::kv_cache::PeerProtocolCompatibility const* validatedCompatibility = nullptr) { uint64_t requestId = llmRequest.getContextPhaseParams().value().getReqId(); auto const& contextState = llmRequest.getDataTransceiverState(); auto const& commState = contextState.getCommState().value(); auto const& destCacheState = contextState.getCacheState().value(); - mCacheTransferLayer.validateSupport(contextState); + std::optional computedCompatibility; + if (validatedCompatibility == nullptr) + { + computedCompatibility = validateRequestSupport(llmRequest); + validatedCompatibility = std::addressof(computedCompatibility.value()); + } + auto const& peerProtocolCompatibility = *validatedCompatibility; + if (!peerProtocolCompatibility.compatible) + { + // Every peer rank explicitly advertised pre-DMA protocol rejection, so complete the existing + // request/ready exchange. CTX will return ready=false and settle its queued response. + TLLM_LOG_WARNING("Completing control-only KV request handshake for request %zu so CTX can reject it: %s", + requestId, peerProtocolCompatibility.reason.c_str()); + } RequestInfo requestInfo(requestId, mSelfState); @@ -1082,6 +1137,12 @@ class CacheReceiver::Impl mSelfState, contextState, resource->mBufferManager, requestInfo.getIndexFromEnd(), requestInfo.getLastBlockKey(), &llmRequest, !common::getEnvKVCacheTimeOutputPath().empty(), mInflightCancelEnabled); + // Bind all later transfer behavior to the negotiated version. Version 1 needs no conditional dispatch. + session.setPeerProtocolVersion(peerProtocolCompatibility.selectedVersion); + if (!peerProtocolCompatibility.compatible) + { + session.rejectPeerProtocol(); + } session.setRecvBufferHolders(std::move(allocatedRecvHolders)); for (size_t ci = 0; ci < allCounterparts.size(); ci++) @@ -1247,6 +1308,26 @@ class CacheReceiver::Impl kCancelled, }; + bool settleLocallyRejectedSession(TransferSession& session, ReadySignalResult const result) + { + if (!session.isPeerProtocolRejected()) + { + return false; + } + if (result == ReadySignalResult::kNotReady) + { + // Every peer confirmed that DMA did not start. + session.releaseRecvBufferHolders(); + } + else + { + // A ready=true or incomplete reply contradicts the local fail-closed decision. The advertised slots may + // be writable by a peer, so they cannot be reused. + session.poisonRecvBufferHolders(); + } + return true; + } + ReadySignalResult receiveReadySignalDetailed(TransferSession& session, std::atomic const& perRequestCancel) { bool isReady = false; @@ -1298,6 +1379,10 @@ class CacheReceiver::Impl bool receiveReadySignal(TransferSession& session) { auto const result = receiveReadySignalDetailed(session, mTerminate); + if (settleLocallyRejectedSession(session, result)) + { + return false; + } if (result == ReadySignalResult::kNotReady) { // An explicit peer rejection proves that no transfer started. @@ -1356,27 +1441,7 @@ class CacheReceiver::Impl 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); - } - } - } - + bool rejectPeerProtocol = false; auto poisonRecvHolders = [&recvHolders]() { for (auto& holder : recvHolders) @@ -1387,10 +1452,46 @@ class CacheReceiver::Impl try { - auto session = sendRequestInfo(llmRequest, &perRequestCancel, std::move(cacheBufferIds)); + // Validate before reserving slots. An incompatible legacy or malformed peer is rejected before any + // destination is advertised, so that safe local failure must not quarantine an unused receive pool. + auto const peerProtocolCompatibility = validateRequestSupport(llmRequest); + + 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); + } + } + } + + rejectPeerProtocol = !peerProtocolCompatibility.compatible; + auto session + = sendRequestInfo(llmRequest, &perRequestCancel, std::move(cacheBufferIds), &peerProtocolCompatibility); session.setRecvBufferHolders(std::move(recvHolders)); session.setTime(TransferSession::kTimeRequestInfo); auto readyResult = receiveReadySignalDetailed(session, perRequestCancel); + if (settleLocallyRejectedSession(session, readyResult)) + { + if (!mInflightCancelEnabled) + { + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + } + llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); + return false; + } if (readyResult == ReadySignalResult::kCancelled) { if (!mInflightCancelEnabled) @@ -1432,7 +1533,7 @@ class CacheReceiver::Impl } catch (...) { - if (mInflightCancelEnabled) + if (mInflightCancelEnabled || rejectPeerProtocol) { poisonRecvHolders(); } diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index 028403b224e8..0512898daa54 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -16,9 +16,11 @@ */ #pragma once +#include #include #include #include +#include #include #include @@ -112,7 +114,7 @@ class TransferSession ~TransferSession() { - if (mEnableInflightCancel) + if (mEnableInflightCancel || mPeerProtocolRejected) { poisonRecvBufferHolders(); } @@ -172,6 +174,26 @@ class TransferSession return mEnableInflightCancel; } + void rejectPeerProtocol() noexcept + { + mPeerProtocolRejected = true; + } + + [[nodiscard]] bool isPeerProtocolRejected() const noexcept + { + return mPeerProtocolRejected; + } + + void setPeerProtocolVersion(std::optional version) noexcept + { + mPeerProtocolVersion = version; + } + + [[nodiscard]] std::optional getPeerProtocolVersion() const noexcept + { + return mPeerProtocolVersion; + } + /// Transfer ownership of preassigned receive-buffer slots to this session. /// Formatters borrow these slots and must not release them independently. void setRecvBufferHolders(std::vector holders) @@ -212,6 +234,8 @@ class TransferSession int32_t mIndexFromEnd{0}; BlockKey mLastBlockKey{}; bool mEnableInflightCancel{false}; + bool mPeerProtocolRejected{false}; + std::optional mPeerProtocolVersion; std::vector mRecvBufferHolders; }; 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 c713eb5ec61e..80b197763912 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include #include @@ -81,20 +82,18 @@ bool isValidBufferKind(uint8_t kind) return false; } -AgentState const* findAgentState(CommState const& commState, std::string const& agentName) +std::uint64_t makeProtocolSessionToken(std::string_view const agentName) { - if (!commState.isAgentState()) + // FNV-1a makes the already-unique registered agent identity available as a compact process-session token. + constexpr std::uint64_t kOffsetBasis{14695981039346656037ULL}; + constexpr std::uint64_t kPrime{1099511628211ULL}; + std::uint64_t token = kOffsetBasis; + for (auto const character : agentName) { - return nullptr; + token ^= static_cast(character); + token *= kPrime; } - for (auto const& agentState : commState.getAgentState()) - { - if (agentState.mAgentName == agentName) - { - return &agentState; - } - } - return nullptr; + return token == 0 ? kOffsetBasis : token; } void validateConnectionIdx( @@ -107,6 +106,29 @@ void validateConnectionIdx( } // namespace +void validateRequestPeerIdentity(batch_manager::RequestInfo const& requestInfo, std::string const& notificationAgent, + std::string const& notificationAddress) +{ + auto const& transState = requestInfo.getTransState(); + TLLM_CHECK_WITH_INFO(transState.getCommState().has_value(), + "AgentConnectionManager received request without comm state from agent '%s'", notificationAgent.c_str()); + auto const& peerCommState = transState.getCommState().value(); + TLLM_CHECK_WITH_INFO(peerCommState.isAgentState(), + "AgentConnectionManager received non-agent comm state from agent '%s'", notificationAgent.c_str()); + auto const& peerAgentStates = peerCommState.getAgentState(); + auto const peerSelfIdx = peerCommState.getSelfIdx(); + TLLM_CHECK_WITH_INFO(peerSelfIdx >= 0 && static_cast(peerSelfIdx) < peerAgentStates.size(), + "AgentConnectionManager received self index %d outside agent-state count %zu from agent '%s'", peerSelfIdx, + peerAgentStates.size(), notificationAgent.c_str()); + auto const& advertisedSelf = peerAgentStates[peerSelfIdx]; + TLLM_CHECK_WITH_INFO(advertisedSelf.mAgentName == notificationAgent, + "AgentConnectionManager received request from '%s' whose comm state identifies self as '%s'", + notificationAgent.c_str(), advertisedSelf.mAgentName.c_str()); + TLLM_CHECK_WITH_INFO(advertisedSelf.mConnectionInfo == notificationAddress, + "AgentConnectionManager received request from '%s' whose comm-state address differs from the live sender", + notificationAgent.c_str()); +} + AgentConnection::AgentConnection( std::string mAgentName, std::string mRemoteAgentName, AgentConnectionManager* mAgentConnectionManager) : mAgentName(mAgentName) @@ -372,16 +394,24 @@ std::optional AgentConnection::getPreAssignedBufferId(uint8_t kind) cons AgentConnectionManager::AgentConnectionManager( std::vector cacheTransBufferManagers, CacheState cacheState, - std::string const& backendType, std::optional rnnCacheState) + std::string const& backendType, std::optional rnnCacheState, + std::optional peerCancellationMode) : mCacheState(std::move(cacheState)) , mRnnCacheState(std::move(rnnCacheState)) , mCacheTransBufferManagers(std::move(cacheTransBufferManagers)) , mRegMemDescs(MemoryType::kVRAM, {}) + , mPeerProtocolAware(peerCancellationMode.has_value()) { TLLM_CUDA_CHECK(cudaGetDevice(&mDeviceId)); TLLM_CHECK(mDeviceId != -1); mAgentName = genUniqueAgentName(); + if (peerCancellationMode.has_value()) + { + auto const descriptor = makePeerProtocolDescriptor( + peerCancellationMode.value() == PeerCancellationMode::kEnabled, makeProtocolSessionToken(mAgentName)); + mAgentName = appendPeerProtocolDescriptor(std::move(mAgentName), descriptor); + } // Create Agent BaseAgentConfig config{mAgentName, true, false, true}; m_Agent = makeTransferAgent(backendType, &config); @@ -450,6 +480,18 @@ AgentConnectionManager::AgentConnectionManager( agentStates[0] = localAgentState; } mCommState = CommState(agentStates, mpi::MpiComm::session().getRank()); + if (peerCancellationMode.has_value()) + { + std::vector localAgentNames; + localAgentNames.reserve(agentStates.size()); + for (auto const& agentState : agentStates) + { + localAgentNames.emplace_back(agentState.mAgentName); + } + auto const localCompatibility = validateLocalPeerProtocol(peerCancellationMode.value(), localAgentNames); + TLLM_CHECK_WITH_INFO(localCompatibility.compatible, + "Local NIXL ranks advertise inconsistent cache-transfer protocols: %s", localCompatibility.reason.c_str()); + } TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), " ***** AgentConnectionManager::AgentConnectionManager mCommState: %s", mCommState.toString().c_str()); } @@ -491,14 +533,10 @@ AgentConnection const* AgentConnectionManager::recvConnectionAndRequestInfo( TLLM_CHECK_WITH_INFO(agent == remoteAgentName, "AgentConnectionManager received RequestAndBufferInfo from '%s' with embedded agent '%s'", agent.c_str(), remoteAgentName.c_str()); - auto const* expectedAgentState = findAgentState(mCommState, remoteAgentName); - if (expectedAgentState != nullptr) + if (mPeerProtocolAware) { - TLLM_CHECK_WITH_INFO(address == expectedAgentState->mConnectionInfo, - "AgentConnectionManager received mismatched connection info for agent '%s'", - remoteAgentName.c_str()); + validateRequestPeerIdentity(requestInfo, remoteAgentName, address); } - else { std::scoped_lock remoteInfoLock(mRemoteConnectionInfoMutex); auto [remoteInfoIt, inserted] = mRemoteConnectionInfo.emplace(remoteAgentName, address); 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 4b167556cc8d..ee496865ed90 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h @@ -22,6 +22,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/executor/cacheCommunicator.h" +#include "tensorrt_llm/executor/cache_transmission/agent_utils/peerProtocol.h" #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/executor/transferAgent.h" #include @@ -35,6 +36,10 @@ namespace tensorrt_llm::executor::kv_cache // that share hostname (--network host) and PID namespace. std::string genUniqueAgentName(); +//! Bind request-advertised peer state to the live notification sender before trusting protocol metadata. +void validateRequestPeerIdentity(batch_manager::RequestInfo const& requestInfo, std::string const& notificationAgent, + std::string const& notificationAddress); + struct RequestAndBufferInfo { std::string mAgentName; @@ -305,7 +310,8 @@ class AgentConnectionManager : public ConnectionManager public: AgentConnectionManager(std::vector cacheTransBufferManagers, CacheState cacheState, std::string const& backendType, - std::optional rnnCacheState = std::nullopt); + std::optional rnnCacheState = std::nullopt, + std::optional peerCancellationMode = std::nullopt); ~AgentConnectionManager(); AgentConnection* recvConnect(DataContext const& ctx, void* data, size_t size) override; [[nodiscard]] std::vector getConnections(CommState const& state) override; @@ -347,6 +353,7 @@ class AgentConnectionManager : public ConnectionManager int mDeviceId; std::string mAgentName; MemoryDescs mRegMemDescs; + bool mPeerProtocolAware{false}; std::atomic mIsRunning{true}; }; diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/peerProtocol.h b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/peerProtocol.h new file mode 100644 index 000000000000..c636739165ea --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/peerProtocol.h @@ -0,0 +1,388 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations under + * the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::executor::kv_cache +{ + +//! Effective cancellation behavior for a cache-transceiver process. +enum class PeerCancellationMode : std::uint8_t +{ + kBaseline = 0, + kEnabled = 1, +}; + +//! Protocol capabilities advertised through the NIXL agent identity. +enum class PeerProtocolCapability : std::uint64_t +{ + kInflightCancellation = std::uint64_t{1} << 0, + //! The peer can settle a protocol mismatch with ready=false before DMA. + kPreDmaProtocolRejection = std::uint64_t{1} << 1, +}; + +struct PeerProtocolDescriptor +{ + std::uint16_t minVersion{1}; + std::uint16_t maxVersion{1}; + std::uint64_t capabilities{0}; + PeerCancellationMode cancellationMode{PeerCancellationMode::kBaseline}; + std::uint64_t sessionToken{0}; + + [[nodiscard]] bool operator==(PeerProtocolDescriptor const& other) const noexcept + { + return minVersion == other.minVersion && maxVersion == other.maxVersion && capabilities == other.capabilities + && cancellationMode == other.cancellationMode && sessionToken == other.sessionToken; + } +}; + +enum class PeerProtocolParseStatus : std::uint8_t +{ + kMissing, + kValid, + kMalformed, +}; + +struct PeerProtocolParseResult +{ + PeerProtocolParseStatus status{PeerProtocolParseStatus::kMissing}; + std::optional descriptor; + std::string reason; +}; + +struct PeerProtocolCompatibility +{ + bool compatible{false}; + bool hasLegacyPeer{false}; + bool allPeersProtocolAware{false}; + bool allPeersCanRejectBeforeDma{false}; + std::optional selectedVersion; + std::string reason; +}; + +namespace detail +{ + +// The terminal envelope avoids interpreting an operator-controlled hostname fragment as protocol metadata. A future +// descriptor schema must use a new versioned marker rather than adding fields to this v1 envelope. +constexpr std::string_view kPeerProtocolMarker{"__trtllm_protocol_v1_"}; +constexpr std::string_view kPeerProtocolTerminator{"__trtllm_protocol_end"}; +constexpr std::uint16_t kPeerProtocolMinVersion{1}; +constexpr std::uint16_t kPeerProtocolMaxVersion{1}; +constexpr std::uint64_t kInflightCancellationCapability + = static_cast(PeerProtocolCapability::kInflightCancellation); +constexpr std::uint64_t kPreDmaProtocolRejectionCapability + = static_cast(PeerProtocolCapability::kPreDmaProtocolRejection); +constexpr std::uint64_t kAdvertisedCapabilities = kInflightCancellationCapability | kPreDmaProtocolRejectionCapability; + +inline bool parseUnsigned(std::string_view text, int const base, std::uint64_t& value) +{ + if (text.empty()) + { + return false; + } + auto const* begin = text.data(); + auto const* end = begin + text.size(); + auto const result = std::from_chars(begin, end, value, base); + return result.ec == std::errc{} && result.ptr == end; +} + +inline std::optional> splitProtocolFields(std::string_view payload) +{ + std::array fields; + for (size_t field = 0; field < fields.size(); ++field) + { + auto const separator = payload.find('_'); + if (field + 1 == fields.size()) + { + if (separator != std::string_view::npos) + { + return std::nullopt; + } + fields[field] = payload; + return fields; + } + if (separator == std::string_view::npos) + { + return std::nullopt; + } + fields[field] = payload.substr(0, separator); + payload.remove_prefix(separator + 1); + } + return std::nullopt; +} + +inline char const* cancellationModeName(PeerCancellationMode const mode) +{ + switch (mode) + { + case PeerCancellationMode::kBaseline: return "baseline"; + case PeerCancellationMode::kEnabled: return "enabled"; + } + return "unknown"; +} + +} // namespace detail + +//! Construct the descriptor supported by this binary. +inline PeerProtocolDescriptor makePeerProtocolDescriptor( + bool const inflightCancellationEnabled, std::uint64_t const sessionToken) +{ + return PeerProtocolDescriptor{detail::kPeerProtocolMinVersion, detail::kPeerProtocolMaxVersion, + detail::kAdvertisedCapabilities, + inflightCancellationEnabled ? PeerCancellationMode::kEnabled : PeerCancellationMode::kBaseline, sessionToken}; +} + +//! Append an opaque, backward-compatible protocol suffix to the actual NIXL agent identity. +inline std::string appendPeerProtocolDescriptor(std::string agentName, PeerProtocolDescriptor const& descriptor) +{ + std::ostringstream suffix; + suffix << detail::kPeerProtocolMarker << descriptor.minVersion << '_' << descriptor.maxVersion << '_' << std::hex + << descriptor.capabilities << '_' << std::dec << static_cast(descriptor.cancellationMode) << '_' + << std::hex << std::setw(16) << std::setfill('0') << descriptor.sessionToken; + suffix << detail::kPeerProtocolTerminator; + agentName += suffix.str(); + return agentName; +} + +//! Parse a descriptor suffix without changing the transport identity string. +inline PeerProtocolParseResult parsePeerProtocolDescriptor(std::string_view const agentName) +{ + if (agentName.size() < detail::kPeerProtocolTerminator.size() + || agentName.substr(agentName.size() - detail::kPeerProtocolTerminator.size()) + != detail::kPeerProtocolTerminator) + { + return {PeerProtocolParseStatus::kMissing, std::nullopt, "peer does not advertise a protocol descriptor"}; + } + + auto const envelopeEnd = agentName.size() - detail::kPeerProtocolTerminator.size(); + auto const markerPos = agentName.rfind(detail::kPeerProtocolMarker, envelopeEnd); + if (markerPos == std::string_view::npos) + { + return {PeerProtocolParseStatus::kMissing, std::nullopt, "peer does not advertise a protocol descriptor"}; + } + + auto const payloadBegin = markerPos + detail::kPeerProtocolMarker.size(); + auto const fields = detail::splitProtocolFields(agentName.substr(payloadBegin, envelopeEnd - payloadBegin)); + if (!fields.has_value()) + { + return {PeerProtocolParseStatus::kMalformed, std::nullopt, + "protocol descriptor must contain exactly version range, capabilities, mode, and session token"}; + } + + std::uint64_t minVersion = 0; + std::uint64_t maxVersion = 0; + std::uint64_t capabilities = 0; + std::uint64_t mode = 0; + std::uint64_t sessionToken = 0; + bool const parsed = detail::parseUnsigned(fields.value()[0], /*base=*/10, minVersion) + && detail::parseUnsigned(fields.value()[1], /*base=*/10, maxVersion) + && detail::parseUnsigned(fields.value()[2], /*base=*/16, capabilities) + && detail::parseUnsigned(fields.value()[3], /*base=*/10, mode) + && detail::parseUnsigned(fields.value()[4], /*base=*/16, sessionToken); + if (!parsed || minVersion > std::numeric_limits::max() + || maxVersion > std::numeric_limits::max()) + { + return {PeerProtocolParseStatus::kMalformed, std::nullopt, + "protocol descriptor contains a non-numeric or out-of-range field"}; + } + if (minVersion == 0 || minVersion > maxVersion) + { + return {PeerProtocolParseStatus::kMalformed, std::nullopt, "protocol version range is invalid"}; + } + if (mode > static_cast(PeerCancellationMode::kEnabled)) + { + return {PeerProtocolParseStatus::kMalformed, std::nullopt, "effective cancellation mode is invalid"}; + } + if (sessionToken == 0) + { + return {PeerProtocolParseStatus::kMalformed, std::nullopt, "protocol session token must be nonzero"}; + } + + return {PeerProtocolParseStatus::kValid, + PeerProtocolDescriptor{static_cast(minVersion), static_cast(maxVersion), + capabilities, static_cast(mode), sessionToken}, + {}}; +} + +//! Validate every peer rank before a ready signal or destination-buffer advertisement. +inline PeerProtocolCompatibility validatePeerProtocol( + PeerProtocolDescriptor const& localDescriptor, std::vector const& peerAgentNames) +{ + if (peerAgentNames.empty()) + { + return {false, false, false, false, std::nullopt, "peer agent-state list is empty"}; + } + if (localDescriptor.minVersion == 0 || localDescriptor.minVersion > localDescriptor.maxVersion) + { + return {false, false, false, false, std::nullopt, "local protocol version range is invalid"}; + } + if (localDescriptor.sessionToken == 0) + { + return {false, false, false, false, std::nullopt, "local protocol session token must be nonzero"}; + } + if (localDescriptor.cancellationMode == PeerCancellationMode::kEnabled + && (localDescriptor.capabilities & detail::kInflightCancellationCapability) == 0) + { + return {false, false, false, false, std::nullopt, + "local protocol descriptor lacks the in-flight cancellation capability"}; + } + if (localDescriptor.cancellationMode == PeerCancellationMode::kEnabled + && (localDescriptor.capabilities & detail::kPreDmaProtocolRejectionCapability) == 0) + { + return {false, false, false, false, std::nullopt, + "local protocol descriptor lacks the pre-DMA protocol-rejection capability"}; + } + + bool hasLegacyPeer = false; + bool allPeersCanRejectBeforeDma = true; + for (size_t rank = 0; rank < peerAgentNames.size(); ++rank) + { + auto const parsed = parsePeerProtocolDescriptor(peerAgentNames[rank]); + if (parsed.status == PeerProtocolParseStatus::kMissing) + { + hasLegacyPeer = true; + allPeersCanRejectBeforeDma = false; + continue; + } + if (parsed.status == PeerProtocolParseStatus::kMalformed) + { + return {false, hasLegacyPeer, false, false, std::nullopt, + "peer rank " + std::to_string(rank) + " has a malformed protocol descriptor: " + parsed.reason}; + } + allPeersCanRejectBeforeDma + &= (parsed.descriptor->capabilities & detail::kPreDmaProtocolRejectionCapability) != 0; + } + + bool const allPeersProtocolAware = !hasLegacyPeer; + auto const localMode = localDescriptor.cancellationMode; + if (localMode == PeerCancellationMode::kEnabled && hasLegacyPeer) + { + return {false, true, false, false, std::nullopt, + "local cancellation mode is enabled but one or more peer ranks do not advertise compatible " + "cancellation semantics"}; + } + + std::uint16_t commonMinVersion = localDescriptor.minVersion; + std::uint16_t commonMaxVersion = localDescriptor.maxVersion; + for (size_t rank = 0; rank < peerAgentNames.size(); ++rank) + { + auto const parsed = parsePeerProtocolDescriptor(peerAgentNames[rank]); + if (parsed.status == PeerProtocolParseStatus::kMissing) + { + continue; + } + auto const& peer = parsed.descriptor.value(); + if (peer.cancellationMode != localMode) + { + return {false, hasLegacyPeer, allPeersProtocolAware, allPeersCanRejectBeforeDma, std::nullopt, + "effective cancellation mode mismatch at peer rank " + std::to_string(rank) + + ": local=" + detail::cancellationModeName(localMode) + + ", peer=" + detail::cancellationModeName(peer.cancellationMode)}; + } + // Cancellation protocol versions do not affect the baseline wire path. Keeping baseline-to-baseline + // transfers compatible lets deployments roll this protocol independently while cancellation remains off. + if (localMode == PeerCancellationMode::kBaseline) + { + continue; + } + + commonMinVersion = std::max(commonMinVersion, peer.minVersion); + commonMaxVersion = std::min(commonMaxVersion, peer.maxVersion); + if (commonMinVersion > commonMaxVersion) + { + return {false, hasLegacyPeer, allPeersProtocolAware, allPeersCanRejectBeforeDma, std::nullopt, + "peer rank " + std::to_string(rank) + " supports protocol versions " + std::to_string(peer.minVersion) + + "-" + std::to_string(peer.maxVersion) + ", but this binary supports " + + std::to_string(localDescriptor.minVersion) + "-" + std::to_string(localDescriptor.maxVersion) + + " and no version is common to every peer rank"}; + } + if ((peer.capabilities & detail::kInflightCancellationCapability) == 0) + { + return {false, hasLegacyPeer, allPeersProtocolAware, allPeersCanRejectBeforeDma, std::nullopt, + "peer rank " + std::to_string(rank) + " lacks the in-flight cancellation capability"}; + } + if ((peer.capabilities & detail::kPreDmaProtocolRejectionCapability) == 0) + { + return {false, hasLegacyPeer, allPeersProtocolAware, false, std::nullopt, + "peer rank " + std::to_string(rank) + " lacks the pre-DMA protocol-rejection capability"}; + } + } + + return {true, hasLegacyPeer, allPeersProtocolAware, allPeersCanRejectBeforeDma, + localMode == PeerCancellationMode::kEnabled ? std::make_optional(commonMaxVersion) : std::nullopt, + hasLegacyPeer ? "compatible legacy baseline peer" : "compatible protocol and effective mode"}; +} + +inline PeerProtocolCompatibility validatePeerProtocol( + PeerCancellationMode const localMode, std::vector const& peerAgentNames) +{ + auto localDescriptor = makePeerProtocolDescriptor(localMode == PeerCancellationMode::kEnabled, /*sessionToken=*/1); + return validatePeerProtocol(localDescriptor, peerAgentNames); +} + +//! Require one exact protocol profile across the ranks of a service instance. Session tokens intentionally differ. +inline PeerProtocolCompatibility validateLocalPeerProtocol( + PeerCancellationMode const localMode, std::vector const& localAgentNames) +{ + auto compatibility = validatePeerProtocol(localMode, localAgentNames); + if (!compatibility.compatible) + { + return compatibility; + } + + std::optional reference; + for (size_t rank = 0; rank < localAgentNames.size(); ++rank) + { + auto const parsed = parsePeerProtocolDescriptor(localAgentNames[rank]); + if (parsed.status != PeerProtocolParseStatus::kValid) + { + return {false, parsed.status == PeerProtocolParseStatus::kMissing, false, false, std::nullopt, + "local rank " + std::to_string(rank) + " does not advertise a valid protocol descriptor"}; + } + if (!reference.has_value()) + { + reference = parsed.descriptor; + continue; + } + auto const& descriptor = parsed.descriptor.value(); + if (descriptor.minVersion != reference->minVersion || descriptor.maxVersion != reference->maxVersion + || descriptor.capabilities != reference->capabilities + || descriptor.cancellationMode != reference->cancellationMode) + { + return {false, false, true, compatibility.allPeersCanRejectBeforeDma, std::nullopt, + "local rank " + std::to_string(rank) + + " advertises a different protocol version, capability set, or effective mode"}; + } + } + return compatibility; +} + +} // namespace tensorrt_llm::executor::kv_cache diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index b6b2672b2e0d..ab528279bab9 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -18,6 +18,7 @@ add_gtest(blockKeyTest blockKeyTest.cpp) add_gtest(radixBlockTreeTest radixBlockTreeTest.cpp) add_gtest(cacheTransBufferTest cacheTransBufferTest.cpp) add_gtest(cacheTransferStateTest cacheTransferStateTest.cpp) +add_gtest(peerProtocolTest peerProtocolTest.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 ff956caa8c73..96ebc3aa7ce6 100644 --- a/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp @@ -337,6 +337,8 @@ TEST_P(BufferIndexHolderLifecycleTest, TransferSessionReleasesRecvSlotOnSuccess) std::vector connections{nullptr}; TransferSession session(std::move(connections), DataContext{0}, std::vector{0}, selfState, std::move(otherState), bufferManager, 0, BlockKey{}, nullptr, false, false); + session.setPeerProtocolVersion(1); + EXPECT_EQ(session.getPeerProtocolVersion(), 1); std::vector holders; holders.emplace_back(mgr(), idx, /*isRecv=*/true); session.setRecvBufferHolders(std::move(holders)); @@ -369,6 +371,62 @@ TEST_P(BufferIndexHolderLifecycleTest, ActiveTransferSessionPoisonsUnreleasedRec EXPECT_TRUE(mgr().hasPoisonedBuffer()); } +TEST_P(BufferIndexHolderLifecycleTest, RejectedBaselineSessionPoisonsOnExceptionUnwind) +{ + if (!isRecv()) + { + GTEST_SKIP() << "TransferSession owns receive slots only"; + } + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + try + { + executor::DataTransceiverState selfState; + executor::DataTransceiverState otherState; + runtime::BufferManager bufferManager{std::make_shared()}; + std::vector connections{nullptr}; + TransferSession session(std::move(connections), DataContext{0}, std::vector{0}, selfState, + std::move(otherState), bufferManager, 0, BlockKey{}, nullptr, false, false); + session.rejectPeerProtocol(); + std::vector holders; + holders.emplace_back(mgr(), idx, /*isRecv=*/true); + session.setRecvBufferHolders(std::move(holders)); + throw std::runtime_error("ready wait failed after advertisement"); + } + catch (std::runtime_error const&) + { + } + + EXPECT_TRUE(mgr().hasPoisonedBuffer()); +} + +TEST_P(BufferIndexHolderLifecycleTest, RejectedBaselineSessionReleasesAfterAllPeersReject) +{ + if (!isRecv()) + { + GTEST_SKIP() << "TransferSession owns receive slots only"; + } + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + { + executor::DataTransceiverState selfState; + executor::DataTransceiverState otherState; + runtime::BufferManager bufferManager{std::make_shared()}; + std::vector connections{nullptr}; + TransferSession session(std::move(connections), DataContext{0}, std::vector{0}, selfState, + std::move(otherState), bufferManager, 0, BlockKey{}, nullptr, false, false); + session.rejectPeerProtocol(); + std::vector holders; + holders.emplace_back(mgr(), idx, /*isRecv=*/true); + session.setRecvBufferHolders(std::move(holders)); + session.releaseRecvBufferHolders(); + } + + EXPECT_EQ(inUse(), before); + EXPECT_FALSE(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/peerProtocolTest.cpp b/cpp/tests/unit_tests/batch_manager/peerProtocolTest.cpp new file mode 100644 index 000000000000..819e723b47e3 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/peerProtocolTest.cpp @@ -0,0 +1,267 @@ +/* + * 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/executor/cache_transmission/agent_utils/peerProtocol.h" + +#include + +#include +#include +#include + +namespace tensorrt_llm::executor::kv_cache +{ +namespace +{ + +constexpr std::uint64_t kSessionToken{0x123456789abcdef0ULL}; + +std::string makeAgentName(bool const enabled, std::uint64_t const sessionToken = kSessionToken) +{ + return appendPeerProtocolDescriptor("opaque_nixl_agent_identity", + makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/enabled, sessionToken)); +} + +TEST(PeerProtocolTest, DescriptorRoundTripsWithoutChangingIdentityPrefix) +{ + auto const descriptor = makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/true, kSessionToken); + auto const agentName = appendPeerProtocolDescriptor("opaque_nixl_agent_identity", descriptor); + EXPECT_EQ(agentName.find("opaque_nixl_agent_identity"), 0); + + auto const parsed = parsePeerProtocolDescriptor(agentName); + ASSERT_EQ(parsed.status, PeerProtocolParseStatus::kValid); + ASSERT_TRUE(parsed.descriptor.has_value()); + EXPECT_EQ(parsed.descriptor.value(), descriptor); +} + +TEST(PeerProtocolTest, MissingDescriptorIsAllowedOnlyForBaselineMode) +{ + auto const baseline = validatePeerProtocol(PeerCancellationMode::kBaseline, {"legacy_agent"}); + EXPECT_TRUE(baseline.compatible); + EXPECT_TRUE(baseline.hasLegacyPeer); + EXPECT_FALSE(baseline.allPeersProtocolAware); + EXPECT_FALSE(baseline.allPeersCanRejectBeforeDma); + + auto const enabled = validatePeerProtocol(PeerCancellationMode::kEnabled, {"legacy_agent"}); + EXPECT_FALSE(enabled.compatible); + EXPECT_TRUE(enabled.hasLegacyPeer); + EXPECT_FALSE(enabled.allPeersProtocolAware); + EXPECT_FALSE(enabled.allPeersCanRejectBeforeDma); + EXPECT_NE(enabled.reason.find("does not advertise"), std::string::npos); +} + +TEST(PeerProtocolTest, MatchingEnabledPeersAreCompatible) +{ + auto const compatibility = validatePeerProtocol( + PeerCancellationMode::kEnabled, {makeAgentName(/*enabled=*/true), makeAgentName(true, kSessionToken + 1)}); + EXPECT_TRUE(compatibility.compatible); + EXPECT_FALSE(compatibility.hasLegacyPeer); + EXPECT_TRUE(compatibility.allPeersProtocolAware); + EXPECT_TRUE(compatibility.allPeersCanRejectBeforeDma); + EXPECT_EQ(compatibility.selectedVersion, 1); +} + +TEST(PeerProtocolTest, BaselineRolloutAssumesDescriptorlessPeersHaveCancellationOff) +{ + auto const compatibility + = validatePeerProtocol(PeerCancellationMode::kBaseline, {makeAgentName(/*enabled=*/false), "legacy_agent"}); + EXPECT_TRUE(compatibility.compatible); + EXPECT_TRUE(compatibility.hasLegacyPeer); + EXPECT_FALSE(compatibility.allPeersProtocolAware); + EXPECT_FALSE(compatibility.allPeersCanRejectBeforeDma); + EXPECT_FALSE(compatibility.selectedVersion.has_value()); +} + +TEST(PeerProtocolTest, EffectiveModeMismatchIsRejected) +{ + auto const enabledLocal = validatePeerProtocol(PeerCancellationMode::kEnabled, {makeAgentName(/*enabled=*/false)}); + EXPECT_FALSE(enabledLocal.compatible); + EXPECT_TRUE(enabledLocal.allPeersProtocolAware); + EXPECT_TRUE(enabledLocal.allPeersCanRejectBeforeDma); + EXPECT_NE(enabledLocal.reason.find("mode mismatch"), std::string::npos); + + auto const baselineLocal = validatePeerProtocol(PeerCancellationMode::kBaseline, {makeAgentName(/*enabled=*/true)}); + EXPECT_FALSE(baselineLocal.compatible); + EXPECT_TRUE(baselineLocal.allPeersProtocolAware); + EXPECT_TRUE(baselineLocal.allPeersCanRejectBeforeDma); + EXPECT_NE(baselineLocal.reason.find("mode mismatch"), std::string::npos); +} + +TEST(PeerProtocolTest, NonOverlappingVersionIsRejected) +{ + auto descriptor = makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/true, kSessionToken); + descriptor.minVersion = 2; + descriptor.maxVersion = 3; + auto const agentName = appendPeerProtocolDescriptor("future_agent", descriptor); + + auto const compatibility = validatePeerProtocol(PeerCancellationMode::kEnabled, {agentName}); + EXPECT_FALSE(compatibility.compatible); + EXPECT_NE(compatibility.reason.find("supports protocol versions 2-3"), std::string::npos); +} + +TEST(PeerProtocolTest, BaselinePeersIgnoreCancellationProtocolVersion) +{ + auto descriptor = makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/false, kSessionToken); + descriptor.minVersion = 2; + descriptor.maxVersion = 3; + auto const agentName = appendPeerProtocolDescriptor("future_baseline_agent", descriptor); + + auto const compatibility = validatePeerProtocol(PeerCancellationMode::kBaseline, {agentName}); + EXPECT_TRUE(compatibility.compatible); + EXPECT_FALSE(compatibility.selectedVersion.has_value()); +} + +TEST(PeerProtocolTest, EnabledPeerMustAdvertiseCancellationCapability) +{ + auto descriptor = makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/true, kSessionToken); + descriptor.capabilities = static_cast(PeerProtocolCapability::kPreDmaProtocolRejection); + auto const agentName = appendPeerProtocolDescriptor("incapable_agent", descriptor); + + auto const compatibility = validatePeerProtocol(PeerCancellationMode::kEnabled, {agentName}); + EXPECT_FALSE(compatibility.compatible); + EXPECT_NE(compatibility.reason.find("lacks the in-flight cancellation capability"), std::string::npos); +} + +TEST(PeerProtocolTest, MalformedAdvertisedDescriptorIsRejectedEvenInBaselineMode) +{ + auto const compatibility = validatePeerProtocol( + PeerCancellationMode::kBaseline, {"agent__trtllm_protocol_v1_1_not-a-version__trtllm_protocol_end"}); + EXPECT_FALSE(compatibility.compatible); + EXPECT_NE(compatibility.reason.find("malformed"), std::string::npos); +} + +TEST(PeerProtocolTest, EmbeddedHostnameMarkerDoesNotForgeLegacyIdentity) +{ + auto const parsed + = parsePeerProtocolDescriptor("legacy__trtllm_protocol_v1_1_1_3_1_1__trtllm_protocol_end_123_456_0"); + EXPECT_EQ(parsed.status, PeerProtocolParseStatus::kMissing); +} + +TEST(PeerProtocolTest, ZeroSessionTokenIsRejected) +{ + auto const agentName = makeAgentName(/*enabled=*/false, /*sessionToken=*/0); + auto const parsed = parsePeerProtocolDescriptor(agentName); + EXPECT_EQ(parsed.status, PeerProtocolParseStatus::kMalformed); + EXPECT_NE(parsed.reason.find("session token"), std::string::npos); +} + +TEST(PeerProtocolTest, V1EnvelopeRejectsUnversionedFutureFields) +{ + auto const suffixPos = makeAgentName(/*enabled=*/true).find("__trtllm_protocol_end"); + ASSERT_NE(suffixPos, std::string::npos); + auto agentName = makeAgentName(/*enabled=*/true); + agentName.insert(suffixPos, "_future_optional_field"); + auto const parsed = parsePeerProtocolDescriptor(agentName); + EXPECT_EQ(parsed.status, PeerProtocolParseStatus::kMalformed); +} + +TEST(PeerProtocolTest, MixedProtocolAwareAndLegacyPeersCannotUseReadyRejection) +{ + auto const compatibility + = validatePeerProtocol(PeerCancellationMode::kBaseline, {makeAgentName(/*enabled=*/true), "legacy_agent"}); + EXPECT_FALSE(compatibility.compatible); + EXPECT_TRUE(compatibility.hasLegacyPeer); + EXPECT_FALSE(compatibility.allPeersProtocolAware); + EXPECT_FALSE(compatibility.allPeersCanRejectBeforeDma); +} + +TEST(PeerProtocolTest, MismatchRequiresEveryPeerToAdvertisePreDmaRejection) +{ + auto descriptor = makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/false, kSessionToken); + descriptor.capabilities &= ~static_cast(PeerProtocolCapability::kPreDmaProtocolRejection); + auto const compatibility = validatePeerProtocol( + PeerCancellationMode::kEnabled, {appendPeerProtocolDescriptor("old_protocol_peer", descriptor)}); + EXPECT_FALSE(compatibility.compatible); + EXPECT_TRUE(compatibility.allPeersProtocolAware); + EXPECT_FALSE(compatibility.allPeersCanRejectBeforeDma); +} + +TEST(PeerProtocolTest, EnabledPeerMustSupportPreDmaProtocolRejection) +{ + auto descriptor = makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/true, kSessionToken); + descriptor.capabilities = static_cast(PeerProtocolCapability::kInflightCancellation); + auto const compatibility = validatePeerProtocol( + PeerCancellationMode::kEnabled, {appendPeerProtocolDescriptor("old_protocol_peer", descriptor)}); + EXPECT_FALSE(compatibility.compatible); + EXPECT_FALSE(compatibility.allPeersCanRejectBeforeDma); + EXPECT_NE(compatibility.reason.find("pre-DMA protocol-rejection"), std::string::npos); +} + +TEST(PeerProtocolTest, EnabledLocalDescriptorMustAdvertiseRequiredCapabilities) +{ + auto local = makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/true, kSessionToken); + auto const peer = makeAgentName(/*enabled=*/true, kSessionToken + 1); + + local.capabilities = static_cast(PeerProtocolCapability::kPreDmaProtocolRejection); + auto const missingCancellation = validatePeerProtocol(local, {peer}); + EXPECT_FALSE(missingCancellation.compatible); + EXPECT_NE(missingCancellation.reason.find("local protocol descriptor lacks the in-flight"), std::string::npos); + + local.capabilities = static_cast(PeerProtocolCapability::kInflightCancellation); + auto const missingRejection = validatePeerProtocol(local, {peer}); + EXPECT_FALSE(missingRejection.compatible); + EXPECT_NE(missingRejection.reason.find("local protocol descriptor lacks the pre-DMA"), std::string::npos); +} + +TEST(PeerProtocolTest, SelectsOneVersionCommonToEveryPeerRank) +{ + auto local = makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/true, kSessionToken); + local.minVersion = 1; + local.maxVersion = 2; + + auto peer0 = makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/true, kSessionToken + 1); + peer0.minVersion = 1; + peer0.maxVersion = 2; + auto peer1 = makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/true, kSessionToken + 2); + peer1.minVersion = 2; + peer1.maxVersion = 3; + auto const compatible = validatePeerProtocol( + local, {appendPeerProtocolDescriptor("peer0", peer0), appendPeerProtocolDescriptor("peer1", peer1)}); + EXPECT_TRUE(compatible.compatible); + EXPECT_EQ(compatible.selectedVersion, 2); + + peer0.minVersion = 1; + peer0.maxVersion = 1; + auto const noCommonVersion = validatePeerProtocol( + local, {appendPeerProtocolDescriptor("peer0", peer0), appendPeerProtocolDescriptor("peer1", peer1)}); + EXPECT_FALSE(noCommonVersion.compatible); + EXPECT_FALSE(noCommonVersion.selectedVersion.has_value()); + EXPECT_NE(noCommonVersion.reason.find("no version is common"), std::string::npos); +} + +TEST(PeerProtocolTest, LocalRanksMustAdvertiseOneExactProtocolProfile) +{ + auto rank0 = makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/true, kSessionToken); + auto rank1 = makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/true, kSessionToken + 1); + auto const matching = validateLocalPeerProtocol(PeerCancellationMode::kEnabled, + {appendPeerProtocolDescriptor("rank0", rank0), appendPeerProtocolDescriptor("rank1", rank1)}); + EXPECT_TRUE(matching.compatible); + + rank1.maxVersion = 2; + auto const mixedVersions = validateLocalPeerProtocol(PeerCancellationMode::kEnabled, + {appendPeerProtocolDescriptor("rank0", rank0), appendPeerProtocolDescriptor("rank1", rank1)}); + EXPECT_FALSE(mixedVersions.compatible); + EXPECT_NE(mixedVersions.reason.find("different protocol version"), std::string::npos); + + auto const legacyRank + = validateLocalPeerProtocol(PeerCancellationMode::kBaseline, {makeAgentName(/*enabled=*/false), "legacy_rank"}); + EXPECT_FALSE(legacyRank.compatible); + EXPECT_NE(legacyRank.reason.find("does not advertise"), std::string::npos); +} + +} // namespace +} // namespace tensorrt_llm::executor::kv_cache diff --git a/cpp/tests/unit_tests/executor/agentCommTest.cpp b/cpp/tests/unit_tests/executor/agentCommTest.cpp index 45acb130ea54..091b30be8d9b 100644 --- a/cpp/tests/unit_tests/executor/agentCommTest.cpp +++ b/cpp/tests/unit_tests/executor/agentCommTest.cpp @@ -19,6 +19,7 @@ #include #include +#include using namespace tensorrt_llm::batch_manager::kv_cache_manager; using namespace tensorrt_llm::runtime; @@ -157,7 +158,11 @@ class AgentCommTest : public ::testing::TestWithParam TEST_P(AgentCommTest, AgentConnectionManagerBasic) { std::vector bufferManagers{mTransBufferManager.get()}; - auto connectionManager = std::make_unique(bufferManagers, *mCacheState, backend); + std::optional const peerMode = backend == "nixl" + ? std::make_optional(PeerCancellationMode::kBaseline) + : std::optional{}; + auto connectionManager + = std::make_unique(bufferManagers, *mCacheState, backend, std::nullopt, peerMode); ASSERT_TRUE(connectionManager != nullptr); ASSERT_EQ(connectionManager->getCacheTransBufferManagers().size(), bufferManagers.size()); ASSERT_TRUE(connectionManager->getCacheTransBufferManagers().front() != nullptr); @@ -167,18 +172,45 @@ TEST_P(AgentCommTest, AgentConnectionManagerBasic) CommState commState = connectionManager->getCommState(); ASSERT_TRUE(commState.isAgentState()); ASSERT_EQ(commState.getAgentState().size(), 1); + auto const parsed = parsePeerProtocolDescriptor(connectionManager->getAgentName()); + if (backend == "nixl") + { + ASSERT_EQ(parsed.status, PeerProtocolParseStatus::kValid); + ASSERT_TRUE(parsed.descriptor.has_value()); + EXPECT_EQ(parsed.descriptor->cancellationMode, PeerCancellationMode::kBaseline); + EXPECT_NE(parsed.descriptor->sessionToken, 0); + } + else + { + EXPECT_EQ(parsed.status, PeerProtocolParseStatus::kMissing); + } } TEST_P(AgentCommTest, AgentConnectionManagerConnect) { std::vector bufferManagers{mTransBufferManager.get()}; - auto connectionManager0 = std::make_unique(bufferManagers, *mCacheState, backend); - auto connectionManager1 = std::make_unique(bufferManagers, *mCacheState, backend); + std::optional const peerMode = backend == "nixl" + ? std::make_optional(PeerCancellationMode::kBaseline) + : std::optional{}; + auto connectionManager0 + = std::make_unique(bufferManagers, *mCacheState, backend, std::nullopt, peerMode); + auto connectionManager1 + = std::make_unique(bufferManagers, *mCacheState, backend, std::nullopt, peerMode); auto agentName0 = connectionManager0->getAgentName(); auto agentName1 = connectionManager1->getAgentName(); ASSERT_TRUE(!agentName0.empty()); ASSERT_TRUE(!agentName1.empty()); ASSERT_TRUE(agentName0 != agentName1); + if (backend == "nixl") + { + auto const descriptor0 = parsePeerProtocolDescriptor(agentName0); + auto const descriptor1 = parsePeerProtocolDescriptor(agentName1); + ASSERT_EQ(descriptor0.status, PeerProtocolParseStatus::kValid); + ASSERT_EQ(descriptor1.status, PeerProtocolParseStatus::kValid); + ASSERT_TRUE(descriptor0.descriptor.has_value()); + ASSERT_TRUE(descriptor1.descriptor.has_value()); + EXPECT_NE(descriptor0.descriptor->sessionToken, descriptor1.descriptor->sessionToken); + } auto commState0 = connectionManager0->getCommState(); auto commState1 = connectionManager1->getCommState(); diff --git a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp index a39756bf7243..f05fbb602b9f 100644 --- a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp +++ b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp @@ -770,6 +770,25 @@ TEST(SerializeUtilsTest, ContextPhaseParams) EXPECT_EQ(state2, stateCopy); } + // The opaque NIXL identity is the CTX-to-GEN protocol carrier and must survive ContextPhaseParams serialization. + { + auto const descriptor + = texec::kv_cache::makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/true, 0x1234); + auto const agentName = texec::kv_cache::appendPeerProtocolDescriptor("agent", descriptor); + auto state = std::make_unique(); + state->setCommState(texec::kv_cache::CommState{{texec::kv_cache::AgentState{agentName, "127.0.0.1:1234"}}, 0}); + auto contextPhaseParams = texec::ContextPhaseParams({10, 20}, 1, state.release(), VecTokens{10, 20}); + + auto serializedState = contextPhaseParams.getSerializedState(); + auto const stateCopy = texec::Serialization::deserializeDataTransceiverState(serializedState); + ASSERT_TRUE(stateCopy.getCommState().has_value()); + auto const& agentStates = stateCopy.getCommState()->getAgentState(); + ASSERT_EQ(agentStates.size(), 1); + auto const parsed = texec::kv_cache::parsePeerProtocolDescriptor(agentStates.front().mAgentName); + ASSERT_EQ(parsed.status, texec::kv_cache::PeerProtocolParseStatus::kValid); + EXPECT_EQ(parsed.descriptor, descriptor); + } + // Test with ctxDpRank and disaggInfoEndpoint { auto state = std::make_unique(); @@ -1347,6 +1366,51 @@ T serializeDeserializeNotification(T const& val) return T::deserialize(iss); } +tensorrt_llm::batch_manager::RequestInfo makeAgentRequestInfo( + std::vector agentStates, int const selfIdx) +{ + texec::DataTransceiverState state; + state.setCommState(kv_cache::CommState{std::move(agentStates), selfIdx}); + return tensorrt_llm::batch_manager::RequestInfo{42, std::move(state)}; +} + +TEST(SerializeUtilsTest, PeerProtocolIdentitySurvivesRequestNotificationSerialization) +{ + auto const descriptor = kv_cache::makePeerProtocolDescriptor(/*inflightCancellationEnabled=*/true, 0x1234); + auto const agentName = kv_cache::appendPeerProtocolDescriptor("sender", descriptor); + std::string const address{"127.0.0.1:8080"}; + auto const requestInfo + = makeAgentRequestInfo({kv_cache::AgentState{agentName, address}, kv_cache::AgentState{"peer", "address"}}, 0); + kv_cache::RequestAndBufferInfo original{agentName, address, requestInfo, + std::vector{kv_cache::MemoryDesc{nullptr, 1024, 0}}, std::nullopt, 0, {0}}; + + auto const deserialized = serializeDeserializeNotification(original); + EXPECT_EQ(deserialized.mRequestInfo.getTransState(), requestInfo.getTransState()); + EXPECT_NO_THROW(kv_cache::validateRequestPeerIdentity(deserialized.mRequestInfo, agentName, deserialized.mAddress)); + auto const& roundTripName = deserialized.mRequestInfo.getTransState() + .getCommState() + ->getAgentState() + .at(deserialized.mRequestInfo.getTransState().getCommState()->getSelfIdx()) + .mAgentName; + auto const parsed = kv_cache::parsePeerProtocolDescriptor(roundTripName); + ASSERT_EQ(parsed.status, kv_cache::PeerProtocolParseStatus::kValid); + EXPECT_EQ(parsed.descriptor, descriptor); +} + +TEST(SerializeUtilsTest, RequestNotificationIdentityMustMatchAdvertisedSelf) +{ + std::string const agentName{"sender"}; + std::string const address{"127.0.0.1:8080"}; + auto const wrongName = makeAgentRequestInfo({kv_cache::AgentState{"other", address}}, 0); + EXPECT_ANY_THROW(kv_cache::validateRequestPeerIdentity(wrongName, agentName, address)); + + auto const wrongAddress = makeAgentRequestInfo({kv_cache::AgentState{agentName, "127.0.0.1:9090"}}, 0); + EXPECT_ANY_THROW(kv_cache::validateRequestPeerIdentity(wrongAddress, agentName, address)); + + auto const wrongSelfIndex = makeAgentRequestInfo({kv_cache::AgentState{agentName, address}}, 1); + EXPECT_ANY_THROW(kv_cache::validateRequestPeerIdentity(wrongSelfIndex, agentName, address)); +} + TEST(SerializeUtilsTest, RequestAndBufferInfo) { // Test with all fields populated including bufferKinds diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index 77f8ea03fb83..51a04f75d449 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -42,6 +42,8 @@ #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/testing/kvCacheManagerTestUtil.h" +#include +#include #include #include #include @@ -735,6 +737,9 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam baseBufferManagers( bufferManagers.begin(), bufferManagers.end()); mConnectionManager = std::make_unique( - baseBufferManagers, *mCacheState, "nixl"); + baseBufferManagers, *mCacheState, "nixl", std::nullopt, peerCancellationMode); } else if (isMooncake) { @@ -792,13 +797,13 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam( - mConnectionManager.get(), mRankInInstance, CacheTransferLayer(*mCacheState, makeFormatter())); + mSender = std::make_unique(mConnectionManager.get(), mRankInInstance, + CacheTransferLayer(*mCacheState, makeFormatter(), nullptr, enableInflightCancel)); } else { - mRequester = std::make_unique( - mConnectionManager.get(), mRankInInstance, CacheTransferLayer(*mCacheState, makeFormatter())); + mRequester = std::make_unique(mConnectionManager.get(), mRankInInstance, + CacheTransferLayer(*mCacheState, makeFormatter(), nullptr, enableInflightCancel)); } TLLM_LOG_DEBUG("setUpCacheTransceiver mSender"); @@ -1270,6 +1275,108 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam request = makeLlmRequest(/*length=*/10); + bool transferSettled = false; + std::future transferFuture; + + if (mIsContext) + { + transferFuture = addRequestAndTransportCacheForContext(request); + mComm->barrier(); + auto const status = transferFuture.wait_for(std::chrono::seconds(30)); + EXPECT_EQ(status, std::future_status::ready); + transferSettled = status == std::future_status::ready; + if (transferSettled) + { + EXPECT_ANY_THROW(transferFuture.get()); + EXPECT_FALSE(mCacheTransBufferManagers.front()->hasPoisonedBuffer()); + } + } + else + { + constexpr std::uint8_t kRecvSentinel{0xA5}; + auto const recvBuffer = mCacheTransBufferManagers.front()->getRecvBuffer(0); + EXPECT_NE(recvBuffer, nullptr); + if (recvBuffer != nullptr) + { + TLLM_CUDA_CHECK(cudaMemset(recvBuffer->data(), kRecvSentinel, recvBuffer->getSizeInBytes())); + TLLM_CUDA_CHECK(cudaDeviceSynchronize()); + } + + mComm->barrier(); + transferFuture = addRequestAndTransportCacheForGeneration(request); + auto const status = transferFuture.wait_for(std::chrono::seconds(30)); + EXPECT_EQ(status, std::future_status::ready); + transferSettled = status == std::future_status::ready; + if (transferSettled && generationMode == texec::kv_cache::PeerCancellationMode::kEnabled) + { + EXPECT_ANY_THROW(transferFuture.get()); + } + else if (transferSettled) + { + EXPECT_NO_THROW(transferFuture.get()); + EXPECT_EQ(request->mLlmRequest->getState(), LlmRequestState::kDISAGG_TRANS_ERROR); + } + + if (transferSettled && recvBuffer != nullptr) + { + std::vector sentinelPrefix(std::min(recvBuffer->getSizeInBytes(), 256)); + TLLM_CUDA_CHECK(cudaMemcpy( + sentinelPrefix.data(), recvBuffer->data(), sentinelPrefix.size(), cudaMemcpyDeviceToHost)); + EXPECT_TRUE(std::all_of(sentinelPrefix.begin(), sentinelPrefix.end(), + [](std::uint8_t const value) { return value == kRecvSentinel; })); + } + + if (transferSettled) + { + auto& recvManager = *mCacheTransBufferManagers.front(); + EXPECT_FALSE(recvManager.hasPoisonedBuffer()); + auto const reacquired = recvManager.assignBufferIndexForRecv(); + EXPECT_TRUE(reacquired.has_value()); + if (reacquired.has_value()) + { + recvManager.freeBufferIndexForRecv(reacquired); + } + } + } + + if (transferSettled) + { + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*request->mLlmRequest); + mManager->removeSequence(request->mLlmRequest->mRequestId, request->mLlmRequest); + } + if (mIsContext) + { + mSender.reset(); + } + else + { + mRequester.reset(); + } + mComm->barrier(); + mConnectionManager.reset(); + } + tensorrt_llm::mpi::MpiComm::world().barrier(); + } + bool mIsContext{false}; bool mIsGeneration{false}; tensorrt_llm::mpi::MpiComm const* mComm; @@ -1282,6 +1389,9 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam mAttentionLayerNumPerPP; @@ -1297,6 +1407,22 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam