[TRTLLM-12721][fix] Add gated disagg in-flight cancellation

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

Author: chienchunhung

cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h

@@ -226,6 +226,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
@@ -273,6 +278,8 @@ class CacheTransceiver : public BaseCacheTransceiver
 
     virtual bool cancelRequest(std::shared_ptr<LlmRequest> llmRequest) override;
 
+    [[nodiscard]] bool hasPoisonedTransferBuffer() const override;
+
 private:
     void initializeCommState();
 

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;

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<int> index, bool isRecv) noexcept
         : mMgr(&mgr)
         , mIndex(index)
-        , mHeld(index.has_value())
+        , mHeld(true)
         , mIsRecv(isRecv)
     {
     }
@@ -120,13 +120,18 @@ 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<int> mIndex{};
     bool mHeld{false};
     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<int> assignBufferIndexForSend();
+    std::optional<int> assignBufferIndexForSend(
+        std::atomic<bool> 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<int> bufferId);
 
+    /// @brief Poison a send buffer index after an unquiesced transfer exit.
+    void poisonBufferIndexForSend(std::optional<int> bufferId) noexcept;
+
     /// @brief Assign a buffer index for receiving.
     /// @return Assigned buffer index, or nullopt if using dynamic buffers.
-    std::optional<int> assignBufferIndexForRecv();
+    std::optional<int> assignBufferIndexForRecv(
+        std::atomic<bool> 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<int> bufferId);
 
+    /// @brief Poison a receive buffer index after an unquiesced transfer exit.
+    void poisonBufferIndexForRecv(std::optional<int> 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,16 +225,20 @@ class BaseTransBufferManager
         std::mutex mBuffersMutex;
         std::condition_variable mBuffersCV;
         std::atomic<int> mConcurrence{0};
+        std::atomic<bool> mPoisoned{false};
     };
 
     std::tuple<std::vector<runtime::ITensor::SharedPtr>, size_t, bool> getOrAllocateBuffers(std::optional<int> bufferId,
         int targetNum, std::vector<size_t> const& requestedNumberOfElements,
         runtime::BufferManager const& bufferManagerToUse, ConcurrenceResource& concurrenceResource);
 
     void allocateBuffer();
-    std::optional<int> assignBufferIndex(ConcurrenceResource& resource, size_t bufferCount, bool onlyUseDynamicBuffer);
+    std::optional<int> assignBufferIndex(ConcurrenceResource& resource, size_t bufferCount, bool onlyUseDynamicBuffer,
+        std::atomic<bool> const* perRequestCancel = nullptr, int64_t waitSliceMs = kBufferAcquireSliceMs);
     void freeBufferIndex(
         ConcurrenceResource& resource, std::optional<int> bufferId, size_t bufferCount, bool onlyUseDynamicBuffer);
+    void poisonBufferIndex(ConcurrenceResource& resource, std::optional<int> bufferId, size_t bufferCount,
+        bool onlyUseDynamicBuffer, char const* direction) noexcept;
 
     size_t mPreAllocBufferSize;
     size_t mRecvBufferCount;

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 && common::getEnvDisaggEnableInflightCancel())
+            {
+                sendHolder.poison();
+            }
+            throw;
+        }
 
         session.setTime(TransferSession::kTimeTransmissions);