[None][feat] Add native KV cache retention to KVCM V2

cpp/include/tensorrt_llm/batch_manager/common.h

@@ -38,6 +38,8 @@ using RequestIdType = std::uint64_t;
 using RequestVector = std::vector<std::shared_ptr<LlmRequest>>;
 using ReqIdsSet = std::unordered_set<RequestIdType>;
 
+inline constexpr int kDefaultColdTierOffloadMinPriority = 30;
+
 class ScheduledRequests
 {
 public:

cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp

@@ -17,6 +17,8 @@
 
 #include "tensorrt_llm/batch_manager/evictionPolicy.h"
 
+#include "tensorrt_llm/batch_manager/common.h"
+
 using namespace tensorrt_llm::batch_manager::kv_cache_manager;
 
 // This implements priority-based eviction.
@@ -31,7 +33,6 @@ auto const kMaxPriority = executor::KvCacheRetentionConfig::kMaxRetentionPriorit
 auto const kNumPriorities = kMaxPriority - kMinPriority + 1;
 
 auto const kDefaultPriority = executor::KvCacheRetentionConfig::kDefaultRetentionPriority;
-executor::RetentionPriority const kDefaultSecondaryOffloadMinPriority = 30;
 
 int const kNumCacheLevels = 2;
 int const kPlaceholderLevel = kNumCacheLevels; // placeholder blocks live at level 2
@@ -80,7 +81,7 @@ void LRUEvictionPolicy::initialize(std::vector<BlockPtr>& mAllBlocksById, std::v
         startIdx += sizes[cacheLevel];
     }
 
-    mSecondaryOffloadMinPriority = secondaryOffloadMinPriority.value_or(kDefaultSecondaryOffloadMinPriority);
+    mSecondaryOffloadMinPriority = secondaryOffloadMinPriority.value_or(kDefaultColdTierOffloadMinPriority);
 }
 
 void LRUEvictionPolicy::initializePlaceholders(std::vector<BlockPtr>& allPlaceholderBlocksById)

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md

@@ -213,6 +213,36 @@ Eviction controller
   Every `KvCache` must be closed before manager shutdown so this deferred cleanup
   runs while `StorageManager` is still alive.
 
+### Beam-search pages
+
+Prefill normally runs at beam width 1. Expanding the beam width does not
+physically duplicate the whole prompt; ownership splits at the prompt tail:
+
+- `_appendBeams` starts at block `floor(expectedPromptLength / tokensPerBlock)`.
+  Full prompt-prefix blocks before that ordinal stay on beam 0, and other beams
+  read that shared prefix through attention cache indirection.
+- The partially filled prompt tail, if any, gets a private writable page per
+  beam, copied from beam 0. Preallocated uncommitted generation blocks are
+  handled the same way. Newly allocated decode blocks have one row per beam.
+  Already committed blocks remain canonical on beam 0, even when the
+  prompt-length hint is conservative.
+- Raw page-index rows of non-leading beams contain `BAD_PAGE_INDEX` for the
+  shared prefix. `getAggregatedPageIndices()` and `Batch` publication expose the
+  actual storage rows. Cache indirection resolves shared attention history; do
+  not duplicate beam-0 pages into every raw row.
+- `getPagePriorities(layerGroupId, beamIndex)` preserves block ordinals. For a
+  block with one shared row it reads beam 0; otherwise it reads the requested
+  beam's private row. A missing lifecycle page yields `kPriorityDefault`.
+- Commit canonicalizes a reusable block to a single beam-0 row and releases the
+  alternative beam rows.
+- The wrapper disables `enablePartialCommit` when `maxBeamWidth > 1`, and
+  changing the beam width while partial commit is enabled is rejected.
+  Full-block reuse remains enabled. Partial reuse is also allowed: on resume,
+  a matched partial prefix is copied into a private uncommitted page before
+  beams expand.
+- SSM current state lives separately in `mSsmBlocks` and is copied per beam on
+  expansion. It is per-beam writable state, not a shared attention prefix.
+
 ## Correctness invariants
 
 - `ReuseScope::salt` is an optional string used directly by scope equality and

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h

