[TRTLLM-11160][feat] Clean up SWA work-arounds with the new radix sea…

cpp/include/tensorrt_llm/batch_manager/evictionPolicy.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-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.
@@ -54,7 +54,7 @@ class BaseEvictionPolicy
     /// @brief Perform any per-iteration bookkeeping
     virtual void refresh() = 0;
 
-    virtual bool verifyQueueIntegrity() = 0;
+    virtual bool verifyQueueIntegrity() const = 0;
 };
 
 struct ExpiringBlockComparator
@@ -95,7 +95,7 @@ class LRUEvictionPolicy : public BaseEvictionPolicy
     // Making this public and virtual makes it possible to test.
     [[nodiscard]] virtual std::chrono::steady_clock::time_point::duration getTime() const;
 
-    bool verifyQueueIntegrity() override;
+    bool verifyQueueIntegrity() const override;
 
 private:
     /// @brief A fixed-size container supporting both non-negative and negative indexing.
@@ -109,6 +109,8 @@ class LRUEvictionPolicy : public BaseEvictionPolicy
 
         T& operator[](KVCacheBlock::IdType id)
         {
+            TLLM_CHECK_WITH_INFO(id != KVCacheBlock::kPlaceholderBlockId,
+                "SWA sentinel placeholders are not part of the indexed free queues");
             return id >= 0 ? positive[id] : negative[-id];
         }
     };

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -305,6 +305,15 @@ class KVCacheBlock : public std::enable_shared_from_this<KVCacheBlock>
 
     static constexpr IdType kCachedBlocksRootId = -1;
 
+    //! Sentinel block ID used by SWA on-demand placeholder blocks (no-arg createPlaceholder()).
+    //! Chosen near the minimum int32 so it remains safely negatable if accidentally normalized,
+    //! while any accidental mAllBlocksById[id] lookup still produces an obvious out-of-range
+    //! failure rather than silently aliasing a real block ID.
+    //! Linear-attention placeholders continue to use negative per-slot IDs via the 2-arg
+    //! createPlaceholder(IdType, SizeType32) overload; this sentinel is specific to the SWA
+    //! path where placeholders are swapped in-place for evicted OOW blocks without an ID.
+    static constexpr IdType kPlaceholderBlockId = std::numeric_limits<IdType>::min() + 1;
+
     explicit KVCacheBlock(IdType blockId, kernels::KVCacheIndex blockIdx, SizeType32 windowSize = -1);
 
     void startScheduling();
@@ -387,11 +396,21 @@ class KVCacheBlock : public std::enable_shared_from_this<KVCacheBlock>
     //! placeholders. Placeholder blocks are excluded from the eviction pool.
     [[nodiscard]] bool isPlaceholder() const;
 
-    //! \brief Create a placeholder KVCacheBlock with no GPU memory.
+    //! \brief Create a placeholder KVCacheBlock with no GPU memory (linear-attention form).
     //! \details The placeholder holds a block ID for sequence bookkeeping but mIsPlaceholder
     //! is set so that getCacheBlockIndices returns a nil index and the eviction pool ignores it.
+    //! Used by the linear-attention queue path (tryAllocatePlaceholderForLinearAttention).
     static BlockPtr createPlaceholder(IdType blockId, SizeType32 windowSize);
 
+    //! \brief Create a placeholder KVCacheBlock with no GPU memory (SWA on-demand form).
+    //! \details Used by WindowBlockManager::detachFrontBlock when an out-of-window block is
+    //! swapped out of a sequence's allocated-blocks list. mIsPlaceholder is set and
+    //! WindowBlockManager::setOffsets maps kPlaceholderBlockId to a nil index. The block
+    //! ID is set to kPlaceholderBlockId so getBlockById can distinguish it from pooled
+    //! linear-attention placeholders. No windowSize is needed because SWA placeholders are
+    //! never attached to the lookup tree.
+    static BlockPtr createPlaceholder();
+
     void detachDescendantsFromLookupTree();
 
     void freeBlockAndAllDescendants();
