[TRTLLM-12721][fix] Harden disagg cancellation consensus and buffer ownership

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>
@@ -244,6 +245,13 @@ class CacheTransceiver : public BaseCacheTransceiver
         std::optional<executor::CacheTransceiverConfig> cacheTransceiverConfig = std::nullopt,
         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,
+        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,
@@ -257,6 +265,19 @@ 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,
+        std::vector<SizeType32> const& rnnLayerNumPerPP, bool enableInflightCancel)
+        : CacheTransceiver(cacheManager,
+            executor::kv_cache::CacheState::ModelConfig{numKvHeadsPerLayer, sizePerHead, tokensPerBlock}, worldConfig,
+            attentionLayerNumPerPP, dataType, attentionType, cacheTransceiverConfig, rnnLayerNumPerPP,
+            enableInflightCancel)
+    {
+    }
+
     virtual ~CacheTransceiver();
 
     void respondAndSendAsync(std::shared_ptr<LlmRequest> llmRequest) override;
@@ -283,6 +304,11 @@ class CacheTransceiver : public BaseCacheTransceiver
 
     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
@@ -311,6 +337,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/cacheCommunicator.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.
@@ -30,9 +30,11 @@ class CommState;
 struct DataContext
 {
 public:
-    explicit DataContext(int tag, std::atomic<bool> const& transferTerminate = sDefaultTransferTerminate)
+    explicit DataContext(int tag, std::atomic<bool> const& transferTerminate = sDefaultTransferTerminate,
+        bool enableInflightCancel = false)
         : mTag{tag}
         , mTransferTerminate(transferTerminate)
+        , mEnableInflightCancel(enableInflightCancel)
     {
     }
 
@@ -46,10 +48,16 @@ struct DataContext
         return mTransferTerminate;
     }
 
+    [[nodiscard]] bool isInflightCancelEnabled() const noexcept
+    {
+        return mEnableInflightCancel;
+    }
+
 private:
     inline static std::atomic<bool> sDefaultTransferTerminate{false};
     int const mTag;
     std::atomic<bool> const& mTransferTerminate;
+    bool const mEnableInflightCancel;
 };
 
 class Connection

cpp/include/tensorrt_llm/executor/transferAgent.h

@@ -299,13 +299,15 @@ class TransferRequest
     /// @param dstDescs Description of the destination memory region.
     /// @param remoteName Name of the remote counterpart.
     /// @param syncMessage Synchronization information for the end of the transfer.
+    /// @param synchronizeStatusAccess Whether status queries can race explicit handle release.
     TransferRequest(TransferOp op, TransferDescs srcDescs, TransferDescs dstDescs, std::string const& remoteName,
-        std::optional<SyncMessage> syncMessage = std::nullopt)
+        std::optional<SyncMessage> syncMessage = std::nullopt, bool synchronizeStatusAccess = false)
         : mOp{op}
         , mSrcDescs{std::move(srcDescs)}
         , mDstDescs{std::move(dstDescs)}
         , mRemoteName{remoteName}
         , mSyncMessage{std::move(syncMessage)}
+        , mSynchronizeStatusAccess{synchronizeStatusAccess}
     {
     }
 
@@ -334,12 +336,18 @@ class TransferRequest
         return mSyncMessage;
     }
 
+    [[nodiscard]] bool synchronizeStatusAccess() const noexcept
+    {
+        return mSynchronizeStatusAccess;
+    }
+
 private:
     TransferOp mOp;
     TransferDescs mSrcDescs;
     TransferDescs mDstDescs;
     std::string mRemoteName;
     std::optional<SyncMessage> mSyncMessage;
+    bool mSynchronizeStatusAccess{false};
 };
 
 enum class TransferState : uint8_t

cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp

@@ -530,7 +530,7 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio
         // 5. send the buffer to the corresponding target. Ideally, we send only once (one buffer) for each target.
 
         auto const* sendCancelFlag
-            = common::getEnvDisaggEnableInflightCancel() ? &session.getDataContext().getTransferTerminate() : nullptr;
+            = session.isInflightCancelEnabled() ? &session.getDataContext().getTransferTerminate() : nullptr;
         auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend(sendCancelFlag);
         BufferIndexHolder sendHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/false);
         int peerDuplicateHeadFactor = targetInfo.mPeerDupHeadFactor;
@@ -623,7 +623,7 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio
         }
         catch (...)
         {
-            if (agentConnection != nullptr && common::getEnvDisaggEnableInflightCancel())
+            if (agentConnection != nullptr && session.isInflightCancelEnabled())
             {
                 sendHolder.poison();
             }