@@ -450,6 +450,9 @@ struct KVCacheManagerConfig
     // Beam search disables this while retaining full-block reuse.
     bool enablePartialCommit = true;
 
+    // Unreferenced GPU pages below this priority are dropped instead of offloaded.
+    Priority coldTierOffloadMinPriority = kDefaultColdTierOffloadMinPriority;
+
     bool enableSwaScratchReuse() const noexcept
     {
         return swaScratchReuse.has_value();

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp

@@ -311,6 +311,24 @@ void EventManager::addCacheLevelUpdated(
         lifeCycle.value());
 }
 
+void EventManager::addPriorityUpdated(
+    Digest const& blockKey, Priority oldPriority, Priority newPriority, LifeCycleId lifeCycle)
+{
+    if (mMaxKvEventEntries <= 0)
+    {
+        return;
+    }
+    std::lock_guard<std::mutex> lock(mMutex);
+    auto const state = mStoredBlocks.find(blockKey);
+    if (state == mStoredBlocks.end() || state->second.lifeCycleIds.count(lifeCycle.value()) == 0)
+    {
+        return;
+    }
+    addEventUnlocked(
+        KVCacheUpdatedData{state->second.blockHash, std::nullopt, KVCacheEventDiff{oldPriority, newPriority}},
+        lifeCycle.value());
+}
+
 void EventManager::addStoredEventUnlocked(KVCacheStoredData data, EventLayerGroupId layerGroupId)
 {
     bool const hasPendingRemovedEvents = !mLatestRemovedBlockHashes.empty();
@@ -547,7 +565,7 @@ std::optional<KVCacheStoredBlockData> EventManager::storedBlockFromBlock(
         if (pageCoversBlock(page, block))
         {
             cacheLevel = page->cacheLevel;
-            priority = page->priority;
+            priority = page->priority();
             foundPage = true;
             break;
         }

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.h

@@ -191,6 +191,8 @@ class EventManager final : public EventSink
     void addRemovedLifeCycle(Digest const& blockKey, LifeCycleId lifeCycle) override;
     void addCacheLevelUpdated(
         Digest const& blockKey, CacheLevel oldLevel, CacheLevel newLevel, LifeCycleId lifeCycle) override;
+    void addPriorityUpdated(
+        Digest const& blockKey, Priority oldPriority, Priority newPriority, LifeCycleId lifeCycle) override;
 
 private:
     enum class HashAlgorithm

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventSink.h

@@ -44,6 +44,9 @@ class EventSink
     virtual void addCacheLevelUpdated(
         Digest const& blockKey, CacheLevel oldLevel, CacheLevel newLevel, LifeCycleId lifeCycle)
         = 0;
+
+    //! Protocols that expose retention metadata receive native page priority changes.
+    virtual void addPriorityUpdated(Digest const&, Priority, Priority, LifeCycleId) {}
 };
 
 } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.cpp

@@ -78,7 +78,7 @@ LRUEvictionPolicy& PrioritizedEvictionPolicy::getOrCreate(Priority p)
 
 NodeRef PrioritizedEvictionPolicy::push(SharedPtr<Page> page, bool evictFirst)
 {
-    Priority p = page->priority;
+    Priority p = page->priority();
     LRUEvictionPolicy& policy = getOrCreate(p);
     return policy.push(std::move(page), evictFirst);
 }