@@ -799,11 +818,12 @@ class WindowBlockManager
     {
         struct ClaimedBlock
         {
-            BlockPtr block;
-            SizeType32 numMatchedTokens; //!< tokens matched in this block
-            bool isPartialMatch;
-            bool needsCopy;     //!< partial match on block with refs or non-leaf (needs getFreeBlock + copy in Phase 2)
-            bool isPlaceholder; //!< placeholder block (linear attention recurrent states)
+            BlockPtr block{nullptr};
+            SizeType32 numMatchedTokens{0}; //!< tokens matched in this block
+            bool isPartialMatch{false};
+            bool needsCopy{false};       //!< partial match on block with refs or non-leaf (needs getFreeBlock + copy)
+            bool isPlaceholder{false};   //!< placeholder block (linear attention recurrent states)
+            bool isTraversalOnly{false}; //!< SWA OOW anchor with a trie node but no cache block value
             bool shouldReleaseCopySource{false}; //!< last copier releases the claimed source after copy
         };
 
@@ -838,7 +858,7 @@ class WindowBlockManager
     //! \brief Batch add sequences with two-phase claim-then-onboard under a single lock.
     //! \details Phase 1 claims all matching blocks across all requests (protecting from eviction).
     //!          Phase 2 onboards host blocks and allocates non-matching blocks.
-    //!          The mCachedBlocksRootMutex is held for the entire operation.
+    //!          The unified lookup tree's mutex is held for the entire operation.
     //! \param sequences  Per-request GenerationRequest references (parallel with other vectors).
     //! \param inputLengths  Per-request effective input length.
     //! \param numContextBlocksVec  Per-request number of context blocks.
@@ -1061,16 +1081,23 @@ class WindowBlockManager
 
     [[nodiscard]] static bool blockInRadixTree(BlockPtr const& block);
 
-    //! \brief Store blocks in cached blocks.
+    //! \brief Store context blocks in the reuse trie for this window.
+    //! \details Called after context phase for both SWA and non-SWA windows.
+    //!          Must be called before any detachFrontBlock call so that OOW blocks
+    //!          are already in the trie when they are replaced with placeholders.
+    void storeContextBlocks(GenerationRequest& sequence, LlmRequest const& llmRequest);
+
+    //! \brief Store blocks in the reuse trie.
     //! \param blockKeys Key of each block.
-    //! \param blockIds Id of each block.
-    //! \param pinBlocks If true, increment ref count for blocks while storing (pin on store).
+    //! \param blocks Block pointers (beam 0 only). OOW slots contain placeholder blocks
+    //!        (isPlaceholder()==true); storeBlocks advances past them via a trie lookup
+    //!        when the anchor still exists, and continues past missing SWA anchors.
+    //! \param pinBlocks If true, increment ref count for blocks while storing.
     //! \return Pair of (num blocks stored for reuse, vector of pinned block IDs).
     [[nodiscard]] std::pair<SizeType32, std::vector<KVCacheBlock::IdType>> storeBlocks(
-        std::vector<BlockKey> const& blockKeys, std::vector<KVCacheBlock::IdType> const& blockIds,
-        bool pinBlocks = false);
+        std::vector<BlockKey> blockKeys, std::vector<BlockPtr> const& blocks, bool pinBlocks = false);
 
-    [[nodiscard]] bool verifyQueueIntegrity();
+    [[nodiscard]] bool verifyQueueIntegrity() const;
 
     // Only needed when sliding window attention + paged context fmha are used together.
     // In that case, a temporary kv cache buffer with maximum chunk size (maxNumTokens) is needed.
@@ -1109,26 +1136,9 @@ class WindowBlockManager
     //! \brief Unpin blocks by block ids directly
     void unpinBlocksById(std::vector<KVCacheBlock::IdType> const& blockIds);
 
