[None][feat] KV cache manager v2: demote finished sub-agent conversations when the parent returns

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp

@@ -2734,6 +2734,82 @@ std::unique_ptr<PlannedDropHandle> KvCache::planCommittedBlockDrop()
     return std::make_unique<PlannedDropHandle>(*mManager, pagesToDrop);
 }
 
+// ---------------------------------------------------------------------------
+// CommittedRetentionHandle
+// ---------------------------------------------------------------------------
+
+CommittedRetentionHandle::CommittedRetentionHandle(KvCacheManager& manager, Block& leaf, BlockOrdinal ownedBegin)
+    : mManager(manager.shared_from_this())
+    , mLeaf(leaf.sharedFromThis())
+    , mOwnedBegin(ownedBegin)
+{
+}
+
+CommittedRetentionHandle::~CommittedRetentionHandle()
+{
+    KVCM2_POISON_ON_EXCEPT(
+        [this]()
+        {
+            // Runs from tp_dealloc/GC with the GIL held: release it before taking the API lock,
+            // as ~PlannedDropHandle does.
+            OptionalGilRelease const gilRelease;
+            auto const apiLock = mManager->lockExclusive();
+            mLeaf = WeakPtr<Block>{};
+        });
+}
+
+int CommittedRetentionHandle::setPriority(Priority priority)
+{
+    TLLM_CHECK_WITH_INFO(
+        kPriorityMin <= priority && priority <= kPriorityMax, "Retention priority must be between 0 and 100");
+    // Exclusive: re-queues pages in the shared eviction lists and reads the tree's links.
+    auto const apiLock = mManager->lockExclusive();
+    auto const leaf = mLeaf.lock();
+    int numPages = 0;
+    // The leaf must not have been extended; an ancestor on the path continues to exactly one block.
+    size_t expectedNext = 0;
+    NodeBase* node = leaf.get();
+    while (node != nullptr && node->type() == NodeBase::Type::kBLOCK)
+    {
+        auto* blk = static_cast<Block*>(node);
+        if (blk->prev == nullptr || blk->ordinal() < mOwnedBegin || blk->next.size() > expectedNext)
+        {
+            break;
+        }
+        for (LifeCycleId lc{0}; lc < blk->numLifeCycles(); ++lc)
+        {
+            if (CommittedPage* page = blk->getPage(lc))
+            {
+                page->claimRetention(priority, std::nullopt);
+                ++numPages;
+            }
+        }
+        node = blk->prev;
+        expectedNext = 1;
+    }
+    return numPages;
+}
+
+std::unique_ptr<CommittedRetentionHandle> KvCache::planCommittedRetention(int ownedBeginTokens)
+{
+    KVCM2_REJECT_IF_POISONED();
+    auto const apiLock = mManager->lockExclusive();
+    if (mCommitState != CommitState::USER_STOP)
+        throw LogicError("plan_committed_retention() requires stop_committing()");
+    TLLM_CHECK_WITH_INFO(ownedBeginTokens >= 0, "owned_begin_tokens must be non-negative");
+    if (numCommittedTokens() == 0)
+        return nullptr;
+    auto const match = mManager->radixTree().match(
+        mReuseScope, toSpan(mCommittedTokens), textOnly(), mManager->enablePartialMatch());
+    if (match.numTokens != numCommittedTokens() || match.blocks.empty())
+        return nullptr;
+    // A block that starts inside the reused prefix is (partly) someone else's.
+    BlockOrdinal const ownedBegin{(ownedBeginTokens + mTokensPerBlock - 1) / mTokensPerBlock};
+    if (ownedBegin >= match.blocks.size())
+        return nullptr;
+    return std::make_unique<CommittedRetentionHandle>(*mManager, *match.blocks.back(), ownedBegin);
+}
+
 // ---------------------------------------------------------------------------
 // _onStopCommitting: release stale held uncommitted pages for SWA layers.
 // Mirrors Python's _on_stop_committing().

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h

@@ -235,6 +235,38 @@ class PlannedDropHandle
     std::optional<std::vector<WeakPtr<CommittedPage>>> mPageRefs;
 };
 