@@ -99,7 +99,7 @@ SharedPtr<Page> PrioritizedEvictionPolicy::pop()
 SharedPtr<Page> PrioritizedEvictionPolicy::remove(NodeRef node)
 {
     auto page = *node;
-    Priority p = page->priority;
+    Priority p = page->priority();
     auto it = mPolicies.find(p);
     TLLM_CHECK_DEBUG(it != mPolicies.end());
     it->second.remove(node);

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp

@@ -55,11 +55,13 @@ int64_t sumSlotBytes(StorageManager const& storage, CacheLevel level, LifeCycleI
 
 KvCache::KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional<BlockRadixTree::ReuseMatch> reuseMatch,
     std::optional<RequestIdType> mId, PriorityCb priorityCb, std::optional<int> expectedPromptLength,
-    std::optional<bool> textOnly, bool enableRequestStats)
+    std::optional<bool> textOnly, bool enableRequestStats,
+    std::optional<executor::KvCacheRetentionConfig> retentionConfig, bool isPrefetch)
     : id(mId)
     , mManager(manager.shared_from_this())
     , mReuseScope(std::move(reuseScope))
     , mPriorityCb(priorityCb ? std::move(priorityCb) : [](BlockOrdinal, LifeCycleId) { return kPriorityDefault; })
+    , mRetentionConfig(std::move(retentionConfig))
     , mStatus(Status::SUSPENDED)
     , mCommitState(CommitState::ALLOWED)
     , mBeamWidth(BeamIndex{1})