-    void initializeSequenceStorageValidity(LlmRequest::RequestIdType requestId)
-    {
-        mIsValidStoreForReuseSequence[requestId] = true;
-    }
-
-    void releaseSequenceStorageValidity(LlmRequest::RequestIdType requestId)
-    {
-        mIsValidStoreForReuseSequence.erase(requestId);
-    }
-
-    //! \brief Return whether this sequence is valid for store for reuse
-    [[nodiscard]] bool isSequenceValidForStoreForReuse(LlmRequest::RequestIdType requestId) const
-    {
-        TLLM_CHECK_WITH_INFO(mIsValidStoreForReuseSequence.count(requestId) > 0, "Sequence should be bookkeeped");
-        return mIsValidStoreForReuseSequence.at(requestId);
-    }
-
     void resetReuseState()
     {
-        std::lock_guard<std::mutex> lock(mCachedBlocksRootMutex);
+        std::lock_guard<std::recursive_mutex> lock(mLookupTree->getMutex());
         // The shared lookup tree is reset once by BlockManager::resetReuseState() before
         // this method is called.  Here we only need to re-create the per-window root block
         // and wire it into the (already fresh) shared tree.
@@ -1139,9 +1149,30 @@ class WindowBlockManager
     }
 
 private:
-    //! \brief Walk the reuse tree with precomputed per-block keys (no lock; callers must hold mCachedBlocksRootMutex).
+    //! \brief Walk the reuse tree with precomputed per-block keys (no lock; callers must hold mLookupTree->getMutex()).
     [[nodiscard]] std::shared_ptr<KVCacheBlock> searchReuseTree(std::vector<BlockKey> const& blockKeys);
 
+    struct ReuseMatch
+    {
+        BlockPtr block;
+        SizeType32 numMatchedTokens{0};
+        bool isPartialMatch{false};
+        bool isTraversalOnly{false};
+    };
+
+    struct ReuseMatchResult
+    {
+        std::vector<ReuseMatch> matches;
+        SizeType32 totalMatchedTokens{0};
+        std::optional<BlockKey> firstNewBlock{std::nullopt};
+    };
+
+    //! \brief Find the reusable prefix, including SWA traversal-only anchors.
+    //! \details Value-less exact nodes are only accepted for SWA when enough later tokens
+    //!          are matched to place the missing anchor outside the attention window.
+    [[nodiscard]] ReuseMatchResult findReusableBlockMatches(std::vector<BlockKey> const& blockKeys,
+        bool enablePartialReuse, bool copyOnPartialReuse, SizeType32 maxMatchedTokens) const;
+
     bool tryAllocatePlaceholderForLinearAttention(GenerationRequest& sequence, bool shareAmongBeams);
 
     //! \brief Add single block to beam of sequence and mAllocatedBlocksPerSeq.
@@ -1151,7 +1182,7 @@ class WindowBlockManager
     void addBlockToAllBeams(BlockPtr const& block, GenerationRequest& sequence);
 
     //! \brief Phase 1: Walk radix tree and claim matching blocks.
-    //! \details Caller must hold mCachedBlocksRootMutex.
+    //! \details Caller must hold mLookupTree->getMutex().
     //!          Uses \p tracker to coordinate partial-match ownership across requests in
     //!          the same batch. \p claimResults is the full vector so that a previous
     //!          request's ClaimedBlock can be retroactively marked needsCopy.
@@ -1165,13 +1196,10 @@ class WindowBlockManager
         GenerationRequest& sequence, SizeType32 inputLength, SizeType32 numContextBlocks, LlmRequest& llmRequest);
 
     //! \brief Phase 2: Onboard claimed host blocks and allocate non-matching blocks.
-    //! \details Caller must hold mCachedBlocksRootMutex.
+    //! \details Caller must hold mLookupTree->getMutex().
     [[nodiscard]] SizeType32 onboardAndAllocateBlocks(
         GenerationRequest& sequence, LlmRequest& llmRequest, ClaimResult& claimResult, bool isEnableBlockReuse);
 
-    //! \brief Free block and all it's descendants. This makes block a claimed leaf block.
-    void freeChildren(BlockPtr const& block);
-
     //! \brief Find block least likely to be reused, free it if necessary and return.
     //! \param sequence Sequence which the free block is allocated for
     //! \param wantPlaceholder If true, return a pre-allocated placeholder block instead of a normal block
@@ -1288,17 +1316,6 @@ class WindowBlockManager
     SizeType32 mPrevMissedBlocks{0};
     SizeType32 mPrevGenAllocBlocks{0};
 
