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 68a398f1b52d..881dc20119f0 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -522,7 +522,8 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio // cache blocks to the corresponding buffer. // 5. send the buffer to the corresponding target. Ideally, we send only once (one buffer) for each target. - auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend(); + auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); + auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend(sendCancelFlag); BufferIndexHolder sendHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/false); int peerDuplicateHeadFactor = targetInfo.mPeerDupHeadFactor; auto bufferTargetNum = targetNum / peerDuplicateHeadFactor; @@ -602,8 +603,19 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio == inputKvCacheBlocksPerWindow.begin()->second.front()->getDataType()); } - sendAllBuffers(session, deviceId, outputSplitCaches, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, - targetInfo, pickUpConnections); + try + { + sendAllBuffers(session, deviceId, outputSplitCaches, bufferCoverTargetNum, preAllocSendBuffer, + bufferManager, targetInfo, pickUpConnections); + } + catch (...) + { + if (agentConnection != nullptr && session.isInflightCancelEnabled()) + { + sendHolder.poison(); + } + throw; + } session.setTime(TransferSession::kTimeTransmissions); @@ -865,8 +877,8 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess else { cacheBufferId = mCacheTransBufferManager->assignBufferIndexForRecv(); + recvHolder = BufferIndexHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); } - recvHolder = BufferIndexHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); auto [recvSplitCachestmp, bufferCoverTargetNumtmp, onlyUseDynamicBuffer] = mCacheTransBufferManager->getOrAllocateRecvBuffers( cacheBufferId, static_cast(targetNum), bufferEleSizes, bufferManager); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 2c69af0059be..a13be408d46f 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( @@ -466,8 +625,10 @@ 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) @@ -515,8 +676,8 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa return nullptr; }; - auto makeCacheTransferLayer - = [&]() { return CacheTransferLayer(*mCacheState, makeFormatter(), makeRnnFormatter()); }; + auto makeCacheTransferLayer = [&]() + { return CacheTransferLayer(*mCacheState, makeFormatter(), makeRnnFormatter(), mInflightCancelEnabled); }; mCacheSender = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); mCacheReceiver = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); @@ -599,6 +760,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 +885,271 @@ void updateKVCacheTransferBW(std::shared_ptr const& mComm, } } +RequestStatuses CacheTransceiver::checkContextTransferStatusWithInflightCancel( + std::optional const& atLeastRequestNum, bool const markComplete) +{ + TLLM_CHECK_WITH_INFO(atLeastRequestNum.has_value(), + "In-flight cancellation requires a finite context-transfer status poll; pass 0 for a nonblocking poll."); + + auto const needsProgress = atLeastRequestNum.value() > 0; + auto const futureWaitInterval + = getTransferFutureWaitInterval(mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(), needsProgress); + auto const kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs().value(); + auto const syncComm = mGroupTensorParaComm; + + auto countReadyFutures = [this]() + { + return static_cast(std::count_if(mSenderFutures.begin(), mSenderFutures.end(), + [](auto& entry) + { return entry.second.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready; })); + }; + if (needsProgress) + { + auto const deadline = std::chrono::steady_clock::now() + futureWaitInterval; + while (countReadyFutures() < atLeastRequestNum.value() && std::chrono::steady_clock::now() < deadline) + { + auto const remaining + = std::chrono::duration_cast(deadline - std::chrono::steady_clock::now()); + std::this_thread::sleep_for(std::min(std::chrono::milliseconds(kTransferFuturePollIntervalMs), remaining)); + } + } + + for (auto const& [request, future] : mSenderFutures) + { + if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + { + continue; + } + auto const elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + if (elapsed.count() > kvTransferTimeoutMs && mTimedOutSenderIds.insert(request->mRequestId).second) + { + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld exceeded the local timeout; requesting topology-wide " + "cancellation.", + request->mRequestId); + } + } + + for (auto it = mSenderFutures.begin(); it != mSenderFutures.end();) + { + auto& [request, future] = *it; + if (future.wait_for(std::chrono::milliseconds(0)) != std::future_status::ready) + { + ++it; + continue; + } + auto const requestId = request->mRequestId; + try + { + future.get(); + auto const failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR; + recordLocalTransferOutcome(requestId, request, failed, mCompletedSenderRequestIds, mFailedSenderRequestIds, + mSenderRequestsAwaitingConsensus); + } + catch (std::exception const& error) + { + TLLM_LOG_ERROR("Error occurred during context transfer for request %ld: %s", requestId, error.what()); + recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedSenderRequestIds, + mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); + } + it = mSenderFutures.erase(it); + } + + auto const localPoisoned = mTopologyTransferBufferPoisoned.load(std::memory_order_relaxed) + || std::any_of(mCacheTransBufferManagerPtrs.begin(), mCacheTransBufferManagerPtrs.end(), + [](BaseTransBufferManager const* manager) { return manager != nullptr && manager->hasPoisonedBuffer(); }); + auto const outcome = exchangeTransferStates( + syncComm, mCompletedSenderRequestIds, mFailedSenderRequestIds, mTimedOutSenderIds, localPoisoned); + if (outcome.transferBufferPoisoned) + { + mTopologyTransferBufferPoisoned.store(true, std::memory_order_relaxed); + } + + for (auto const requestId : outcome.failedRequestIds) + { + auto const futureIt = std::find_if(mSenderFutures.begin(), mSenderFutures.end(), + [requestId](auto const& entry) { return entry.first->mRequestId == requestId; }); + if (futureIt == mSenderFutures.end() + || futureIt->second.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + { + continue; + } + if (outcome.timedOutRequestIds.find(requestId) != outcome.timedOutRequestIds.end() + && mTimedOutSenderIds.insert(requestId).second) + { + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld exceeded the topology-wide timeout; cancelling it on " + "every participating rank.", + requestId); + } + (void) mCacheSender->cancelRequest(*futureIt->first); + } + + RequestStatuses requestsStatus; + for (auto const requestId : outcome.quiescedRequestIds) + { + auto const requestIt = mSenderRequestsAwaitingConsensus.find(requestId); + if (requestIt == mSenderRequestsAwaitingConsensus.end()) + { + continue; + } + if (outcome.failedRequestIds.find(requestId) != outcome.failedRequestIds.end()) + { + requestIt->second->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + requestsStatus.errorRequestIds.insert(requestId); + } + else + { + requestsStatus.completedRequestIds.insert(requestId); + if (markComplete) + { + requestIt->second->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); + } + } + mTimedOutSenderIds.erase(requestId); + eraseLocalTransferOutcome( + requestId, mCompletedSenderRequestIds, mFailedSenderRequestIds, mSenderRequestsAwaitingConsensus); + } + return requestsStatus; +} + +void CacheTransceiver::checkGenTransferStatusWithInflightCancel(std::optional const& atLeastRequestNum) +{ + TLLM_CHECK_WITH_INFO(atLeastRequestNum.has_value(), + "In-flight cancellation requires a finite generation-transfer status poll; pass 0 for a nonblocking poll."); + + auto const needsProgress = atLeastRequestNum.value() > 0; + auto const futureWaitInterval + = getTransferFutureWaitInterval(mCacheTransceiverConfig->getKvTransferPollIntervalMs(), needsProgress); + auto const kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs().value(); + auto const syncComm = mGroupComm; + + auto countReadyFutures = [this]() + { + return static_cast(std::count_if(mRequesterFutures.begin(), mRequesterFutures.end(), + [](auto& entry) + { return entry.second.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready; })); + }; + if (needsProgress) + { + auto const deadline = std::chrono::steady_clock::now() + futureWaitInterval; + while (countReadyFutures() < atLeastRequestNum.value() && std::chrono::steady_clock::now() < deadline) + { + auto const remaining + = std::chrono::duration_cast(deadline - std::chrono::steady_clock::now()); + std::this_thread::sleep_for(std::min(std::chrono::milliseconds(kTransferFuturePollIntervalMs), remaining)); + } + } + + for (auto const& [request, future] : mRequesterFutures) + { + if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + { + continue; + } + auto const elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + if (elapsed.count() > kvTransferTimeoutMs && mTimedOutRequesterIds.insert(request->mRequestId).second) + { + TLLM_LOG_WARNING( + "Generation KV cache transfer for request %ld exceeded the local timeout; requesting topology-wide " + "cancellation.", + request->mRequestId); + } + } + + for (auto it = mRequesterFutures.begin(); it != mRequesterFutures.end();) + { + auto& [request, future] = *it; + if (future.wait_for(std::chrono::milliseconds(0)) != std::future_status::ready) + { + ++it; + continue; + } + auto const requestId = request->mRequestId; + try + { + future.get(); + auto const failed = request->getState() == LlmRequestState::kDISAGG_TRANS_ERROR; + if (failed) + { + request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); + } + recordLocalTransferOutcome(requestId, request, failed, mCompletedRequesterRequestIds, + mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); + } + catch (std::exception const& error) + { + TLLM_LOG_ERROR("Error occurred during generation transfer for request %ld: %s", requestId, error.what()); + recordLocalTransferOutcome(requestId, request, /*failed=*/true, mCompletedRequesterRequestIds, + mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); + } + it = mRequesterFutures.erase(it); + } + + auto const localPoisoned = mTopologyTransferBufferPoisoned.load(std::memory_order_relaxed) + || std::any_of(mCacheTransBufferManagerPtrs.begin(), mCacheTransBufferManagerPtrs.end(), + [](BaseTransBufferManager const* manager) { return manager != nullptr && manager->hasPoisonedBuffer(); }); + auto const outcome = exchangeTransferStates( + syncComm, mCompletedRequesterRequestIds, mFailedRequesterRequestIds, mTimedOutRequesterIds, localPoisoned); + if (outcome.transferBufferPoisoned) + { + mTopologyTransferBufferPoisoned.store(true, std::memory_order_relaxed); + } + + for (auto const requestId : outcome.failedRequestIds) + { + auto const futureIt = std::find_if(mRequesterFutures.begin(), mRequesterFutures.end(), + [requestId](auto const& entry) { return entry.first->mRequestId == requestId; }); + if (futureIt == mRequesterFutures.end() + || futureIt->second.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready) + { + continue; + } + if (outcome.timedOutRequestIds.find(requestId) != outcome.timedOutRequestIds.end() + && mTimedOutRequesterIds.insert(requestId).second) + { + TLLM_LOG_WARNING( + "Generation KV cache transfer for request %ld exceeded the topology-wide timeout; cancelling it on " + "every participating rank.", + requestId); + } + (void) mCacheReceiver->cancelRequest(*futureIt->first); + } + + for (auto const requestId : outcome.quiescedRequestIds) + { + auto const requestIt = mRequesterRequestsAwaitingConsensus.find(requestId); + if (requestIt == mRequesterRequestsAwaitingConsensus.end()) + { + continue; + } + if (outcome.failedRequestIds.find(requestId) != outcome.failedRequestIds.end()) + { + requestIt->second->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + } + else + { + requestIt->second->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); + if (!common::getEnvKVCacheTimeOutputPath().empty()) + { + updateKVCacheTransferBW(syncComm, requestIt->second.get()); + } + } + mTimedOutRequesterIds.erase(requestId); + eraseLocalTransferOutcome( + requestId, mCompletedRequesterRequestIds, mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); + } +} + RequestStatuses CacheTransceiver::checkContextTransferStatus( std::optional const& atLeastRequestNum, bool markComplete) { + if (mInflightCancelEnabled) + { + return checkContextTransferStatusWithInflightCancel(atLeastRequestNum, markComplete); + } bool const blockAll = !atLeastRequestNum.has_value(); std::optional senderFutureTimeoutMs = std::nullopt; if (mCacheTransceiverConfig.has_value()) @@ -892,6 +1318,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 +1535,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..a61a8651439a 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp @@ -25,15 +25,18 @@ #include "tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h" #include +#include 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); } @@ -43,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 " @@ -73,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 { @@ -143,4 +167,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..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" @@ -49,15 +50,19 @@ 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; 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. @@ -86,10 +91,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..b6ccb89bb95d --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/cacheTransferState.h @@ -0,0 +1,137 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/assert.h" + +#include +#include +#include +#include +#include + +namespace tensorrt_llm::batch_manager::detail +{ + +using TransferRequestId = std::uint64_t; + +enum TransferStateFlag : std::uint64_t +{ + kTransferCompleted = 1U << 0, + kTransferFailed = 1U << 1, + kTransferTimedOut = 1U << 2, +}; + +enum TransferGlobalFlag : std::uint64_t +{ + kTransferBufferPoisoned = 1U << 0, +}; + +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; + // Quiescence is an intersection: every rank has reached a local terminal state. + std::unordered_set quiescedRequestIds; + 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.quiescedRequestIds.insert(requestId); + if (!logicallyFailed && counts.completed == rankCount) + { + outcome.completedRequestIds.insert(requestId); + } + } + } + return outcome; +} + +} // namespace tensorrt_llm::batch_manager::detail diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 5aac98450be9..5f39b91ea167 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,40 @@ 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); + mPeerProtocolErrors.erase(requestId); + } + { + std::lock_guard lg(mInFlightCancelMutex); + mInFlightCancelFlags.erase(requestId); + } + } + + void discardTransferState(LlmRequest::RequestIdType requestId) noexcept + { + try + { + std::unique_lock lk(mMtxForMap); + mRequestToSession.erase(requestId); + mPeerProtocolErrors.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() @@ -381,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()); @@ -392,19 +443,32 @@ 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); + // 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; } @@ -426,16 +490,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 its topology can reach quiescence 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 +626,246 @@ 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) + std::optional peerProtocolError; { - mRemainSendCount.erase(reqId); + 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); + 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 && !peerProtocolError.has_value(); + } + + // 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); - // Check if the request is cancelled - bool isReady = true; + 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)); } } - mCurrentRequest = std::nullopt; + else + { + 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, + "%s for request %zu", errorMessage.c_str(), reqId))); + } + discardTransferState(reqId); + } } void response() noexcept { + std::exception_ptr responseException; try { tensorrt_llm::common::setThreadName("dataTransResp"); TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); - while (!mTerminate || !mAnyReady) + while (true) { - if (!mAnyReady) { - std::unique_lock lk(mCondMutex); - mSenderCv.wait(lk, [this]() { return (mAnyReady || mTerminate); }); + std::unique_lock lock(mSenderMutex); + mSenderCv.wait(lock, + [this]() { return mTerminate || !mReadyResponses.empty() || !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 + void failPendingResponses(std::exception_ptr const& exception) noexcept { - return mCurrentRequest.value(); - } - - [[nodiscard]] std::map::iterator getCurrentResponse() - { - std::scoped_lock lk(mSenderMutex); - return mReadyResponses.find(getCurrentRequestId()); + std::map pendingResponses; + { + std::scoped_lock lock(mSenderMutex); + pendingResponses.swap(mReadyResponses); + mCurrentRequest = std::nullopt; + mCancelledRequests.clear(); + mRemainSendCount.clear(); + } + for (auto& entry : pendingResponses) + { + failResponse(entry.second, exception); + } } public: @@ -749,11 +879,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; @@ -762,11 +894,15 @@ class CacheSender::Impl executor::kv_cache::ConnectionManager* mManager; std::map mRequestToSession; + std::unordered_map mPeerProtocolErrors; 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 +911,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 +926,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 +959,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,28 +979,74 @@ 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); } - session.exportMeasure(mMeasuresFile, false); + } + catch (...) + { + if (session.isInflightCancelEnabled()) + { + session.poisonRecvBufferHolders(); + } + else + { + session.releaseRecvBufferHolders(); + } + throw; } } - TransferSession sendRequestInfo(LlmRequest const& llmRequest) + 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 = {}, + 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); @@ -880,12 +1076,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()); } @@ -914,6 +1125,20 @@ 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); + // 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++) { auto rank = allCounterparts[ci]; @@ -966,18 +1191,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) @@ -1021,6 +1242,7 @@ class CacheReceiver::Impl bool isCancelled = false; auto& asyncResource = mInstanceToAsyncResource.at(processInfo); + std::optional queuedCancelledReqId; { std::unique_lock lck(asyncResource->mMtxForQueue); auto it = std::find_if(asyncResource->mRequestsQueue.begin(), asyncResource->mRequestsQueue.end(), @@ -1047,45 +1269,138 @@ 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, + }; + + 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 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; } - return isReadyFinal; + 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"); + } + + 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. + 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; @@ -1098,29 +1413,141 @@ 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; + bool rejectPeerProtocol = false; + auto poisonRecvHolders = [&recvHolders]() + { + for (auto& holder : recvHolders) + { + holder.poison(); + } + }; + + try { - // Reuse the error state for the cancelled request. - llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + // 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) + { + 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()); - return; } - receiveSync(session); - llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); + catch (...) + { + if (mInflightCancelEnabled || rejectPeerProtocol) + { + poisonRecvHolders(); + } + if (!mInflightCancelEnabled) + { + llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR); + } + llmRequest.setKvCacheTransferEnd(LlmRequest::getSteadyClockNow()); + throw; + } TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End calling requestSync for request ID: %zu, context request ID: %zu.", llmRequest.mRequestId, llmRequest.getContextPhaseParams().value().getReqId()); + return true; + } + + [[nodiscard]] bool requestSync(LlmRequest& llmRequest) + { + return requestSync(llmRequest, mTerminate); } struct RequestAndPromise @@ -1129,16 +1556,20 @@ class CacheReceiver::Impl // protects worker-side dereferences and the promise itself from premature destruction. std::shared_ptr mRequest; std::unique_ptr> mPromise; + std::shared_ptr> mCancelFlag; RequestAndPromise() : mRequest(nullptr) , mPromise(nullptr) + , mCancelFlag(nullptr) { } - RequestAndPromise(std::shared_ptr request, std::unique_ptr>&& promise) + RequestAndPromise(std::shared_ptr request, std::unique_ptr>&& promise, + std::shared_ptr> cancelFlag) : mRequest(std::move(request)) , mPromise(std::move(promise)) + , mCancelFlag(std::move(cancelFlag)) { } @@ -1147,6 +1578,7 @@ class CacheReceiver::Impl RequestAndPromise(RequestAndPromise&& other) noexcept : mRequest(std::move(other.mRequest)) , mPromise(std::move(other.mPromise)) + , mCancelFlag(std::move(other.mCancelFlag)) { } @@ -1162,6 +1594,7 @@ class CacheReceiver::Impl mRequest = std::move(other.mRequest); mPromise = std::move(other.mPromise); + mCancelFlag = std::move(other.mCancelFlag); } return *this; } @@ -1205,8 +1638,20 @@ class CacheReceiver::Impl try { TLLM_CHECK_WITH_INFO(requestAndPromise.mRequest != nullptr, "requestAndPromise.mRequest is null"); - requestSync(*requestAndPromise.mRequest); - requestAndPromise.mPromise->set_value(); + TLLM_CHECK_WITH_INFO( + requestAndPromise.mCancelFlag != nullptr, "requestAndPromise.mCancelFlag is null"); + if (requestSync(*requestAndPromise.mRequest, *requestAndPromise.mCancelFlag) + || !mInflightCancelEnabled) + { + requestAndPromise.mPromise->set_value(); + } + else + { + requestAndPromise.mPromise->set_exception(std::make_exception_ptr(TLLM_REQUEST_EXCEPTION( + requestAndPromise.mRequest->mRequestId, common::RequestErrorCode::kNETWORK_ERROR, + "Generation KV cache transfer failed for request %zu", + requestAndPromise.mRequest->mRequestId))); + } } catch (tensorrt_llm::common::RequestSpecificException const& err) { @@ -1224,6 +1669,19 @@ class CacheReceiver::Impl requestAndPromise.mRequest->getContextPhaseParams().value().getReqId(), err.what()); requestAndPromise.mPromise->set_exception(std::current_exception()); } + catch (...) + { + TLLM_LOG_ERROR("Unknown exception in CacheReceiver request() loop"); + if (requestAndPromise.mPromise) + { + requestAndPromise.mPromise->set_exception(std::current_exception()); + } + } + if (requestAndPromise.mRequest != nullptr) + { + std::lock_guard lg(mInFlightCancelMutex); + mInFlightCancelFlags.erase(requestAndPromise.mRequest->mRequestId); + } } } } @@ -1244,6 +1702,7 @@ class CacheReceiver::Impl std::unordered_map> mInstanceToAsyncResource; executor::kv_cache::ConnectionManager* mManager; executor::DataTransceiverState mSelfState; + bool mInflightCancelEnabled{false}; CacheTransferLayer mCacheTransferLayer; std::unordered_map> mProcessToResources; std::mutex mProcessIoResouceMutex; @@ -1251,6 +1710,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..0512898daa54 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"); @@ -16,12 +16,15 @@ */ #pragma once +#include #include #include #include +#include #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 +87,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 +98,7 @@ class TransferSession , mRequest(llmRequest) , mIndexFromEnd(indexFromEnd) , mLastBlockKey(lastBlockKey) + , mEnableInflightCancel(enableInflightCancel) { TLLM_CHECK(!mConnections.empty()); if (recordTiming) @@ -102,6 +107,19 @@ class TransferSession } } + TransferSession(TransferSession const&) = delete; + TransferSession& operator=(TransferSession const&) = delete; + TransferSession(TransferSession&&) noexcept = default; + TransferSession& operator=(TransferSession&&) noexcept = delete; + + ~TransferSession() + { + if (mEnableInflightCancel || mPeerProtocolRejected) + { + poisonRecvBufferHolders(); + } + } + [[nodiscard]] std::vector const& getConnections() const; // should be called only during the initialization of the TransferSession @@ -151,6 +169,59 @@ class TransferSession mCounterPartRanks = std::move(ranks); } + [[nodiscard]] bool isInflightCancelEnabled() const noexcept + { + 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) + { + 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 +233,10 @@ class TransferSession std::unique_ptr mTimes; int32_t mIndexFromEnd{0}; BlockKey mLastBlockKey{}; + bool mEnableInflightCancel{false}; + bool mPeerProtocolRejected{false}; + std::optional mPeerProtocolVersion; + std::vector mRecvBufferHolders; }; using UniqueToken = tensorrt_llm::runtime::UniqueToken; diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index 812323e92b81..304db228425b 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -253,7 +253,8 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses return bufferSizeForTarget; }; auto bufferEleSizes = getBufferSizeForTarget(); - auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend(); + auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); + auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend(sendCancelFlag); BufferIndexHolder sendHolder( *mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/false); auto result = mCacheTransBufferManagers[transferIndexerKCache]->getOrAllocateSendBuffers( @@ -339,48 +340,59 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses session.appendMeasure(startTime, endTime, outputSplitCaches.at(cacheIdx)->getSizeInBytes()); }; - if (pickUpConnections.size() > 1) + try { - if (!common::getEnvEnableReceiveKVCacheParallel()) + if (pickUpConnections.size() > 1) { - TLLM_LOG_DEBUG("Disable parallel receiving of the KV cache."); - for (size_t i = 0; i < pickUpConnections.size(); i++) + if (!common::getEnvEnableReceiveKVCacheParallel()) { - sendBufferFun(deviceId, pickUpConnections[i]); + TLLM_LOG_DEBUG("Disable parallel receiving of the KV cache."); + for (size_t i = 0; i < pickUpConnections.size(); i++) + { + sendBufferFun(deviceId, pickUpConnections[i]); + } } - } - else - { - // concurrency num - auto concurrencyNum - = std::min(std::max(static_cast(1), bufferCoverTargetNum), pPDomainSize * cPDomainSize); + else + { + // concurrency num + auto concurrencyNum + = std::min(std::max(static_cast(1), bufferCoverTargetNum), pPDomainSize * cPDomainSize); - auto remainSendNum = pickUpConnections.size(); + auto remainSendNum = pickUpConnections.size(); - while (remainSendNum > 0) - { - auto sendConcurrencyNum = std::min(remainSendNum, concurrencyNum); - std::vector> futures; - futures.reserve(sendConcurrencyNum); - for (size_t i = 0; i < sendConcurrencyNum; i++) - { - size_t idx = i + (pickUpConnections.size() - remainSendNum); - size_t connIdx = pickUpConnections[idx]; - TLLM_CHECK(idx < pickUpConnections.size()); - TLLM_CHECK(connIdx < session.getConnections().size()); - futures.push_back(std::async(std::launch::async, sendBufferFun, deviceId, connIdx)); - } - for (auto& future : futures) + while (remainSendNum > 0) { - future.get(); + auto sendConcurrencyNum = std::min(remainSendNum, concurrencyNum); + std::vector> futures; + futures.reserve(sendConcurrencyNum); + for (size_t i = 0; i < sendConcurrencyNum; i++) + { + size_t idx = i + (pickUpConnections.size() - remainSendNum); + size_t connIdx = pickUpConnections[idx]; + TLLM_CHECK(idx < pickUpConnections.size()); + TLLM_CHECK(connIdx < session.getConnections().size()); + futures.push_back(std::async(std::launch::async, sendBufferFun, deviceId, connIdx)); + } + for (auto& future : futures) + { + future.get(); + } + remainSendNum -= sendConcurrencyNum; } - remainSendNum -= sendConcurrencyNum; } } + else + { + sendBufferFun(deviceId, pickUpConnections[0]); + } } - else + catch (...) { - sendBufferFun(deviceId, pickUpConnections[0]); + if (agentConnection != nullptr && session.isInflightCancelEnabled()) + { + sendHolder.poison(); + } + throw; } sendHolder.release(); } @@ -497,9 +509,9 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s else { cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForRecv(); + recvHolder = BufferIndexHolder( + *mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/true); } - recvHolder - = BufferIndexHolder(*mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/true); auto targetNum = pickUpConnections.size(); diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp index 05742af17a28..af8943d8c055 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -23,6 +23,7 @@ #include "tensorrt_llm/batch_manager/rnnStateManager.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/dataType.h" +#include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/nvtxUtils.h" #include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h" @@ -184,7 +185,9 @@ void RnnCacheFormatter::formatSlotMode(TransferSession& session) / targetInfo.mDomainTPSize; } - auto cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForSend(); + auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); + auto cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForSend(sendCancelFlag); + BufferIndexHolder sendHolder(*mRnnCacheTransBufferManager, cacheBufferId, /*isRecv=*/false); auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateSendBuffers( cacheBufferId, static_cast(bufferTargetNum), bufferSizesPerTarget, bufferManager); auto& outputBuffers = std::get<0>(allocationResult); @@ -222,12 +225,23 @@ void RnnCacheFormatter::formatSlotMode(TransferSession& session) auto preAllocSendBuffer = mRnnCacheTransBufferManager->getSendBuffer(cacheBufferId); - sendAllBuffers(session, deviceId, outputBuffers, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, - targetInfo, pickUpConnections); + try + { + sendAllBuffers(session, deviceId, outputBuffers, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, + targetInfo, pickUpConnections); + } + catch (...) + { + if (agentConnection != nullptr && session.isInflightCancelEnabled()) + { + sendHolder.poison(); + } + throw; + } session.setTime(TransferSession::kTimeTransmissions); - mRnnCacheTransBufferManager->freeBufferIndexForSend(cacheBufferId); + sendHolder.release(); session.setTime(TransferSession::kTimePostprocess); TLLM_LOG_DEBUG( @@ -348,6 +362,7 @@ void RnnCacheFormatter::unformatSlotMode(TransferSession& session) size_t remainNoCoverSourceNum = 0; size_t bufferCoverSourceNum = 0; std::optional cacheBufferId = std::nullopt; + BufferIndexHolder recvHolder; auto preAssignedRnnId = connections[pickUpConnections[0]]->getPreAssignedBufferId(static_cast(BufferKind::kRNN)); @@ -358,6 +373,7 @@ void RnnCacheFormatter::unformatSlotMode(TransferSession& session) else { cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForRecv(); + recvHolder = BufferIndexHolder(*mRnnCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); } auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateRecvBuffers( @@ -505,10 +521,7 @@ void RnnCacheFormatter::unformatSlotMode(TransferSession& session) bufferManager.getStream().synchronize(); - if (cacheBufferId.has_value()) - { - mRnnCacheTransBufferManager->freeBufferIndexForRecv(cacheBufferId); - } + recvHolder.release(); session.setTime(TransferSession::kTimePostprocess); TLLM_LOG_DEBUG( @@ -643,11 +656,16 @@ void RnnCacheFormatter::formatUnifiedPoolMode(TransferSession& session) bufferSizesPerTarget[t] = ssmBufBytes + convBufBytes; } - auto cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForSend(); + auto const* sendCancelFlag = &session.getDataContext().getTransferTerminate(); + auto cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForSend(sendCancelFlag); + BufferIndexHolder sendHolder(*mRnnCacheTransBufferManager, cacheBufferId, /*isRecv=*/false); auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateSendBuffers( cacheBufferId, static_cast(bufferTargetNum), bufferSizesPerTarget, bufferManager); auto& outputBuffers = std::get<0>(allocationResult); auto& bufferCoverTargetNum = std::get<1>(allocationResult); + auto const* agentConnection = rnnSendConns.empty() + ? nullptr + : dynamic_cast(connections[rnnSendConns[0]]); // Split each outputBuffer into SSM and conv portions. std::vector ssmOutputBuffers(numTargets); @@ -677,11 +695,22 @@ void RnnCacheFormatter::formatUnifiedPoolMode(TransferSession& session) int deviceId; TLLM_CUDA_CHECK(cudaGetDevice(&deviceId)); auto preAllocSendBuffer = mRnnCacheTransBufferManager->getSendBuffer(cacheBufferId); - sendAllBuffers(session, deviceId, outputBuffers, bufferCoverTargetNum, preAllocSendBuffer, bufferManager, - targetInfo, rnnSendConns); + try + { + sendAllBuffers(session, deviceId, outputBuffers, bufferCoverTargetNum, preAllocSendBuffer, + bufferManager, targetInfo, rnnSendConns); + } + catch (...) + { + if (agentConnection != nullptr && session.isInflightCancelEnabled()) + { + sendHolder.poison(); + } + throw; + } session.setTime(TransferSession::kTimeTransmissions); - mRnnCacheTransBufferManager->freeBufferIndexForSend(cacheBufferId); + sendHolder.release(); } } @@ -821,6 +850,7 @@ void RnnCacheFormatter::unformatUnifiedPoolMode(TransferSession& session) // Use pre-assigned buffer ID from NIXL connection if available (same as slot mode). std::optional cacheBufferId = std::nullopt; + BufferIndexHolder recvHolder; auto preAssignedRnnId = connections[rnnRecvConns[0]]->getPreAssignedBufferId(static_cast(BufferKind::kRNN)); if (preAssignedRnnId.has_value()) @@ -830,6 +860,7 @@ void RnnCacheFormatter::unformatUnifiedPoolMode(TransferSession& session) else { cacheBufferId = mRnnCacheTransBufferManager->assignBufferIndexForRecv(); + recvHolder = BufferIndexHolder(*mRnnCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); } auto allocationResult = mRnnCacheTransBufferManager->getOrAllocateRecvBuffers( @@ -876,7 +907,7 @@ void RnnCacheFormatter::unformatUnifiedPoolMode(TransferSession& session) bufferManager.getStream().synchronize(); - mRnnCacheTransBufferManager->freeBufferIndexForRecv(cacheBufferId); + recvHolder.release(); } } diff --git a/cpp/tensorrt_llm/common/envUtils.cpp b/cpp/tensorrt_llm/common/envUtils.cpp index 8987e5048c90..3a48bb602b48 100644 --- a/cpp/tensorrt_llm/common/envUtils.cpp +++ b/cpp/tensorrt_llm/common/envUtils.cpp @@ -395,6 +395,12 @@ bool getEnvTryZCopyForKVCacheTransfer() return zcopyForSysmmetricKVCache; } +bool getEnvDisaggEnableInflightCancel() +{ + static bool const enabled = getBoolEnv("TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL"); + return enabled; +} + bool getEnvForceDeterministic() { static bool const forceDeterministic = getBoolEnv("FORCE_DETERMINISTIC"); diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h index d8e72b7b02cc..b1950194321c 100644 --- a/cpp/tensorrt_llm/common/envUtils.h +++ b/cpp/tensorrt_llm/common/envUtils.h @@ -110,6 +110,9 @@ std::string const& getEnvKVCacheTimeOutputPath(); bool getEnvTryZCopyForKVCacheTransfer(); +// Opt-in for disaggregated KV transfer in-flight cancellation and fail-closed transfer-buffer quarantine. +bool getEnvDisaggEnableInflightCancel(); + // Force deterministic behavior for all kernels. bool getEnvForceDeterministic(); diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp index a21d470b898a..cfb9abb20d34 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) @@ -189,8 +211,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 +248,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()); @@ -255,6 +309,10 @@ 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()); } @@ -291,9 +349,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; } @@ -323,16 +389,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); @@ -401,6 +475,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()); } @@ -442,14 +528,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); @@ -692,7 +774,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 +782,7 @@ void AgentConnectionManager::waitForNotification( if (!mIsRunning) { - return; + return false; } updateUnhandledNotifications(); std::scoped_lock lock(mNotificationMutex); @@ -733,7 +815,7 @@ void AgentConnectionManager::waitForNotification( { it = mUnhandledNotifications.erase(it); } - return; + return true; } } } @@ -753,7 +835,7 @@ void AgentConnectionManager::waitForNotification( { it = mUnhandledNotifications.erase(it); } - return; + return true; } } } @@ -773,24 +855,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..ee496865ed90 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"); @@ -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; @@ -260,13 +265,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; @@ -303,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; @@ -320,11 +328,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; @@ -345,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/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..ab528279bab9 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -17,6 +17,8 @@ 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(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 d4d561f55d07..eb3386dd55c2 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 @@ -146,13 +147,14 @@ TEST_P(BufferIndexHolderLifecycleTest, DefaultConstructedExplicitReleaseIsNoOp) EXPECT_EQ(inUse(), before); } -// Holder built from a nullopt index is disarmed (mHeld == false). +// A nullopt index still represents ownership for dynamic-buffer mode: an +// interrupted transfer must be able to poison the whole dynamic pool. 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); } @@ -286,6 +288,112 @@ 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); + session.setPeerProtocolVersion(1); + EXPECT_EQ(session.getPeerProtocolVersion(), 1); + 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()); +} + +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/cacheTransferStateTest.cpp b/cpp/tests/unit_tests/batch_manager/cacheTransferStateTest.cpp new file mode 100644 index 000000000000..a2234c6ea748 --- /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.quiescedRequestIds.empty()); + + auto const secondPoll = summarizePackedTransferStates({{0, 17, kTransferCompleted}, {0, 17, kTransferCompleted}}); + EXPECT_EQ(secondPoll.completedRequestIds, std::unordered_set{17}); + EXPECT_EQ(secondPoll.quiescedRequestIds, std::unordered_set{17}); +} + +TEST(CacheTransferStateTest, FailureUnionPrecedesQuiescenceIntersection) +{ + auto const firstPoll = summarizePackedTransferStates({{0, 29, kTransferFailed}, {0}}); + EXPECT_EQ(firstPoll.failedRequestIds, std::unordered_set{29}); + EXPECT_TRUE(firstPoll.quiescedRequestIds.empty()); + + auto const secondPoll = summarizePackedTransferStates({{0, 29, kTransferFailed}, {0, 29, kTransferCompleted}}); + EXPECT_EQ(secondPoll.failedRequestIds, std::unordered_set{29}); + EXPECT_EQ(secondPoll.quiescedRequestIds, std::unordered_set{29}); + EXPECT_TRUE(secondPoll.completedRequestIds.empty()); +} + +TEST(CacheTransferStateTest, TimeoutVoteWinsOverLateSuccess) +{ + auto const outcome + = summarizePackedTransferStates({{0, 41, kTransferCompleted | kTransferTimedOut}, {0, 41, kTransferCompleted}}); + EXPECT_EQ(outcome.timedOutRequestIds, std::unordered_set{41}); + EXPECT_EQ(outcome.failedRequestIds, std::unordered_set{41}); + EXPECT_EQ(outcome.quiescedRequestIds, std::unordered_set{41}); + EXPECT_TRUE(outcome.completedRequestIds.empty()); +} + +TEST(CacheTransferStateTest, LargeRequestIdIsPreserved) +{ + auto const requestId = std::numeric_limits::max(); + auto const outcome + = summarizePackedTransferStates({{0, requestId, kTransferCompleted}, {0, requestId, kTransferCompleted}}); + EXPECT_EQ(outcome.completedRequestIds, std::unordered_set{requestId}); + EXPECT_EQ(outcome.quiescedRequestIds, std::unordered_set{requestId}); +} + +TEST(CacheTransferStateTest, MalformedPayloadIsRejected) +{ + EXPECT_THROW((void) summarizePackedTransferStates({{}}), std::exception); + EXPECT_THROW((void) summarizePackedTransferStates({{0, 17}}), std::exception); +} + +TEST(CacheTransferStateTest, DuplicateRequestIdOnOneRankIsRejected) +{ + EXPECT_THROW( + (void) summarizePackedTransferStates({{0, 17, kTransferCompleted, 17, kTransferFailed}}), std::exception); +} + +TEST(CacheTransferStateTest, PoisonIsTopologyWideIndependentOfRequestState) +{ + auto const outcome = summarizePackedTransferStates({{0}, {kTransferBufferPoisoned}}); + EXPECT_TRUE(outcome.transferBufferPoisoned); + EXPECT_TRUE(outcome.failedRequestIds.empty()); + EXPECT_TRUE(outcome.quiescedRequestIds.empty()); +} + +TEST(CacheTransferStateTest, 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/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 194d5267c6b3..c59f0f526f7f 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; @@ -155,7 +157,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); @@ -165,18 +171,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 c280f8293f4a..80fa1c24ed7f 100644 --- a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp +++ b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp @@ -752,6 +752,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(); @@ -1329,6 +1348,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 d1ca104cca1a..1f415cbc7b5e 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 + +# Disaggregated Peer Cancellation and Terminal Protocol + +| | | +| --- | --- | +| **Status as of** | 2026-06-30 | +| **Status** | Design scaffold; this draft implements no runtime behavior | +| **JIRA** | TRTLLM-12721 | +| **Scope** | C++ transceiver; `CacheTransceiver` NIXL backend over NIXL-UCX; preallocated buffer; positive transfer timeout; overlap enabled; PP=1; CP=1; attention-DP off; layerwise and zero-copy off | +| **Prerequisite** | #15798: CTX-GEN protocol and effective-mode negotiation | +| **Dependent rollout** | #15738: qualified default-on policy | +| **Default behavior** | Unchanged; cancellation remains default-off in this PR | +| **Unsupported/unqualified paths** | Python transceiver, direct UCX, MPI, Mooncake, NIXL-libfabric, non-preallocated transfer, nonpositive timeout, overlap disabled, PP>1, CP>1, attention-DP, layerwise, zero-copy, and unqualified EP/topology combinations | + +## Purpose + +#15798 provides a fixed v1 descriptor, effective-mode and local-profile checks, per-view highest-common-version +selection, and pre-DMA protocol rejection. This phase defines the next layer: generation-safe transfer identity, +asynchronous peer cancellation, terminal acknowledgement, and bounded retirement of cancellation metadata. Phase B2 +extends #15798's per-view result into one topology-wide cancellation-control version, capability set, and serialization +ABI selection. + +CTX and GEN do not synchronously vote on each cancellation. Either endpoint may initiate cancellation. Local +participants still reach a rank-consistent outcome, and cancellation-related physical resource reuse waits for a +terminal result that is safe for every expected peer. The ordinary, uncontended success path keeps its existing +authoritative local-completion cleanup and does not wait for a terminal-control exchange. + +This draft intentionally adds no live cancellation messages. In particular, it does not claim that a cancelled NIXL +transfer is quiescent and does not return a possibly written buffer to the reusable pool. + +## Safety invariants + +1. Request ID alone never identifies a transfer. +2. A transfer attempt is valid only in the process epochs that created it. +3. Only an exact active-registry match may authorize `ready=true` or DMA. +4. Cancellation is idempotent and may be initiated by either endpoint. +5. An endpoint may report `LOCAL_NOT_STARTED` only when its rank-pair did not cross the ready/DMA commit boundary; + topology-wide `NOT_STARTED` requires matching evidence from both endpoints of every expected rank-pair and every + participating local rank. +6. A source reports `LOCAL_SOURCE_COMPLETE` only after backend source access ended. A destination reports + `LOCAL_DESTINATION_COMPLETE` only after transfer completion and destination postprocessing. Topology-wide + `COMPLETED` requires both reports for every expected rank-pair and agreement on every local rank. +7. Missing, conflicting, stale, malformed, or otherwise unproven state becomes `UNKNOWN`. +8. A timeout, cache expiry, message receipt, or backend-handle release is never by itself proof that memory is reusable. +9. No cancellation-related result other than topology-wide `NOT_STARTED` permits early release of source/send or + destination/receive leases reserved for the attempt. +10. `UNKNOWN` makes the logical attempt unsafe, quarantines each still-owned and unproven local lease, and follows the + bounded fail/restart policy. It never retroactively mutates a lease generation already released after authoritative + local completion. +11. Cancellation-related early cleanup and logical request completion occur only after rank-consistent aggregation. + Physical endpoint-local resources may follow the existing success fast path immediately after authoritative local + completion; bounded evidence is retained for a later cancellation race. +12. A delayed message cannot affect a newer attempt that reused the same request ID. +13. A control receipt acknowledges exact bytes only. It is neither quiescence evidence nor a topology-wide decision. +14. Network input never directly names memory to release. Cleanup resolves an exact active lease generation from the + registry and runs through one cleanup-once token. +15. A rank-local timeout or transport callback records cancellation intent only. It does not call a cancellation + primitive, complete a promise, or release a lease before the topology-specific local-rank decision. +16. Completion evidence is attempt-bundle-wide: every KV, KV-indexer, RNN, registration, status handle, and required + postprocessing step covered by the parent lease must be terminal. One completed component never releases a partial + parent lease or certifies the edge. + +## Transfer identity + +### Process epoch + +#15798 advertises a nonzero process-session token for every registered NIXL agent. This token is the endpoint's +`instance_epoch` and changes whenever that process or registered agent restarts. + +An epoch is an identity discriminator, not a timestamp. It must not be reused by a replacement process. + +### Context ticket + +CTX creates a context ticket before publishing `ContextPhaseParams`: + +```text +ContextTicket { + uint64 ctx_request_id; // full-width ID; never packed or truncated + uint64 ctx_instance_epoch; // nonzero logical CTX publication epoch + uint64 ctx_attempt; // nonzero, unique in that publication epoch + Direction direction; // initially CTX_TO_GEN +} +``` + +`ctx_instance_epoch` is generated by the authoritative CTX rank for the logical CTX publication session; it is distinct +from the per-rank agent epochs advertised by #15798. `ctx_attempt` is generated once per exported context result. +Reusing a request ID or restarting the CTX publication session creates a new ticket. + +The same logical ticket must be visible on every participating CTX rank. The implementation must generate it at one +authoritative rank and distribute it at an existing rank-lockstep point, or provide an equivalent proof that every rank +derives the same value. + +### Full transfer key and topology identity + +GEN creates one nonzero `gen_attempt` when it begins one logical consumption of a context result. A control transport +retransmission to the same endpoint reuses that attempt; it does not create another protocol identity. An application +retry that can select another worker or allocate different destination buffers is a new logical attempt. It is permitted +only after proving the previous attempt never opened or after obtaining a new context result and `ContextTicket`. +Version 1 does not authorize sequential rebinding of one ticket after an unresolved attempt. + +Each rank-pair transfer is identified by: + +```text +TransferKey { + uint64 ctx_request_id; + uint64 ctx_instance_epoch; + uint64 ctx_attempt; + uint64 gen_attempt; + Direction direction; + + uint32 source_rank; + uint32 destination_rank; + uint64 source_instance_epoch; + uint64 destination_instance_epoch; + byte expected_topology_digest[32]; +} +``` + +The full `ContextTicket`, including `ctx_instance_epoch`, is embedded in every key and control payload. The source and +destination epochs come from #15798's peer descriptors. `source_rank` and `destination_rank` are zero-based indices in +the ordered source and destination `CommState` agent-state vectors used for this transfer. They are not assumed to be +MPI world ranks, TP ranks, or interchangeable with another communicator's indices. + +`expected_topology_digest` is SHA-256 over the following padding-free, network-order canonical bytes: + +```text +bytes domain[8]="TRTLTOP1" | u16 topology_schema_version=1 | u8 direction | u16 selected_control_version | +u64 selected_capabilities | u64 serialization_abi_id | +u32 source_count | repeated(source_count, u32 agent_name_length | bytes agent_name | u64 instance_epoch) | +u32 destination_count | repeated(destination_count, u32 agent_name_length | bytes agent_name | u64 instance_epoch) | +u32 edge_count | repeated(edge_count, u32 source_rank | u32 destination_rank) +``` + +Agent lists retain their `CommState` order; edges are sorted lexicographically and duplicates are invalid. Names are the +exact NIXL agent-identity bytes, not locale-normalized text. Counts and name lengths use the same topology/message bounds +as decoding. Every rank derives the same digest at a lockstep point. The digest freezes the exact participant set for the +attempt; world size is never substituted for the topology-specific set. The outer NIXL notification sender, embedded +agent identity, rank indices, instance epochs, and digest must all agree before connection, buffer, ready, or DMA state +changes. + +Golden digest fixtures cover asymmetric source/destination counts, multiple edges, and identical numeric rank indices +under different ordered agent namespaces. + +No field may be compressed into request-ID high bits. Counters must fail closed before wraparound. + +## Active-context registry + +CTX owns an active-context registry keyed by `ContextTicket`. It is the authorization source for new transfer work. GEN +owns a corresponding local attempt registry keyed by `(ContextTicket, gen_attempt)`; the services do not share memory or +registry ownership. On each endpoint, the parent entry owns binding and topology aggregation, and each expected +rank-pair has a child entry keyed by its exact `TransferKey`. The following shape is logical and exists independently at +both endpoints, with role-specific local leases and evidence. + +An entry contains at least: + +```text +ContextEntry { + ContextTicket ticket; + ContextState state; + optional binding; + ParentLease parent_lease; // CTX publication lease until bind; then role-specific attempt lease + ExpectedPeerSet expected_peers; + map edges; + optional topology_terminal; + monotonic_time created_at; + monotonic_time updated_at; +} + +TransferEdgeEntry { + TransferKey key; + EdgeState state; + ImmutableBufferBinding buffers; + EndpointLeaseClaim local_claim; + optional local_evidence; + optional remote_evidence; + DedupeWindow dedupe; + optional terminal_message; + optional terminal_receipt; +} +``` + +`GenBinding` contains one logical generation attempt plus the destination rank/epoch information needed to construct +every expected rank-pair key. It is installed atomically and never replaced by a conflicting attempt. Per-edge state is +required because one logical transfer can fan out to multiple peers and observe partial progress. + +State is partitioned by owner; one mixed enum must not be shared accidentally across parent and edge records. The +diagrams below are normative. An implementation may split them further but must preserve at least these distinctions: + +```text +CTX parent: PUBLISHED, BOUND, REDUCING, SUCCESS_COMPLETE, + TERMINAL_NOT_STARTED, TERMINAL_COMPLETED, TERMINAL_UNKNOWN, RETIRED +CTX edge: BOUND, READY_PREPARED, CANCEL_PENDING, READY_DECIDED, IN_FLIGHT, + EVIDENCE_NOT_STARTED, EVIDENCE_SOURCE_COMPLETE, EVIDENCE_UNKNOWN, RETIRED +GEN parent: CONTEXT_IMPORTED, ADMISSION_PREPARED, ADMISSION_DECIDED, OPEN_PUBLISHING, OPEN_SENT, + CANCEL_PENDING, CANCEL_SENT, + IN_FLIGHT, REDUCING, UNINSTALLED_UNKNOWN, SUCCESS_COMPLETE, LOCAL_CANCELLED, + TERMINAL_NOT_STARTED, TERMINAL_COMPLETED, TERMINAL_UNKNOWN, RETIRED +GEN edge: OPEN_SENT, CANCEL_SENT, READY_ACCEPTED, IN_FLIGHT, + EVIDENCE_NOT_STARTED, EVIDENCE_DESTINATION_COMPLETE, EVIDENCE_UNKNOWN, RETIRED +``` + +Required behavior: + +- Create `PUBLISHED` before the context result becomes externally visible. +- Accept a new generation open/request only when its context ticket exactly matches a `PUBLISHED` entry. An exact + duplicate after binding returns the cached decision without allocating or submitting again. +- An authenticated but incompatible `OPEN` proposal consumes the ticket: CTX removes authorization, commits + `TERMINAL_UNKNOWN`, quarantines its still-owned publication lease, caches the exact negative READY keyed by the + proposal key/fingerprint, and enters bounded fail/restart. A duplicate receives the byte-equivalent rejection; a + different `gen_attempt` can never rebind that ticket. +- Bind one generation attempt atomically. A conflicting attempt or endpoint epoch produces `UNKNOWN`; it never replaces + the existing binding. A retransmission reuses the exact `gen_attempt`, key, payload, and message ID. +- Install the topology-wide `READY_DECIDED` state on every participating local rank before sending or making + `ready=true` visible. +- Remove transfer authorization when entering any terminal state. +- Retain only bounded terminal metadata for duplicate replies and diagnostics. +- Treat an absent or retired ticket as unauthorized. An exact no-peer terminal-cache hit may reproduce + `READY=NOT_READY + LOCAL_NOT_STARTED`; every miss or mismatch receives `READY=NOT_READY + LOCAL_UNKNOWN`. Neither + case recreates an active entry or authorizes DMA. + +This default-deny rule is what makes bounded tombstone retirement safe: evicting diagnostic terminal metadata can +reduce a late peer's answer from `NOT_STARTED` to `UNKNOWN`, but can never reopen DMA. + +### Resource leases and ABA protection + +CTX creates a `ContextPublicationLease` before publishing a context result. On successful bind, it atomically transfers +that ownership into a parent `TransferAttemptLease`; GEN creates its parent attempt lease before advertising destination +memory. Together the CTX and GEN attempt halves form the logical request-level `TransferLease`. A parent local half +owns, as applicable: + +- the pinned request object; +- KV, KV-indexer, and RNN send/receive `BufferIndexHolder` resources; +- dynamic and zero-copy registrations used by a future supported cell; +- backend status handles and their connection, agent, and plugin lifetimes; +- immutable buffer bindings; and +- one cleanup-once token shared by every completion, cancellation, timeout, and teardown path. + +A no-peer CTX publication owns its source request and KV resources through the publication lease before the context +result becomes visible. Physical holders shared by multiple peer connections are owned once by the parent attempt, with +one cleanup token. Child edge entries hold non-owning segment claims and evidence; they cannot independently free a +shared holder. Resources are never partially released from a request lease because one edge or buffer kind appeared +idle. + +Every pool-backed binding records `(pool_id, slot_id, slot_generation, buffer_kind)`. Slot generation increments before +a freed slot is leased again. A handler may release or quarantine memory only after resolving the exact active +`TransferKey` and matching all recorded lease generations. A terminal-cache entry contains evidence and fingerprints; +it never owns memory and cannot release a slot. Minimum slot states are `FREE`, `LEASED`, +`QUARANTINED_PENDING_RESTART`, and `PERMANENTLY_QUARANTINED`. + +`UNKNOWN` atomically removes every still-owned, unproven local lease generation from reusable capacity and marks the +worker unhealthy. A slot already released after authoritative local completion remains released; cached evidence and +generation matching prevent a late topology outcome from touching its newer occupant. The logical attempt still becomes +`UNKNOWN` and triggers restart. The executor that owns the transceiver must stop admission, reject or settle outstanding +work, preserve all quarantined allocations, and request process exit within a configured +`unknown_restart_deadline`; failure to exit gracefully escalates to forced termination. +The deployment supervisor owns restart. This is the bounded fail/restart policy, not an optional logging action. + +Shutdown order is: stop admission; freeze new bindings; settle cancellation/terminal evidence for a bounded interval; +permanently isolate unresolved leases; join response and tracker workers; deregister memory while backing storage still +exists; and destroy the NIXL agent/plugin last. No destructor may silently return a quarantined slot to its pool. + +Response and background status workers may validate input, record immutable backend evidence, and wake the executor. +They do not release resources, fulfill user promises, run a collective, or invoke cancellation primitives. The executor +applies rank-consistent state transitions and cleanup at its ordered lockstep point. + +### Active-registry bounds + +The active registry and terminal cache have different capacity rules: + +- Reaching the active-entry limit applies admission backpressure. CTX must fail or defer publication before exposing a + new context result; it must not evict an authorized active entry to make room. +- A `PUBLISHED`, `BOUND`, or `READY_PREPARED` entry that expires before topology-wide ready commit is cancelled through + the normal rank-consensus path and may become `NOT_STARTED` only with complete local evidence. +- A `READY_DECIDED` or `IN_FLIGHT` entry that exceeds its deadline becomes `UNKNOWN`; each endpoint quarantines its + still-owned, unproven local lease generations and the fail/restart policy runs. +- Capacity and TTL never silently remove active authorization. Only an explicit terminal transition removes it. +- Terminal duplicate-response metadata may use bounded age/capacity eviction only after its mandatory control lifetime + (or no-peer publication validity) expires. Capacity pressure before that point applies admission backpressure; it does + not evict an in-contract response. A later miss remains default-denied and returns `UNKNOWN`. + +## Control protocol + +Phase B2 extends #15798's descriptor and compatibility result with a distinct +`PEER_CANCELLATION_CONTROL` capability, a `serialization_abi_id`, and a selected control version. ABI ID 1 means the +version-1 carrier and legacy payload frozen below on a 64-bit little-endian host. Every participating local rank +contributes its own range/capabilities/ABI plus the descriptors for every remote edge it can reach, including an empty +edge list. One topology-wide reduction first computes `global_min=max(all min_version)`, +`global_max=min(all max_version)`, and +`global_capabilities=bitwise_and(all capabilities)` across every local and remote endpoint in the frozen topology. The +topology rejects if the range is empty or the required control bit is absent; otherwise every rank selects +`global_max`. All ABI IDs must be identical. Selection is never performed independently over rank-local peer subsets. +Every participant installs the same capability set, selected version, ABI, peer epochs, and topology digest at an +existing ordered lockstep point. + +The capability value is `uint64(1) << 2`. #15798 already owns bit 0 for in-flight cancellation and bit 1 for pre-DMA +protocol rejection. Its v1 descriptor is a fixed five-field envelope and must not be extended in place. B2 therefore +uses a new descriptor-schema marker with one required ABI field: + +```text +__trtllm_protocol_v2________trtllm_protocol_end +``` + +A #15798 v1 parser treats the v2 marker as no descriptor and never advertises bit 2, so it cannot receive the active +carrier. A B2 parser may accept the fixed v1 descriptor only for baseline compatibility; active control requires a valid +six-field v2 descriptor, bit 2, and a nonzero ABI ID. ABI ID 1 is encoded as hexadecimal `1`; future incompatible native +schemas allocate a new nonzero ID and control version rather than changing either descriptor schema in place. + +Fixed-peer deployments preflight every peer before accepting traffic. Dynamic CTX routing cannot know the GEN +`CommState` before the first request. In that case, `OPEN` is the sole capability-gated proposal: GEN first installs a +candidate selection across all GEN ranks using its local descriptors and the CTX descriptors carried by the context +result, then sends `OPEN` only because every target CTX agent advertises the control capability and compatible ABI. CTX +queues the proposal, recomputes and installs the exact topology-wide result across all CTX ranks, and compares the +version, ABI, epochs, and topology digest before mutating transfer/data connections, buffers, readiness, or DMA. A +mismatch returns `READY=NOT_READY` plus `LOCAL_UNKNOWN`. Establishing or looking up the NIXL control-plane connection +needed to receive `OPEN` is allowed; it does not authorize transfer state. + +`OPEN`, a proposal-scoped `READY=NOT_READY + LOCAL_UNKNOWN` rejection, and an exact no-peer terminal-cache response of +`READY=NOT_READY + LOCAL_NOT_STARTED` are the only pre-install messages; they use a carrier version/ABI already proven +common by the advertised descriptors. Positive READY, CANCEL, TERMINAL, and receipt traffic requires both sides to +install the exact selection. A local GEN cancellation after `OPEN` but before READY remains pending and sends no CANCEL; +READY/timeout resolution below determines its safe outcome. An incompatible previously unseen peer fails that request +without silently falling back. + +GEN tracks `peer_installation_proven` per edge. It becomes true only after validating `READY=READY` or a bound-session +`TERMINAL` carrying the exact selected version/key/digest; an `OPEN` proposal or negative READY alone does not set it. +GEN may parse a candidate-version TERMINAL to validate this proof because its own topology installed the candidate +before `OPEN`; a valid TERMINAL demonstrates that CTX installed the same selection. Outgoing CANCEL/TERMINAL is forbidden +on every edge whose proof remains false. +During a rolling upgrade, active cancellation remains disabled for a routing pool until every possible peer supports +the selected control protocol. + +### Notification carrier and dispatch + +Version 1 appends `ControlMessageInfo` as variant index 3 of the existing NIXL `NotificationInfo` variant; existing +variant indices 0 through 2 and their bytes remain unchanged. `ControlMessageInfo` contains one bounded +`ControlEnvelope`. It is emitted only after the descriptor precheck proves a common control carrier/ABI; full dynamic +selection may still be pending for the `OPEN` proposal and its negative rejection. An old peer never receives the new +variant. A raw `TRTLCAN1` envelope is not placed directly into the notification queue because current receive paths would +parse its first bytes as the native legacy variant index. + +The exact carrier is `[legacy native-size_t variant_index=3][32-byte ControlEnvelope][payload]`. There is no generic +string/vector length between the variant index and envelope. The receiver parses directly from the length already +reported by NIXL, validates that the span contains the complete fixed header and declared payload, and only then copies +or allocates. Because the outer index and embedded `REQUEST_AND_BUFFER_INFO` remain legacy-native, peers must have a +matching `serialization_abi_id`. Control v1/ABI 1 freezes the exact field sequence, fundamental widths, and native +encoding defined here; any incompatible `RequestAndBufferInfo` schema/layout or architecture requires a new ABI ID and +selected control version. + +The notification dispatcher first decodes the bounded outer variant, preserves the NIXL-reported sender identity, and +routes control messages to one serialized control handler. Request, ready, sync, and control notifications may +interleave in one queue; unrelated variants remain queued for their existing consumers. Before registry mutation, the +control handler verifies that the outer sender equals the embedded agent identity and the endpoint named by the exact +key. The outer decoder checks the notification length before reading the native variant index and requires exact +consumption for `ControlMessageInfo`; it does not call the current unchecked stream decoder on malformed input. Golden +tests freeze indices 0 through 3, and interleaving tests cover every legacy/control ordering. + +All new envelope, TLV-header, transfer-key, and scalar-TLV integers use unsigned network byte order. The opaque embedded +legacy `REQUEST_AND_BUFFER_INFO` value retains its current native representation. The fixed 32-byte length-delimited +envelope is: + +```text +ControlEnvelope { + byte magic[8]; // ASCII "TRTLCAN1" + uint16 protocol_version; + uint16 message_type; + uint32 flags; // must be zero in version 1 + uint32 payload_length; // at most 1048576 bytes + uint32 payload_crc32c; // CRC-32C (Castagnoli) over payload + uint64 message_id; // nonzero; scoped by sender epoch and key + bytes payload; +} +``` + +The 1 MiB value is a protocol ceiling, not proof that the NIXL notification backend accepts that size. At admission, +the implementation computes the effective ceiling as the minimum of the protocol limit and the configured or queried +backend notification limit, including outer-variant overhead. Version 1 does not fragment control messages: a request +that cannot fit is rejected before allocation or advertisement. Silent fallback after active negotiation is forbidden. + +Version 1 rejects unknown flag bits, an unknown message type, a zero message ID, a payload above the effective ceiling, +a length mismatch, trailing bytes, or a CRC mismatch before registry mutation. Message type values are fixed as `OPEN=1`, +`READY=2`, `CANCEL=3`, `TERMINAL=4`, and `TERMINAL_RECEIPT=5`. Golden byte fixtures must freeze every message type and +both endian conversion and malformed-input behavior. + +Initial message types are: + +| Message | Direction | Purpose | +| --- | --- | --- | +| `OPEN` | GEN to CTX | Atomically validate/bind the context ticket and carry complete request-and-buffer information | +| `READY` | CTX to GEN | Return the ready decision for the exact transfer key | +| `CANCEL` | Either | Request logical cancellation idempotently | +| `TERMINAL` | Either | Report sender-local terminal evidence for the exact rank-pair key | +| `TERMINAL_RECEIPT` | Either | Confirm the exact key, message ID, evidence, and digest received | + +Control payloads use the same network-order `uint16 type`, `uint16 flags`, `uint32 length`, and value TLV header defined +for the context trailer. The version-1 types are `TRANSFER_KEY=1`, `REQUEST_AND_BUFFER_INFO=2`, `READY_DECISION=3`, +`CANCEL_REASON=4`, `LOCAL_EVIDENCE=5`, `TERMINAL_MESSAGE_ID=6`, and `PAYLOAD_SHA256=7`. Known fields may appear only +once. `TRANSFER_KEY` is always required and has this exact 89-byte, padding-free value layout: + +```text +u64 ctx_request_id | u64 ctx_instance_epoch | u64 ctx_attempt | u64 gen_attempt | +u8 direction | u32 source_rank | u32 destination_rank | +u64 source_instance_epoch | u64 destination_instance_epoch | bytes expected_topology_digest[32] +``` + +`OPEN` additionally requires `REQUEST_AND_BUFFER_INFO`; `READY` requires `READY_DECISION`; `CANCEL` requires a 32-bit +`CANCEL_REASON`; `TERMINAL` requires a one-byte `LOCAL_EVIDENCE`; and `TERMINAL_RECEIPT` requires the original 64-bit +terminal message ID, one-byte evidence, and 32-byte SHA-256 terminal-payload digest. `READY=NOT_READY` also requires +exactly one `LOCAL_EVIDENCE`, whose value is `LOCAL_NOT_STARTED` for a bound pre-commit rejection or exact cached +no-peer terminal hit, or `LOCAL_UNKNOWN` for an unbound/proposal rejection or cache miss. `READY=READY` forbids +`LOCAL_EVIDENCE`. All new scalar TLV values use network byte order; +the opaque legacy value is the stated exception. Unknown required TLVs, duplicates, bad lengths, and missing or +forbidden message-specific fields reject before state mutation. Golden fixtures cover positive READY and both negative +evidence cases. + +Version-1 enum values are fixed: + +| Field | Values | +| --- | --- | +| `Direction` | `INVALID=0`, `CTX_TO_GEN=1` | +| `READY_DECISION` | `INVALID=0`, `READY=1`, `NOT_READY=2` | +| `LOCAL_EVIDENCE` | `INVALID=0`, `LOCAL_NOT_STARTED=1`, `LOCAL_SOURCE_COMPLETE=2`, `LOCAL_DESTINATION_COMPLETE=3`, `LOCAL_UNKNOWN=4` | +| `CANCEL_REASON` | `INVALID=0`, `USER=1`, `DEADLINE=2`, `SHUTDOWN=3`, `PEER_FAILURE=4`, `PROTOCOL_ERROR=5` | + +Unknown required enum values reject active mode. `INVALID` is never valid on the wire. + +The new envelope does not add positional fields to existing `RequestInfo`, `RequestAndBufferInfo`, or +`DataTransceiverState` serialization. Active mode uses one `OPEN` envelope containing a length-delimited copy of the +complete current `RequestAndBufferInfo`: agent name, connection address, nested `RequestInfo`, destination memory +descriptors, optional metadata, valid connection index, and buffer kinds. Carrying only `RequestInfo` is invalid because +the source needs the other fields to establish the connection and DMA destination. + +Before calling any existing legacy deserializer, active mode must run a checked structural preflight (or use a new +checked bounded reader) over every nested string, optional, descriptor, vector, and scalar. Counts are bounded by the +negotiated topology and configured buffer kinds; lengths and arithmetic are checked against remaining input and the +effective message budget; no `reserve`, allocation, connection, or registry mutation occurs first; and exact byte +consumption is required. A valid envelope around an unchecked legacy blob is not considered bounded decoding. + +Dispatch validates the envelope, outer sender, embedded agent/address, full key, connection index, nested request ID, +mode, and active-registry entry before the decoded request can reach any ready or DMA handler. Each endpoint separately +resolves its immutable local binding and slot generations from that registry; wire-provided bytes never directly assert +a trusted local generation. A bare legacy request can authorize work only in baseline mode. A separately reordered +sidecar is not permitted. The baseline path continues using its existing messages unchanged. + +### Cancellation + +`CANCEL` contains the exact `TransferKey`, a reason code, and the sender's monotonically increasing message ID. + +Receiving `CANCEL`: + +- is idempotent; +- never blocks on a cross-service vote; +- returns the already committed terminal result when one exists; +- contributes `LOCAL_NOT_STARTED` evidence only after the pending rank-wide ready decision proves that no participant + committed; +- contributes `LOCAL_UNKNOWN` evidence for ambiguous post-commit work; +- may return the already committed role-specific `LOCAL_SOURCE_COMPLETE` or `LOCAL_DESTINATION_COMPLETE` evidence only + when that local completion was already authoritative. + +Cancellation is serialized with the edge state under one local lock. Once cancellation wins that lock, a late or +duplicate `READY` cannot authorize DMA. If cancellation arrives while a ready decision is being reduced, it is marked +pending until that decision resolves: a rejected decision may prove `LOCAL_NOT_STARTED`; a committed decision is +post-commit and therefore `LOCAL_UNKNOWN` unless role-specific completion was already authoritative. + +The initial implementation has no `DRAINING` or `QUIESCED` terminal claim. If a rank-pair cannot prove +`LOCAL_NOT_STARTED`, it reports `LOCAL_UNKNOWN`. + +### Terminal acknowledgement + +`TERMINAL` carries sender-local evidence for one exact rank-pair key: + +- `LOCAL_NOT_STARTED`: the sender proves that rank-pair did not cross the topology-wide ready/DMA commit boundary; +- `LOCAL_SOURCE_COMPLETE`: the source proves its backend will no longer access source/send memory; +- `LOCAL_DESTINATION_COMPLETE`: the destination proves transfer completion and destination postprocessing; +- `LOCAL_UNKNOWN`: any other state. + +A rank-pair report is evidence, not a global terminal decision. Only the all-peer and local-rank reduction may commit +one of these topology-wide terminal dispositions: + +| Disposition | Meaning | Resource consequence | +| --- | --- | --- | +| `NOT_STARTED` | Every expected rank-pair and participating local rank proved it did not cross ready/DMA commit | Release both source/send and destination/receive leases reserved for the unstarted attempt | +| `COMPLETED` | Every rank-pair has matching source-complete and destination-complete evidence, and every local rank agrees | Follow normal successful cleanup for both source and destination leases | +| `UNKNOWN` | Any source/destination quiescence or exact topology state is unproven | Quarantine still-owned, unproven local lease generations; poison and fail/restart; never mutate an authoritatively released generation | + +`QUIESCED` is deliberately absent. It may be added only in a later PR that cites a backend contract proving no further +source or destination memory access after cancellation. + +Terminal state is monotonic. Once an attempt commits `UNKNOWN`, a late success notification does not make its +quarantined source or destination memory reusable in this phase. A backend-handle `release()` alone is not +`LOCAL_SOURCE_COMPLETE` proof. + +When cancellation or ambiguity makes local evidence immutable at the ordered local decision point, the executor +constructs and sends its canonical `TERMINAL` promptly on each edge with installation proof, outside the state lock, +without waiting for peer evidence. A +`CANCEL` is answered by this `TERMINAL`; there is no separate cancellation acknowledgement. This rule prevents both endpoints from waiting in +`REDUCING` for evidence that neither has published. A `TERMINAL_RECEIPT` echoes the exact key, terminal message ID, +sender-local evidence, and payload digest. It acknowledges those bytes only, never a topology disposition, and gates +only duplicate-response metadata retirement. + +Uncontended success does not require this exchange. Each endpoint follows its existing authoritative successful +completion and endpoint-local cleanup, records immutable completion evidence in the bounded terminal cache, and sends a +`TERMINAL` only if a later cancellation or ambiguity requires cross-service reduction. + +Physical cleanup on this fast path is local and bundle-wide: an endpoint releases its parent attempt lease only after +all local edge claims and every owned buffer kind/status/postprocessing step are authoritatively complete. Logical +request success still waits for topology-local rank agreement. If another rank later becomes `UNKNOWN`, the already +proven and released generation remains safe; only still-owned/unproven generations quarantine, while the logical +request and worker follow fail/restart. + +### Topology-local intent coordination + +Local admission, cancellation, and logical completion decisions follow one unconditional ordered protocol. Once per +executor decision iteration, every participant enters the same topology-scoped batched operation, including ranks with +no records. The version-1 contribution is a fixed header `(schema_version, decision_round_id, record_count)` plus +records sorted by `(ContextTicket, key_present, source_rank, destination_rank)`, using `UINT32_MAX` ranks when no key +exists. Each record contains the full available +key, intent (`ADMIT`, `CANCEL`, `COMPLETE`, or `REJECT`), effective mode/version/topology digest, and local evidence. +Empty contributions are explicit; a request never gates whether a rank enters the collective. + +The internal version-1 contribution uses network order and this bounded encoding: + +```text +u16 schema_version=1 | u16 flags=0 | u64 decision_round_id | u32 record_count | +repeated(record_count, + u32 record_length | + ContextTicket[25] | u64 proposed_gen_attempt | u8 key_present | + optional(TransferKey[89]) | u8 intent | u8 local_evidence | + u8 open_publication_state | u8 effective_mode | u16 selected_version | u64 selected_capabilities | + u64 serialization_abi_id | bytes expected_topology_digest[32]) +``` + +Intent values are `ADMIT=1`, `CANCEL=2`, `COMPLETE=3`, and `REJECT=4`; zero is invalid. Before a key exists, the +authoritative GEN participant contributes the one nonzero proposed attempt and every other participant contributes zero; +the reduction rejects multiple distinct nonzero proposals and returns the agreed attempt to all ranks. Counts, record +lengths, key presence, enum ranges, ordering, duplicates, and exact consumption are validated before applying a record. +Internal `local_evidence=0` means `NONE` and is allowed only for admission; nonzero values match the control enum. +`open_publication_state` is `UNSET=0`, `MAYBE_OBSERVABLE=1`, `ACKED=2`, or `FAILED_INDETERMINATE=3`. + +The reduction applies this state-dependent precedence; implementations may not choose a different tie-breaker: + +| State at decision round | Contributions | Decision | +| --- | --- | --- | +| Before every edge is `READY_DECIDED` | all `ADMIT` and identical identity/mode/topology | Continue admission/ready preparation | +| Before every edge is `READY_DECIDED` | any `CANCEL` or `REJECT` | Prevent READY; `NOT_STARTED` only if every local edge/rank proves no commit, otherwise `UNKNOWN` | +| At or after any edge is `READY_DECIDED` | any `CANCEL` or `REJECT` without prior bundle completion | `UNKNOWN` | +| Any state | bundle-wide endpoint-local completion was committed on every required local edge/rank in an earlier decision round | Preserve role-specific local completion evidence and enter peer reduction; topology `COMPLETED` still requires matching remote evidence | +| Any state | `CANCEL` races `COMPLETE` in the same round, or completion coverage/order differs by rank | `UNKNOWN` | +| Any state | conflicting identity, topology, mode, version, or unrecognized intent | `UNKNOWN` | + +Thus cancellation never loses to a same-round ready/admit proposal. Local completion evidence is preserved only when +its bundle-wide evidence point is strictly earlier than the cancellation decision epoch; it does not become topology +completion without remote evidence. Tests cover same-round +`CANCEL+COMPLETE`, `CANCEL+ADMIT`, and `CANCEL+READY_PREPARED` on different ranks. + +Request cancellation is parent-scoped to `(ContextTicket, gen_attempt)`. Any authenticated child-edge `CANCEL` is +promoted to one parent intent and atomically marks every sibling edge `CANCEL_PENDING`. After the local decision, the +endpoint sends CANCEL, negative READY, or TERMINAL as applicable on every expected edge. A fault that is truly confined +to one edge still makes the parent `UNKNOWN`; version 1 never lets sibling edges continue to READY/DMA after one edge is +cancelled or indeterminate. + +Pre-`OPEN` local release requires `open_publication_state=UNSET` from every participant in the same decision round. If +any rank reports `MAYBE_OBSERVABLE`, `ACKED`, or `FAILED_INDETERMINATE`, the parent remains cancel-pending, waits for all +known notify outcomes, and follows installed `OPEN_SENT` cancellation or `UNKNOWN` semantics. One rank can never release +because only its local notify had not started. + +The reduction returns one decision per record. Before that decision commits, no rank may send a locally initiated +`CANCEL` or `TERMINAL`, invoke a backend cancellation primitive, complete a user promise, drop the request from +scheduling, or release a cancellation-related lease. Response/status workers only authenticate and enqueue peer intent +or immutable evidence. An authenticated peer `CANCEL` is queued immediately, but the executor publishes the response +after the corresponding local decision round. The ready prepare/install protocol below uses the same ordered round +framework and participant set, rather than adding a request-gated collective. + +GEN performs a pre-`OPEN` admission round before destination allocation or advertisement. That round agrees on the +`ContextTicket`, one shared `gen_attempt`, selected mode/version, expected topology and digest, and whether cancellation +is already pending. Only `ADMISSION_DECIDED` may create the parent receive lease and send `OPEN`. GEN's later aggregation +of CTX `READY` replies governs its local DMA/outcome and quarantine; it does not authorize CTX source DMA, which remains +guarded by CTX's own ready prepare/install decision. + +## Endpoint state machine + +```text +CTX: + +ABSENT + publish ContextTicket + -> PUBLISHED + +PUBLISHED + matching OPEN after CTX topology-local proposal decision + -> BOUND + authenticated OPEN proposal rejected by CTX topology-local decision + -> TERMINAL_UNKNOWN (remove authorization, cache READY=NOT_READY + LOCAL_UNKNOWN, fail/restart) + local CTX cancellation + -> topology-local decision -> TERMINAL_NOT_STARTED (cache and settle; no peer reduction) + unbound peer CANCEL + -> reject without mutation; READY=NOT_READY + LOCAL_UNKNOWN + +BOUND + local/peer CANCEL before ready intent is sealed + -> CANCEL_PENDING + topology-local cancel decision commits while still pre-ready + -> LOCAL_NOT_STARTED -> REDUCING (publish READY=NOT_READY and TERMINAL) + validate exact binding and prepare ready intent + -> READY_PREPARED + +READY_PREPARED + cancel arrives while rank decision is pending + -> CANCEL_PENDING (do not publish READY or decide evidence yet) + rank-wide prepare/install protocol rejects + -> LOCAL_NOT_STARTED -> REDUCING (publish READY=NOT_READY and TERMINAL) + rank-wide prepare/install protocol commits on every participant + -> READY_DECIDED + +READY_DECIDED + make READY=true visible, then submit/progress DMA + -> IN_FLIGHT + pending/new cancel, timeout, or transport uncertainty + -> CANCEL_PENDING -> topology-local decision -> LOCAL_UNKNOWN -> REDUCING + +IN_FLIGHT + source backend access ended + -> SUCCESS_COMPLETE (physical local cleanup; cache LOCAL_SOURCE_COMPLETE) + cancel, timeout, failure, conflicting evidence + -> CANCEL_PENDING -> topology-local decision -> LOCAL_UNKNOWN -> REDUCING + +REDUCING + enter with immutable evidence committed by the local decision and publish/retry TERMINAL + all expected peer evidence plus local-rank consensus + -> TERMINAL_NOT_STARTED | TERMINAL_COMPLETED | TERMINAL_UNKNOWN + +SUCCESS_COMPLETE + late CANCEL within retained evidence window + -> CANCEL_PENDING -> topology-local decision -> REDUCING (publish cached LOCAL_SOURCE_COMPLETE) + authoritative cleanup plus maximum accepted control lifetime + -> RETIRED + +TERMINAL_* + reply idempotently to duplicate messages + fix local disposition/cleanup, receive required receipts, and pass retention horizon + -> RETIRED + +RETIRED / ABSENT + stale OPEN, CANCEL, or TERMINAL + -> reject; report UNKNOWN; never reopen +``` + +GEN follows the corresponding states: + +```text +CONTEXT_IMPORTED + enqueue admission or local cancellation intent + -> ADMISSION_PREPARED + +ADMISSION_PREPARED + topology-local admission decision succeeds + -> ADMISSION_DECIDED + topology-local cancellation/rejection decision succeeds + -> LOCAL_CANCELLED (release unadvertised GEN resources; send no peer CANCEL) + +ADMISSION_DECIDED + create parent attempt lease, allocate destination buffers, and seal open_publish_committed + -> OPEN_PUBLISHING + local cancel wins before open_publish_committed + -> topology-local decision; all participants report publication UNSET -> LOCAL_CANCELLED + local cancel, but any participant may have published + -> CANCEL_PENDING + +OPEN_PUBLISHING + NIXL notification publication succeeds + -> OPEN_SENT + NIXL notification publication succeeds with cancellation pending + -> CANCEL_PENDING + cancel while publication may be observable + -> CANCEL_PENDING (wait for publication result; do not release) + notification result is failed or indeterminate + -> UNINSTALLED_UNKNOWN + +OPEN_SENT + local cancel wins the edge-state lock before any READY=true is accepted + -> CANCEL_PENDING (late READY cannot authorize DMA) + every edge returns exact cached no-peer READY=NOT_READY + LOCAL_NOT_STARTED + -> topology-local consensus -> TERMINAL_NOT_STARTED (no TERMINAL message required) + every installed/bound edge returns READY=NOT_READY + LOCAL_NOT_STARTED and matching TERMINAL + -> TERMINAL_NOT_STARTED + all expected peers READY=true, then local-rank binding/ready prepare/install succeeds + -> IN_FLIGHT + cancellation pending when any READY=true is accepted + -> LOCAL_UNKNOWN -> REDUCING + mixed, missing, conflicting, or invalid replies with any installed-edge proof + -> LOCAL_UNKNOWN -> REDUCING (send control only on proven installed edges) + missing/conflicting/invalid replies with no installed-edge proof + -> UNINSTALLED_UNKNOWN + +CANCEL_PENDING + every edge returns exact cached no-peer READY=NOT_READY + LOCAL_NOT_STARTED + -> topology-local consensus -> TERMINAL_NOT_STARTED (no TERMINAL message required) + every installed/bound edge returns READY=NOT_READY + LOCAL_NOT_STARTED and matching TERMINAL + -> TERMINAL_NOT_STARTED + any READY=READY + -> topology-local post-commit decision -> CANCEL_SENT + LOCAL_UNKNOWN -> REDUCING + READY=NOT_READY + LOCAL_UNKNOWN, missing/conflicting response, or deadline, with any installed-edge proof + -> LOCAL_UNKNOWN -> REDUCING (send control only on proven installed edges) + same outcomes with no installed-edge proof + -> UNINSTALLED_UNKNOWN + +CANCEL_SENT + retry CANCEL/TERMINAL while reducing peer evidence; exhaustion remains UNKNOWN + +IN_FLIGHT + local destination transfer and postprocessing complete + -> SUCCESS_COMPLETE (physical local cleanup; cache LOCAL_DESTINATION_COMPLETE) + cancel, timeout, failure, or uncertainty + -> CANCEL_PENDING -> topology-local decision -> LOCAL_UNKNOWN -> REDUCING + +REDUCING + enter with immutable evidence committed by the local decision + publish/retry CANCEL or TERMINAL only for edges with installation proof + all expected peer evidence plus local-rank consensus + -> TERMINAL_NOT_STARTED | TERMINAL_COMPLETED | TERMINAL_UNKNOWN + +UNINSTALLED_UNKNOWN + quarantine still-owned local leases, mark worker unhealthy, and begin bounded fail/restart + send no CANCEL or TERMINAL because peer installation is unproven + +SUCCESS_COMPLETE + late CANCEL within retained evidence window + -> CANCEL_PENDING -> topology-local decision -> REDUCING (publish cached LOCAL_DESTINATION_COMPLETE) + authoritative cleanup plus maximum accepted control lifetime + -> RETIRED +``` + +`LOCAL_CANCELLED` is an endpoint-local outcome, not a topology disposition. It cannot mutate the still-unbound CTX +publication. The service/orchestrator may cancel the CTX request independently; otherwise CTX's bounded publication +validity/expiry path performs the no-peer decision. + +`open_publish_committed` is set under the attempt lock before entering the NIXL notify API and is never cleared. The +payload and parent lease already exist at that point. Cancellation before the bit is set may take the local path after +the topology decision only when every participant reports `UNSET`; cancellation at or after any rank's bit follows +`OPEN_SENT`/indeterminate semantics for the entire parent. Version 1 treats notify failure as +indeterminate and `UNKNOWN` unless the transport provides an explicit proof that no peer could observe any bytes. + +### Ready-decision linearization + +No rank may make `READY=true` visible or start destination DMA until every topology-specific local participant has +installed the same irreversible decision for the logical ticket, generation binding, effective mode, selected protocol +version, expected-peer set, and ready intent. The response thread records input and wakes the executor; it never enters a +collective. + +The executor uses a bounded, divergence-safe two-phase protocol at its existing ordered status-poll lockstep point: + +1. Under the edge-state lock, each participant either seals `READY_PREPARED` or records rejection. A participant that + has contributed a prepared intent cannot independently roll it back; a concurrent cancellation becomes pending. +2. A topology-scoped reduction computes one prepare decision. A conflicting or rejected participant selects reject. If + the collective returns that decision, all prepared ranks can prove no READY was published and resolve pending + cancellation as `LOCAL_NOT_STARTED`. +3. On unanimous prepare, each participant installs `READY_DECIDED` under its state lock, then enters an ordered + installation barrier for that decision epoch. A bounded collective failure yields `UNKNOWN` and no rank publishes + READY. +4. Only after the installation barrier returns may a rank publish `READY=true` or authorize DMA. Cancellation pending + at or after `READY_DECIDED` is post-commit and yields `LOCAL_UNKNOWN` unless authoritative completion already exists. + +The implementation should piggyback or batch these phases with existing topology-aware status polling; it must not add +an independently ordered per-request collective in the response worker. Network, collectives, CUDA, promise completion, +and callbacks run outside registry locks. Deterministic tests inject cancellation before prepare, after prepare input, +after decision installation, after the installation barrier, and before/after READY publication. + +The current blocking `allgather`/`allgatherv` status path is not a bounded primitive and cannot turn a rank that never +enters into an in-process `UNKNOWN`. B3 must make every intent/ready decision round use a nonblocking/timeout-capable +topology collective or arm a process/communicator watchdog that terminates the affected local workers by the decision +deadline. A watchdog exit is the bound; the surviving service observes changed peer epochs and follows restart +handling. Qualification injects one rank skipping admission, cancellation, ready-prepare, ready-install, and completion +rounds and proves bounded termination with no asymmetric side effect or READY publication. + +## Multi-peer and local-rank reduction + +First reduce the two endpoint reports for every exact rank-pair key: + +```text +source LOCAL_NOT_STARTED + destination LOCAL_NOT_STARTED + -> EDGE_NOT_STARTED +source LOCAL_SOURCE_COMPLETE + destination LOCAL_DESTINATION_COMPLETE + -> EDGE_COMPLETED +anything else + -> EDGE_UNKNOWN +``` + +"Anything else" includes: + +- mixed not-started and source/destination-complete evidence; +- any explicit `LOCAL_UNKNOWN`; +- a missing reply at the bounded control deadline; +- conflicting transfer keys or process epochs; +- one peer reporting ready while another rejects; +- a semantically invalid message whose envelope, sender, and exact key were first validated; or +- a duplicate message with conflicting contents. + +An invalid frame, length, or CRC has no trustworthy key and must not mutate the edge named by untrusted bytes. It is a +connection-scoped protocol fault: stop admission to that sender, mark every active binding on the authenticated sender +connection `UNKNOWN` through the local-rank decision path, and begin bounded fail/restart. A valid frame with an +authenticated sender and exact key but conflicting semantics affects only that key. + +All expected edges must then agree: all `EDGE_NOT_STARTED` becomes topology-wide `NOT_STARTED`, all `EDGE_COMPLETED` +becomes topology-wide `COMPLETED`, and anything else becomes `UNKNOWN`. The peer-edge reduction and local-rank reduction +form one topology decision; neither intermediate result authorizes resource release. Release requires the same safe +disposition on every participating local rank. Any rank uncertainty makes the topology-wide result `UNKNOWN`. + +The parent `ContextEntry` aggregates child `TransferEdgeEntry` results; child progress is never stored in one shared +ready bit or evidence field. Partial edge progress therefore remains representable. If an implementation cannot prove +the exact per-edge state, it must collapse the entire topology to `UNKNOWN`, not infer a common state from one edge. + +## No-peer cancellation + +CTX cancellation can occur before GEN consumes the context result and before CTX knows a GEN endpoint. + +The required behavior is: + +1. CTX finds the exact active `ContextTicket` and starts the local-rank cancellation decision. +2. If every participating CTX rank proves the ticket is still `PUBLISHED` and uncommitted, the ranks atomically remove + transfer authorization and commit topology-wide `NOT_STARTED`. +3. A missing, bound, committed, or conflicting local rank instead commits `UNKNOWN` and triggers fail/restart. +4. CTX settles the local future without waiting for a GEN peer after the local-rank decision completes. +5. A bounded terminal cache remembers `NOT_STARTED` for an exact late ticket. +6. A late matching `OPEN` receives: + - `READY=NOT_READY + LOCAL_NOT_STARTED` while the terminal cache entry remains; + - `READY=NOT_READY + LOCAL_UNKNOWN` after that cache entry is evicted. +7. An active-registry miss without an exact terminal-cache hit, or any mismatched ticket, receives + `READY=NOT_READY + LOCAL_UNKNOWN`; it can never start DMA. + +No cross-service acknowledgement is possible when no peer ever exists. Memory safety therefore comes from the +active-registry allowlist and non-reusable attempt identity, not from retaining an unbounded request-ID tombstone. + +Terminal-cache age or capacity eviction is allowed only because the fallback is `UNKNOWN` plus rejection. TTL expiry +never authorizes transfer or memory reuse. A stale cache miss does not quarantine or release an unrelated current slot: +the receiver has no active exact lease generation to mutate, so it rejects and reports `UNKNOWN` without changing its +health. The requester that still owns the exact unproven lease quarantines it and enters bounded restart. The receiver +restarts only when it has a matching active lease or the fault is connection-scoped corruption. + +## Retry, deduplication, and reconnect + +### Deduplication + +- Every new control message has a nonzero `message_id`, strictly increasing within + `(sender_instance_epoch, TransferKey)`; a byte-identical retransmission reuses its original ID. +- The receiver keeps a bounded window of `(message_id, payload_sha256, reply)` entries plus the highest accepted ID and + current terminal reply for that scope. +- Repeated `OPEN` for the same exact key returns the same ready/terminal result and never allocates or submits twice. +- Repeated `CANCEL` returns the current terminal result. +- Repeated `TERMINAL` returns a `TERMINAL_RECEIPT` that names the exact key, terminal message ID, sender-local evidence, + and terminal-payload digest. The receipt does not assert a topology disposition. +- The same in-window message ID with different content is a protocol violation and contributes `LOCAL_UNKNOWN`. +- An ID at or below the evicted-window floor never mutates state. A terminal scope returns its cached terminal reply; a + nonterminal scope returns `LOCAL_UNKNOWN`. It does not claim to compare content that is no longer retained. + +### Retry + +- `OPEN`, `CANCEL`, and `TERMINAL` use bounded retry with backoff. A retransmission reuses the byte-identical payload, + `message_id`, `gen_attempt`, and `TransferKey`; creating a new message ID is a new protocol message, not a retry. +- The retry deadline is no longer than the configured transfer/control deadline. +- Exhausting `OPEN` or any cancellation/ambiguity exchange contributes `LOCAL_UNKNOWN` and therefore commits + topology-wide `UNKNOWN`; it never assumes cancellation succeeded. Ordinary successful completion does not start a + terminal exchange and is not made unavailable by loss of the control channel. +- Promise completion and resource release must not depend on an unbounded control retry loop. +- The OpenAI client's regenerated `disagg_request_id` is worker-request identity only. In active mode, + `ContextPhaseParams.req_id` remains the original `ctx_request_id`. A blind generation-only `_post_with_retry` may + select a different worker or buffer binding and therefore cannot rebind the old ticket. It must either prove that the + first attempt never reached `OPEN` and reuse the exact binding, or rerun context to obtain a new ticket. The current + retry path needs an end-to-end guard and test before active mode is supported. + +### Reconnect and restart + +- Reconnection with unchanged process epochs may replay an idempotent message. +- A changed source or destination epoch invalidates every old `TransferKey`. +- In-flight work is not resumed across a process-epoch change. +- The surviving endpoint detects the epoch change, transitions every active binding to the old peer epoch to `UNKNOWN`, + quarantines its exact local leases, stops admission, and enters bounded fail/restart. A stale message never creates a + new entry. +- A replacement process starts with an empty active registry, rejects old keys, and does not attempt to contribute + evidence for state it never owned. Recovery requires a new context result and new attempts. +- State reconstruction across process restart is out of scope. + +## Serialization constraints + +### Context ticket trailer + +The context ticket crosses the service boundary in the existing length-delimited `ContextPhaseParams.opaque_state`. +The legacy `DataTransceiverState` bytes remain an exact prefix: + +```text +[legacy DataTransceiverState bytes] +[versioned TLV payload] +[fixed trailer footer] +``` + +All trailer integers use unsigned network byte order. The fixed 20-byte footer is: + +```text +TrailerFooter { + byte magic[8]; // ASCII "TRTLCTX1" + uint16 schema_version; // 1 + uint16 flags; // must be zero in version 1 + uint32 payload_length; // at most 4096 bytes + uint32 payload_crc32c; // CRC-32C (Castagnoli) over TLV payload +} +``` + +Each payload field uses an 8-byte TLV header: network-order `uint16 type`, `uint16 flags`, and `uint32 length`, followed +by exactly `length` bytes. TLV flag bit 0 means required; all other version-1 bits are zero. Version-1 type IDs are +`CTX_REQUEST_ID=1`, `CTX_INSTANCE_EPOCH=2`, `CTX_ATTEMPT=3`, `DIRECTION=4`, `REQUIRED_CAPABILITIES=5`, and +`SERIALIZATION_ABI_ID=6`; each appears exactly once and is required. Integer values are network-order 64-bit values, +except `DIRECTION`, which is a one-byte enum. Schema version 1 requires ABI ID 1. Duplicate known fields, unknown +required fields, bad scalar lengths, nonzero reserved bits, trailing payload bytes, or a CRC mismatch reject active mode +before registry mutation. Unknown optional TLVs are skipped. + +The payload contains bounded TLVs for: + +- full `ctx_request_id`; +- `ctx_instance_epoch`; +- `ctx_attempt`; +- direction; +- required protocol capability bits; and +- serialization ABI ID. + +Rules: + +- The footer is located from the end of the opaque byte vector. +- An active `opaque_state` is at most 16 MiB in total, including the legacy prefix, trailer payload, and footer. The + bound is checked before parsing or allocation. +- `payload_length` is bounds-checked before allocation or parsing. +- The complete TLV payload is at most 4 KiB. +- Unknown optional TLVs are skipped. +- Unknown required TLVs reject active mode. +- The request ID in the trailer must equal `ContextPhaseParams.req_id`. The active C++ path constructs this field from + `ctx_request_id`, even when Python assigns a different `disagg_request_id` to a retried worker request. +- The network-order trailer and ABI ID are validated from the end before structurally parsing the native legacy prefix. + A mismatched or unknown ABI is rejected without interpreting that prefix. +- With a ticket present, `setReqId` accepts only the same value; C++ and Nanobind mutation APIs cannot create a second + in-memory identity. +- A malformed or truncated trailer is never treated as an active ticket. +- A legacy opaque state with no trailer remains valid only for baseline mode; enabled mode rejects it before buffer + allocation. + +The generic positional serialization of `DataTransceiverState` must not change. + +### Selected preservation strategy + +`ContextPhaseParams` owns a private typed `optional` alongside its typed `DataTransceiverState`. + +- The `vector` constructor applies the total bound, locates and validates the fixed footer from the end, and runs a + no-allocation checked structural preflight over the complete legacy `DataTransceiverState` prefix. Only after nested + counts, lengths, checked arithmetic, and exact prefix consumption pass may it deserialize the legacy state and store + the parsed ticket separately. +- `getSerializedState()` reserializes the typed legacy state and appends the canonical TLV payload/footer when a ticket + is present. +- Copy, move, equality, Nanobind pickle, and `opaque_state` access preserve the typed ticket explicitly; they do not + depend on unparsed trailing bytes. +- CTX result construction moves the complete typed `ContextPhaseParams` through an internal factory that updates first + tokens and optional draft tokens while preserving the legacy state and ticket. Reconstructing a new object from + `releaseState()` alone is forbidden because the current result path would drop the sidecar ticket. +- The class-layout change requires the normal C++ API/ABI review before implementation. No public constructor argument + changes are required. + +Generic C++ `ContextPhaseParams` serialization is currently positional and embeds state inside larger `Request` and +`Result` objects, so an unmarked trailer would corrupt following fields. Version 1 therefore selects an active-only +tagged state encoding after the existing `hasState=true` byte: + +```text +ACTIVE_STATE_ENCODING { + byte tag = 0xD7; // outside canonical legacy optional values 0 and 1 + byte magic[8] = "TRTLCTX1"; + uint32 opaque_state_length; // network order; at most 16 MiB in version 1 + bytes opaque_state; // legacy state prefix plus validated ticket trailer +} +``` + +A new decoder peeks the next byte: canonical `0` or `1` selects the byte-identical legacy `DataTransceiverState` path; +`0xD7` selects the bounded active envelope; every other value is malformed. Baseline serialization emits the exact old +bytes. Active encoding is deliberately unreadable by an old worker and may be emitted only after a local-topology +capability check proves every serializer/deserializer is new. If that proof is unavailable, enabled mode is rejected at +request admission. Frozen old `ContextPhaseParams`, `Request`, and `Result` fixtures plus active golden fixtures are +required. + +The implementation must preserve the opaque trailer through: + +- CTX result construction; +- Nanobind copy, move, pickle, and `opaque_state` access; +- base64 encode/decode in the OpenAI disaggregated service; +- generation request construction; +- worker/RPC/MPI handoff used by the qualified executor path. + +It must not assume that the current `ContextPhaseParams(std::vector)` constructor preserves trailing bytes; +explicit storage and round-trip tests are required. A dedicated `createResult()` round trip must begin with a typed +ticket and compare the returned opaque trailer byte-for-byte. If a qualified handoff cannot carry a length-delimited +opaque extension, the feature must reject that configuration rather than modify a generic positional wire layout +silently. + +The OpenAI request layer enforces the bound before C++: `encoded_opaque_state` is ASCII base64 with at most 22,369,624 +characters (the encoding ceiling for 16 MiB), uses strict alphabet/padding validation, and is decoded only with strict +validation enabled. The decoded value is then checked at 16 MiB before constructing `ContextPhaseParams`. Malformed or +oversized input is rejected before allocating an unbounded decoded buffer. + +Malformed-prefix fuzzing uses an otherwise valid ticket/footer around truncated scalars and oversized nested +string/vector counts. Every case must reject before allocation and without falling back to baseline mode. + +### Control messages + +- New control payloads are length-delimited and versioned. +- Existing positional `RequestInfo` serialization is unchanged. +- Existing baseline notification variants and indices are unchanged; the active-only control variant is appended. +- Unknown message types or required fields are rejected before state mutation. +- Decoders validate the authenticated sender, lengths, effective backend limit, enum ranges, nested counts, duplicate + fields, checked arithmetic, and exact payload consumption before allocation or mutation. +- Frozen legacy byte fixtures protect old baseline decoding. + +## Tombstone and terminal-cache retirement + +The registry separates authorization from duplicate-response metadata: + +`maximum_accepted_control_lifetime` is the configured maximum control retry deadline plus the maximum notification queue +delay and allowed same-epoch reconnect replay delay. A receiver emits `TERMINAL_RECEIPT` only after atomically storing +the canonical terminal bytes, sender-local evidence, and dedupe fingerprint. + +- The active registry is an allowlist. +- Terminal entries cannot authorize DMA. +- The terminal cache may be bounded by count and age. +- A cache hit may reproduce `NOT_STARTED` or `COMPLETED`. +- A cache miss returns `UNKNOWN` and rejects the operation. +- Request-ID reuse is safe because a new attempt has a different ticket. +- Known-peer terminal metadata may retire only after the local topology disposition and cleanup action are fixed, every + expected peer returns a matching `TERMINAL_RECEIPT`, and the maximum accepted control lifetime has elapsed since the + last send. Receipt alone is insufficient. +- Normal-success evidence that never needed a terminal exchange may retire only after its configured duplicate/control + horizon, which is at least the maximum accepted control lifetime, and authoritative local cleanup; a later stale + message is rejected and cannot touch a newer slot generation. +- No-peer metadata may retire by bounded eviction because an unknown late ticket is default-denied. +- Eviction never releases a poisoned transfer buffer. +- `UNKNOWN` buffer ownership remains governed by the process fail/restart policy. + +## Performance requirements + +The protocol is not qualified merely because it is safe. Its no-fault path must preserve the current completion fast +path and avoid a terminal round trip. Rank agreement must reuse or batch with the existing topology-aware status-poll +sequence. Response workers do not launch collectives; one unconditional bounded batch per existing executor decision +iteration replaces request-gated/per-request control collectives. + +Qualification reports, separately for payload counts 0, 1, 128, and 1024 and for local NVLink and cross-node IB: + +- number, payload bytes, and p50/p95/p99 time of added local-rank collective operations per transfer phase; +- control-message encode/decode, queue, retry, and end-to-end latency; +- notification size at and above the effective carrier limit; +- throughput, TTFT, and TPOT p50/p95/p99 with cancellation disabled, enabled without faults, and under cancellation; and +- active-registry, terminal-cache, and quarantined-capacity growth under a marathon run. + +Initial no-fault guardrails are at most 3% throughput regression and at most 5% TTFT/TPOT p99 regression relative to +the same C++ NIXL configuration with the feature off. Any exception requires a measured, reviewed budget before +default-on rollout. The ready-decision collective count and latency are blocking metrics, not diagnostic-only counters. + +## Observability + +Expose per-worker counters or gauges for: + +- active context tickets; +- bound and in-flight transfers; +- terminal-cache entries and oldest age; +- terminal-cache capacity/age evictions; +- sent, received, retried, and duplicate control messages; +- `NOT_STARTED`, `COMPLETED`, and `UNKNOWN` outcomes; +- rejected unknown tickets; +- stale process epochs; +- attempt and direction mismatches; +- malformed envelopes and trailers; +- control retry exhaustion; +- mixed or missing peer dispositions; +- poisoned buffers and fail/restart escalation. + +Logs should include the request ID, attempts, endpoint epochs, direction, peer rank, and one stable reason code. Repeated +retries must be rate-limited. Do not log opaque payload bytes. + +Required alerts for default-on qualification include: + +- nonzero or increasing `UNKNOWN` rate; +- terminal-cache exhaustion; +- oldest active/terminal age above the configured bound; +- repeated stale-epoch traffic; +- poisoned pool capacity; +- fail/restart events. + +## Rollout and rollback + +1. `OFF`: land this design scaffold and implementation pieces with no active control behavior. +2. `OBSERVE_ONLY`: create identities and record would-be decisions without changing existing transfer or cleanup. +3. `OPT_IN_FAIL_CLOSED`: enable the default-deny registry, typed ticket, exact lease generations, and bounded + fail/restart behavior for explicit homogeneous C++ NIXL deployments. +4. `OPT_IN_PROTOCOL`: enable the negotiated control carrier and lifecycle after deterministic TP=2 fault tests, + serialization fuzzing, and performance qualification pass. +5. Canary 1%, 10%, and 50% of exact supported deployment cells, with rollback at each step on `UNKNOWN`, restart, + latency, throughput, or capacity thresholds. +6. Run the disaggregated stress suite and full post-merge-stage CI at every rollout-changing commit. +7. Make #15738 default-on only for an exact qualified cell after this protocol or an explicitly approved bounded + fail/restart alternative satisfies the same safety requirements. Unsupported cells remain baseline. + +The existing environment flag remains the kill switch. Rollback stops new active-mode sessions first. Existing transfer +keys finish under the protocol that created them; rollback must not reinterpret an `UNKNOWN` transfer as safe. If safe +drain cannot be proven, restart. + +Baseline mode and unsupported transports remain unchanged. + +Rollback tests must cover active sessions, a changed peer epoch, control loss, and mixed-version routing. The kill +switch stops only new active admissions; an already active key finishes under the version and mode frozen at admission. + +## Test matrix + +| Area | Case | Required result | +| --- | --- | --- | +| Legacy compatibility | Legacy opaque state, baseline mode | Existing path unchanged | +| Legacy rejection | Missing trailer, enabled mode | Reject before allocation/advertisement | +| Version selection | Different overlapping ranges across local ranks and peers | Every rank freezes the same highest common control version | +| Serialization ABI | Peer omits or mismatches ABI ID 1 | Active proposal rejected; old peer receives no control variant | +| Dynamic precheck | Advertised CTX descriptor has incompatible mode/version/ABI | GEN rejects before allocation/`OPEN`; no silent fallback | +| Dynamic proposal | CTX recomputes a different topology/version from `OPEN` | After bounded advertisement but before READY/DMA, retire ticket as `UNKNOWN` and quarantine/restart exact leases | +| Rejected proposal replay | Duplicate rejected `OPEN` or different attempt for its ticket | Cached negative READY for duplicate; new attempt denied; authorization never reopens | +| Active dispatch | Bare or reordered legacy `RequestAndBufferInfo` after active negotiation | Reject; only atomic `OPEN` can authorize work | +| Notification carrier | Control mixed with request, ready, and sync variants | Stable variant indices; correct sender-preserving dispatch | +| Complete `OPEN` | Agent, address, request, descriptors, metadata, connection index, and kinds | Exact bounded round trip and binding validation | +| Bounded decode | Tiny payload with huge nested string/vector/count | Reject before reserve, allocation, connection, or registry mutation | +| Trailer | New round trip through Nanobind and HTTP base64 | Exact ticket preserved | +| Result construction | Typed ticket passes through `createResult()` | Returned opaque trailer preserved byte-for-byte | +| Ticket mutation | C++ or Nanobind `setReqId` changes a ticketed object | Different ID rejected; identical ID remains valid | +| HTTP retry | GEN request retries with a regenerated `disagg_request_id` | `ctx_request_id` stays stable; exact unstarted binding is replayed or context is rerun with a new ticket | +| Trailer | Truncated, oversized, bad CRC, unknown required TLV | Reject without registry mutation | +| Trailer ABI | Valid footer with unknown/mismatched serialization ABI | Reject before parsing the native legacy prefix | +| HTTP opaque state | Oversized or malformed base64 input | Reject by encoded/decoded bound and strict validation before C++ construction | +| Fixed serialization | Frozen baseline `ContextPhaseParams`, `Request`, `Result`, `DataTransceiverState`, and `RequestInfo` fixtures | Byte-identical | +| Active serialization | Tagged active state plus every control message | Golden portable-field bytes, frozen ABI-1 legacy blob, and bounded decode | +| Identity | Same request ID, new CTX attempt | Old messages cannot bind | +| Identity | Same CTX ticket, new GEN attempt | Conflict becomes `UNKNOWN` | +| Identity | Source/destination epoch mismatch | Stale reject and `UNKNOWN` | +| Identity | Direction or rank mismatch | Reject without DMA | +| Identity | Same rank numbers in a different `CommState` topology | Topology digest mismatch; reject without mutation | +| No peer | Cancel before `OPEN` | Rank-consistent local completion and cached `NOT_STARTED` | +| No peer | Late `OPEN` before cache eviction | `READY=NOT_READY + LOCAL_NOT_STARTED`; no TERMINAL required | +| No peer | Late `OPEN` after cache eviction | `READY=NOT_READY + LOCAL_UNKNOWN` | +| GEN admission | One local rank cancels or proposes a different attempt before `OPEN` | No destination allocation/advertisement; rank-consistent local rejection | +| OPEN publication | Cancel/failure immediately before, during, and after notify | Pre-commit local cancel only; possibly observable/indeterminate publication retains or quarantines lease | +| OPEN rank race | One rank is `UNSET` while a sibling is `MAYBE_OBSERVABLE` | Parent cannot take local-cancel release; wait/resolve or become `UNKNOWN` | +| Uninstalled timeout | Dropped `OPEN`/READY or proposal rejection while cancellation is pending | Local quarantine and bounded restart; no CANCEL/TERMINAL sent without installation proof | +| Intent batch | Empty rank contribution or one rank skips a decision round | No asymmetric side effect; bounded collective failure or watchdog exit | +| Pre-commit | Cancel after `OPEN`, before ready commit | Topology-wide `NOT_STARTED`; reserved source/destination leases releasable | +| Negative READY | Bound pre-commit rejection with dropped/reordered READY or TERMINAL | `READY=NOT_READY` carries `LOCAL_NOT_STARTED`; retry stays idempotent and bounded | +| Ready prepare | Cancel after one rank contributes ready intent | Pending until common decision; never asymmetric `NOT_STARTED`/READY | +| Ready install | Rank timeout before installation barrier completes | No READY is published; topology-wide `UNKNOWN` | +| Commit race | Cancel races with `ready=true` | Exactly one linearized result; post-commit uncertainty is `UNKNOWN` | +| Reorder | `CANCEL` crosses delayed or duplicate `READY=true` | Cancellation winner blocks DMA; committed READY makes result `UNKNOWN` | +| Completion race | Cancel races with successful completion | `COMPLETED` only with authoritative all-peer success; otherwise `UNKNOWN` | +| Terminal progress | Both endpoints cancel simultaneously | Each publishes local TERMINAL without waiting; no circular wait | +| Success fast path | Control channel drops after uncontended completion | Existing authoritative cleanup completes; no mandatory terminal round trip | +| Duplicate | Repeated `OPEN`, `CANCEL`, `TERMINAL`, receipt | Idempotent; no double allocation, response, or release | +| Conflict | Same message ID, different payload | Protocol error and `UNKNOWN` | +| Dedupe eviction | Old ID below retained fingerprint window | No mutation; cached terminal or `UNKNOWN` response | +| Retry | Dropped control replies | Bounded retry; exhaustion becomes `UNKNOWN` | +| Reconnect | Same epoch reconnect | Idempotent replay | +| Restart | Changed peer epoch | No resume; stale reject and `UNKNOWN` | +| Restart ownership | `UNKNOWN` does not drain within configured bound | Admission stops, leases remain isolated, process exits or is terminated | +| Multi-peer | Both endpoints report `LOCAL_NOT_STARTED` on every edge | `NOT_STARTED` | +| Multi-peer | Source-complete plus destination-complete on every edge | `COMPLETED` | +| Multi-peer | Mixed, missing, or conflicting peers | `UNKNOWN` | +| Multi-peer cancel | One child edge receives cancel before ready | Parent marks every sibling pending; no edge continues to READY/DMA | +| Intent race | CANCEL vs COMPLETE/ADMIT/READY_PREPARED in one round | Normative precedence table; same-round completion race is `UNKNOWN` | +| TP=2 | One local rank sees mismatch/timeout | Rank-consistent `UNKNOWN`; no asymmetric release | +| TP=2 binding | Local ranks bind different GEN attempt/epoch | No `READY=true`; topology-wide `UNKNOWN` | +| Resource safety | `UNKNOWN` after advertised buffer | Still-owned, unproven local leases remain quarantined; authoritatively released generations are untouched; fail/restart | +| ABA | Delayed terminal arrives after slot reuse | Old slot generation cannot release or quarantine the new lease | +| Source safety | Cancel after NIXL submission while source may still be read | Send staging slot is not reused or unpoisoned | +| Bundle completion | KV completes while indexer or RNN remains active | No completion claim; entire parent lease becomes `UNKNOWN`/quarantined on deadline | +| Teardown | Shutdown with active, completed, and unknown leases | Workers join and memory deregisters before agent/plugin destruction; no silent reuse | +| Capacity | Repeated no-peer cancellations | Registry/cache remain bounded; late work never reopens | +| Active capacity | Active registry limit or deadline | Admission backpressure or explicit terminal transition; no silent eviction | +| Default-off | Existing synchronous/asynchronous suites | No behavior regression | +| Initial cell | C++ NIXL-over-UCX, preallocated, timeout>0, overlap, PP=CP=1, no attention-DP/layerwise/zero-copy | Active protocol may be admitted after all other gates pass | +| Unsupported path | Python/direct-UCX/MPI/Mooncake/libfabric or any initial-cell gate mismatch | No new messages; existing baseline policy preserved | +| GPU integration | Deterministic before-ready and after-ready races on NIXL | Expected disposition, no hang, no unsafe reuse | +| Stress | B200/B300 cancellation and restart campaign | Bounded metadata, exact counters, no silent reuse | +| Performance | No-fault active mode across payload/topology matrix | Throughput regression <=3%; TTFT/TPOT p99 regression <=5% | + +## Implementation phases + +### Phase B0 - Design scaffold + +- Land this document. +- Review and approve identity, disposition, ownership, and serialization decisions. +- Add no runtime behavior. + +### Phase B1 - Identity and registry + +- Add `ContextTicket` and `TransferKey` value types. +- Add rank-consistent CTX and GEN attempt generation. +- Add typed `ContextPhaseParams` ticket storage, the opaque-state trailer, active-only tagged internal serialization, + frozen baseline fixtures, and round-trip tests. +- Add the active-context allowlist and bounded terminal cache. +- Add endpoint transfer leases, buffer slot generations, cleanup-once ownership, and no-peer publication pinning. +- Reject unknown or stale tickets before `ready=true`. +- Keep cancellation default-off. + +### Phase B2 - Control protocol + +- Extend #15798 with capability bit 2, the v2 descriptor schema, serialization ABI ID 1, deterministic selected + version, and immutable per-session compatibility result. +- Add the atomic length-delimited `OPEN`, `READY`, `CANCEL`, `TERMINAL`, and `TERMINAL_RECEIPT` envelopes and golden + byte fixtures. +- Append the active-only control notification variant and add checked bounded decoding for complete + `RequestAndBufferInfo` payloads. +- Implement exact-key deduplication and bounded retry. +- Preserve the legacy baseline path. + +### Phase B3 - Lifecycle integration + +- Connect pre-commit cancellation to `LOCAL_NOT_STARTED` evidence. +- Connect source-backend completion to `LOCAL_SOURCE_COMPLETE` and destination postprocessing to + `LOCAL_DESTINATION_COMPLETE` evidence. +- Map every post-commit ambiguity to `LOCAL_UNKNOWN` evidence. +- Add the ready prepare/decision/installation linearization protocol without running collectives in response threads. +- Add all-peer and local-rank reduction. +- Retire old request-ID-only tombstones after the allowlist is authoritative. +- Integrate source/destination lease quarantine, poison, and bounded fail/restart behavior. + +### Phase B4 - Qualification + +- Add deterministic TP=2 NIXL race/fault tests. +- Run synchronous and asynchronous disaggregated integration suites. +- Run the original gen-only reproducer and cancellation stress suite on B200/B300. +- Run full pre-merge and post-merge-stage CI. +- Validate counters, alert thresholds, carrier limits, and performance guardrails under load. + +## Draft blockers + +The draft must remain non-runtime until these are resolved: + +- prove rank-consistent CTX and GEN attempt generation plus binding/ready commitment without adding a divergent + collective; +- replace or watchdog-bound the current blocking status collectives and prove unconditional empty participation; +- complete C++ API/ABI review for typed `ContextPhaseParams` ticket storage; +- prove the opaque trailer and active-only tagged state survive every qualified service and worker handoff while all + baseline `ContextPhaseParams`, `Request`, and `Result` bytes remain unchanged; +- prove local-topology capability before any active tagged state is sent to a worker; +- prove process/publication epoch generation does not reuse an identity after restart; +- define the control-message dispatch thread and lock ordering; +- implement and fuzz a checked bounded decoder for every nested `RequestAndBufferInfo` field; +- validate the selected protocol version and active-only control carrier across fixed and dynamically routed peers; +- freeze and test descriptor capability bit 2, the v2 marker, serialization ABI ID 1, and proposal-scoped rejection; +- keep network, CUDA, promise, and user callbacks outside registry locks; +- define connection teardown and same-epoch reconnect ownership; +- prove the exact ready-commit linearization point; +- prove the exact existing events that justify `LOCAL_SOURCE_COMPLETE` and `LOCAL_DESTINATION_COMPLETE` independently; +- keep every cancelled post-commit path at `UNKNOWN` until a backend quiescence contract exists; +- add deterministic TP=2 NIXL cancel/complete/ready race tests; +- demonstrate active-registry admission backpressure, explicit expiry transitions, and bounded terminal-cache behavior; +- derive and enforce the maximum accepted control-message lifetime and mandatory retention horizon; +- define and test the process fail/restart trigger for unresolved `UNKNOWN` state; +- demonstrate normal successful cleanup remains independent of the terminal-control channel; +- reject every configuration outside the exact initial supported cell before active admission; +- meet the no-fault throughput, TTFT, and TPOT guardrails. + +## Non-goals + +This phase does not: + +- add a synchronous CTX-GEN cancellation vote; +- claim `QUIESCED` after NIXL cancellation; +- unpoison or reuse a buffer after ambiguous mid-transfer cancellation; +- add a reusable cross-transport recovery framework beyond the endpoint leases required for this C++ NIXL cell; +- support Python transceiver, direct UCX, MPI, Mooncake, or NIXL-libfabric; +- resume an in-flight transfer across process restart; +- change default-on policy; +- change generic positional `DataTransceiverState` or `RequestInfo` encoding. + +## Alternatives considered + +- **Request-ID tombstones with a TTL:** rejected because request IDs are reusable and expiry is not proof of quiescence. +- **Backend-handle release or timeout as cancellation completion:** rejected because neither proves that source and + destination memory are no longer accessed. +- **A synchronous cross-service cancellation vote:** rejected because it cannot handle no-peer publication cancellation + and would put network latency on every cancel path. +- **A raw envelope or independently reordered sidecar notification:** rejected because current notification receivers + unconditionally decode the legacy variant and because reordering can separate authorization from the buffer payload. +- **Appending fields to generic positional serialization:** rejected because it would corrupt older nested + `Request`/`Result` decoders and change the baseline wire path. +- **Mandatory terminal exchange for every successful transfer:** rejected to preserve current availability and latency; + immutable completion evidence is retained and exchanged only if cancellation or ambiguity requires it. +- **Bounded fail/restart without peer protocol:** remains an acceptable #15738 dependency alternative if it provides the + same exact lease quarantine, admission stop, and bounded supervisor-restart guarantees. + +## Implementation ownership + +| Concern | Owning component | Phase | +| --- | --- | --- | +| Descriptor capability and selected version | `peerProtocol` / connection admission | B2 | +| Context ticket and result preservation | `ContextPhaseParams`, `LlmRequest::createResult`, Python disaggregated params | B1 | +| Active registry and per-edge lifecycle | C++ cache transceiver / executor state owner | B1-B3 | +| Notification carrier and bounded decode | NIXL `AgentConnectionManager` | B2 | +| Ready prepare/install agreement | Executor topology status-poll path | B3 | +| Lease generation, quarantine, and cleanup-once | Transfer-buffer managers and request lifecycle | B1-B3 | +| Unhealthy-worker exit and supervisor restart | Executor/server health owner and deployment supervisor | B3-B4 | +| Fault, integration, stress, and performance qualification | Disaggregated-serving CI suites | B4 | + +## Dependency graph + +Arrows point from prerequisite to dependent. #15798 and #15799 are default-on/production-hardening follow-ups; they do +not block the default-off merges of #15238, #15794, or #15795. + +```mermaid +flowchart LR + PR15238["#15238
gated cancellation
(default off)"] + PR15737["#15737
CacheSender lost-wakeup fix"] + PR15794["#15794
C++ protocol and buffer safety
(default off)"] + PR15795["#15795
PyExecutor lifecycle safety
(default off)"] + PRA["#15798
CTX-GEN protocol/mode negotiation"] + PRB["#15799
peer cancel/terminal protocol
(design scaffold)"] + PR15738["#15738
qualified default-on policy"] + + PR15238 --> PR15794 + PR15737 --> PR15794 + PR15794 --> PR15795 + PR15794 --> PRA + PRA --> PRB + PR15795 --> PR15738 + PRA --> PR15738 + PRB -.->|or bounded fail/restart policy| PR15738 +``` diff --git a/docs/source/index.rst b/docs/source/index.rst index 9d92a8770148..3dacdd8bb2a4 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -98,6 +98,7 @@ Welcome to TensorRT LLM's Documentation! developer-guide/dev-containers.md developer-guide/api-change.md developer-guide/kv-transfer.md + developer-guide/disagg-peer-cancellation-control-protocol.md developer-guide/telemetry.md developer-guide/sparse-attention-development-guide.md diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index b272740deab4..561aa4169130 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -1282,6 +1282,7 @@ def create_autodeploy_executor( attention_type_cpp, cache_transceiver_config, mamba_cache_manager=None, + inflight_cancel_supported_by_executor=False, ) # Guided (structured) decoding. diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index ac0feab0347d..4ec94412e076 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import copy import dataclasses import os @@ -2193,8 +2208,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..9963dca7a0bc 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import concurrent import contextlib import functools @@ -54,6 +69,8 @@ def result(self): DEFAULT_SERVER_WAITING_TIMEOUT = 2100 # Timeout for the accuracy evaluation DEFAULT_ACC_EVALUATION_TIMEOUT = 1500 +# Timeout for each request sent to the disaggregated server +DEFAULT_REQUEST_TIMEOUT_S = 1800000 @functools.lru_cache(maxsize=1) @@ -152,6 +169,8 @@ def launch_disaggregated_llm( enable_perf=False, extra_env: Optional[Dict[str, str]] = None, gen_extra_env: Optional[Dict[str, str]] = None, + request_timeout_s: float = DEFAULT_REQUEST_TIMEOUT_S, + request_max_retries: Optional[int] = None, ): temp_dir = tempfile.TemporaryDirectory() disaggregated_serving_config_path = os.path.join( @@ -405,9 +424,14 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): f"Server is not ready after {server_waiting_timeout} seconds. Please check the logs for more details." ) - client = openai.OpenAI(api_key="1234567890", - base_url=f"http://localhost:{serve_port}/v1", - timeout=1800000) + client_kwargs: Dict[str, Any] = { + "api_key": "1234567890", + "base_url": f"http://localhost:{serve_port}/v1", + "timeout": request_timeout_s, + } + if request_max_retries is not None: + client_kwargs["max_retries"] = request_max_retries + client = openai.OpenAI(**client_kwargs) def send_request(prompt: str, sampling_params: SamplingParams, streaming: bool): @@ -1077,17 +1101,18 @@ def test_nixl_backend(self): @pytest.mark.skip_less_device_memory(60000) @skip_no_hopper def test_gen_only_sync(self): - """Test gen-only synchronous KV transfer path with NIXL Python transceiver. + """Test gen-only synchronous KV transfer path with NIXL C++ transceiver. Sets TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1 so the gen worker calls - request_and_receive_sync instead of the async path. Accuracy must be - identical to the standard async path. + the blocking request_and_receive_sync path. A bounded client timeout + makes a stuck transfer fail this test instead of waiting for its outer + one-hour timeout. """ ctx_server_config = { "disable_overlap_scheduler": True, "cache_transceiver_config": { "backend": "NIXL", - "transceiver_runtime": "PYTHON", + "transceiver_runtime": "CPP", "max_tokens_in_buffer": 4096, }, } @@ -1095,7 +1120,7 @@ def test_gen_only_sync(self): "disable_overlap_scheduler": True, "cache_transceiver_config": { "backend": "NIXL", - "transceiver_runtime": "PYTHON", + "transceiver_runtime": "CPP", "max_tokens_in_buffer": 4096, }, } @@ -1116,6 +1141,8 @@ def test_gen_only_sync(self): self.MODEL_PATH, # Apply to both servers: gen worker uses sync receive path. extra_env={"TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP": "1"}, + request_timeout_s=120, + request_max_retries=0, ) as llm: run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a958f8b24951..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 2bd377836f93..7c06c07c042b 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -17,6 +17,7 @@ l0_dgx_b200: tests: - unittest/_torch/misc/test_autotuner.py::test_autotuner_distributed_strategy - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp2-TRTLLM] + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_only_sync # ------------- KV Cache V2 Scheduler IT (multi-GPU) --------------- - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2DSv3Lite::test_mtp_draft_tokens - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2DSv3Lite::test_mtp_chunked_draft_tokens diff --git a/tests/integration/test_lists/test-db/l0_sanity_check.yml b/tests/integration/test_lists/test-db/l0_sanity_check.yml index bc947305f276..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/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index e6931ca348d9..2c996c1f6bcb 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -378,10 +378,6 @@ perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_kimi-k25-thinking-fp4 perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6368078) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_128k8k_con64_ctx1_pp8_gen1_dep32_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6302903) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_1k1k_con1024_ctx1_dep4_gen1_dep32_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6302903) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_1k1k_con1024_ctx1_dep4_gen1_dep8_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6323889) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_1k1k_con2048_ctx2_dep4_gen1_dep16_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6323889) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_1k1k_con2048_ctx2_dep4_gen1_dep16_eplb288_mtp3_ccb-NIXL] SKIP (https://nvbugs/6323889) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_1k1k_con3072_ctx1_dep4_gen1_dep4_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6323889) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6302903) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_1k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6280721) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_1k1k_con2048_ctx1_dep4_gen1_dep4_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6374905) @@ -391,18 +387,11 @@ perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4 perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_32k4k_con256_ctx1_dep8_gen1_dep8_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6374893) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_8k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6280721) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_8k1k_con4096_ctx1_dep4_gen1_dep32_eplb256_mtp0_ccb-NIXL] SKIP (https://nvbugs/6302903) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_gpt-oss-120b-fp4_1k1k_con2048_ctx1_tp1_gen1_dep2_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6324123) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_gpt-oss-120b-fp4_1k1k_con512_ctx1_tp1_gen1_dep2_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6324123) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6324123) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_kimi-k25-thinking-fp4_1k1k_con4096_ctx1_dep4_gen1_dep8_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6323074) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6388238) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_deepseek-r1-fp4_1k1k_con3072_ctx1_dep4_gen1_dep4_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6323889) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_glm-5-fp4_1k1k_con4096_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL] SKIP (https://nvbugs/6324131) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL] SKIP (https://nvbugs/6324131) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_glm-5-fp4_8k1k_con1_ctx1_dep2_gen1_tep8_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6324131) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_kimi-k25-thinking-fp4_1k1k_con4096_ctx1_dep4_gen1_dep8_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6323074) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6323074) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL] SKIP (https://nvbugs/6323074) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6374893) perf/test_visual_gen_perf_sanity.py::test_visual_gen_e2e[vg_upload-ltx2_blackwell-ltx2_2stage_bf16_i2v_cfg2_ulysses4_compile_on] SKIP (https://nvbugs/6294413) perf/test_visual_gen_perf_sanity.py::test_visual_gen_e2e[vg_upload-ltx2_blackwell-ltx2_2stage_bf16_t2v_cfg2_ulysses4_compile_on] SKIP (https://nvbugs/6294413) diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py new file mode 100644 index 000000000000..1020126e800a --- /dev/null +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -0,0 +1,572 @@ +# 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_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.