[None][fix] harden Python native KV transfer ownership

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -999,7 +999,8 @@ class WindowBlockManager
 
     void storeNewBlock(GenerationRequest& sequence, OptionalRef<LlmRequest const> llmRequest);
 
-    //! \brief Pin blocks associated with a sequence to prevent eviction.
+    //! \brief Pin each distinct physical block associated with a sequence once.
+    //! \details Blocks shared by multiple beams acquire one pin reference, not one per beam.
     void pinBlocks(GenerationRequest& sequence);
 
     //! \brief Release blocks of the sequence.
@@ -1256,7 +1257,8 @@ class WindowBlockManager
     [[nodiscard]] std::shared_ptr<KVCacheBlock> findBlocksInReuseTreeByBlockKeys(
         std::vector<BlockKey> const& blockKeys);
 
-    //! \brief Unpin blocks by block ids directly
+    //! \brief Unpin blocks by block ids directly.
+    //! \details Every id must be unique and identify an outstanding pin reference.
     void unpinBlocksById(std::vector<KVCacheBlock::IdType> const& blockIds);
 
     void truncateBlocks(LlmRequest::VecTokens const& targetTokens, SizeType32 numTokensToKeep);
@@ -1555,8 +1557,9 @@ class BlockManager
 
     void schedulingReleaseBlocks(LlmRequest::RequestIdType requestId);
 
-    /// @brief Pin all blocks associated with a sequence across all window managers.
+    /// @brief Pin each distinct physical block associated with a sequence once across all window managers.
     /// @param sequence The generation request whose blocks should be pinned.
+    /// @note Block-ID unpinning currently supports single-window managers only.
     void pinBlocks(GenerationRequest& sequence);
 
     void unpinBlocksById(std::vector<KVCacheBlock::IdType> const& blockIds);
@@ -2214,6 +2217,8 @@ class BaseKVCacheManager
         std::vector<BlockKey> const& blockKeys, SizeType32 windowSize)
         = 0;
 
+    /// @brief Consume one outstanding pin reference for each unique block id.
+    /// @note Callers must not pass duplicate ids or ids without an outstanding pin.
     virtual void unpinBlocksById(std::vector<KVCacheBlock::IdType> const& blockIds) = 0;
 
     /// @brief Release cached blocks for a token sequence beyond a given prefix length.

cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp

@@ -42,6 +42,7 @@
 #include <map>
 #include <optional>
 #include <sstream>
+#include <unordered_set>
 #include <utility>
 
 namespace tc = tensorrt_llm::common;