-    // Mutex for the cached blocks root
-    mutable std::mutex mCachedBlocksRootMutex;
-
-    // Record which sequence is using the block
-    std::map<KVCacheBlock::IdType, LlmRequest::RequestIdType> mBlockToSequence;
-    // Record whether a sequence has all blocks held valid.
-    // The boolean value is set to true upon first encounter of a new sequence.
-    // It may be invalidated to false when other sequence acquires a block that
-    // is used by another sequence.
-    std::map<LlmRequest::RequestIdType, bool> mIsValidStoreForReuseSequence;
-
     // Whether to enable indexer K cache
     bool mEnableIndexerKCache;
     // Quant block size for indexer K cache
@@ -1413,14 +1430,13 @@ class BlockManager
     void offloadBlock(BlockPtr const& block, SizeType32 windowSize,
         executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, std::string const& directory = "");
 
-    [[nodiscard]] std::pair<SizeType32, std::vector<KVCacheBlock::IdType>> storeBlocks(
-        std::vector<BlockKey> const& blockKeys, std::vector<KVCacheBlock::IdType> const& blockIds,
-        SizeType32 windowSize, bool pinBlocks = false)
+    [[nodiscard]] std::pair<SizeType32, std::vector<KVCacheBlock::IdType>> storeBlocks(std::vector<BlockKey> blockKeys,
+        std::vector<BlockPtr> const& blocks, SizeType32 windowSize, bool pinBlocks = false)
     {
-        return mWindowBlockManagers.at(windowSize).storeBlocks(blockKeys, blockIds, pinBlocks);
+        return mWindowBlockManagers.at(windowSize).storeBlocks(std::move(blockKeys), blocks, pinBlocks);
     }
 
-    [[nodiscard]] bool verifyQueueIntegrity(SizeType32 windowSize);
+    [[nodiscard]] bool verifyQueueIntegrity(SizeType32 windowSize) const;
 
     void releasePools();
 
@@ -1695,48 +1711,6 @@ class BlockManager
     //! context block that goes OOW.
     void adjustBlocksIfNeeded(GenerationRequest& sequence);
 