@@ -74,6 +76,26 @@ KvCache::KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional<B
     , mNumCommittedBlocks(0)
     , mTokensPerBlock(manager.tokensPerBlock())
 {
+    if (mRetentionConfig.has_value())
+    {
+        if (mRetentionConfig->getTransferMode() != executor::KvCacheTransferMode::DRAM)
+        {
+            throw std::invalid_argument(
+                "KVCacheManagerV2 does not support per-request GDS/POSIX transfer_mode or directory; "
+                "its offload storage is configured by the manager's cache tiers");
+        }
+        if (!expectedPromptLength.has_value())
+        {
+            throw std::invalid_argument("Retention requires the complete prompt length");
+        }
+        auto const priority = mRetentionConfig->getDecodeRetentionPriority();
+        if (priority < kPriorityMin || priority > kPriorityMax)
+        {
+            throw std::invalid_argument("Decode retention priority must be between 0 and 100");
+        }
+        mPromptRetentions
+            = mRetentionConfig->getPerBlockRetentionPriorityDuration(mTokensPerBlock, *expectedPromptLength);
+    }
     LifeCycleId numLc = manager.storage().numLifeCycles();
 
     // Initialise page index buffers: [beamIdx][lcId] = empty vector
@@ -99,7 +121,7 @@ KvCache::KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional<B
 
     if (reuseMatch.has_value())
     {
-        _setupForReuse(*reuseMatch);
+        _setupForReuse(*reuseMatch, !isPrefetch);
     }
 
     _collectReuseCandidates();
@@ -143,9 +165,56 @@ CachedCudaEvent KvCache::finishEvent() const
 
 Priority KvCache::getPriority(BlockOrdinal ordinal, LifeCycleId lc) const
 {
+    auto const retention = getRetention(ordinal);
+    if (retention.retentionPriority.has_value())
+    {
+        return *retention.retentionPriority;
+    }
     return mPriorityCb(ordinal, lc);
 }
 
+executor::RetentionPriorityAndDuration KvCache::getRetention(BlockOrdinal ordinal) const
+{
+    if (!mRetentionConfig.has_value() || ordinal.value() < 0)
+    {
+        return {std::nullopt, std::nullopt};
+    }
+    if (static_cast<size_t>(ordinal.value()) < mPromptRetentions.size())
+    {
+        return mPromptRetentions.at(ordinal.value());
+    }
+    return {mRetentionConfig->getDecodeRetentionPriority(), mRetentionConfig->getDecodeDurationMs()};
+}
+
+std::vector<Priority> KvCache::getPagePriorities(LayerGroupId layerGroupId, BeamIndex beamIndex) const
+{
+    KVCM2_REJECT_IF_POISONED();
+    auto const apiLock = mManager->lockShared();
+    if (beamIndex < BeamIndex{0} || beamIndex >= mBeamWidth)
+    {
+        throw std::out_of_range("Beam index is outside the request's beam width");
+    }
+    if (layerGroupId.value() < 0 || LifeCycleId{layerGroupId.value()} >= mManager->storage().numLifeCycles())
+    {
+        throw std::out_of_range("Layer group is outside the configured life cycles");
+    }
+    std::vector<Priority> result;
+    result.reserve(mBlocks.stdSize());
+    bool const isSsm
+        = std::holds_alternative<SsmLifeCycle>(mManager->storage().getLifeCycle(LifeCycleId{layerGroupId.value()}));
+    for (auto const& block : mBlocks)
+    {
+        // Shared prompt blocks keep only beam 0; writable tails have a row per beam.
+        auto const pageBeamIndex = block.pages.size() == BeamIndex{1} ? kDefaultBeamIndex : beamIndex;
+        auto const* page = blockPageGetPage(block.pages[pageBeamIndex].at(layerGroupId)).get();
+        if (isSsm && block.treeBlock)
+            page = block.treeBlock->getPage(LifeCycleId{layerGroupId.value()});
+        // Preserve ordinals, including absent pages, just like getAggregatedPageIndices(validOnly=false).
+        result.push_back(page ? page->priority() : kPriorityDefault);
+    }
+    return result;
+}
+
 StorageManager* KvCache::storageManager() const
 {
     return &mManager->storage();
@@ -248,11 +317,18 @@ void KvCache::_selectOptionalGroups(std::vector<LifeCycleId> const& groups)
     }
     for (auto lc : groups)
     {
+        bool const newlySelected = std::find(mSelectedOptionalGroups.begin(), mSelectedOptionalGroups.end(), lc)
+            == mSelectedOptionalGroups.end();
         auto const stale = _getStaleRange(mHistoryLength, mManager->lifeCycles()[lc]);
         for (auto const& [ordinal, holder] : mOptionalCandidates[lc])
         {
             if (!stale.contains(ordinal))
             {
+                if (newlySelected)
+                {
+                    auto const retention = getRetention(ordinal);
+                    holder->page->claimRetention(retention.retentionPriority, retention.durationMs);
+                }
                 mBlocks[ordinal].pages[kDefaultBeamIndex][lc] = holder;
             }
         }
@@ -772,6 +848,7 @@ bool KvCache::resume(std::optional<CUstream> stream, std::optional<bool> isDecod
         // Phase 1: Copy locked sources into private GPU slots. A shared sparse
         // prefix may be locked on host by another request and must stay there.
         std::vector<SharedPageLock*> srcLocks;
+        TypedVec<LifeCycleId, SharedPtr<Page>> copiedSources(numLc);
         for (LifeCycleId lcIdx{0}; lcIdx < numLc; ++lcIdx)
         {
             if (!deferredSlots[lcIdx].has_value())
@@ -793,6 +870,7 @@ bool KvCache::resume(std::optional<CUstream> stream, std::optional<bool> isDecod
             TLLM_CHECK_DEBUG(lock && lock->isValid());
             bool const hasPartialReuseSource = _hasReuseSource(*sourcePage);
             srcLocks.push_back(lock);
+            copiedSources[lcIdx] = lock->page();
 
             CacheLevel const sourceLevel = lock->page()->cacheLevel;
             deferredCopiesStarted = true;
@@ -851,6 +929,12 @@ bool KvCache::resume(std::optional<CUstream> stream, std::optional<bool> isDecod
             }
 
             auto newPage = makeShared<UncommittedPage>(*this, blockOrdinal, lcIdx, kHotLevel, beamIdx);
+            if (blockOrdinal != kBadBlockOrdinal)
+            {
+                auto const& source = copiedSources[lcIdx];
+                TLLM_CHECK_DEBUG(source);
+                newPage->inheritRetention(*source);
+            }
             newPage->setSlot(newSlot);
             auto newLock = newPage->lock(*this, beamIdx, blockOrdinal, lcIdx, /*skipWait=*/true);
             *targetBp = std::move(newLock);
@@ -875,6 +959,7 @@ bool KvCache::resume(std::optional<CUstream> stream, std::optional<bool> isDecod
     }
     mOptionalCandidates.clear();
     // Deferred copies survive a failed decode admission and must not be repeated on retry.
+    mPartialCopyRetentionHolds.clear();
     mNeverResumed = false;
     if (mIsDecoding)
     {
@@ -1266,6 +1351,7 @@ bool KvCache::_hasReuseSource(BlockPage const& page)
 
 void KvCache::_clearBlocks()
 {
+    mPartialCopyRetentionHolds.clear();
     // Drop last block first (mirrors Python: while self._blocks: self._blocks.pop()).
     while (!mBlocks.empty())
         mBlocks.pop_back();
@@ -1315,6 +1401,15 @@ CommittedPage* KvCache::_copyPageToTreeBlock(
         newSlot.readyEvent = CachedCudaEvent(reinterpret_cast<CudaStream>(stream));
         auto committed = makeShared<CommittedPage>(
             &storageMgr, treeBlock, lcIdx, lvl, numTokensInBlock, getPriority(treeBlock->ordinal(), lcIdx));
+        auto const retention = getRetention(treeBlock->ordinal());
+        if (mManager->lifeCycles().ssmLifeCycleId() == lcIdx)
+        {
+            committed->claimRetention(retention.retentionPriority, retention.durationMs);
+        }
+        else
+        {
+            committed->inheritRetention(*srcPage);
+        }
         committed->setSlot(newSlot);
         // Drops the superseded page, deferred until the copy is issued: an
         // OutOfPagesError above must not destroy a usable shorter snapshot.
@@ -1819,6 +1914,11 @@ bool KvCache::resize(std::optional<int> capacity, std::optional<int> historyLeng
         TLLM_CHECK_DEBUG(std::all_of(slots.begin(), slots.end(), [](auto const& vec) { return vec.empty(); }));
     }
 
+    for (auto const& backup : backupHolders)
+    {
+        // Leaving the attention window ends the retention interval without changing priority.
+        backup.holder->page->claimRetention(std::nullopt, std::nullopt);
+    }
     mCapacity = newCap;
     _publishHistoryLength(newHist);
     _refreshGenerationAllocReady();
@@ -2787,7 +2887,7 @@ void KvCache::dropPartialBlockCachedTokenAttribution()
 // _setupForReuse: find existing blocks in radix tree matching input tokens.
 // ---------------------------------------------------------------------------
 
-void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match)
+void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match, bool claimRetention)
 {
     auto const& matched = match.blocks;
     int const numTokens = match.numTokens;
@@ -2866,7 +2966,18 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match)
             TLLM_CHECK_WITH_INFO(page, "Expected page in non-stale block");
             CacheLevel const level = page->cacheLevel;
             auto& bpSlot = mBlocks[ordinal].pages[beamIdx][lcId];
+            bool const partialCopy = hasPartialMatch && ordinal == fullReusedEnd;
+            bool const borrowPartial = partialCopy && page->hasRetentionConsumers();
             bpSlot = page->hold();
+            if (claimRetention && !borrowPartial)
+            {
+                auto const retention = getRetention(ordinal);
+                page->claimRetention(retention.retentionPriority, retention.durationMs);
+            }
+            if (partialCopy)
+            {
+                mPartialCopyRetentionHolds.push_back(page->borrowRetentionForCopy());
+            }
             if (!isAttention)
             {
                 return;
@@ -2937,6 +3048,11 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match)
         TLLM_CHECK_WITH_INFO(snapshotPage, "Last matched block must have SSM snapshot after truncation");
         ssmLevel = snapshotPage->cacheLevel;
         mSsmBlocks[kDefaultBeamIndex][*ssmLcId] = snapshotPage->hold();
+        if (claimRetention)
+        {
+            auto const retention = getRetention(snapshotBlock.ordinal());
+            snapshotPage->claimRetention(retention.retentionPriority, retention.durationMs);
+        }
     }
 
     _finalizeCachedTokensByLevel(numTokens, attentionLevels, ssmLevel);