@@ -2978,11 +2979,22 @@ void BlockManager::unpinBlocksById(std::vector<KVCacheBlock::IdType> const& bloc
 
 void WindowBlockManager::pinBlocks(GenerationRequest& sequence)
 {
+    std::lock_guard<std::recursive_mutex> lock(mLookupTree->getMutex());
     auto const requestId = sequence.getRequestId();
     auto& allocatedBlocks = mAllocatedBlocksPerSeq.at(requestId);
+    // Shared beams repeat the same BlockPtr in mAllocatedBlocksPerSeq. One
+    // reference is sufficient to pin that physical block. The matching unpin
+    // operation must likewise provide that physical block ID only once.
+    std::unordered_set<KVCacheBlock*> pinnedBlocks;
+    pinnedBlocks.reserve(allocatedBlocks.size());
     for (auto& block : allocatedBlocks)
     {
-        block->incRefCount();
+        // Placeholders have no addressable pool slot or non-negative block ID,
+        // so they cannot produce a token accepted by unpinBlocksById.
+        if (!block->isPlaceholder() && pinnedBlocks.insert(block.get()).second)
+        {
+            block->incRefCount();
+        }
     }
 }
 
@@ -2993,18 +3005,36 @@ void WindowBlockManager::unpinBlocksById(std::vector<KVCacheBlock::IdType> const
         return;
     }
 
+    std::lock_guard<std::recursive_mutex> lock(mLookupTree->getMutex());
+
+    // Validate the complete operation before changing any refcount. Raw block
+    // IDs do not identify which refcount contribution is a pin token, so a
+    // duplicate could consume a live sequence or another owner's reference.
+    std::unordered_set<KVCacheBlock::IdType> seenBlockIds;
+    std::vector<BlockPtr> blocksToUnpin;
+    seenBlockIds.reserve(blockIds.size());
+    blocksToUnpin.reserve(blockIds.size());
     for (auto const& blockId : blockIds)
     {
         TLLM_CHECK_WITH_INFO(blockId >= 0 && static_cast<size_t>(blockId) < mAllBlocksById.size(),
             "Block id %d is out of range", blockId);
+        auto const inserted = seenBlockIds.insert(blockId).second;
+        TLLM_CHECK_WITH_INFO(inserted, "Duplicate block id %d in one unpin operation", blockId);
         auto block = mAllBlocksById[blockId];
         if (block && block->getBlockId() != KVCacheBlock::kCachedBlocksRootId)
         {
-            block->decRefCount();
-            if (!block->hasRefs())
-            {
-                mEvictionPolicy->releaseBlock(block);
-            }
+            TLLM_CHECK_WITH_INFO(
+                block->hasRefs(), "Can't unpin block (id=%d) that is not allocated", static_cast<int>(blockId));
+            blocksToUnpin.push_back(block);
+        }
+    }
+
+    for (auto const& block : blocksToUnpin)
+    {
+        block->decRefCount();
+        if (!block->hasRefs())
+        {
+            mEvictionPolicy->releaseBlock(block);
         }
     }
 }

cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp

@@ -4903,16 +4903,86 @@ TEST_F(KVCacheManagerTest, PinAndUnpinBlocksById)
     ASSERT_TRUE(lastBlockIdOpt.has_value());
     auto const& allBlockIds = kvCacheManager.getCacheBlockIds(requestId, maxAttentionWindow)[0];
     std::vector<SizeType32> pinnedBlockIds(allBlockIds.begin(), allBlockIds.end());
+    ASSERT_FALSE(pinnedBlockIds.empty());
+
+    // Raw block IDs are not owner-scoped pin handles. A duplicate must fail
+    // before it can consume the live sequence reference.
+    EXPECT_THROW(kvCacheManager.unpinBlocksById({pinnedBlockIds.front(), pinnedBlockIds.front()}), std::runtime_error);
+
     tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*llmRequest);
     (void) kvCacheManager.removeSequence(requestId, llmRequest);
     auto const freeAfterRemovePinned = kvCacheManager.getNumFreeBlocks();
     EXPECT_LT(freeAfterRemovePinned, totalBlocks);
 
-    kvCacheManager.unpinBlocksById(pinnedBlockIds);
+    auto invalidBlockIds = pinnedBlockIds;
+    invalidBlockIds.push_back(totalBlocks);
+    EXPECT_THROW(kvCacheManager.unpinBlocksById(invalidBlockIds), std::runtime_error);
+
+    auto const freeBlockId = [&]() -> SizeType32
+    {
+        for (SizeType32 blockId = 0; blockId < totalBlocks; ++blockId)
+        {
+            if (std::find(pinnedBlockIds.begin(), pinnedBlockIds.end(), blockId) == pinnedBlockIds.end())
+            {
+                return blockId;
+            }
+        }
+        return totalBlocks;
+    }();
+    ASSERT_LT(freeBlockId, totalBlocks);
+    EXPECT_THROW(kvCacheManager.unpinBlocksById({pinnedBlockIds.front(), freeBlockId}), std::runtime_error);
+    EXPECT_NO_THROW(kvCacheManager.unpinBlocksById(pinnedBlockIds));
     auto const freeAfterUnpin = kvCacheManager.getNumFreeBlocks();
     EXPECT_EQ(freeAfterUnpin, totalBlocks);
 }
 