+// ---------------------------------------------------------------------------
+// CommittedRetentionHandle — re-prioritizes the committed blocks a closed
+// request computed itself, without owning them.
+//
+// Keeps a weak reference to the last block of the request's committed path.
+// setPriority() walks toward the root over the blocks at or after `ownedBegin`
+// (the first block the conversation computed rather than reused) and stops at a
+// block another path extends or branches from: from there up the prefix is shared.
+// ---------------------------------------------------------------------------
+class CommittedRetentionHandle
+{
+public:
+    // Holds a strong reference to the manager (setPriority() takes its API lock), like
+    // PlannedDropHandle: release every handle before shutting the manager down.
+    CommittedRetentionHandle(KvCacheManager& manager, Block& leaf, BlockOrdinal ownedBegin);
+
+    // Releases the weak block reference under the manager's API lock (its count is not atomic).
+    ~CommittedRetentionHandle();
+
+    CommittedRetentionHandle(CommittedRetentionHandle const&) = delete;
+    CommittedRetentionHandle& operator=(CommittedRetentionHandle const&) = delete;
+
+    // Returns the number of pages whose priority was set: zero once the leaf block left the
+    // tree or the path became shared.
+    int setPriority(Priority priority);
+
+private:
+    std::shared_ptr<KvCacheManager> mManager;
+    WeakPtr<Block> mLeaf;
+    BlockOrdinal mOwnedBegin;
+};
+
 // Availability at the claimed prefix, indexed by LayerGroupId. OPTIONAL
 // candidates are held until the first successful resume or close.
 // Describes the initial prefix claim, not subsequent request progress.
@@ -500,6 +532,13 @@ class KvCache : public std::enable_shared_from_this<KvCache>
     // shut down: applying a plan takes the manager's API lock and mutates its state.
     std::unique_ptr<PlannedDropHandle> planCommittedBlockDrop();
 
+    // Handle for lowering the retention priority of this request's own committed blocks later
+    // (e.g. once its conversation is known to be finished). `ownedBeginTokens` is the prefix the
+    // conversation reused rather than computed; blocks that start inside it are never touched.
+    // Must be called after stopCommitting(). Returns nullptr when the committed path is not fully
+    // present in the tree or the request owns no block.
+    std::unique_ptr<CommittedRetentionHandle> planCommittedRetention(int ownedBeginTokens);
+
     int historyLength() const noexcept
     {
         return mHistoryLength;

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp

@@ -1060,22 +1060,27 @@ void StorageManager::_prepareFreeSlots(TypedVec<CacheLevel, TypedVec<PoolGroupIn
                 }
             }
         });
+    // One stable pass per life cycle, not an erase per moved page (quadratic for a whole demoted
+    // prefix): kept and moved pages both keep their order.
     for (auto& pages : fallenPages)
     {
-        for (auto iter = pages.begin(); iter != pages.end();)
+        auto kept = pages.begin();
+        for (auto iter = pages.begin(); iter != pages.end(); ++iter)
         {
             auto const& page = *iter;
             if (page->cacheLevel == kHotLevel && page->status() == PageStatus::DROPPABLE
                 && page->priority() < mColdTierOffloadMinPriority)
             {
                 belowOffloadThreshold.at(page->lifeCycle).push_back(std::move(*iter));
-                iter = pages.erase(iter);
             }
             else
             {
-                ++iter;
+                if (kept != iter)
+                    *kept = std::move(*iter);
+                ++kept;
             }
         }
+        pages.erase(kept, pages.end());
     }
     if (TLLM_UNLIKELY(gDebug))
     {

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp

@@ -920,6 +920,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m)
     m.attr("NDEBUG") = nb::bool_(!kv::gDebug);
     m.attr("DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY")
         = nb::int_(tensorrt_llm::batch_manager::kDefaultColdTierOffloadMinPriority);
+    m.attr("DEFAULT_RETENTION_PRIORITY") = nb::int_(kv::kPriorityDefault);
 
     // ---- Exceptions --------------------------------------------------------
     static nb::object sOutOfMemoryError = nb::exception<kv::OutOfMemoryError>(m, "OutOfMemoryError");
@@ -1993,6 +1994,10 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m)
     nb::class_<kv::PlannedDropHandle>(m, "PlannedDropHandle")
         .def("drop", &kv::PlannedDropHandle::drop, nb::call_guard<nb::gil_scoped_release>());
 
+    nb::class_<kv::CommittedRetentionHandle>(m, "CommittedRetentionHandle")
+        .def("set_priority", &kv::CommittedRetentionHandle::setPriority, nb::arg("priority"),
+            nb::call_guard<nb::gil_scoped_release>());
+
     nb::class_<kv::ReuseGroupStatus>(m, "ReuseGroupStatus")
         .def_prop_ro("group_id", [](kv::ReuseGroupStatus const& s) { return s.groupId.value(); })
         .def_ro("policy", &kv::ReuseGroupStatus::policy)
@@ -2202,6 +2207,10 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m)
         .def(
             "plan_committed_block_drop", [](kv::KvCache& self) { return self.planCommittedBlockDrop(); },
             nb::call_guard<nb::gil_scoped_release>())
+        .def(
+            "plan_committed_retention",
+            [](kv::KvCache& self, int ownedBeginTokens) { return self.planCommittedRetention(ownedBeginTokens); },
+            nb::arg("owned_begin_tokens") = 0, nb::call_guard<nb::gil_scoped_release>())
         .def(
             "get_aggregated_page_indices",
             [](kv::KvCache const& self, int layerGroupId, int beamIdx, bool validOnly) {