-    //! \brief Return whether the sequence is already managed by the block manager
-    [[nodiscard]] bool isSequenceHeld(LlmRequest::RequestIdType requestId) const
-    {
-        return mManagedSequences.count(requestId) > 0;
-    }
-
-    //! \brief Add a sequence to the managed sequences
-    //! \details Take the sequence into account for the manager. Initialize
-    //! sequence storage validity under all window sizes.
-    void holdSequence(LlmRequest::RequestIdType requestId)
-    {
-        mManagedSequences.insert(requestId);
-        for (auto const& [windowSize, metadata] : mWindowSizeToMetadata)
-        {
-            mWindowBlockManagers.at(windowSize).initializeSequenceStorageValidity(requestId);
-        }
-    }
-
-    //! \brief Remove a sequence from the managed sequences.
-    //! \details Remove sequence from the managed sequences and remove sequence
-    //! storage
-    void releaseSequence(LlmRequest::RequestIdType requestId)
-    {
-        mManagedSequences.erase(requestId);
-        for (auto const& [windowSize, metadata] : mWindowSizeToMetadata)
-        {
-            mWindowBlockManagers.at(windowSize).releaseSequenceStorageValidity(requestId);
-        }
-    }
-
-    //! \brief Return whether the sequence is still valid for store-for-reuse
-    //! regarding the specific window size.
-    //! \details Currently this utility function is only used under
-    //! kvCacheManagerTest.cpp. Checking for store-for-reuse for each window
-    //! size is done in an iterating fashion under BlockManager::releaseBlocks.
-    bool isSequenceValidForStoreForReuse(LlmRequest::RequestIdType requestId, SizeType32 windowSize) const
-    {
-        TLLM_CHECK_WITH_INFO(
-            mWindowBlockManagers.count(windowSize) > 0, "Querying window size is not found under mWindowBlockManager");
-        return mWindowBlockManagers.at(windowSize).isSequenceValidForStoreForReuse(requestId);
-    }
-
     void resetReuseState()
     {
         // Reset the shared tree once; all blocks' LookupNodePtr references to the old
@@ -1786,8 +1760,6 @@ class BlockManager
     std::vector<SizeType32> mLayerToWindowSize;
     std::vector<SizeType32> mAbsolutePoolToWindowSize;
     std::vector<SizeType32> mAbsolutePoolToRelativePoolIndex;
-    // Record what sequences are currently managed by the block manager
-    std::set<LlmRequest::RequestIdType> mManagedSequences;
 
     bool mIsEnableIndexerKCache{false};
     SizeType32 mIndexerKCacheQuantBlockSize{0};

cpp/include/tensorrt_llm/batch_manager/radixBlockTree.h

@@ -22,6 +22,7 @@
 #include "tensorrt_llm/common/assert.h"
 #include "tensorrt_llm/common/logger.h"
 
+#include <mutex>
 #include <optional>
 #include <vector>
 
@@ -68,8 +69,37 @@ inline constexpr int kRecurrentStates = -1;
 class UnifiedBlockTree : public templated_trie::Trie<BlockKey, BlockKeyHasher, int, std::hash<int>, BlockPtr, true>
 {
 public:
+    using Base = templated_trie::Trie<BlockKey, BlockKeyHasher, int, std::hash<int>, BlockPtr, true>;
+
     UnifiedBlockTree() = default;
 
+    // std::mutex is not movable, so define move operations explicitly.
+    // The trie contents (parent class data) are moved; each instance keeps its own mutex.
+    UnifiedBlockTree(UnifiedBlockTree&& other) noexcept
+        : Base(std::move(other))
+    {
+    }
+
+    UnifiedBlockTree& operator=(UnifiedBlockTree&& other) noexcept
+    {
+        if (this != &other)
+        {
+            Base::operator=(std::move(other));
+        }
+        return *this;
+    }
+
+    //! \brief Returns the shared mutex that guards all trie operations.
+    //! All reads and writes to the trie (insertions, lookups, detachFromLookupNode)
+    //! must be performed while holding this mutex.  A recursive mutex is used so
+    //! that call paths which hold the tree lock (e.g. BlockManager::addSequenceBatch
+    //! two-phase claim) can safely invoke helpers (e.g. getFreeBlock / detachFromLookupNode)
+    //! that also need to acquire it without deadlocking.
+    [[nodiscard]] std::recursive_mutex& getMutex() noexcept
+    {
+        return mMutex;
+    }
+
     //! \brief Insert a block into the tree at the given prefix position for a specific window size.
     //! \details This is a tree-only insertion: it does NOT set block->mLookupNode. The block is
     //! stored as a value in the trie node but carries no back-reference to that node. Use this for testing. For
@@ -203,6 +233,9 @@ class UnifiedBlockTree : public templated_trie::Trie<BlockKey, BlockKeyHasher, i
             }
         }
     }
+
+private:
+    std::recursive_mutex mMutex;
 };
 
 } // namespace tensorrt_llm::batch_manager::radix_block_tree

cpp/tensorrt_llm/batch_manager/evictionPolicy.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");
@@ -103,7 +103,7 @@ void LRUEvictionPolicy::initializePlaceholders(std::vector<BlockPtr>& allPlaceho
     }
 }
 
-bool LRUEvictionPolicy::verifyQueueIntegrity()
+bool LRUEvictionPolicy::verifyQueueIntegrity() const
 {
     static char const* const levelToStr[] = {"primary", "secondary", "placeholder"};
     static const std::function<bool(BlockPtr const&)> levelValidators[]
@@ -172,6 +172,14 @@ void LRUEvictionPolicy::releaseBlock(BlockPtr block, bool toFront)
     TLLM_CHECK_WITH_INFO(
         block->getBlockId() != tensorrt_llm::batch_manager::kv_cache_manager::KVCacheBlock::kCachedBlocksRootId,
         "Attempted to release the cached-blocks root into the eviction queue");
+    // SWA on-demand placeholders are transient sentinels created by createPlaceholder() and
+    // are not part of the pooled placeholder free queues. Skip re-inserting only those
+    // sentinels; pooled linear-attention placeholders must fall through and return to the
+    // placeholder queue at kPlaceholderLevel.
+    if (block->isPlaceholder() && block->getBlockId() == KVCacheBlock::kPlaceholderBlockId)
+    {
+        return;
+    }
     SizeType32 const cacheLevel = getCacheLevel(block);
     SizeType32 const id = block->getBlockId();