+TEST_F(KVCacheManagerTest, PinAndUnpinSharedBeamBlocksOncePerPhysicalBlock)
+{
+    using namespace tensorrt_llm::batch_manager::kv_cache_manager;
+    auto constexpr numLayers = 2;
+    auto constexpr numKvHeads = 2;
+    auto constexpr sizePerHead = 16;
+    auto constexpr tokensPerBlock = 4;
+    auto constexpr blocksInPrimaryPool = 8;
+    auto constexpr blocksInSecondaryPool = 0;
+    auto constexpr maxNumSequences = 8;
+    auto const stream = std::make_shared<tr::CudaStream>();
+
+    auto constexpr beamWidth = 2;
+    auto const maxAttentionWindow = tokensPerBlock * blocksInPrimaryPool;
+    BlocksPerWindow const blocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}};
+    KVCacheManager kvCacheManager(numLayers, numKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences,
+        beamWidth, std::vector<BlockManager::SizeType32>{maxAttentionWindow}, tensorrt_llm::DataType::kHALF, 0, stream,
+        maxAttentionWindow, maxAttentionWindow, true);
+    kvCacheManager.allocatePools(false);
+
+    LlmRequest::RequestIdType requestId{0};
+    auto inputTokens = std::make_shared<VecTokens>(VecTokens{0, 1, 2, 3, 4, 5, 6, 7});
+    tr::SamplingConfig const samplingConfig{beamWidth};
+    auto llmRequest = std::make_shared<LlmRequest>(requestId, 0, inputTokens, samplingConfig, false);
+    kvCacheManager.addSequenceBatch(
+        {{{requestId, static_cast<SizeType32>(inputTokens->size()), beamWidth}}}, {std::ref(*llmRequest)});
+
+    kvCacheManager.pinBlocks(requestId);
+    auto const& perBeamBlockIds = kvCacheManager.getCacheBlockIds(requestId, maxAttentionWindow);
+    std::vector<SizeType32> perBeamPinnedBlockIds;
+    for (auto const& beamBlockIds : perBeamBlockIds)
+    {
+        perBeamPinnedBlockIds.insert(perBeamPinnedBlockIds.end(), beamBlockIds.begin(), beamBlockIds.end());
+    }
+    ASSERT_FALSE(perBeamPinnedBlockIds.empty());
+    std::set<SizeType32> const uniquePinnedBlockIds(perBeamPinnedBlockIds.begin(), perBeamPinnedBlockIds.end());
+    EXPECT_LT(uniquePinnedBlockIds.size(), perBeamPinnedBlockIds.size());
+
+    tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*llmRequest);
+    (void) kvCacheManager.removeSequence(requestId, llmRequest);
+    EXPECT_LT(kvCacheManager.getNumFreeBlocks(), kvCacheManager.getMaxNumBlocks());
+
+    EXPECT_NO_THROW(kvCacheManager.unpinBlocksById(
+        std::vector<SizeType32>(uniquePinnedBlockIds.begin(), uniquePinnedBlockIds.end())));
+    EXPECT_EQ(kvCacheManager.getNumFreeBlocks(), kvCacheManager.getMaxNumBlocks());
+}
+
 // Regression test for NVBug 6018647: storeBlocks(pin=true) on a zero-ref block
 // that sits in the eviction free queue must call claimBlock() before incRefCount().
 // Without the fix, unpinBlocksById inserts the block into the free queue a second

docs/design/README.md

@@ -0,0 +1,13 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+-->
+
+# Design Notes
+
+This directory contains draft engineering design notes and implementation
+plans. They are review artifacts, not part of the published Sphinx user
+documentation under `docs/source`. Each note records its own status and target
+scope.
+
+- [Python native disaggregated KV transfer ownership and lifecycle](disagg-kv-transfer-ownership/README.md)