[TRTLLM-12721][docs] Design disaggregated peer cancellation protocol

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 <atomic>
 #include <future>
 #include <memory>
 #include <mutex>
@@ -226,6 +227,11 @@ class BaseCacheTransceiver
     [[nodiscard]] virtual bool checkGenTransferComplete() const = 0;
 
     virtual bool cancelRequest(std::shared_ptr<LlmRequest> 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<SizeType32> const& rnnLayerNumPerPP = {});
 
+    CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager,
+        executor::kv_cache::CacheState::ModelConfig const& cacheStateModelCfg, runtime::WorldConfig const& worldConfig,
+        std::vector<SizeType32> const& attentionLayerNumPerPP, nvinfer1::DataType dataType,
+        executor::kv_cache::CacheState::AttentionType attentionType,
+        std::optional<executor::CacheTransceiverConfig> cacheTransceiverConfig,
+        rnn_state_manager::RnnStateManager* rnnStateManager, std::vector<SizeType32> const& rnnLayerNumPerPP,
+        bool enableInflightCancel);
+
     CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, std::vector<SizeType32> numKvHeadsPerLayer,
         SizeType32 sizePerHead, SizeType32 tokensPerBlock, runtime::WorldConfig const& worldConfig,
         std::vector<SizeType32> const& attentionLayerNumPerPP, nvinfer1::DataType dataType,
@@ -254,6 +268,20 @@ class CacheTransceiver : public BaseCacheTransceiver
     {
     }
 
+    CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, std::vector<SizeType32> numKvHeadsPerLayer,
+        SizeType32 sizePerHead, SizeType32 tokensPerBlock, runtime::WorldConfig const& worldConfig,
+        std::vector<SizeType32> const& attentionLayerNumPerPP, nvinfer1::DataType dataType,
+        executor::kv_cache::CacheState::AttentionType attentionType,
+        std::optional<executor::CacheTransceiverConfig> cacheTransceiverConfig,
+        rnn_state_manager::RnnStateManager* rnnStateManager, std::vector<SizeType32> 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> llmRequest) override;
@@ -273,11 +301,18 @@ class CacheTransceiver : public BaseCacheTransceiver
 
     virtual bool cancelRequest(std::shared_ptr<LlmRequest> llmRequest) override;
 
+    [[nodiscard]] bool hasPoisonedTransferBuffer() const override;
+
 private:
     void initializeCommState();
 
     void setContextState(LlmRequest* llmRequest);
 
+    RequestStatuses checkContextTransferStatusWithInflightCancel(
+        std::optional<int> const& atLeastRequestNum, bool markComplete);
+
+    void checkGenTransferStatusWithInflightCancel(std::optional<int> const& atLeastRequestNum);
+
     std::unique_ptr<CacheSender> mCacheSender;
     std::unique_ptr<CacheReceiver> 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<executor::kv_cache::CacheState> mCacheState;
     std::unique_ptr<executor::kv_cache::ConnectionManager> mManager;
     std::optional<executor::CacheTransceiverConfig> mCacheTransceiverConfig;
+    bool mInflightCancelEnabled{false};
+    std::atomic<bool> mTopologyTransferBufferPoisoned{false};
     std::vector<std::unique_ptr<kv_cache_manager::CacheTransBufferManager>> mCacheTransBufferManagers;
     std::vector<BaseTransBufferManager*> mCacheTransBufferManagerPtrs;
 

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

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 <chrono>
+#include <exception>
 #include <mutex>
 
 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<size_t> maxNumTokens)
     : mDataType{dataType}
@@ -78,26 +120,40 @@ BaseTransBufferManager::BaseTransBufferManager(
     allocateBuffer();
 }
 
-std::optional<int> BaseTransBufferManager::assignBufferIndexForSend()
+std::optional<int> BaseTransBufferManager::assignBufferIndexForSend(
+    std::atomic<bool> const* perRequestCancel, int64_t waitSliceMs)
 {
-    return assignBufferIndex(mConcurrenceSendResource, mSendBufferCount, mOnlyUseDynamicBuffer);
+    return assignBufferIndex(
+        mConcurrenceSendResource, mSendBufferCount, mOnlyUseDynamicBuffer, perRequestCancel, waitSliceMs);
 }
 
 void BaseTransBufferManager::freeBufferIndexForSend(std::optional<int> bufferId)
 {
     freeBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, mOnlyUseDynamicBuffer);
 }
 
-std::optional<int> BaseTransBufferManager::assignBufferIndexForRecv()
+void BaseTransBufferManager::poisonBufferIndexForSend(std::optional<int> bufferId) noexcept
+{
+    poisonBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, mOnlyUseDynamicBuffer, "send");
+}
+
+std::optional<int> BaseTransBufferManager::assignBufferIndexForRecv(
+    std::atomic<bool> const* perRequestCancel, int64_t waitSliceMs)
 {
-    return assignBufferIndex(mConcurrenceRecvResource, mRecvBufferCount, mOnlyUseDynamicBuffer);
+    return assignBufferIndex(
+        mConcurrenceRecvResource, mRecvBufferCount, mOnlyUseDynamicBuffer, perRequestCancel, waitSliceMs);
 }
 
 void BaseTransBufferManager::freeBufferIndexForRecv(std::optional<int> bufferId)
 {
     freeBufferIndex(mConcurrenceRecvResource, bufferId, mRecvBufferCount, mOnlyUseDynamicBuffer);
 }
 
+void BaseTransBufferManager::poisonBufferIndexForRecv(std::optional<int> bufferId) noexcept
+{
+    poisonBufferIndex(mConcurrenceRecvResource, bufferId, mRecvBufferCount, mOnlyUseDynamicBuffer, "recv");
+}
+
 std::tuple<std::vector<runtime::ITensor::SharedPtr>, size_t, bool> BaseTransBufferManager::getOrAllocateSendBuffers(
     std::optional<int> bufferId, int targetNum, std::vector<size_t> const& requestedNumberOfElements,
     runtime::BufferManager const& bufferManagerToUse)
@@ -249,16 +305,44 @@ void BaseTransBufferManager::allocateBuffer()
     }
 }
 
-std::optional<int> BaseTransBufferManager::assignBufferIndex(
-    ConcurrenceResource& resource, size_t bufferCount, bool onlyUseDynamicBuffer)
+std::optional<int> BaseTransBufferManager::assignBufferIndex(ConcurrenceResource& resource, size_t bufferCount,
+    bool onlyUseDynamicBuffer, std::atomic<bool> 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<size_t>(resource.mConcurrence) < bufferCount; });
+    auto const predicate = [&resource, bufferCount]()
+    {
+        return resource.mPoisoned.load(std::memory_order_relaxed)
+            || static_cast<size_t>(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,13 +372,61 @@ void BaseTransBufferManager::freeBufferIndex(
         TLLM_CHECK(static_cast<size_t>(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--;
         resource.mBuffersCV.notify_one();
     }
 }
 
+void BaseTransBufferManager::poisonBufferIndex(ConcurrenceResource& resource, std::optional<int> 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<size_t>(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;