From e3b4392437d29aba5f75d484ec417cb78a489354 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 10 Oct 2026 02:59:15 -0700 Subject: [PATCH 01/23] feat: implement native V2 KV cache retention Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../kv_cache_manager_v2/config.cpp | 4 + .../kv_cache_manager_v2/config.h | 3 + .../kv_cache_manager_v2/eventManager.cpp | 8 +- .../kv_cache_manager_v2/eventManager.h | 2 + .../kv_cache_manager_v2/eventSink.h | 3 + .../evictionController.cpp | 4 +- .../kv_cache_manager_v2/kvCache.cpp | 67 ++++++- .../kv_cache_manager_v2/kvCache.h | 12 +- .../kv_cache_manager_v2/kvCacheManager.cpp | 14 +- .../kv_cache_manager_v2/kvCacheManager.h | 6 +- .../kv_cache_manager_v2/page.cpp | 21 ++- .../batch_manager/kv_cache_manager_v2/page.h | 21 ++- .../kv_cache_manager_v2/storageManager.cpp | 110 ++++++++++- .../kv_cache_manager_v2/storageManager.h | 12 +- .../batch_manager/kvCacheManagerV2.cpp | 24 ++- .../source/examples/kvcacheretentionconfig.md | 26 ++- .../connectors/kv_cache_connector.py | 24 ++- .../kv_cache/kv_cache_manager_v2.py | 12 ++ .../runtime/kv_cache_manager_v2/__init__.py | 1 + .../runtime/kv_cache_manager_v2/__init__.pyi | 8 + .../defs/llmapi/test_llm_api_connector.py | 141 +++----------- .../integration/test_lists/test-db/l0_a10.yml | 5 +- .../test_lists/test-db/l0_b200.yml | 5 + .../test_kv_cache_manager_v2.py | 177 ++++++++++++++++++ 24 files changed, 560 insertions(+), 150 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp index 8a6e20701d73..425a25bddd02 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp @@ -40,6 +40,10 @@ void DiskCacheTierConfig::validate() const void KVCacheManagerConfig::validate() const { + if (secondaryOffloadMinPriority < kPriorityMin || secondaryOffloadMinPriority > kPriorityMax) + { + throw AssertionError("secondary_offload_min_priority must be between 0 and 100"); + } if (swaScratchReuse.has_value()) { swaScratchReuse->validate(); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h index 355f6adde0ba..715742299fdc 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h +++ b/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 secondaryOffloadMinPriority = 30; + bool enableSwaScratchReuse() const noexcept { return swaScratchReuse.has_value(); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp index c17071a55de5..da9f2b96796e 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp @@ -311,6 +311,12 @@ void EventManager::addCacheLevelUpdated( lifeCycle.value()); } +void EventManager::addPriorityUpdated( + Digest const& blockKey, Priority oldPriority, Priority newPriority, LifeCycleId lifeCycle) +{ + addUpdatedEvent(blockKey, std::nullopt, KVCacheEventDiff{oldPriority, newPriority}, lifeCycle.value()); +} + void EventManager::addStoredEventUnlocked(KVCacheStoredData data, EventLayerGroupId layerGroupId) { bool const hasPendingRemovedEvents = !mLatestRemovedBlockHashes.empty(); @@ -547,7 +553,7 @@ std::optional EventManager::storedBlockFromBlock( if (pageCoversBlock(page, block)) { cacheLevel = page->cacheLevel; - priority = page->priority; + priority = page->priority(); foundPage = true; break; } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.h index 8f00c24b0a9f..c4bffdc0333f 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.h +++ b/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 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventSink.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventSink.h index fbdce560b871..eace2073f4b0 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventSink.h +++ b/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 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.cpp index 015271306fe0..7c12e436040d 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/evictionController.cpp +++ b/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, 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 PrioritizedEvictionPolicy::pop() SharedPtr 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); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index 4bfc0d244370..62ca6dc08ab0 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/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 reuseMatch, std::optional mId, PriorityCb priorityCb, std::optional expectedPromptLength, - std::optional textOnly, bool enableRequestStats) + std::optional textOnly, bool enableRequestStats, + std::optional retentionConfig) : 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::optionalgetTransferMode() != 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 @@ -143,9 +165,44 @@ 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(ordinal.value()) < mPromptRetentions.size()) + { + return mPromptRetentions.at(ordinal.value()); + } + return {mRetentionConfig->getDecodeRetentionPriority(), mRetentionConfig->getDecodeDurationMs()}; +} + +std::vector KvCache::getPagePriorities(LayerGroupId layerGroupId, BeamIndex beamIndex) const +{ + KVCM2_REJECT_IF_POISONED(); + auto const apiLock = mManager->lockShared(); + std::vector result; + result.reserve(mBlocks.stdSize()); + for (auto const& block : mBlocks) + { + auto const* page = beamIndex < block.pages.size() + ? blockPageGetPage(block.pages[beamIndex].at(layerGroupId)).get() + : nullptr; + // 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(); @@ -1315,6 +1372,8 @@ CommittedPage* KvCache::_copyPageToTreeBlock( newSlot.readyEvent = CachedCudaEvent(reinterpret_cast(stream)); auto committed = makeShared( &storageMgr, treeBlock, lcIdx, lvl, numTokensInBlock, getPriority(treeBlock->ordinal(), lcIdx)); + auto const retention = getRetention(treeBlock->ordinal()); + committed->claimRetention(retention.retentionPriority, retention.durationMs); committed->setSlot(newSlot); // Drops the superseded page, deferred until the copy is issued: an // OutOfPagesError above must not destroy a usable shorter snapshot. @@ -2333,6 +2392,8 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) for (size_t ri = 0; ri < reuseTasks.size(); ++ri) { LifeCycleId lc = reuseTasks[ri].lifeCycle; + auto const retention = getRetention(BlockOrdinal{ord}); + locks[ri].page()->claimRetention(retention.retentionPriority, retention.durationMs); sb.pages[kDefaultBeamIndex][lc] = std::move(locks[ri]); } if (mIsDecoding) @@ -2867,6 +2928,8 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match) CacheLevel const level = page->cacheLevel; auto& bpSlot = mBlocks[ordinal].pages[beamIdx][lcId]; bpSlot = page->hold(); + auto const retention = getRetention(ordinal); + page->claimRetention(retention.retentionPriority, retention.durationMs); if (!isAttention) { return; @@ -2937,6 +3000,8 @@ 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(); + auto const retention = getRetention(snapshotBlock.ordinal()); + snapshotPage->claimRetention(retention.retentionPriority, retention.durationMs); } _finalizeCachedTokensByLevel(numTokens, attentionLevels, ssmLevel); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h index 8444d3049f72..b42556ab2c7c 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h @@ -27,6 +27,7 @@ #include "kv_cache_manager_v2/utils/funcGuard.h" #include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/executor/executor.h" #include #include #include @@ -248,7 +249,8 @@ class KvCache : public std::enable_shared_from_this KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional reuseMatch, std::optional id, PriorityCb priorityCb, std::optional expectedPromptLength = std::nullopt, - std::optional textOnly = std::nullopt, bool enableRequestStats = false); + std::optional textOnly = std::nullopt, bool enableRequestStats = false, + std::optional retentionConfig = std::nullopt); ~KvCache(); @@ -537,6 +539,12 @@ class KvCache : public std::enable_shared_from_this // Priority for (blockOrdinal, lifeCycleId) based on the callback. Priority getPriority(BlockOrdinal ordinal, LifeCycleId lc) const; + //! Retention is resolved by block start, with decode beginning after the final prompt block. + executor::RetentionPriorityAndDuration getRetention(BlockOrdinal ordinal) const; + + //! Current priorities for the request's allocated pages in one layer group. + std::vector getPagePriorities(LayerGroupId layerGroupId, BeamIndex beamIndex = kDefaultBeamIndex) const; + // Reference to StorageManager (for page acquisition/release). StorageManager* storageManager() const; @@ -760,6 +768,8 @@ class KvCache : public std::enable_shared_from_this std::shared_ptr mManager; ReuseScope mReuseScope; PriorityCb mPriorityCb; + std::optional mRetentionConfig; + std::vector mPromptRetentions; std::optional mCudaStream; Status mStatus; CommitState mCommitState; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp index f91f62018dee..65faccbef8a3 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp @@ -125,7 +125,7 @@ KvCacheManager::KvCacheManager(KVCacheManagerConfig const& config, std::shared_p StorageConfig storageConfig = createStorageConfig(mConfig); mStorage = std::make_shared(mLifeCycles, storageConfig, mConfig.tokensPerBlock, std::move(coldPageCodec), mConfig.swaScratchReuse, mConfig.typicalStep, mConfig.constraints, - mConfig.initialPoolRatio, mEventSink, mConfig.maxUtilForResume); + mConfig.initialPoolRatio, mEventSink, mConfig.maxUtilForResume, mConfig.secondaryOffloadMinPriority); mTargetRatioListHot = _currentHotRatio(); mTargetRatioListCold = _currentColdRatios(); @@ -187,7 +187,8 @@ void KvCacheManager::clearReusableBlocks() std::shared_ptr KvCacheManager::createKvCache(ReuseScope reuseScope, TokenSpan inputTokens, std::optional id, KvCache::PriorityCb priorityCb, std::optional expectedPromptLength, - std::optional textOnly, bool enableRequestStats) + std::optional textOnly, bool enableRequestStats, + std::optional retentionConfig) { KVCM2_API_GUARD(); auto const apiLock = lockExclusive(); @@ -211,7 +212,14 @@ std::shared_ptr KvCacheManager::createKvCache(ReuseScope reuseScope, To } return std::make_shared(*this, std::move(reuseScope), std::move(reuseMatch), std::move(id), - std::move(priorityCb), expectedPromptLength, textOnly, enableRequestStats); + std::move(priorityCb), expectedPromptLength, textOnly, enableRequestStats, std::move(retentionConfig)); +} + +void KvCacheManager::refreshRetention() +{ + KVCM2_API_GUARD(); + auto const apiLock = lockExclusive(); + mStorage->refreshRetention(); } BlockRadixTree::ReuseMatch KvCacheManager::matchReuse( diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h index 0adb1976874d..d011bdf233ac 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h @@ -147,7 +147,11 @@ class KvCacheManager : public std::enable_shared_from_this std::shared_ptr createKvCache(ReuseScope reuseScope = {}, TokenSpan inputTokens = {}, std::optional id = std::nullopt, KvCache::PriorityCb priorityCb = {}, std::optional expectedPromptLength = std::nullopt, std::optional textOnly = std::nullopt, - bool enableRequestStats = false); + bool enableRequestStats = false, + std::optional retentionConfig = std::nullopt); + + //! Restore expired, unreferenced pages to the default retention priority. + void refreshRetention(); // knownNoDigest: from external text_only knowledge, never a scan (see Hasher::update). // Defaults false (safe: the scanning path is taken). diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp index 2fb89b143ea6..a7db0f0ac043 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp @@ -37,8 +37,8 @@ Page::Page(StorageManager* mgr, LifeCycleId lc, CacheLevel level, Priority prio) : manager(mgr) , lifeCycle(lc) , cacheLevel(level) - , priority(prio) , nodeRef(std::nullopt) + , mPriority(prio) { } @@ -52,6 +52,7 @@ Page::~Page() KVCM2_POISON_ON_EXCEPT( [this]() { + manager->cancelRetentionExpiry(*this); TLLM_CHECK_DEBUG_WITH_INFO(status() == PageStatus::DROPPABLE && !scheduledForEviction(), "Page destroyed while still held or scheduled for eviction"); if (hasValidSlot()) @@ -77,6 +78,7 @@ PageStatus Page::status() const noexcept SharedPtr Page::hold() { + manager->cancelRetentionExpiry(*this); // Return existing holder if any. auto h = holder.lock(); if (h) @@ -98,6 +100,16 @@ SharedPtr Page::hold() return h; } +void Page::claimRetention(std::optional priority, std::optional duration) +{ + manager->cancelRetentionExpiry(*this); + if (priority.has_value()) + { + manager->updatePriority(*this, *priority); + } + mRetentionDuration = duration; +} + SharedPageLock Page::lock(KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal ordinal, LifeCycleId lc, bool skipWait) { return hold()->lock(kvCache, beamIndex, ordinal, lc, skipWait); @@ -165,6 +177,8 @@ UncommittedPage::UncommittedPage(KvCache& kvc, BlockOrdinal ord, LifeCycleId lc, , ordinal(ord) , beamIndex(bi) { + auto const retention = kvc.getRetention(ord); + claimRetention(retention.retentionPriority, retention.durationMs); } UncommittedPage::~UncommittedPage() @@ -219,7 +233,9 @@ SharedPtr UncommittedPage::convertToCommitted( // Set the ready event before transfer (matches Python: self.ready_event = ready_event). this->readyEvent = std::move(readyEv); - auto committed = makeShared(manager, blk, lifeCycle, cacheLevel, numTokensInBlock, priority); + auto committed = makeShared(manager, blk, lifeCycle, cacheLevel, numTokensInBlock, priority()); + auto const retention = kvCache->getRetention(blk->ordinal()); + committed->claimRetention(retention.retentionPriority, retention.durationMs); // Move slot id to the committed page; invalidate our slot. committed->setSlotId(slotId()); // asserts valid committed->readyEvent = std::move(readyEvent); @@ -257,6 +273,7 @@ PageHolder::~PageHolder() // If it's a committed page, schedule for eviction (if evictable). if (page->isCommitted()) { + manager->scheduleRetentionExpiry(*page); if (!page->scheduledForEviction()) manager->scheduleForEviction(*page); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h index 9ecd538fe079..08ef0a5e8532 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h @@ -25,7 +25,9 @@ #include "kv_cache_manager_v2/utils/cudaEvent.h" #include "kv_cache_manager_v2/utils/sharedPtr.h" +#include #include +#include #include #include @@ -38,6 +40,8 @@ class KvCache; class PageHolder; class UniqPageLock; class SharedPageLock; +class Page; +using RetentionExpiryQueue = std::multimap>; // --------------------------------------------------------------------------- // Page — base class for all KV-cache pages. @@ -50,9 +54,6 @@ class Page : public Slot, public EnableSharedFromThis StorageManager* manager; LifeCycleId lifeCycle; CacheLevel cacheLevel; - // Immutable: PrioritizedEvictionPolicy locates a scheduled page's sub-queue by this value, - // so changing it while the page is scheduled would erase from the wrong list. - Priority const priority; WeakPtr holder; // empty → DROPPABLE std::optional nodeRef; // present → scheduled for eviction @@ -64,6 +65,14 @@ class Page : public Slot, public EnableSharedFromThis bool isPrivate() const; + Priority priority() const noexcept + { + return mPriority; + } + + //! Apply a request claim, preserving priority when unspecified and replacing its duration. + void claimRetention(std::optional priority, std::optional duration); + PageStatus status() const noexcept; bool scheduledForEviction() const noexcept @@ -82,6 +91,12 @@ class Page : public Slot, public EnableSharedFromThis // skip_wait: caller guarantees the page is ready on kvCache's stream. SharedPageLock lock( KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal ordinal, LifeCycleId lifeCycle, bool skipWait = false); + +private: + friend class StorageManager; + Priority mPriority; + std::optional mRetentionDuration; + std::optional mRetentionExpiry; }; // --------------------------------------------------------------------------- diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp index 0f19a38841a6..f5072e53f3cb 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp @@ -162,8 +162,8 @@ void sortFallenPagesByPriority(TypedVec> { for (auto& pages : fallenPages) { - std::stable_sort( - pages.begin(), pages.end(), [](auto const& lhs, auto const& rhs) { return lhs->priority < rhs->priority; }); + std::stable_sort(pages.begin(), pages.end(), + [](auto const& lhs, auto const& rhs) { return lhs->priority() < rhs->priority(); }); } } @@ -214,9 +214,10 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi std::unique_ptr coldPageCodec, std::optional swaScratchReuse, std::optional const& typicalBatch, std::vector const& constraints, std::optional> const& initialPoolRatio, std::shared_ptr eventSink, - float maxUtilForResume) + float maxUtilForResume, Priority secondaryOffloadMinPriority) : mLifeCycles(lifeCycles) , mEventSink(std::move(eventSink)) + , mSecondaryOffloadMinPriority(secondaryOffloadMinPriority) , mHotPoolGroupMapping(config.lifeCycleGrouping()) , mSwaScratchReuse(std::move(swaScratchReuse)) , mColdPageCodec(coldPageCodec ? std::move(coldPageCodec) : createDefaultKvCacheColdPageCodec()) @@ -770,7 +771,76 @@ bool StorageManager::isEvictable(Page const& page, std::optional lev void StorageManager::scheduleForEviction(Page& page) { if (isEvictable(page)) + { + mLevels.at(page.cacheLevel).controller.scheduleForEviction(page); + scheduleRetentionExpiry(page); + } +} + +void StorageManager::updatePriority(Page& page, Priority priority) +{ + TLLM_CHECK_WITH_INFO( + kPriorityMin <= priority && priority <= kPriorityMax, "Retention priority must be between 0 and 100"); + if (priority == page.mPriority) + { + return; + } + // Keep ownership across removal: a droppable page may be owned only by its queue. + auto const keepAlive = page.sharedFromThis(); + bool const queued = page.scheduledForEviction(); + if (queued) + { + excludeFromEviction(page); + } + auto const oldPriority = page.mPriority; + page.mPriority = priority; + if (queued) + { mLevels.at(page.cacheLevel).controller.scheduleForEviction(page); + } + if (mEventSink && page.isCommitted()) + { + auto const& committed = static_cast(page); + if (committed.block != nullptr && committed.block->holdsPage(committed)) + { + mEventSink->addPriorityUpdated(committed.block->key, oldPriority, priority, page.lifeCycle); + } + } +} + +void StorageManager::cancelRetentionExpiry(Page& page) +{ + if (page.mRetentionExpiry.has_value()) + { + mRetentionExpiries.erase(*page.mRetentionExpiry); + page.mRetentionExpiry.reset(); + } +} + +void StorageManager::scheduleRetentionExpiry(Page& page) +{ + if (page.status() == PageStatus::DROPPABLE && page.mRetentionDuration.has_value() + && page.priority() != kPriorityDefault && !page.mRetentionExpiry.has_value()) + { + page.mRetentionExpiry = mRetentionExpiries.emplace( + std::chrono::steady_clock::now() + *page.mRetentionDuration, page.sharedFromThis()); + } +} + +void StorageManager::refreshRetention() +{ + auto const now = std::chrono::steady_clock::now(); + while (!mRetentionExpiries.empty() && mRetentionExpiries.begin()->first <= now) + { + auto const page = mRetentionExpiries.begin()->second.lock(); + mRetentionExpiries.erase(mRetentionExpiries.begin()); + if (page) + { + page->mRetentionExpiry.reset(); + TLLM_CHECK_DEBUG(page->status() == PageStatus::DROPPABLE); + updatePriority(*page, kPriorityDefault); + } + } } void StorageManager::excludeFromEviction(Page& page) @@ -810,6 +880,7 @@ void StorageManager::excludeFromEviction(Page& page) void StorageManager::prepareFreeSlots(CacheLevel level, TypedVec const& requirements, MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) { + refreshRetention(); TypedVec> goals(numCacheLevels()); for (CacheLevel lvl{0}; lvl < goals.size(); ++lvl) { @@ -827,6 +898,7 @@ void StorageManager::prepareFreeSlots(CacheLevel level, TypedVec const& minNumPages, DropRecorder const& dropRecorder) { + refreshRetention(); auto evicted = mLevels.at(level).controller.evict(minNumPages); auto rescheduleEvictedPagesOnFailure = makeEvictionRollbackGuard(evicted); @@ -883,6 +955,28 @@ void StorageManager::_prepareFreeSlots(TypedVec> dropped; + for (auto iter = pages.begin(); iter != pages.end();) + { + auto const& page = *iter; + if (page->cacheLevel == kHotLevel && page->status() == PageStatus::DROPPABLE + && page->priority() < mSecondaryOffloadMinPriority) + { + dropped.push_back(std::move(*iter)); + iter = pages.erase(iter); + } + else + { + ++iter; + } + } + if (dropRecorder && !dropped.empty()) + { + dropRecorder(dropped, kHotLevel); + } + } if (TLLM_UNLIKELY(gDebug)) { TLLM_CHECK_WITH_INFO(goals.size() == numCacheLevels(), "goals.rows must equal numCacheLevels"); @@ -976,7 +1070,7 @@ void StorageManager::_prepareFreeSlots(TypedVec std::optional { auto const& fallen = fallenPages.at(lifeCycle); - return fallen.empty() ? std::nullopt : std::optional{fallen.back()->priority}; + return fallen.empty() ? std::nullopt : std::optional{fallen.back()->priority()}; }; while (count > 0) { @@ -988,7 +1082,7 @@ void StorageManager::_prepareFreeSlots(TypedVecpriority < priority; }); + [](auto const& page, Priority priority) { return page->priority() < priority; }); SlotCount const numAccepted = std::min(count, slotCountValueFromSize(std::distance(matchingBegin, fallen.end()))); matchingBegin = fallen.end() - numAccepted; @@ -1095,6 +1189,11 @@ void StorageManager::_prepareFreeSlots(TypedVecstatus() == PageStatus::DROPPABLE) + { + // Offloading an unreferenced page ends its released retention interval. + page->claimRetention(std::nullopt, std::nullopt); + } scheduleForEviction(*page); } } @@ -1719,6 +1818,7 @@ void StorageManager::expandPoolGroup(CacheLevel level, PoolGroupIndex pgIdx, Slo void StorageManager::shrinkPoolGroup( CacheLevel level, PoolGroupIndex pgIdx, SlotCount newNumSlots, std::vector> const& persistentPages) { + refreshRetention(); auto& pg = poolGroup(level, pgIdx); auto& allocator = pg.slotAllocator(); auto& ctrl = mLevels.at(level).controller; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h index d9c9b64abef9..816a6eaa330f 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h @@ -23,6 +23,7 @@ #include "kv_cache_manager_v2/eventSink.h" #include "kv_cache_manager_v2/evictionController.h" #include "kv_cache_manager_v2/lifeCycleRegistry.h" +#include "kv_cache_manager_v2/page.h" #include "kv_cache_manager_v2/storage/config.h" #include "kv_cache_manager_v2/storage/core.h" #include "tensorrt_llm/common/assert.h" @@ -128,7 +129,8 @@ class StorageManager : public std::enable_shared_from_this std::optional swaScratchReuse = std::nullopt, std::optional const& typicalBatch = std::nullopt, std::vector const& constraints = {}, std::optional> const& initialPoolRatio = std::nullopt, - std::shared_ptr eventSink = nullptr, float maxUtilForResume = 1.0f); + std::shared_ptr eventSink = nullptr, float maxUtilForResume = 1.0f, + Priority secondaryOffloadMinPriority = 30); ~StorageManager(); StorageManager(StorageManager const&) = delete; @@ -164,6 +166,12 @@ class StorageManager : public std::enable_shared_from_this // Remove a page from the eviction queue. void excludeFromEviction(Page& page); + //! Change priority only while detached from its current priority queue. + void updatePriority(Page& page, Priority priority); + void cancelRetentionExpiry(Page& page); + void scheduleRetentionExpiry(Page& page); + void refreshRetention(); + // Check if a page is evictable (optionally at a target level). bool isEvictable(Page const& page, std::optional level = std::nullopt) const noexcept; @@ -416,6 +424,7 @@ class StorageManager : public std::enable_shared_from_this LifeCycleRegistry const& mLifeCycles; std::shared_ptr mEventSink; + Priority mSecondaryOffloadMinPriority; LifeCyclePoolGroupMapping mHotPoolGroupMapping; LifeCyclePoolGroupMapping mColdPoolGroupMapping; // Codec batching representative for each lifecycle. @@ -462,6 +471,7 @@ class StorageManager : public std::enable_shared_from_this TypedVec mMinSlots; // All GPU cache levels borrow this allocator. It must outlive mLevels. std::unique_ptr mGpuPhysMemAllocator; + RetentionExpiryQueue mRetentionExpiries; TypedVec mLevels; static constexpr size_t kDefaultPageStagingBytes = 64u << 20u; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index f9a5c5db276d..1bbc5d1ba78c 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1748,7 +1748,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) std::vector constraints, std::optional typicalStep, std::optional> initialPoolRatio, std::optional swaScratchReuse, bool commitMinSnapshot, bool enableStats, - bool textOnly, bool enablePartialCommit, kv::PoolRebalanceConfig poolRebalance) + bool textOnly, bool enablePartialCommit, kv::PoolRebalanceConfig poolRebalance, + kv::Priority secondaryOffloadMinPriority) { new (cfg) kv::KVCacheManagerConfig(); cfg->tokensPerBlock = tokensPerBlock; @@ -1773,6 +1774,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) cfg->textOnly = textOnly; cfg->enablePartialCommit = enablePartialCommit; cfg->poolRebalance = std::move(poolRebalance); + cfg->secondaryOffloadMinPriority = secondaryOffloadMinPriority; // Mirror Python's __post_init__: validate at construction. Config-integrity // failures raise AssertionError (translated below). cfg->validate(); @@ -1783,7 +1785,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::arg("typical_step") = std::nullopt, nb::arg("initial_pool_ratio").none() = std::nullopt, nb::arg("swa_scratch_reuse").none() = std::nullopt, nb::arg("commit_min_snapshot") = false, nb::arg("enable_stats") = true, nb::arg("text_only") = false, nb::arg("enable_partial_commit") = true, - nb::arg("pool_rebalance") = kv::PoolRebalanceConfig{}) + nb::arg("pool_rebalance") = kv::PoolRebalanceConfig{}, + nb::arg("secondary_offload_min_priority") = 30) .def_rw("tokens_per_block", &kv::KVCacheManagerConfig::tokensPerBlock) .def_rw("cache_tiers", &kv::KVCacheManagerConfig::cacheTiers) .def_rw("layers", &kv::KVCacheManagerConfig::layers) @@ -1791,6 +1794,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def_rw("enable_partial_reuse", &kv::KVCacheManagerConfig::enablePartialReuse) .def_rw("reuse_match_backoff", &kv::KVCacheManagerConfig::reuseMatchBackoff) .def_rw("enable_partial_commit", &kv::KVCacheManagerConfig::enablePartialCommit) + .def_rw("secondary_offload_min_priority", &kv::KVCacheManagerConfig::secondaryOffloadMinPriority) .def_rw("typical_step", &kv::KVCacheManagerConfig::typicalStep) .def_rw("constraints", &kv::KVCacheManagerConfig::constraints) .def_rw("initial_pool_ratio", &kv::KVCacheManagerConfig::initialPoolRatio, @@ -1954,6 +1958,11 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::call_guard()) .def("acknowledge_page_storage", &kv::KvCache::acknowledgePageStorage, nb::arg("version"), nb::call_guard()) + .def( + "get_page_priorities", + [](kv::KvCache const& self, int layerGroupId, int beamIndex) + { return self.getPagePriorities(kv::LayerGroupId{layerGroupId}, kv::BeamIndex{beamIndex}); }, + nb::arg("layer_group_id"), nb::arg("beam_index") = 0, nb::call_guard()) .def( "get_page_storage_snapshot", [](kv::KvCache const& self, int layerGroupId, int beamIdx) @@ -2640,6 +2649,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::arg("config"), nb::arg("event_manager").none() = nb::none(), nb::arg("cold_page_codec").none() = nb::none()) .def("shutdown", &kv::KvCacheManager::shutdown, nb::call_guard()) + .def("refresh_retention", &kv::KvCacheManager::refreshRetention, nb::call_guard()) .def("is_sparse", &kv::KvCacheManager::isSparse, nb::arg("layer_id"), nb::arg("data_role")) .def( "clear_reusable_blocks", &kv::KvCacheManager::clearReusableBlocks, nb::call_guard()) @@ -2647,7 +2657,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) "create_kv_cache", [](std::shared_ptr self, nb::object reuseScopeObj, nb::object inputTokens, std::optional id, nb::object customPriorityCallback, - std::optional expectedPromptLength, std::optional textOnly, bool enableRequestStats) + std::optional expectedPromptLength, std::optional textOnly, bool enableRequestStats, + std::optional retentionConfig) { kv::ReuseScope reuseScope = castReuseScope(std::move(reuseScopeObj)); kv::KvCache::PriorityCb priorityCb = castPriorityCallback(*self, std::move(customPriorityCallback)); @@ -2655,7 +2666,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) { nb::gil_scoped_release release; return self->createKvCache(std::move(reuseScope), kv::TokenSpan{}, id, std::move(priorityCb), - expectedPromptLength, textOnly, enableRequestStats); + expectedPromptLength, textOnly, enableRequestStats, std::move(retentionConfig)); } return withTokens(inputTokens, [&](kv::TokenSpan view, bool knownNoDigest) @@ -2672,12 +2683,13 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) } nb::gil_scoped_release release; return self->createKvCache(std::move(reuseScope), view, id, std::move(priorityCb), promptLen, - textOnly, enableRequestStats); + textOnly, enableRequestStats, std::move(retentionConfig)); }); }, nb::arg("reuse_scope") = nb::none(), nb::arg("input_tokens") = nb::none(), nb::arg("id") = std::nullopt, nb::arg("custom_priority_callback") = nb::none(), nb::arg("expected_prompt_length") = std::nullopt, - nb::arg("text_only") = std::nullopt, nb::arg("enable_request_stats") = false) + nb::arg("text_only") = std::nullopt, nb::arg("enable_request_stats") = false, + nb::arg("kv_cache_retention_config") = std::nullopt) .def( "probe_reuse", [](std::shared_ptr self, nb::object reuseScopeObj, nb::object inputTokens) diff --git a/docs/source/examples/kvcacheretentionconfig.md b/docs/source/examples/kvcacheretentionconfig.md index e490f223798f..4651f725d5dd 100644 --- a/docs/source/examples/kvcacheretentionconfig.md +++ b/docs/source/examples/kvcacheretentionconfig.md @@ -28,7 +28,7 @@ if __name__ == '__main__': main() ``` -The blocks from the prompts are stored for reuse with the default priority of 35 on a scale from 1 to 100, where 100 is highest priority and 1 is lowest priority. Assume you know that the first four tokens of each prompt represent a system prompt that should be stored with high priority (100). You can achieve this by providing a KV cache retention config object when you submit the prompts for generation: +The blocks from the prompts are stored for reuse with the default priority of 35 on a scale from 0 to 100, where 100 is highest priority and 0 is lowest priority. Assume you know that the first four tokens of each prompt represent a system prompt that should be stored with high priority (100). You can achieve this by providing a KV cache retention config object when you submit the prompts for generation: ```python from tensorrt_llm import LLM, SamplingParams @@ -66,3 +66,27 @@ if __name__ == '__main__': ``` This example uses a single ```kv_cache_retention_config``` object for all the prompts. You can also provide a list that must have the same length as the list of prompts. + +KV cache manager V2 applies this configuration to native pages, including reused +prefix pages. A token range selects blocks whose first token falls within the +range. The final partial prompt block uses the prompt policy; subsequent blocks +use the decode policy. + +When requests share a page, the latest claim with an explicit priority sets its +priority. Releasing an older request does not restore that request's priority. +A claim without an explicit priority preserves the existing priority, but replaces +the duration, including clearing it when no duration is supplied. The duration +starts when the last reference is released. Expired pages return to priority 35; +pages still referenced by a request cannot be permanently evicted. + +On GPU eviction, unreferenced pages below `secondary_offload_min_priority` +(default 30) are discarded. Pages at or above it can be offloaded. Offloading +such a page preserves its priority and clears its duration, as in V1; a later +claim can supply a new duration. Referenced pages may move between tiers but +are never discarded by this threshold. + +V2 supports the default `transfer_mode=DRAM`. Its host and disk offload tiers are +configured at manager creation. The per-request `GDS` and `POSIX_DEBUG_FALLBACK` +transfer modes, whose `directory` selects V1 block files, are rejected by V2. +They are not aliases for V2's manager-wide disk tier. With `DRAM`, `directory` +does not select storage. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 3eb1f191cd3c..48d555679f57 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -575,8 +575,7 @@ def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManag new_block_ids_by_layer_group: List[List[int]] = [] if is_v2: - # Block hashes and retention priorities have no V2 accessor yet, so - # they are reported empty rather than guessed at. + # Block hashes have no V2 accessor yet. block_hashes = [] indices_by_group = kv_cache_manager.get_page_indices_by_layer_group(req) while len(self.block_ids_by_layer_group) < len(indices_by_group): @@ -620,13 +619,22 @@ def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManag req ) # Specdec with draft tokens is not supported yet. - # V2 does not honour KvCacheRetentionConfig, so it cannot provide - # per-block retention priorities for connector offload filtering. + # Report the native pages' current priorities, including shared-page claims. priorities = None - if not is_v2 and req.kv_cache_retention_config is not None: - priorities = [ - kv_cache_manager.get_priority_by_block_id(block_id) for block_id in new_block_ids - ] + if req.kv_cache_retention_config is not None: + if is_v2: + priorities = ( + kv_cache_manager.kv_cache_map[req.py_request_id].get_page_priorities(0)[ + len(self.block_ids) - len(new_block_ids) : len(self.block_ids) + ] + if len(new_block_ids_by_layer_group) == 1 + else [] + ) + else: + priorities = [ + kv_cache_manager.get_priority_by_block_id(block_id) + for block_id in new_block_ids + ] return RequestData( req.request_id, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 601407d0d3af..105993450743 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -49,6 +49,7 @@ prefer_pinned, str_dtype_to_torch, ) +from tensorrt_llm.bindings.executor import KvCacheRetentionConfig from tensorrt_llm.bindings.internal.batch_manager import KvCacheIterationStats, KvCacheStats from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( IndexMapper, @@ -3048,6 +3049,11 @@ def _build_base_config( ratio_tolerance=rebalance_cfg.ratio_tolerance, moving_average_decay=rebalance_cfg.moving_average_decay, ), + secondary_offload_min_priority=( + 30 + if kv_cache_config.secondary_offload_min_priority is None + else kv_cache_config.secondary_offload_min_priority + ), ) def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: @@ -3779,6 +3785,7 @@ def prepare_context_cache(self, req: LlmRequest, reuse_limit: int | None = None) req.lora_task_id, tokens, cache_salt=req.cache_salt, + kv_cache_retention_config=req.kv_cache_retention_config, is_dummy=req.is_dummy, enable_request_stats=req.return_perf_metrics, expected_prompt_length=( @@ -4245,6 +4252,7 @@ def get_page_indices_by_layer_group(self, request: LlmRequest) -> List[List[int] @nvtx_range("prepare_resources_kv_cache_manager_v2") def prepare_resources(self, scheduled_batch: ScheduledRequests): + self.impl.refresh_retention() for request in scheduled_batch.context_requests: ready = self._disagg_receive_ready.get(request.py_request_id) if ready is not None: @@ -4619,6 +4627,8 @@ def _mirror_draft_kv_cache(self, req: LlmRequest): None, cache_salt=req.cache_salt, is_dummy=req.is_dummy, + expected_prompt_length=req.prompt_len, + kv_cache_retention_config=req.kv_cache_retention_config, ) if kv_cache is None: return None @@ -6535,6 +6545,7 @@ def _create_kv_cache( is_dummy: bool = False, enable_request_stats: bool = False, expected_prompt_length: int | None = None, + kv_cache_retention_config: KvCacheRetentionConfig | None = None, ): assert request_id not in self.kv_cache_map, ( f"KV cache for request {request_id} already exists" @@ -6575,6 +6586,7 @@ def _create_kv_cache( id=request_id, enable_request_stats=enable_request_stats, expected_prompt_length=expected_prompt_length, + kv_cache_retention_config=kv_cache_retention_config, **priority_kwargs, ) self.kv_cache_map[request_id] = kv_cache diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 903458676ba8..722ec078d2f2 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -133,6 +133,7 @@ class _KVCacheManagerConfigFieldSpec: text_only: bool = False enable_partial_commit: bool = True pool_rebalance: object = None + secondary_offload_min_priority: int = 30 KVCacheManagerConfig.__dataclass_fields__ = _KVCacheManagerConfigFieldSpec.__dataclass_fields__ diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index c1e3ed04e75d..3d58accb38a3 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -32,6 +32,8 @@ from typing import ( Union, ) +from tensorrt_llm.bindings.executor import KvCacheRetentionConfig + # From _common.py NDEBUG: Final[int] DEFAULT_BEAM_INDEX: Final[BeamIndex] @@ -328,6 +330,7 @@ class KVCacheManagerConfig: text_only: bool = False enable_partial_commit: bool = True pool_rebalance: PoolRebalanceConfig = ... + secondary_offload_min_priority: Priority = 30 @property def enable_swa_scratch_reuse(self) -> bool: ... @@ -686,6 +689,9 @@ class _KVCache: """Bind a standalone consumer's row; Batch members must use Batch.add/remove.""" def acknowledge_page_storage(self, version: int) -> bool: """Clear dirty state after all groups/beams use this same version, if it is still current.""" + def get_page_priorities( + self, layer_group_id: LayerGroupId, beam_index: BeamIndex = ... + ) -> list[Priority]: ... def get_page_storage_snapshot( self, layer_group_id: LayerGroupId, beam_id: BeamIndex = DEFAULT_BEAM_INDEX ) -> PageStorageSnapshot: @@ -841,6 +847,7 @@ class KVCacheManager: expected_prompt_length: int | None = None, text_only: bool | None = None, enable_request_stats: bool = False, + kv_cache_retention_config: KvCacheRetentionConfig | None = None, ) -> _KVCache: """Create a suspended cache with a prefill-to-generation boundary. @@ -852,6 +859,7 @@ class KVCacheManager: it also marks generation-phase allocation stats. """ ... + def refresh_retention(self) -> None: ... def probe_reuse( self, reuse_scope: ReuseScope | None = None, diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index b7921433db26..febc997b5cd3 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -42,7 +42,6 @@ KvCacheConfig, KvCacheConnectorConfig, NGramDecodingConfig, SchedulerConfig) from tensorrt_llm.llmapi.llm_utils import KvCacheRetentionConfig -from tensorrt_llm.logger import logger as trtllm_logger_singleton from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX from ..conftest import get_sm_version, llm_models_root @@ -1237,39 +1236,12 @@ def test_connector_multi_request(enforce_single_worker, model_with_connector): @pytest.mark.threadleak(enabled=False) -@pytest.mark.parametrize("use_kv_cache_manager_v2", [ - pytest.param(False, id="kv_cache_manager_v1"), - pytest.param( - True, - id="kv_cache_manager_v2", - marks=pytest.mark.xfail( - strict=True, - reason= - "KvCacheRetentionConfig does not reach KVCacheManagerV2 at all " - "(per-page priority comes from custom_priority_callback, which " - "V2 never overrides), so the connector reports priorities=None."), - ), -], +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], + ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], indirect=True) -def test_connector_priorities(enforce_single_worker, model_with_connector): - """Test that retention priorities flow through the connector correctly. - - This test verifies that when KvCacheRetentionConfig is provided, - the RequestData.priorities field is populated with the correct - per-block priorities based on the token ranges. - - KNOWN GAP -- `xfail(strict=True)` on `kv_cache_manager_v2`. - `KvCacheRetentionConfig` does not reach KVCacheManagerV2 at all: V2's - per-page priority comes from `custom_priority_callback` - (kv_cache_manager_v2/_core/_kv_cache_manager.py), which KVCacheManagerV2 - never overrides, so every page carries the default priority and the - connector reports `priorities=None`. A user who sets a retention config on - V2 silently gets none of it -- not only through the connector. The - assertions below stay the correct expectation for both managers rather than - being relaxed per manager, so wiring retention into V2 turns this green - instead of needing the test rewritten; `strict=True` is what makes it fail - loudly on that day rather than passing silently. - """ +def test_connector_priorities(enforce_single_worker, model_with_connector, + use_kv_cache_manager_v2): + """Native prompt/decode priorities reach the connector's offload filter.""" BLOCK_SIZE = 32 NUM_INPUT_TOKENS = 64 # 2 blocks NUM_TOKENS = 4 @@ -1279,6 +1251,7 @@ def test_connector_priorities(enforce_single_worker, model_with_connector): model_fn, scheduler, worker = model_with_connector model = model_fn(disable_overlap_scheduler=True) + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False worker.get_finished.return_value = [], [] @@ -1334,83 +1307,23 @@ def test_connector_priorities(enforce_single_worker, model_with_connector): assert request.priorities[ 1] == LOW_PRIORITY, f"Expected priority {LOW_PRIORITY} for block 1, got {request.priorities[1]}" - -@pytest.mark.threadleak(enabled=False) -@pytest.mark.parametrize("use_kv_cache_manager_v2", [True], - ids=["kv_cache_manager_v2"], - indirect=True) -def test_connector_warns_that_retention_is_ignored_on_v2( - enforce_single_worker, model_with_connector, caplog): - """A retention config that has no effect must say so. - - `KvCacheRetentionConfig` does not reach KVCacheManagerV2 at all, so the - connector reports `priorities=None` there. `test_connector_priorities` - pins that gap as `xfail(strict=True)`; this pins the diagnostic, which is - the only thing standing between a user's retention config and it being - dropped in silence. - - V2 only: on V1 the config is honoured and there is nothing to warn about. - """ - model_fn, scheduler, worker = model_with_connector - model = model_fn(disable_overlap_scheduler=True) - - scheduler.get_num_new_matched_tokens.return_value = 0, False - worker.get_finished.return_value = [], [] - - retention_config = KvCacheRetentionConfig(token_range_retention_configs=[ - KvCacheRetentionConfig.TokenRangeRetentionConfig(token_start=0, - token_end=32, - priority=80) - ], - decode_retention_priority=10) - - # Two gates sit between `logger.warning_once` and `caplog`, and both are - # process-global rather than per-test: - # - # 1. The TensorRT-LLM logger defaults to `error` (tensorrt_llm/logger.py:163), - # and `Logger.log` drops anything below that before it reaches Python - # logging at all. The LLM constructor raises the level to `info` only - # while it parses arguments and restores it afterwards, so a warning - # emitted during `generate` is discarded unless the level is raised here. - # 2. `log_once` records its key before consulting the level, so the key is - # consumed by whichever test drove this path first in the process -- - # `test_connector_priorities[kv_cache_manager_v2]` above sets a - # retention config on V2 too, and it runs first. - # - # Raise the level and clear the key so the assertion below observes this - # test's own emission instead of the leftovers of test ordering. - previous_level = trtllm_logger_singleton.level - trtllm_logger_singleton.set_level("warning") - trtllm_logger_singleton._appeared_keys.discard( - _V2_RETENTION_IGNORED_LOG_KEY) - - # The TensorRT-LLM logger sets `propagate = False`, so caplog only sees its - # records once its handler is attached to that logger by name. - trtllm_logger = logging.getLogger(TRTLLM_LOGGER_NAME) - trtllm_logger.addHandler(caplog.handler) - try: - with caplog.at_level(logging.WARNING, logger=TRTLLM_LOGGER_NAME): - generate_and_wait(model, - scheduler, - worker, [0] * 64, - sampling_params=SamplingParams(max_tokens=4, - ignore_eos=True), - kv_cache_retention_config=retention_config) - finally: - trtllm_logger.removeHandler(caplog.handler) - trtllm_logger_singleton.set_level(previous_level) - - assert "KvCacheRetentionConfig has no effect" in caplog.text, ( - "A retention config was set on KVCacheManagerV2 and nothing said it " - "would be ignored. The user's configuration is silently dropped:\n" - f"{caplog.text}") - - # The other half: the warning describes what actually happened. - request = scheduler.build_connector_meta.call_args_list[0].args[ - 0].new_requests[0] - assert request.priorities is None, ( - "priorities are populated on V2 after all, so the warning is now " - "wrong -- revisit it together with `test_connector_priorities`.") + # A connector can select blocks from actual scheduler metadata by priority. + selected = [ + block_id + for block_id, priority in zip(request.new_block_ids, request.priorities) + if priority >= HIGH_PRIORITY + ] + assert selected == [request.new_block_ids[0]] + decode_requests = [ + request_data + for meta_call in scheduler.build_connector_meta.call_args_list[1:] + for request_data in meta_call.args[0].cached_requests + if request_data.new_block_ids + ] + assert decode_requests, "generation must allocate a new decode block" + for request_data in decode_requests: + assert request_data.priorities == [LOW_PRIORITY] * len( + request_data.new_block_ids) @pytest.mark.threadleak(enabled=False) @@ -2140,7 +2053,13 @@ def test_connector_cancel_holds_pages_until_the_save_reads_them( cancelled = model.generate_async( [100] * CANCELLED_SAVE_PROMPT_TOKENS, SamplingParams(max_tokens=CANCELLED_SAVE_MAX_TOKENS, - ignore_eos=True)) + ignore_eos=True), + kv_cache_retention_config=KvCacheRetentionConfig( + token_range_retention_configs=[ + KvCacheRetentionConfig.TokenRangeRetentionConfig( + 0, None, 0) + ], + decode_retention_priority=0)) assert control.decoding.wait(60), "Request never reached generation" cancelled.abort() diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 2b5b6f887860..1fd54e4d1171 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -165,9 +165,7 @@ l0_a10: # by hand. test_connector_runs_on_kv_cache_manager_v2 is what makes the V2 # entries meaningful -- the creator silently falls back to the V1 manager for # combinations it cannot serve, so without it every V2 id below could pass - # while running V1. test_connector_priorities[kv_cache_manager_v2] is - # xfail(strict=True) in the suite, not omitted here, so the retention gap - # stays counted. The two e2e cases below are the only multi-layer-group ones + # while running V1. The two e2e cases below are the only multi-layer-group ones # that move bytes -- the rest drive a MagicMock and assert routing. # test_connector_vswa_e2e_gemma3 is the interleaved-attention run (Gemma-3 # asks for KVCacheManagerV2 itself, and its layer groups come out different @@ -216,7 +214,6 @@ l0_a10: - llmapi/test_llm_api_connector.py::test_connector_disagg_prefill[kv_cache_manager_v2-True] - llmapi/test_llm_api_connector.py::test_connector_multi_request[kv_cache_manager_v2] - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v2] - - llmapi/test_llm_api_connector.py::test_connector_warns_that_retention_is_ignored_on_v2[kv_cache_manager_v2] - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v2] - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v2-host_offloading] - llmapi/test_llm_api_connector.py::test_connector_rejects_unsupported_config[kv_cache_manager_v2-beam_search] diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 0061bf2389ad..7407cc9a149e 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -289,6 +289,11 @@ l0_b200: - unittest/tools/test_layer_wise_benchmarks.py::test_kimi_k3_gen_dep[1] - unittest/tools/test_layer_wise_benchmarks.py::test_performance_alignment[1] - unittest/kv_cache_manager_v2_tests + - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v1] + - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_cancel_holds_pages_until_the_save_reads_them # ------------- KV Cache V2 Scheduler IT --------------- - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2Llama::test_v2_vs_v1_basic - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2Llama::test_token_budget_limited diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 6573a6c5a8c2..838e29e4cc26 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -25,6 +25,7 @@ import unittest from contextlib import contextmanager from dataclasses import dataclass, replace +from datetime import timedelta from importlib.util import find_spec from random import randbytes from statistics import median @@ -82,7 +83,9 @@ MemToMemTask, copy_device_to_device, ) + from bindings.executor import KvCacheRetentionConfig, KvCacheTransferMode # isort: skip else: + from tensorrt_llm.bindings.executor import KvCacheRetentionConfig, KvCacheTransferMode from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( MemToMemTask, copy_device_to_device, @@ -409,6 +412,180 @@ def prepare( self.manager = KVCacheManager(self.cfg) +class TestRetention(TestKVCacheManagerV2): + def setUp(self) -> None: + super().setUp() + self.caches = [] + self.stream = CachedCudaStream() + + def tearDown(self) -> None: + for cache in reversed(self.caches): + cache.close() + self.caches.clear() + super().tearDown() + + def _cache(self, tokens, config=None, prompt_length=None): + cache = self.manager.create_kv_cache( + input_tokens=tokens, + expected_prompt_length=len(tokens) if prompt_length is None else prompt_length, + kv_cache_retention_config=config, + ) + self.caches.append(cache) + self.assertTrue(cache.resume(self.stream.handle)) + self.assertTrue(cache.resize(len(tokens))) + remaining = tokens[cache.num_committed_tokens :] + if tokens: + self.engine.execute( + [Step(cache, remaining, tokens[: cache.num_committed_tokens])], self.stream.handle + ) + if remaining: + cache.commit(remaining) + return cache + + @staticmethod + def _config(priority, duration=None): + if priority is None: + return None + return KvCacheRetentionConfig( + [KvCacheRetentionConfig.TokenRangeRetentionConfig(0, None, priority, duration)] + ) + + @parameterized.expand( + [(80, 10, False), (80, 10, True), (10, 80, False), (10, 80, True), (80, None, False)] + ) + def test_shared_claim_order_controls_eviction(self, first, second, close_newer_first): + self.prepare(2 << 20, 0, 0, 1, None, 0, tokens_per_block=4, kv_buf_size=65536) + prompt = [1, 2, 3, 4] + older = self._cache(prompt, self._config(first)) + newer = self._cache(prompt, self._config(second)) + priority = first if second is None else second + self.assertEqual(older.get_page_priorities(0), [priority]) + self.assertEqual(newer.get_aggregated_page_indices(0), older.get_aggregated_page_indices(0)) + requests = [newer, older] if close_newer_first else [older, newer] + requests[0].close() + self.assertEqual(requests[1].get_page_priorities(0), [priority]) + competitor = self._cache([5, 6, 7, 8], self._config(50)) + competitor.close() + # Referenced pages cannot be permanently evicted, irrespective of priority. + free = self.manager.get_storage_statistics()[0].free + pin = self._cache([]) + self.assertTrue(pin.resize((free + 1) * 4)) + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 4) + self.assertEqual(self.manager.probe_reuse(input_tokens=[5, 6, 7, 8]), 0) + pin.close() + requests[1].close() + competitor = self._cache([5, 6, 7, 8], self._config(50)) + competitor.close() + free = self.manager.get_storage_statistics()[0].free + pressure = self._cache([]) + self.assertTrue(pressure.resize((free + 1) * 4)) + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 4 if priority > 50 else 0) + self.assertEqual( + self.manager.probe_reuse(input_tokens=[5, 6, 7, 8]), 0 if priority > 50 else 4 + ) + + def test_prompt_decode_duration_and_unconfigured_claim(self): + self.prepare(2 << 20, 0, 0, 1, None, 0, tokens_per_block=4, kv_buf_size=65536) + duration = timedelta(milliseconds=20) + config = KvCacheRetentionConfig( + [KvCacheRetentionConfig.TokenRangeRetentionConfig(1, None, 80, duration)], + decode_retention_priority=10, + decode_duration_ms=duration, + ) + tokens = list(range(12)) + # A five-token prompt occupies two blocks; its partial second block stays context. + first = self._cache(tokens[:5], config) + self.assertTrue(first.resize(len(tokens))) + self.engine.execute([Step(first, tokens[5:], tokens[:5])], self.stream.handle) + first.commit(tokens[5:]) + second = self._cache(tokens, config, prompt_length=5) + self.assertEqual(first.get_page_priorities(0), [35, 80, 10]) + first.close() + time.sleep(0.04) + self.manager.refresh_retention() + self.assertEqual(second.get_page_priorities(0), [35, 80, 10]) + second.close() + time.sleep(0.04) + self.manager.refresh_retention() + reclaimed = self._cache(tokens) + self.assertEqual(reclaimed.get_page_priorities(0), [35, 35, 35]) + reclaimed.close() + long_duration = timedelta(seconds=1) + config = KvCacheRetentionConfig( + [KvCacheRetentionConfig.TokenRangeRetentionConfig(1, None, 80, long_duration)], + decode_retention_priority=10, + decode_duration_ms=long_duration, + ) + timed = self._cache(tokens, config, prompt_length=5) + timed.close() + untimed = self._cache(tokens) + untimed.close() + time.sleep(1.05) + self.manager.refresh_retention() + self.assertEqual(self._cache(tokens).get_page_priorities(0), [35, 80, 10]) + + def test_offload_reload_and_error_cleanup(self): + self.prepare(2 << 20, 2 << 20, 0, 1, None, 0, tokens_per_block=4, kv_buf_size=65536) + tokens = [1, 2, 3, 4] + cached = self._cache(tokens, self._config(80)) + cached.suspend() + free = self.manager.get_storage_statistics()[0].free + pressure = self._cache([]) + self.assertTrue(pressure.resize((free + 1) * 4)) + snapshot = cached.get_page_storage_snapshot(0) + self.assertEqual(snapshot.cache_levels, [CacheLevel(1)]) + pressure.close() + self.assertTrue(cached.resume(self.stream.handle)) + self.assertEqual(cached.get_page_priorities(0), [80]) + self.assertEqual(cached.get_page_storage_snapshot(0).cache_levels, [GPU_LEVEL]) + self.engine.execute([Step(cached, [], tokens)], self.stream.handle) + self.stream.synchronize() + capacity = cached.capacity + total = self.manager.get_storage_statistics()[0].total + self.assertFalse(cached.resize((total + 1) * 4)) + self.assertEqual(cached.capacity, capacity) + self.assertEqual(cached.get_page_priorities(0), [80]) + self.engine.execute([Step(cached, [], tokens)], self.stream.handle) + self.stream.synchronize() + cached.close() + # Unsupported file-transfer semantics fail before claiming any reusable pages. + for mode in (KvCacheTransferMode.GDS, KvCacheTransferMode.POSIX_DEBUG_FALLBACK): + config = KvCacheRetentionConfig([], transfer_mode=mode, directory="/tmp") + with self.assertRaisesRegex(ValueError, "per-request GDS/POSIX"): + self.manager.create_kv_cache(input_tokens=tokens, kv_cache_retention_config=config) + self.assertEqual(self._cache(tokens).get_page_priorities(0), [80]) + for cache in self.caches: + cache.close() + self.manager.clear_reusable_blocks() + for level in (GPU_LEVEL, CacheLevel(1)): + for stats in self.manager.get_storage_statistics(level): + self.assertEqual(stats.free, stats.total) + self.assertEqual(stats.evictable, 0) + + @parameterized.expand([(29, 30), (30, 30), (49, 50), (50, 50)]) + def test_released_offload_threshold_and_duration(self, priority, threshold): + self.cfg = create_config(4, 2 << 20, 2 << 20, 0, 1, None, 0, 65536) + self.cfg.secondary_offload_min_priority = threshold + self.engine = FakeEngine(self.cfg) + self.manager = KVCacheManager(self.cfg) + tokens = [1, 2, 3, 4] + cached = self._cache(tokens, self._config(priority, timedelta(seconds=2))) + cached.close() + free = self.manager.get_storage_statistics()[0].free + pressure = self._cache([]) + self.assertTrue(pressure.resize((free + 1) * 4)) + retained = priority >= threshold + self.assertEqual(self.manager.probe_reuse(input_tokens=tokens), 4 if retained else 0) + self.assertEqual( + self.manager.get_storage_statistics(CacheLevel(1))[0].evictable, 1 if retained else 0 + ) + pressure.close() + time.sleep(2.05) + self.manager.refresh_retention() + # GPU eviction-driven offload clears the old duration, retaining the priority. + self.assertEqual(self._cache(tokens).get_page_priorities(0), [priority if retained else 35]) + + class TestStorageStatistics(TestKVCacheManagerV2): def test_statistics_are_detached_snapshots(self) -> None: self.prepare(8 << 20, 8 << 20, 0, 2, 16, 0, tokens_per_block=4, kv_buf_size=4096) From eccc995205d38489d5dd237b5af19c596bd6cc25 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:18:12 -0700 Subject: [PATCH 02/23] fix: preserve retention across partial reuse and lifecycle boundaries Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../kv_cache_manager_v2/eventManager.cpp | 14 ++++- .../kv_cache_manager_v2/kvCache.cpp | 32 +++++++++- .../kv_cache_manager_v2/kvCache.h | 1 + .../kv_cache_manager_v2/page.cpp | 22 ++++++- .../batch_manager/kv_cache_manager_v2/page.h | 5 ++ .../kv_cache_manager_v2/storageManager.cpp | 10 +++- .../kv_cache_manager_v2/storageManager.h | 2 + .../kv_cache/mamba_cache_manager.py | 3 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 20 ++++--- .../defs/llmapi/test_llm_api_connector.py | 16 +++++ .../test_lists/test-db/l0_b200.yml | 1 + .../kv_cache/test_mamba_cache_manager.py | 18 +++++- .../test_kv_cache_manager_v2.py | 58 +++++++++++++++++-- 13 files changed, 180 insertions(+), 22 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp index da9f2b96796e..295f10d0b77b 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp @@ -314,7 +314,19 @@ void EventManager::addCacheLevelUpdated( void EventManager::addPriorityUpdated( Digest const& blockKey, Priority oldPriority, Priority newPriority, LifeCycleId lifeCycle) { - addUpdatedEvent(blockKey, std::nullopt, KVCacheEventDiff{oldPriority, newPriority}, lifeCycle.value()); + if (mMaxKvEventEntries <= 0) + { + return; + } + std::lock_guard 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) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index 62ca6dc08ab0..87963715d8d2 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -829,6 +829,7 @@ bool KvCache::resume(std::optional stream, std::optional 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 srcLocks; + TypedVec> copiedSources(numLc); for (LifeCycleId lcIdx{0}; lcIdx < numLc; ++lcIdx) { if (!deferredSlots[lcIdx].has_value()) @@ -850,6 +851,7 @@ bool KvCache::resume(std::optional stream, std::optional 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; @@ -908,6 +910,12 @@ bool KvCache::resume(std::optional stream, std::optional isDecod } auto newPage = makeShared(*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); @@ -932,6 +940,7 @@ bool KvCache::resume(std::optional stream, std::optional isDecod } mOptionalCandidates.clear(); // Deferred copies survive a failed decode admission and must not be repeated on retry. + mPartialCopyRetentionHolds.clear(); mNeverResumed = false; if (mIsDecoding) { @@ -1323,6 +1332,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(); @@ -1373,7 +1383,14 @@ CommittedPage* KvCache::_copyPageToTreeBlock( auto committed = makeShared( &storageMgr, treeBlock, lcIdx, lvl, numTokensInBlock, getPriority(treeBlock->ordinal(), lcIdx)); auto const retention = getRetention(treeBlock->ordinal()); - committed->claimRetention(retention.retentionPriority, retention.durationMs); + 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. @@ -2927,9 +2944,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(); - auto const retention = getRetention(ordinal); - page->claimRetention(retention.retentionPriority, retention.durationMs); + if (!borrowPartial) + { + auto const retention = getRetention(ordinal); + page->claimRetention(retention.retentionPriority, retention.durationMs); + } + if (partialCopy) + { + mPartialCopyRetentionHolds.push_back(page->borrowRetentionForCopy()); + } if (!isAttention) { return; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h index b42556ab2c7c..76a635793b70 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h @@ -794,6 +794,7 @@ class KvCache : public std::enable_shared_from_this std::optional mPageStorageRow; TypedVec mBlocks; + std::vector>> mPartialCopyRetentionHolds; std::vector mCommittedTokens; // Initial current-residency provenance, observed before reused pages are held or promoted. diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp index a7db0f0ac043..0f6b120ee070 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp @@ -110,6 +110,25 @@ void Page::claimRetention(std::optional priority, std::optionaluniqLock.expired() || h.useCount() > mPartialCopyHolders + 1); +} + +FuncGuard> Page::borrowRetentionForCopy() +{ + FuncGuard> guard([self = sharedFromThis()]() { --self->mPartialCopyHolders; }); + ++mPartialCopyHolders; + return guard; +} + SharedPageLock Page::lock(KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal ordinal, LifeCycleId lc, bool skipWait) { return hold()->lock(kvCache, beamIndex, ordinal, lc, skipWait); @@ -234,8 +253,7 @@ SharedPtr UncommittedPage::convertToCommitted( this->readyEvent = std::move(readyEv); auto committed = makeShared(manager, blk, lifeCycle, cacheLevel, numTokensInBlock, priority()); - auto const retention = kvCache->getRetention(blk->ordinal()); - committed->claimRetention(retention.retentionPriority, retention.durationMs); + committed->inheritRetention(*this); // Move slot id to the committed page; invalidate our slot. committed->setSlotId(slotId()); // asserts valid committed->readyEvent = std::move(readyEvent); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h index 08ef0a5e8532..b10075de042f 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h @@ -23,6 +23,7 @@ #include "kv_cache_manager_v2/lifeCycleRegistry.h" #include "kv_cache_manager_v2/storage/core.h" #include "kv_cache_manager_v2/utils/cudaEvent.h" +#include "kv_cache_manager_v2/utils/funcGuard.h" #include "kv_cache_manager_v2/utils/sharedPtr.h" #include @@ -72,6 +73,9 @@ class Page : public Slot, public EnableSharedFromThis //! Apply a request claim, preserving priority when unspecified and replacing its duration. void claimRetention(std::optional priority, std::optional duration); + void inheritRetention(Page const& source); + bool hasRetentionConsumers() const; + FuncGuard> borrowRetentionForCopy(); PageStatus status() const noexcept; @@ -95,6 +99,7 @@ class Page : public Slot, public EnableSharedFromThis private: friend class StorageManager; Priority mPriority; + int mPartialCopyHolders = 0; std::optional mRetentionDuration; std::optional mRetentionExpiry; }; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp index f5072e53f3cb..7ff04e74d5fd 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp @@ -881,6 +881,12 @@ void StorageManager::prepareFreeSlots(CacheLevel level, TypedVec const& requirements, + MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) +{ TypedVec> goals(numCacheLevels()); for (CacheLevel lvl{0}; lvl < goals.size(); ++lvl) { @@ -898,7 +904,6 @@ void StorageManager::prepareFreeSlots(CacheLevel level, TypedVec const& minNumPages, DropRecorder const& dropRecorder) { - refreshRetention(); auto evicted = mLevels.at(level).controller.evict(minNumPages); auto rescheduleEvictedPagesOnFailure = makeEvictionRollbackGuard(evicted); @@ -1904,7 +1909,8 @@ void StorageManager::shrinkPoolGroup( // Ensure free slots for the overflow pages. TypedVec reqs(numPoolGroups(level), 0); reqs[pgIdx] = slotCountValueFromSize(overflowPages.size()); - prepareFreeSlots(level, reqs); + // The eviction order captured above must remain stable through defragmentation. + _prepareFreeSlots(level, reqs, MigrationRecorder{}, DropRecorder{}); // A17: all overflow pages must be at the expected cache level. TLLM_CHECK_DEBUG_WITH_INFO(std::all_of(overflowPages.begin(), overflowPages.end(), diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h index 816a6eaa330f..569cbff98145 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h @@ -395,6 +395,8 @@ class StorageManager : public std::enable_shared_from_this [[nodiscard]] auto makeEvictionRollbackGuard(TypedVec>> const& evicted); + void _prepareFreeSlots(CacheLevel level, TypedVec const& requirements, + MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder); void _prepareFreeSlots(TypedVec>& goals, CacheLevel lvlId, PagesByLifeCycle& fallenPages, MigrationRecorder const& migrationRecorder = {}, DropRecorder const& dropRecorder = {}); diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index f612de0660c9..1bc6779f814e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -46,6 +46,7 @@ from tensorrt_llm._utils import (TensorWrapper, convert_to_torch_tensor, nvtx_range, prefer_pinned, torch_dtype_to_binding) +from tensorrt_llm.bindings.executor import KvCacheRetentionConfig from tensorrt_llm.bindings.internal.batch_manager import ( LinearAttentionMetadata, LinearCacheType) from tensorrt_llm.llmapi.llm_args import KvCacheConfig @@ -3624,6 +3625,7 @@ def _create_kv_cache( is_dummy: bool = False, enable_request_stats: bool = False, expected_prompt_length: Optional[int] = None, + kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None, ) -> Optional[_KVCache]: kv_cache = super()._create_kv_cache( request_id, @@ -3633,6 +3635,7 @@ def _create_kv_cache( is_dummy=is_dummy, enable_request_stats=enable_request_stats, expected_prompt_length=expected_prompt_length, + kv_cache_retention_config=kv_cache_retention_config, ) if (self.mapping.rank == 0 and kv_cache is not None and input_tokens is not None and not is_dummy diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a582900ec098..4ffb3e1275d4 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -33,6 +33,7 @@ global_mpi_size, is_trace_enabled, mpi_comm, mpi_disabled, nvtx_range, set_thread_local_mpi_comm, trace_func) +from tensorrt_llm.bindings import executor as executor_bindings from tensorrt_llm.bindings.executor import (DisServingRequestStats, FinishReason, InflightBatchingStats, IterationStats, KvCacheStats, @@ -6363,15 +6364,16 @@ def _configure_csa2_remote_tail(self, request: LlmRequest) -> None: def _validate_request(self, request: LlmRequest): self._configure_csa2_remote_tail(request) - if (isinstance(self.kv_cache_manager, KVCacheManagerV2) - and request.kv_cache_retention_config is not None): - logger.warning_once( - "KvCacheRetentionConfig has no effect with KV cache manager V2; " - "retention priorities and durations are ignored. Set " - "kv_cache_config.use_kv_cache_manager_v2=False to use retention " - "with KV cache manager V1.", - key=_V2_RETENTION_IGNORED_LOG_KEY) - + retention = request.kv_cache_retention_config + if self._is_kv_manager_v2 and retention is not None: + if retention.transfer_mode != executor_bindings.KvCacheTransferMode.DRAM: + raise ValueError( + "KVCacheManagerV2 does not support per-request GDS/POSIX " + "transfer_mode or directory; configure the manager's cache tiers instead." + ) + if not 0 <= retention.decode_retention_priority <= 100: + raise ValueError( + "Decode retention priority must be between 0 and 100") # Validate context-side pipelined-transfer constraints. disagg_params = request.py_disaggregated_params if (self.kv_cache_transceiver is not None diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index febc997b5cd3..dc6ae8045328 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -37,6 +37,8 @@ from tensorrt_llm._torch.pyexecutor.py_executor import \ _V2_RETENTION_IGNORED_LOG_KEY from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.bindings.executor import KvCacheTransferMode +from tensorrt_llm.llmapi import RequestError from tensorrt_llm.llmapi.llm_args import (CacheTransceiverConfig, CapacitySchedulerPolicy, KvCacheConfig, KvCacheConnectorConfig, @@ -1277,6 +1279,20 @@ def test_connector_priorities(enforce_single_worker, model_with_connector, sampling_params = SamplingParams(max_tokens=NUM_TOKENS, ignore_eos=True) + if use_kv_cache_manager_v2: + for mode in (KvCacheTransferMode.GDS, + KvCacheTransferMode.POSIX_DEBUG_FALLBACK): + with pytest.raises(RequestError, match="per-request GDS/POSIX"): + model.generate([0] * NUM_INPUT_TOKENS, + sampling_params=sampling_params, + kv_cache_retention_config=KvCacheRetentionConfig( + [], transfer_mode=mode, directory="/tmp")) + with pytest.raises(RequestError, match="between 0 and 100"): + model.generate([0] * NUM_INPUT_TOKENS, + sampling_params=sampling_params, + kv_cache_retention_config=KvCacheRetentionConfig( + [], decode_retention_priority=101)) + generate_and_wait(model, scheduler, worker, [0] * NUM_INPUT_TOKENS, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 7407cc9a149e..ec98130ff647 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -290,6 +290,7 @@ l0_b200: - unittest/tools/test_layer_wise_benchmarks.py::test_performance_alignment[1] - unittest/kv_cache_manager_v2_tests - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v1] + - unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_v2_hybrid_retains_configured_number_of_conversation_turns - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v2] - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v1] - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v2] diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 0387894bf82c..ac7feea518a3 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -62,6 +62,7 @@ ) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import torch_dtype_to_binding +from tensorrt_llm.bindings.executor import KvCacheRetentionConfig from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.llmapi.llm_args import ( @@ -2425,6 +2426,7 @@ def _build_v2_hybrid_with_mamba_layer( mamba_ssm_cache_dtype=torch.float16, kda_replay_num_spec=None, enable_kv_pool_rebalance=False, + secondary_offload_min_priority=None, ): """Construct a real MambaHybridCacheManagerV2.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -2453,6 +2455,7 @@ def _build_v2_hybrid_with_mamba_layer( ), dtype=kv_cache_dtype, enable_kv_pool_rebalance=enable_kv_pool_rebalance, + secondary_offload_min_priority=secondary_offload_min_priority, ) return MambaHybridCacheManagerV2( mamba_d_state=8, @@ -2501,6 +2504,7 @@ def _make_v2_conversation_request( request_id: int, tokens: list[int], conversation_id: str, + retention: KvCacheRetentionConfig | None = None, ) -> LlmRequest: request = LlmRequest( request_id=request_id, @@ -2508,6 +2512,7 @@ def _make_v2_conversation_request( input_tokens=tokens, sampling_config=SamplingConfig(), is_streaming=False, + kv_cache_retention_config=retention, ) request.py_conversation_params = ConversationParams(conversation_id=conversation_id) return request @@ -2948,8 +2953,14 @@ def test_v2_hybrid_retains_configured_number_of_conversation_turns(): block_reuse_policy="per_conversation", max_num_turns=2, additional_snapshot_offsets_from_end=[0], + secondary_offload_min_priority=50, + ) + request_a = _make_v2_conversation_request( + 1, + list(range(64)), + "conv-1", + KvCacheRetentionConfig([KvCacheRetentionConfig.TokenRangeRetentionConfig(0, None, 80)]), ) - request_a = _make_v2_conversation_request(1, list(range(64)), "conv-1") request_b = _make_v2_conversation_request(2, list(range(100, 164)), "conv-1") # Use fresh conversation IDs so probes query the shared prefix cache # without altering conv-1's retained-turn accounting. @@ -2971,6 +2982,11 @@ def test_v2_hybrid_retains_configured_number_of_conversation_turns(): try: _run_v2_hybrid_context(mgr, request_a) + cache_a = mgr.kv_cache_map[request_a.py_request_id] + attention_group = mgr.impl.get_layer_group_id(1) + priorities = cache_a.get_page_priorities(attention_group) + assert priorities and set(priorities) == {80} + assert mgr.impl.init_config.secondary_offload_min_priority == 50 request_a_state_index = mgr.get_state_indices([request_a.py_request_id], [False])[0] mgr.free_resources(request_a) _run_v2_hybrid_context(mgr, request_b) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 838e29e4cc26..697572dcf44a 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -51,8 +51,10 @@ GpuCacheTierConfig, HostCacheTierConfig, KVCacheDesc, + KVCacheEventManager, KVCacheManager, KVCacheManagerConfig, + KVCacheUpdatedData, LayerGroupId, LayerId, LogicError, @@ -107,8 +109,10 @@ GpuCacheTierConfig, HostCacheTierConfig, KVCacheDesc, + KVCacheEventManager, KVCacheManager, KVCacheManagerConfig, + KVCacheUpdatedData, LayerGroupId, LayerId, LogicError, @@ -532,14 +536,21 @@ def test_offload_reload_and_error_cleanup(self): free = self.manager.get_storage_statistics()[0].free pressure = self._cache([]) self.assertTrue(pressure.resize((free + 1) * 4)) - snapshot = cached.get_page_storage_snapshot(0) - self.assertEqual(snapshot.cache_levels, [CacheLevel(1)]) + self.assertEqual(_introspection.active_page_stats(cached)[0], [0, 1]) pressure.close() self.assertTrue(cached.resume(self.stream.handle)) self.assertEqual(cached.get_page_priorities(0), [80]) self.assertEqual(cached.get_page_storage_snapshot(0).cache_levels, [GPU_LEVEL]) self.engine.execute([Step(cached, [], tokens)], self.stream.handle) self.stream.synchronize() + for config in (None, self._config(10)): + partial = self._cache(tokens[:3], config) + self.assertEqual(partial.num_committed_tokens, 3) + self.assertEqual(partial.get_page_priorities(0), [80]) + self.assertEqual(cached.get_page_priorities(0), [80]) + self.assertNotEqual( + partial.get_aggregated_page_indices(0), cached.get_aggregated_page_indices(0) + ) capacity = cached.capacity total = self.manager.get_storage_statistics()[0].total self.assertFalse(cached.resize((total + 1) * 4)) @@ -562,6 +573,23 @@ def test_offload_reload_and_error_cleanup(self): self.assertEqual(stats.free, stats.total) self.assertEqual(stats.evictable, 0) + donor = self._cache(list(range(8)), self._config(80)) + donor.close() + pending = [] + for priority in (10, 90): + copy = self.manager.create_kv_cache( + input_tokens=list(range(3)), + expected_prompt_length=3, + kv_cache_retention_config=self._config(priority), + ) + self.caches.append(copy) + pending.append(copy) + for copy in pending: + self.assertTrue(copy.resume(self.stream.handle)) + self.assertEqual(copy.get_page_priorities(0), [90]) + self.engine.execute([Step(copy, [], list(range(3)))], self.stream.handle) + self.stream.synchronize() + @parameterized.expand([(29, 30), (30, 30), (49, 50), (50, 50)]) def test_released_offload_threshold_and_duration(self, priority, threshold): self.cfg = create_config(4, 2 << 20, 2 << 20, 0, 1, None, 0, 65536) @@ -4990,6 +5018,7 @@ def prepare_partial( gpu_quota: int = 64 << 20, window_size: int | None = None, reuse_match_backoff: int = 0, + event_manager=None, ) -> None: kv_buf_size = 8192 window_size = self.WINDOW_SIZE if window_size is None else window_size @@ -5018,7 +5047,7 @@ def prepare_partial( reuse_match_backoff=reuse_match_backoff, ) self.engine = FakeEngine(self.cfg) - self.manager = KVCacheManager(self.cfg) + self.manager = KVCacheManager(self.cfg, event_manager=event_manager) @property def _full_attn_lc_id(self) -> int: @@ -5161,7 +5190,8 @@ def test_replaced_reused_page_does_not_outlive_its_block(self) -> None: s.take_finish_event().synchronize() def test_rewind_endpoint_survives_longer_sibling_created_after(self) -> None: - self.prepare_partial() + events = KVCacheEventManager(max_kv_event_entries=128) + self.prepare_partial(event_manager=events) base = [TokenId(i) for i in range(80)] extended = base + [TokenId(i) for i in range(1000, 1080)] rewind = base + [TokenId(2000)] @@ -5180,6 +5210,26 @@ def test_rewind_endpoint_survives_longer_sibling_created_after(self) -> None: # The partial SWA page is stale at the longer endpoint and must not constrain # the full-attention lifecycle's reusable prefix. self.assertEqual(self.manager.probe_reuse(input_tokens=extended), len(extended)) + events.flush_iteration_events() + events.get_latest_events(0) + claimant = self.manager.create_kv_cache( + input_tokens=rewind, + expected_prompt_length=len(rewind), + kv_cache_retention_config=KvCacheRetentionConfig( + [KvCacheRetentionConfig.TokenRangeRetentionConfig(0, None, 80)] + ), + ) + try: + events.flush_iteration_events() + updates = [ + event + for event in events.get_latest_events(0) + if isinstance(event.data, KVCacheUpdatedData) + ] + self.assertTrue(updates) + self.assertEqual({event.layer_group_id for event in updates}, {self._full_attn_lc_id}) + finally: + claimant.close() def test_rewind_endpoint_attaches_to_longer_sibling_created_before(self) -> None: self.prepare_partial() From 5240193fd7b350b1c075c7daaeb6a156d25e6354 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:26:05 -0700 Subject: [PATCH 03/23] fix: retain V1 policy across prefetch and cache lifecycle transitions Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../kv_cache_manager_v2/kvCache.cpp | 19 ++++-- .../kv_cache_manager_v2/kvCache.h | 4 +- .../kv_cache_manager_v2/kvCacheManager.cpp | 11 ++++ .../kv_cache_manager_v2/kvCacheManager.h | 3 + .../kv_cache_manager_v2/page.cpp | 16 ++++- .../kv_cache_manager_v2/storageManager.cpp | 28 ++++++--- .../batch_manager/kvCacheManagerV2.cpp | 13 ++++ .../kv_cache/kv_cache_manager_v2.py | 62 +++++++++++++------ .../runtime/kv_cache_manager_v2/__init__.pyi | 4 ++ .../test_kv_cache_manager_v2.py | 22 ++++++- 10 files changed, 145 insertions(+), 37 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index 87963715d8d2..f51a4c83b383 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -56,7 +56,7 @@ int64_t sumSlotBytes(StorageManager const& storage, CacheLevel level, LifeCycleI KvCache::KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional reuseMatch, std::optional mId, PriorityCb priorityCb, std::optional expectedPromptLength, std::optional textOnly, bool enableRequestStats, - std::optional retentionConfig) + std::optional retentionConfig, bool isPrefetch) : id(mId) , mManager(manager.shared_from_this()) , mReuseScope(std::move(reuseScope)) @@ -121,7 +121,7 @@ KvCache::KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional KvCache::getPagePriorities(LayerGroupId layerGroupId, Beam auto const apiLock = mManager->lockShared(); std::vector result; result.reserve(mBlocks.stdSize()); + bool const isSsm + = std::holds_alternative(mManager->storage().getLifeCycle(LifeCycleId{layerGroupId.value()})); for (auto const& block : mBlocks) { auto const* page = beamIndex < block.pages.size() ? blockPageGetPage(block.pages[beamIndex].at(layerGroupId)).get() : nullptr; + 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); } @@ -1895,6 +1899,11 @@ bool KvCache::resize(std::optional capacity, std::optional 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(); @@ -2409,8 +2418,6 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) for (size_t ri = 0; ri < reuseTasks.size(); ++ri) { LifeCycleId lc = reuseTasks[ri].lifeCycle; - auto const retention = getRetention(BlockOrdinal{ord}); - locks[ri].page()->claimRetention(retention.retentionPriority, retention.durationMs); sb.pages[kDefaultBeamIndex][lc] = std::move(locks[ri]); } if (mIsDecoding) @@ -2865,7 +2872,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; @@ -2947,7 +2954,7 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match) bool const partialCopy = hasPartialMatch && ordinal == fullReusedEnd; bool const borrowPartial = partialCopy && page->hasRetentionConsumers(); bpSlot = page->hold(); - if (!borrowPartial) + if (claimRetention && !borrowPartial) { auto const retention = getRetention(ordinal); page->claimRetention(retention.retentionPriority, retention.durationMs); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h index 76a635793b70..f6f248e372b0 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h @@ -250,7 +250,7 @@ class KvCache : public std::enable_shared_from_this KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional reuseMatch, std::optional id, PriorityCb priorityCb, std::optional expectedPromptLength = std::nullopt, std::optional textOnly = std::nullopt, bool enableRequestStats = false, - std::optional retentionConfig = std::nullopt); + std::optional retentionConfig = std::nullopt, bool isPrefetch = false); ~KvCache(); @@ -623,7 +623,7 @@ class KvCache : public std::enable_shared_from_this // counts. Called at the end of _setupForReuse, which collects them in the same walk. void _finalizeCachedTokensByLevel( int numTokens, TypedVec const& attentionLevels, std::optional ssmLevel); - void _setupForReuse(BlockRadixTree::ReuseMatch const& match); + void _setupForReuse(BlockRadixTree::ReuseMatch const& match, bool claimRetention); // Reconstruct the committed token sequence from a match's blocks (mirrors // Python's _get_matched_tokens); used when reuse-matching no longer has the // raw input tokens in scope. diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp index 65faccbef8a3..54ad290f5a3d 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp @@ -222,6 +222,17 @@ void KvCacheManager::refreshRetention() mStorage->refreshRetention(); } +bool KvCacheManager::prefetchReuse(ReuseScope reuseScope, TokenSpan inputTokens, CacheLevel target, bool knownNoDigest) +{ + KVCM2_API_GUARD(); + auto const apiLock = lockExclusive(); + auto reuseMatch = matchReuse(reuseScope, inputTokens, knownNoDigest); + auto cache = std::make_shared(*this, std::move(reuseScope), std::move(reuseMatch), std::nullopt, + KvCache::PriorityCb{}, std::nullopt, knownNoDigest, false, std::nullopt, /*isPrefetch=*/true); + auto close = FuncGuard([&cache]() { cache->close(); }); + return cache->prefetch(target); +} + BlockRadixTree::ReuseMatch KvCacheManager::matchReuse( ReuseScope const& reuseScope, TokenSpan inputTokens, bool knownNoDigest) const { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h index d011bdf233ac..8a90a273472e 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h @@ -153,6 +153,9 @@ class KvCacheManager : public std::enable_shared_from_this //! Restore expired, unreferenced pages to the default retention priority. void refreshRetention(); + // Move reuse-matched pages without claiming or extending their request retention policy. + bool prefetchReuse(ReuseScope reuseScope, TokenSpan inputTokens, CacheLevel target, bool knownNoDigest); + // knownNoDigest: from external text_only knowledge, never a scan (see Hasher::update). // Defaults false (safe: the scanning path is taken). BlockRadixTree::ReuseMatch matchReuse( diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp index 0f6b120ee070..d13e1af69ee7 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp @@ -78,7 +78,6 @@ PageStatus Page::status() const noexcept SharedPtr Page::hold() { - manager->cancelRetentionExpiry(*this); // Return existing holder if any. auto h = holder.lock(); if (h) @@ -252,8 +251,19 @@ SharedPtr UncommittedPage::convertToCommitted( // Set the ready event before transfer (matches Python: self.ready_event = ready_event). this->readyEvent = std::move(readyEv); - auto committed = makeShared(manager, blk, lifeCycle, cacheLevel, numTokensInBlock, priority()); - committed->inheritRetention(*this); + bool const isSsm = std::holds_alternative(manager->getLifeCycle(lifeCycle)); + auto const committedPriority = isSsm ? kvCache->getPriority(blk->ordinal(), lifeCycle) : priority(); + auto committed + = makeShared(manager, blk, lifeCycle, cacheLevel, numTokensInBlock, committedPriority); + if (isSsm) + { + auto const retention = kvCache->getRetention(blk->ordinal()); + committed->claimRetention(retention.retentionPriority, retention.durationMs); + } + else + { + committed->inheritRetention(*this); + } // Move slot id to the committed page; invalidate our slot. committed->setSlotId(slotId()); // asserts valid committed->readyEvent = std::move(readyEvent); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp index 7ff04e74d5fd..3eed1b3a2229 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp @@ -837,7 +837,7 @@ void StorageManager::refreshRetention() if (page) { page->mRetentionExpiry.reset(); - TLLM_CHECK_DEBUG(page->status() == PageStatus::DROPPABLE); + // A prefetch may hold storage without claiming or extending its retention interval. updatePriority(*page, kPriorityDefault); } } @@ -960,16 +960,29 @@ void StorageManager::_prepareFreeSlots(TypedVec>> belowOffloadThreshold(numLifeCycles()); + auto restoreBelowThreshold = FuncGuard( + [&]() + { + for (auto const& pages : belowOffloadThreshold) + { + for (auto iter = pages.rbegin(); iter != pages.rend(); ++iter) + { + auto const& page = *iter; + if (!page->scheduledForEviction()) + mLevels.at(page->cacheLevel).controller.scheduleForEviction(*page, /*evictFirst=*/true); + } + } + }); for (auto& pages : fallenPages) { - std::vector> dropped; for (auto iter = pages.begin(); iter != pages.end();) { auto const& page = *iter; if (page->cacheLevel == kHotLevel && page->status() == PageStatus::DROPPABLE && page->priority() < mSecondaryOffloadMinPriority) { - dropped.push_back(std::move(*iter)); + belowOffloadThreshold.at(page->lifeCycle).push_back(std::move(*iter)); iter = pages.erase(iter); } else @@ -977,10 +990,6 @@ void StorageManager::_prepareFreeSlots(TypedVec()) .def("refresh_retention", &kv::KvCacheManager::refreshRetention, nb::call_guard()) + .def( + "prefetch_reuse", + [](kv::KvCacheManager& self, nb::object reuseScopeObj, nb::object inputTokens, kv::CacheLevel target) + { + auto reuseScope = castReuseScope(std::move(reuseScopeObj)); + return withTokens(inputTokens, + [&](kv::TokenSpan view, bool knownNoDigest) + { + nb::gil_scoped_release release; + return self.prefetchReuse(std::move(reuseScope), view, target, knownNoDigest); + }); + }, + nb::arg("reuse_scope"), nb::arg("input_tokens"), nb::arg("target")) .def("is_sparse", &kv::KvCacheManager::isSparse, nb::arg("layer_id"), nb::arg("data_role")) .def( "clear_reusable_blocks", &kv::KvCacheManager::clearReusableBlocks, nb::call_guard()) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 105993450743..a3592e58eaee 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3757,6 +3757,42 @@ def probe_context_reuse(self, req: LlmRequest) -> int | None: self._context_reuse_tokens(req), ) + def _retention_config_for_request(self, req: LlmRequest) -> KvCacheRetentionConfig | None: + config = req.kv_cache_retention_config + if config is None or not self._has_cp_helix: + return config + # Retention ranges use rank-local physical tokens; native ordinals use CP super-blocks. + prompt_blocks = (req.prompt_len + self.tokens_per_block - 1) // self.tokens_per_block + prompt_end = prompt_blocks * self.tokens_per_block + ranges = [] + for token_range in config.token_range_retention_configs: + end = min(prompt_end, token_range.token_end or prompt_end) + if token_range.token_start < end: + ranges.append( + KvCacheRetentionConfig.TokenRangeRetentionConfig( + token_range.token_start * self._helix_cp_size, + end * self._helix_cp_size, + token_range.priority, + token_range.duration_ms, + ) + ) + # A short or empty rank can allocate a decode page within the final global prompt block. + ranges.append( + KvCacheRetentionConfig.TokenRangeRetentionConfig( + prompt_end * self._helix_cp_size, + None, + config.decode_retention_priority, + config.decode_duration_ms, + ) + ) + return KvCacheRetentionConfig( + ranges, + decode_retention_priority=config.decode_retention_priority, + decode_duration_ms=config.decode_duration_ms, + transfer_mode=config.transfer_mode, + directory=config.directory, + ) + def prepare_context_cache(self, req: LlmRequest, reuse_limit: int | None = None) -> int | None: """Create/resume a context cache without touching request cursors. ``reuse_limit`` caps the claimed depth while retaining its lookahead evidence; the returned @@ -3785,7 +3821,7 @@ def prepare_context_cache(self, req: LlmRequest, reuse_limit: int | None = None) req.lora_task_id, tokens, cache_salt=req.cache_salt, - kv_cache_retention_config=req.kv_cache_retention_config, + kv_cache_retention_config=self._retention_config_for_request(req), is_dummy=req.is_dummy, enable_request_stats=req.return_perf_metrics, expected_prompt_length=( @@ -4628,7 +4664,7 @@ def _mirror_draft_kv_cache(self, req: LlmRequest): cache_salt=req.cache_salt, is_dummy=req.is_dummy, expected_prompt_length=req.prompt_len, - kv_cache_retention_config=req.kv_cache_retention_config, + kv_cache_retention_config=self._retention_config_for_request(req), ) if kv_cache is None: return None @@ -6674,29 +6710,19 @@ def prefetch_for_context_tokens(self, requests: list) -> bool: """ if not self.enable_block_reuse: return False - # Prefetch via a transient KV cache that holds the reuse-matched blocks, - # prefetches disk->host, then closes. Holding blocks costs no GPU space - # (never resumed) and close() needs no stream sync. The transient cache - # is NOT registered in kv_cache_map / IndexMapper. + # Prefetch holds storage without claiming a request's retention interval. success = True for req in requests: tokens = self._context_reuse_tokens(req) # Use the same salt derivation as _create_kv_cache so the transient # cache hits the same radix-tree blocks. reuse_salt = self._derive_reuse_salt(req.cache_salt) - kv_cache = self.impl.create_kv_cache( - ReuseScope(lora_id=req.lora_task_id, salt=reuse_salt), - tokens, + prefetched = self.impl.prefetch_reuse( + ReuseScope(lora_id=req.lora_task_id, salt=reuse_salt), tokens, CACHE_LEVEL1 ) - try: - # Prefetch to the first tier below GPU (host if present, - # otherwise disk). prefetch() is a best-effort hint. - prefetched = kv_cache.prefetch(CACHE_LEVEL1) - if not prefetched: - logger.warning("prefetch failed for request %s", req.py_request_id) - success = False - finally: - kv_cache.close() + if not prefetched: + logger.warning("prefetch failed for request %s", req.py_request_id) + success = False return success def reset_reuse_state(self): diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 3d58accb38a3..5a30d89e74dc 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -860,6 +860,10 @@ class KVCacheManager: """ ... def refresh_retention(self) -> None: ... + def prefetch_reuse( + self, reuse_scope: ReuseScope, input_tokens: Sequence[TokenIdExt], target: CacheLevel + ) -> bool: + """Prefetch matched storage while preserving request retention policies and deadlines.""" def probe_reuse( self, reuse_scope: ReuseScope | None = None, diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 697572dcf44a..f8272c068574 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -460,9 +460,20 @@ def _config(priority, duration=None): def test_shared_claim_order_controls_eviction(self, first, second, close_newer_first): self.prepare(2 << 20, 0, 0, 1, None, 0, tokens_per_block=4, kv_buf_size=65536) prompt = [1, 2, 3, 4] + late = self.manager.create_kv_cache( + input_tokens=prompt, kv_cache_retention_config=self._config(100) + ) + self.caches.append(late) + self.assertTrue(late.resume(self.stream.handle)) + self.assertTrue(late.resize(len(prompt))) + self.engine.execute([Step(late, prompt, [])], self.stream.handle) older = self._cache(prompt, self._config(first)) newer = self._cache(prompt, self._config(second)) priority = first if second is None else second + # Deduplicating an already computed page is not a new request claim. + late.commit(prompt) + self.assertEqual(late.get_page_priorities(0), [priority]) + late.close() self.assertEqual(older.get_page_priorities(0), [priority]) self.assertEqual(newer.get_aggregated_page_indices(0), older.get_aggregated_page_indices(0)) requests = [newer, older] if close_newer_first else [older, newer] @@ -510,6 +521,7 @@ def test_prompt_decode_duration_and_unconfigured_claim(self): self.assertEqual(second.get_page_priorities(0), [35, 80, 10]) second.close() time.sleep(0.04) + self.assertTrue(self.manager.prefetch_reuse(ReuseScope(), tokens, GPU_LEVEL)) self.manager.refresh_retention() reclaimed = self._cache(tokens) self.assertEqual(reclaimed.get_page_priorities(0), [35, 35, 35]) @@ -3607,7 +3619,12 @@ def test_commit_is_end_moves_partial_attention_and_ssm_pages(self) -> None: stream = cast(CudaStream, stream_holder.handle) prompt = [self.next_token() for _ in range(48)] - kv_cache = self.manager.create_kv_cache() + kv_cache = self.manager.create_kv_cache( + expected_prompt_length=len(prompt), + kv_cache_retention_config=KvCacheRetentionConfig( + [KvCacheRetentionConfig.TokenRangeRetentionConfig(0, None, 80)] + ), + ) kv_cache.resume(stream) kv_cache.capacity = len(prompt) kv_cache.history_length = len(prompt) @@ -3636,6 +3653,9 @@ def test_commit_is_end_moves_partial_attention_and_ssm_pages(self) -> None: self.assertEqual(attn_page[0], attn_tail_slot) self.assertEqual(ssm_page[0], ssm_slot) self.assertEqual(ssm_page[1], 16) + probe = self.manager.create_kv_cache(input_tokens=prompt) + self.assertEqual(probe.get_page_priorities(LayerGroupId(ssm_lc_id))[-1], 80) + probe.close() def test_ssm_snapshot_moves_to_covering_block(self) -> None: """A snapshot on a partial block survives the full sibling that replaces it.""" From 173f4572edc9e55adccce1a0317dea0272590d09 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:31:17 -0700 Subject: [PATCH 04/23] test: cover retention lifecycle and rank-local policy boundaries Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../kv_cache_manager_v2/kvCache.cpp | 9 ++- .../kv_cache_manager_v2/kvCacheManager.cpp | 2 + .../kvCacheManagerV2ColdPageTest.cpp | 28 ++++++++-- .../test_lists/test-db/l0_b200.yml | 1 + .../test_kv_cache_v2_first_new_block_probe.py | 1 + .../kv_cache/test_kvcm2_integration.py | 56 +++++++++++++++++++ .../test_kv_cache_manager_v2.py | 44 +++++++++++++-- 7 files changed, 130 insertions(+), 11 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index f51a4c83b383..16162179f452 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -190,6 +190,8 @@ std::vector KvCache::getPagePriorities(LayerGroupId layerGroupId, Beam { 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"); std::vector result; result.reserve(mBlocks.stdSize()); bool const isSsm @@ -3033,8 +3035,11 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match, bool claim TLLM_CHECK_WITH_INFO(snapshotPage, "Last matched block must have SSM snapshot after truncation"); ssmLevel = snapshotPage->cacheLevel; mSsmBlocks[kDefaultBeamIndex][*ssmLcId] = snapshotPage->hold(); - auto const retention = getRetention(snapshotBlock.ordinal()); - snapshotPage->claimRetention(retention.retentionPriority, retention.durationMs); + if (claimRetention) + { + auto const retention = getRetention(snapshotBlock.ordinal()); + snapshotPage->claimRetention(retention.retentionPriority, retention.durationMs); + } } _finalizeCachedTokensByLevel(numTokens, attentionLevels, ssmLevel); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp index 54ad290f5a3d..35d2a62a8558 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp @@ -226,6 +226,8 @@ bool KvCacheManager::prefetchReuse(ReuseScope reuseScope, TokenSpan inputTokens, { KVCM2_API_GUARD(); auto const apiLock = lockExclusive(); + if (target < kHotLevel || target >= mStorage->numCacheLevels()) + throw std::out_of_range("Prefetch target is outside the configured cache tiers"); auto reuseMatch = matchReuse(reuseScope, inputTokens, knownNoDigest); auto cache = std::make_shared(*this, std::move(reuseScope), std::move(reuseMatch), std::nullopt, KvCache::PriorityCb{}, std::nullopt, knownNoDigest, false, std::nullopt, /*isPrefetch=*/true); diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp index d3a22708aaf7..011f8b9dec98 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp @@ -813,24 +813,42 @@ TEST(KvCacheManagerV2ColdPageTest, ForceEvictFailureReschedulesFallenPage) ASSERT_EQ(cudaSetDevice(0), cudaSuccess); auto codec = std::make_unique(AsyncRejectingColdPageCodec::Operation::kEncode); auto* codecPtr = codec.get(); - auto manager = std::make_shared(makeTieredConfig(), nullptr, std::move(codec)); + auto config = makeSplitColdGroupingConfig(); + std::get(config.layers[1]).buffers.front().size *= 2; + auto manager = std::make_shared(std::move(config), nullptr, std::move(codec)); auto releaseCodecGuard = FuncGuard([codecPtr]() { codecPtr->release(); }); auto& storage = manager->storage(); - LifeCycleId const lifeCycle{0}; - TypedVec oneSlot(LifeCycleId{1}, 1); + LifeCycleId const lifeCycle{1}; + TypedVec oneSlot(LifeCycleId{2}, 1); auto hotSlots = storage.newGpuSlots(oneSlot); - auto sourcePage = makeCommittedPage(*manager, storage, kHotLevel, hotSlots[lifeCycle].front()); + auto sourcePage = makeCommittedPage(*manager, storage, kHotLevel, hotSlots[lifeCycle].front(), lifeCycle, 80); + auto lowPage + = makeCommittedPage(*manager, storage, kHotLevel, hotSlots[LifeCycleId{0}].front(), LifeCycleId{0}, 10); + SlotId const lowSlotId = lowPage->slotId(); + WeakPtr lowPageRef = lowPage; + lowPage.reset(); SlotId const sourceSlotId = sourcePage->slotId(); TypedVec evictOne(storage.numPoolGroups(kHotLevel), 0); evictOne[storage.getPoolGroupIndex(kHotLevel, lifeCycle)] = 1; + evictOne[storage.getPoolGroupIndex(kHotLevel, LifeCycleId{0})] = 1; + size_t dropped = 0; + auto recordDrop = [&](std::vector> const& pages, CacheLevel) { dropped += pages.size(); }; - EXPECT_THROW(storage.forceEvict(kHotLevel, evictOne), TllmException); + EXPECT_THROW(storage.forceEvict(kHotLevel, evictOne, recordDrop), TllmException); ASSERT_TRUE(codecPtr->launched()); EXPECT_EQ(sourcePage->cacheLevel, kHotLevel); EXPECT_EQ(sourcePage->slotId(), sourceSlotId); EXPECT_TRUE(sourcePage->scheduledForEviction()); EXPECT_FALSE(sourcePage->queryReady()); + EXPECT_EQ(sourcePage->priority(), 80); + EXPECT_EQ(dropped, 0); + lowPage = lowPageRef.lock(); + ASSERT_TRUE(lowPage); + EXPECT_EQ(lowPage->cacheLevel, kHotLevel); + EXPECT_EQ(lowPage->slotId(), lowSlotId); + EXPECT_TRUE(lowPage->scheduledForEviction()); + EXPECT_EQ(lowPage->priority(), 10); codecPtr->release(); sourcePage->readyEvent.synchronize(); diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index ec98130ff647..34f3e5b83f0f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -291,6 +291,7 @@ l0_b200: - unittest/kv_cache_manager_v2_tests - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v1] - unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_v2_hybrid_retains_configured_number_of_conversation_turns + - unittest/_torch/executor/kv_cache/test_kvcm2_integration.py::test_helix_retention_uses_rank_local_ranges - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v2] - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v1] - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v2] diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py index 90402c3c0ae1..e0cc136f8c47 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_first_new_block_probe.py @@ -100,6 +100,7 @@ def __init__( self.is_dummy = False self.lora_task_id = lora_task_id self.cache_salt = cache_salt + self.kv_cache_retention_config = None # LlmRequest's own default; a truthy value here turns on the # per-request stats branch of _create_kv_cache. self.return_perf_metrics = return_perf_metrics diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index 2a1d353bb9b1..12ffd70bef12 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -41,6 +41,7 @@ from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType, SamplingConfig from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE +from tensorrt_llm.bindings.executor import KvCacheRetentionConfig from tensorrt_llm.bindings.internal.batch_manager import CacheType, LinearCacheType from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.llmapi.llm_args import ( @@ -974,6 +975,7 @@ def test_prepare_context_cache_records_lookup_without_mutating_cursor( py_request_id=7, lora_task_id=3, cache_salt="11", + kv_cache_retention_config=None, is_dummy=False, return_perf_metrics=False, prompt_len=8, @@ -1524,6 +1526,7 @@ class _ContextRequest: use_conversation_params: bool = True lora_task_id: int | None = None cache_salt: str | None = None + kv_cache_retention_config: KvCacheRetentionConfig | None = None is_first_context_chunk: bool = True is_last_context_chunk: bool = True is_disagg_generation_init_state: bool = False @@ -1739,6 +1742,59 @@ def test_generation_dummy_does_not_reuse_real_request_prefix() -> None: manager.shutdown() +@pytest.mark.parametrize( + "global_prompt,local_prompt,rank,range_end,expected", + [(256, 128, 0, 128, [35, 80]), (64, 0, 1, 128, [10]), (130, 65, 0, 65, [35, 80])], +) +def test_helix_retention_uses_rank_local_ranges( + monkeypatch: pytest.MonkeyPatch, + global_prompt: int, + local_prompt: int, + rank: int, + range_end: int, + expected: list[int], +) -> None: + # Only cross-process coordination is replaced; allocation and policy run in native V2 on GPU. + dist = SimpleNamespace(allreduce=lambda value, op: value, local_world_size=1) + monkeypatch.setattr(Distributed, "get", lambda mapping: dist) + manager = KVCacheManagerV2( + KvCacheConfig(enable_block_reuse=False, max_gpu_total_bytes=4 << 20), + CacheType.SELF, + num_layers=1, + num_kv_heads=2, + head_dim=128, + tokens_per_block=64, + max_seq_len=512, + max_batch_size=1, + max_num_tokens=256, + mapping=Mapping( + world_size=2, cp_size=2, tp_size=1, rank=rank, cp_config={"cp_type": "HELIX"} + ), + dtype=DataType.HALF, + ) + request = _ContextRequest( + 1, list(range(local_prompt)), local_prompt, "retention", use_conversation_params=False + ) + request.total_input_len_cp = global_prompt + request.kv_cache_retention_config = KvCacheRetentionConfig( + [KvCacheRetentionConfig.TokenRangeRetentionConfig(64, range_end, 80)], + decode_retention_priority=10, + ) + try: + assert isinstance(manager.impl, KVCacheManager) + assert manager.impl.init_config.tokens_per_block == 128 + assert manager.prepare_context_cache(request) is not None + cache = manager.kv_cache_map[request.py_request_id] + assert cache.resize(global_prompt) + group = manager.impl.get_layer_group_id(0) + assert cache.get_page_priorities(group) == expected + assert cache.resize((len(expected) + 1) * 128) + assert cache.get_page_priorities(group) == expected + [10] + finally: + manager.free_resources(request) + manager.shutdown() + + @pytest.fixture def max_num_turns() -> int: return 1 diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index f8272c068574..486d0963f93d 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -515,6 +515,10 @@ def test_prompt_decode_duration_and_unconfigured_claim(self): first.commit(tokens[5:]) second = self._cache(tokens, config, prompt_length=5) self.assertEqual(first.get_page_priorities(0), [35, 80, 10]) + with self.assertRaises(IndexError): + first.get_page_priorities(0, -1) + with self.assertRaises(IndexError): + self.manager.prefetch_reuse(ReuseScope(), tokens, CacheLevel(-1)) first.close() time.sleep(0.04) self.manager.refresh_retention() @@ -540,6 +544,26 @@ def test_prompt_decode_duration_and_unconfigured_claim(self): self.manager.refresh_retention() self.assertEqual(self._cache(tokens).get_page_priorities(0), [35, 80, 10]) + for shared in (False, True): + with self.subTest(shared_swa_page=shared): + for cache in self.caches: + cache.close() + self.caches.clear() + self.manager.shutdown() + self.prepare(2 << 20, 0, 0, 1, 4, 0, tokens_per_block=4, kv_buf_size=65536) + config = self._config(80, duration) + source = self._cache(tokens[:4], config) + sharer = self._cache(tokens[:4], config) if shared else None + self.assertTrue(source.resize(8)) + self.engine.execute([Step(source, tokens[4:8], tokens[:4])], self.stream.handle) + source.commit(tokens[4:8]) + source.close() + if sharer is not None: + sharer.close() + time.sleep(0.04) + self.manager.refresh_retention() + self.assertEqual(self._cache(tokens[:4]).get_page_priorities(0), [80]) + def test_offload_reload_and_error_cleanup(self): self.prepare(2 << 20, 2 << 20, 0, 1, None, 0, tokens_per_block=4, kv_buf_size=65536) tokens = [1, 2, 3, 4] @@ -3622,7 +3646,11 @@ def test_commit_is_end_moves_partial_attention_and_ssm_pages(self) -> None: kv_cache = self.manager.create_kv_cache( expected_prompt_length=len(prompt), kv_cache_retention_config=KvCacheRetentionConfig( - [KvCacheRetentionConfig.TokenRangeRetentionConfig(0, None, 80)] + [ + KvCacheRetentionConfig.TokenRangeRetentionConfig( + 0, None, 80, timedelta(milliseconds=20) + ) + ] ), ) kv_cache.resume(stream) @@ -3636,7 +3664,10 @@ def test_commit_is_end_moves_partial_attention_and_ssm_pages(self) -> None: ssm_slot = kv_cache.get_ssm_block_base_index(LayerGroupId(ssm_lc_id)) kv_cache.commit(prompt, is_end=True) - kv_cache.close() + try: + self.assertEqual(kv_cache.get_page_priorities(LayerGroupId(ssm_lc_id))[-1], 80) + finally: + kv_cache.close() _, attn_pages = _introspection.reuse_match_pages( self.manager, ReuseScope(), prompt, attn_lc_id, self.manager.enable_partial_match @@ -3653,9 +3684,14 @@ def test_commit_is_end_moves_partial_attention_and_ssm_pages(self) -> None: self.assertEqual(attn_page[0], attn_tail_slot) self.assertEqual(ssm_page[0], ssm_slot) self.assertEqual(ssm_page[1], 16) + time.sleep(0.04) + self.assertTrue(self.manager.prefetch_reuse(ReuseScope(), prompt, GPU_LEVEL)) + self.manager.refresh_retention() probe = self.manager.create_kv_cache(input_tokens=prompt) - self.assertEqual(probe.get_page_priorities(LayerGroupId(ssm_lc_id))[-1], 80) - probe.close() + try: + self.assertEqual(probe.get_page_priorities(LayerGroupId(ssm_lc_id))[-1], 35) + finally: + probe.close() def test_ssm_snapshot_moves_to_covering_block(self) -> None: """A snapshot on a partial block survives the full sibling that replaces it.""" From f77f698b9d591b730b430b6a8f1f322fc31d9b80 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:33:19 -0700 Subject: [PATCH 05/23] fix: keep prefetch a no-op for GPU-only caches Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index a3592e58eaee..eafdc7a70de2 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -6710,6 +6710,8 @@ def prefetch_for_context_tokens(self, requests: list) -> bool: """ if not self.enable_block_reuse: return False + if not self.can_evict: + return True # Prefetch holds storage without claiming a request's retention interval. success = True for req in requests: From b081744e02c8fae2dd66d0f9398bca6357f549c7 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:43:27 -0700 Subject: [PATCH 06/23] test: cover GPU-only retention prefetch Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../_torch/executor/kv_cache/test_mamba_cache_manager.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index ac7feea518a3..71151f1b5441 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -2987,6 +2987,8 @@ def test_v2_hybrid_retains_configured_number_of_conversation_turns(): priorities = cache_a.get_page_priorities(attention_group) assert priorities and set(priorities) == {80} assert mgr.impl.init_config.secondary_offload_min_priority == 50 + assert not mgr.can_evict + assert mgr.prefetch_for_context_tokens([request_a]) is True request_a_state_index = mgr.get_state_indices([request_a.py_request_id], [False])[0] mgr.free_resources(request_a) _run_v2_hybrid_context(mgr, request_b) From f9251507dcbc5b1e76dbe7ad54662024f6bd3a94 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:45:17 -0700 Subject: [PATCH 07/23] test: explicitly configure GPU-only prefetch coverage Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../_torch/executor/kv_cache/test_mamba_cache_manager.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 71151f1b5441..3355ebd7d408 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -2427,6 +2427,7 @@ def _build_v2_hybrid_with_mamba_layer( kda_replay_num_spec=None, enable_kv_pool_rebalance=False, secondary_offload_min_priority=None, + host_cache_size=None, ): """Construct a real MambaHybridCacheManagerV2.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -2456,6 +2457,7 @@ def _build_v2_hybrid_with_mamba_layer( dtype=kv_cache_dtype, enable_kv_pool_rebalance=enable_kv_pool_rebalance, secondary_offload_min_priority=secondary_offload_min_priority, + host_cache_size=host_cache_size, ) return MambaHybridCacheManagerV2( mamba_d_state=8, @@ -2954,6 +2956,7 @@ def test_v2_hybrid_retains_configured_number_of_conversation_turns(): max_num_turns=2, additional_snapshot_offsets_from_end=[0], secondary_offload_min_priority=50, + host_cache_size=0, ) request_a = _make_v2_conversation_request( 1, From 6f911b73c32f833526b2ee27be4304071d4616a0 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:53:26 -0700 Subject: [PATCH 08/23] fix: align retention type declarations and test fixtures Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../source/examples/kvcacheretentionconfig.md | 6 +++-- .../runtime/kv_cache_manager_v2/__init__.pyi | 2 +- .../defs/llmapi/test_llm_api_connector.py | 10 +++----- .../kv_cache/test_kv_cache_v2_scheduler.py | 1 + .../kv_cache/test_kvcm2_integration.py | 2 +- .../_torch/sampler/test_beam_search.py | 1 + .../disaggregated/test_chunked_transfer.py | 24 +++++++++---------- 7 files changed, 23 insertions(+), 23 deletions(-) diff --git a/docs/source/examples/kvcacheretentionconfig.md b/docs/source/examples/kvcacheretentionconfig.md index 4651f725d5dd..4bde2e29decd 100644 --- a/docs/source/examples/kvcacheretentionconfig.md +++ b/docs/source/examples/kvcacheretentionconfig.md @@ -76,8 +76,10 @@ When requests share a page, the latest claim with an explicit priority sets its priority. Releasing an older request does not restore that request's priority. A claim without an explicit priority preserves the existing priority, but replaces the duration, including clearing it when no duration is supplied. The duration -starts when the last reference is released. Expired pages return to priority 35; -pages still referenced by a request cannot be permanently evicted. +starts when a page has no request references. An independently published +snapshot can reach this state before its originating request finishes. Prefetch +preserves the existing expiry. Expired pages return to priority 35; pages still +referenced by a request cannot be permanently evicted. On GPU eviction, unreferenced pages below `secondary_offload_min_priority` (default 30) are discarded. Pages at or above it can be offloaded. Offloading diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 5a30d89e74dc..81517522e770 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -330,7 +330,7 @@ class KVCacheManagerConfig: text_only: bool = False enable_partial_commit: bool = True pool_rebalance: PoolRebalanceConfig = ... - secondary_offload_min_priority: Priority = 30 + secondary_offload_min_priority: int = 30 @property def enable_swa_scratch_reuse(self) -> bool: ... diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index dc6ae8045328..49b0ae8a24da 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -265,13 +265,9 @@ def generate_and_wait(model, scheduler, worker, *args, **kwargs): def test_v2_connector_contract_does_not_reuse_the_v1_methods(): """The V2 connector path implements none of the V1 accessors, by design. - Something depends on that, and it does not ask: `update_and_build_data` - reports `block_hashes` and `priorities` empty on V2 by branching on - `isinstance(manager, KVCacheManagerV2)`, not on `hasattr`. Those - short-circuits are only correct while V2 genuinely has no such accessor - - the day one is added (retention priorities are a known gap; see - `test_connector_priorities`) the branch keeps reporting nothing while the - data exists, and this is what says so. + The V2 branch of `update_and_build_data` reads native page priorities and + reports page indices per layer group. It must not call the V1 accessors, + which assume one flat block-id space across all layers. A static check rather than an end-to-end run: under V2 the connector would die at the *first* method it reached, so no run can report more than one at diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index 2efd68e5096e..bd902ea09770 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -93,6 +93,7 @@ def test_resume_restores_page_buffers_and_sparse_metadata_row(sparse: bool) -> N @pytest.mark.parametrize("is_draft", [False, True]) def test_sparse_metadata_publishes_after_preparation(is_draft: bool) -> None: manager = object.__new__(KVCacheManagerV2) + manager.impl = Mock() order = Mock() manager._disagg_receive_ready = {} manager.is_draft = is_draft diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index 12ffd70bef12..685648f1c1e2 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -1744,7 +1744,7 @@ def test_generation_dummy_does_not_reuse_real_request_prefix() -> None: @pytest.mark.parametrize( "global_prompt,local_prompt,rank,range_end,expected", - [(256, 128, 0, 128, [35, 80]), (64, 0, 1, 128, [10]), (130, 65, 0, 65, [35, 80])], + [(256, 128, 0, 128, [35, 80]), (64, 0, 1, 128, [10]), (130, 66, 0, 65, [35, 80])], ) def test_helix_retention_uses_rank_local_ranges( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 72fe06068076..28e251e1f0e1 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -2016,6 +2016,7 @@ def test_vbws_rejects_decreasing_beam_width_array(beam_width_array: list[int], # Exercise the real beam admission checks with unrelated cache, transport, # token-range, sampler, and budget validation disabled in the executor stub. executor = types.SimpleNamespace( + _is_kv_manager_v2=False, max_beam_width=request.py_beam_width, kv_cache_manager=None, kv_cache_transceiver=None, diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index d75c51a6ca77..7d67ecec6fed 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -630,14 +630,14 @@ def test_pipelined_transfer_requires_gen_first_flow(): """ValueError when a real request is not using gen-first flow.""" from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - executor = MagicMock() + executor = MagicMock(_is_kv_manager_v2=False) executor.is_warmup = False executor.kv_cache_transceiver.pipeline_transfer_enabled = True executor.model_engine.attn_runtime_features.chunked_prefill = True executor._validate_token_id_range = MagicMock() executor.sampler.validate_request = MagicMock() - request = MagicMock() + request = MagicMock(kv_cache_retention_config=None) request.sampling_config = None request.py_beam_width = 1 request.llm_request_type = LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY @@ -657,13 +657,13 @@ def test_pipelined_transfer_requires_chunked_prefill(): """Pipelined context transfer requires chunked prefill.""" from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - executor = MagicMock() + executor = MagicMock(_is_kv_manager_v2=False) executor.kv_cache_transceiver.pipeline_transfer_enabled = True executor.model_engine = SimpleNamespace( attn_runtime_features=SimpleNamespace(chunked_prefill=False) ) - request = MagicMock() + request = MagicMock(kv_cache_retention_config=None) request.llm_request_type = LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY request.py_disaggregated_params = SimpleNamespace( schedule_style=DisaggScheduleStyle.GENERATION_FIRST @@ -680,12 +680,12 @@ def test_pipelined_transfer_rejects_pipeline_parallelism_for_context_request(): """Context workers require all layers on one pipeline rank.""" from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - executor = MagicMock() + executor = MagicMock(_is_kv_manager_v2=False) executor.kv_cache_transceiver.pipeline_transfer_enabled = True executor.model_engine.attn_runtime_features.chunked_prefill = True executor.dist.pp_size = 2 - request = MagicMock() + request = MagicMock(kv_cache_retention_config=None) request.py_beam_width = 1 request.llm_request_type = LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY request.py_disaggregated_params = SimpleNamespace( @@ -714,11 +714,11 @@ def test_pipelined_transfer_requires_single_beam_for_context_request(): """Context-side pipelining rejects beam search before scheduling.""" from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - executor = MagicMock() + executor = MagicMock(_is_kv_manager_v2=False) executor.kv_cache_transceiver.pipeline_transfer_enabled = True executor.model_engine.attn_runtime_features.chunked_prefill = True - request = MagicMock() + request = MagicMock(kv_cache_retention_config=None) request.py_beam_width = 2 request.llm_request_type = LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY request.is_context_only_request = True @@ -734,14 +734,14 @@ def test_pipelined_transfer_allows_non_disaggregated_request(): """Requests without disaggregated parameters do not transfer KV cache.""" from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - executor = MagicMock() + executor = MagicMock(_is_kv_manager_v2=False) executor.is_warmup = False executor.max_beam_width = 1 executor.kv_cache_transceiver.pipeline_transfer_enabled = True executor._validate_token_id_range = MagicMock() executor.sampler.validate_request = MagicMock() - request = MagicMock() + request = MagicMock(kv_cache_retention_config=None) request.sampling_config = None request.py_beam_width = 1 request.llm_request_type = LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY @@ -757,7 +757,7 @@ def test_pipelined_transfer_allows_generation_only_request(): """Generation workers do not build prefill chunks or enforce sender-only limits.""" from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - executor = MagicMock() + executor = MagicMock(_is_kv_manager_v2=False) executor.is_warmup = False executor.max_beam_width = 2 executor.kv_cache_transceiver.pipeline_transfer_enabled = True @@ -766,7 +766,7 @@ def test_pipelined_transfer_allows_generation_only_request(): executor._validate_token_id_range = MagicMock() executor.sampler.validate_request = MagicMock() - request = MagicMock() + request = MagicMock(kv_cache_retention_config=None) request.sampling_config = None request.py_beam_width = 2 request.llm_request_type = LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY From ea9c3f8326a47cea0ee0fae1020a3dd22b73bff6 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 10 Oct 2026 04:00:21 -0700 Subject: [PATCH 09/23] test: provide native manager in disaggregated ordering fixture Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- tests/unittest/_torch/executor/test_disagg_receive_ordering.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/_torch/executor/test_disagg_receive_ordering.py b/tests/unittest/_torch/executor/test_disagg_receive_ordering.py index f55be54cfd6c..f25ca9339d1a 100644 --- a/tests/unittest/_torch/executor/test_disagg_receive_ordering.py +++ b/tests/unittest/_torch/executor/test_disagg_receive_ordering.py @@ -51,6 +51,7 @@ def _manager() -> Mock: """Run the real admission and resource hooks with mocked cache allocation.""" manager = Mock(spec=KVCacheManagerV2) manager._sharing = None # no lender: free_resources closes and removes the cache itself + manager.impl = Mock() manager._stream = Mock() manager._disagg_receive_ready = {} manager._pending_preemption = {} From edfe5d5e8bbdbad2efb87bf157a869e2d8f8be28 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 10 Oct 2026 04:31:52 -0700 Subject: [PATCH 10/23] fix: convert prefetch cache levels at the native binding Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp | 4 ++-- .../kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 7a31a3c6c26b..349df3d1c15d 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -2652,14 +2652,14 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def("refresh_retention", &kv::KvCacheManager::refreshRetention, nb::call_guard()) .def( "prefetch_reuse", - [](kv::KvCacheManager& self, nb::object reuseScopeObj, nb::object inputTokens, kv::CacheLevel target) + [](kv::KvCacheManager& self, nb::object reuseScopeObj, nb::object inputTokens, int target) { auto reuseScope = castReuseScope(std::move(reuseScopeObj)); return withTokens(inputTokens, [&](kv::TokenSpan view, bool knownNoDigest) { nb::gil_scoped_release release; - return self.prefetchReuse(std::move(reuseScope), view, target, knownNoDigest); + return self.prefetchReuse(std::move(reuseScope), view, kv::CacheLevel{target}, knownNoDigest); }); }, nb::arg("reuse_scope"), nb::arg("input_tokens"), nb::arg("target")) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 486d0963f93d..be76b50ef358 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -519,6 +519,8 @@ def test_prompt_decode_duration_and_unconfigured_claim(self): first.get_page_priorities(0, -1) with self.assertRaises(IndexError): self.manager.prefetch_reuse(ReuseScope(), tokens, CacheLevel(-1)) + with self.assertRaises(IndexError): + self.manager.prefetch_reuse(ReuseScope(), tokens, CacheLevel(1)) first.close() time.sleep(0.04) self.manager.refresh_retention() From 00635f2610aed47cf8ca64cf2a68106974ac6c3c Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sat, 10 Oct 2026 04:42:34 -0700 Subject: [PATCH 11/23] test: assert the manager used by default retention controls Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- tests/integration/defs/llmapi/test_llm_api_connector.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index 49b0ae8a24da..dac65f7001ef 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -1343,11 +1343,13 @@ def test_connector_priorities(enforce_single_worker, model_with_connector, ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], indirect=True) def test_connector_priorities_default(enforce_single_worker, - model_with_connector): + model_with_connector, + use_kv_cache_manager_v2): """Test that priorities are None when no retention config is provided.""" model_fn, scheduler, worker = model_with_connector model = model_fn(disable_overlap_scheduler=True) + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) scheduler.get_num_new_matched_tokens.return_value = 0, False worker.get_finished.return_value = [], [] From a396fea29c96ab78b827e3b74fc81b029dcea4c6 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sun, 11 Oct 2026 08:05:45 -0700 Subject: [PATCH 12/23] Rename the cache offload threshold for cold tiers Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../kv_cache_manager_v2/config.cpp | 4 -- .../kv_cache_manager_v2/config.h | 2 +- .../kv_cache_manager_v2/kvCache.cpp | 7 +++ .../kv_cache_manager_v2/kvCacheManager.cpp | 2 +- .../kv_cache_manager_v2/storageManager.cpp | 8 ++-- .../kv_cache_manager_v2/storageManager.h | 4 +- .../batch_manager/kvCacheManagerV2.cpp | 8 ++-- docs/source/developer-guide/telemetry.md | 2 +- .../source/examples/kvcacheretentionconfig.md | 8 +++- docs/source/features/kvcache.md | 2 +- .../run_cache_transceiver_test.py | 2 +- .../kv_cache/kv_cache_manager_v2.py | 6 +-- .../_torch/pyexecutor/resource_manager.py | 7 ++- tensorrt_llm/llmapi/llm_args.py | 38 ++++++++++++++-- .../runtime/kv_cache_manager_v2/__init__.py | 2 +- .../runtime/kv_cache_manager_v2/__init__.pyi | 2 +- .../usage/llm_args_golden_manifest.json | 10 ++--- .../kv_cache/test_mamba_cache_manager.py | 8 ++-- .../_torch/executor/test_py_executor.py | 1 + .../_torch/executor/test_resource_manager.py | 9 +++- .../test_cache_transceiver_single_process.py | 4 +- .../disaggregated/test_kv_transfer.py | 2 +- .../test_kv_cache_manager_v2.py | 4 +- .../test_optional_checkpoint.py | 44 ++++++++++++++++--- tests/unittest/llmapi/test_llm_args.py | 43 +++++++++++++++++- 25 files changed, 174 insertions(+), 55 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp index 425a25bddd02..8a6e20701d73 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp @@ -40,10 +40,6 @@ void DiskCacheTierConfig::validate() const void KVCacheManagerConfig::validate() const { - if (secondaryOffloadMinPriority < kPriorityMin || secondaryOffloadMinPriority > kPriorityMax) - { - throw AssertionError("secondary_offload_min_priority must be between 0 and 100"); - } if (swaScratchReuse.has_value()) { swaScratchReuse->validate(); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h index 715742299fdc..0ea731080f58 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h @@ -451,7 +451,7 @@ struct KVCacheManagerConfig bool enablePartialCommit = true; // Unreferenced GPU pages below this priority are dropped instead of offloaded. - Priority secondaryOffloadMinPriority = 30; + Priority coldTierOffloadMinPriority = 30; bool enableSwaScratchReuse() const noexcept { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index 16162179f452..1b81d1a1804c 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -311,11 +311,18 @@ void KvCache::_selectOptionalGroups(std::vector 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; } } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp index 35d2a62a8558..e1fb3d21b89a 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp @@ -125,7 +125,7 @@ KvCacheManager::KvCacheManager(KVCacheManagerConfig const& config, std::shared_p StorageConfig storageConfig = createStorageConfig(mConfig); mStorage = std::make_shared(mLifeCycles, storageConfig, mConfig.tokensPerBlock, std::move(coldPageCodec), mConfig.swaScratchReuse, mConfig.typicalStep, mConfig.constraints, - mConfig.initialPoolRatio, mEventSink, mConfig.maxUtilForResume, mConfig.secondaryOffloadMinPriority); + mConfig.initialPoolRatio, mEventSink, mConfig.maxUtilForResume, mConfig.coldTierOffloadMinPriority); mTargetRatioListHot = _currentHotRatio(); mTargetRatioListCold = _currentColdRatios(); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp index 3eed1b3a2229..ccfad1c559e2 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp @@ -214,10 +214,10 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi std::unique_ptr coldPageCodec, std::optional swaScratchReuse, std::optional const& typicalBatch, std::vector const& constraints, std::optional> const& initialPoolRatio, std::shared_ptr eventSink, - float maxUtilForResume, Priority secondaryOffloadMinPriority) + float maxUtilForResume, Priority coldTierOffloadMinPriority) : mLifeCycles(lifeCycles) , mEventSink(std::move(eventSink)) - , mSecondaryOffloadMinPriority(secondaryOffloadMinPriority) + , mColdTierOffloadMinPriority(coldTierOffloadMinPriority) , mHotPoolGroupMapping(config.lifeCycleGrouping()) , mSwaScratchReuse(std::move(swaScratchReuse)) , mColdPageCodec(coldPageCodec ? std::move(coldPageCodec) : createDefaultKvCacheColdPageCodec()) @@ -969,7 +969,7 @@ void StorageManager::_prepareFreeSlots(TypedVecscheduledForEviction()) + if (page->hasValidSlot() && !page->scheduledForEviction() && isEvictable(*page)) mLevels.at(page->cacheLevel).controller.scheduleForEviction(*page, /*evictFirst=*/true); } } @@ -980,7 +980,7 @@ void StorageManager::_prepareFreeSlots(TypedVeccacheLevel == kHotLevel && page->status() == PageStatus::DROPPABLE - && page->priority() < mSecondaryOffloadMinPriority) + && page->priority() < mColdTierOffloadMinPriority) { belowOffloadThreshold.at(page->lifeCycle).push_back(std::move(*iter)); iter = pages.erase(iter); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h index 569cbff98145..81c24b91f9aa 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h @@ -130,7 +130,7 @@ class StorageManager : public std::enable_shared_from_this std::optional const& typicalBatch = std::nullopt, std::vector const& constraints = {}, std::optional> const& initialPoolRatio = std::nullopt, std::shared_ptr eventSink = nullptr, float maxUtilForResume = 1.0f, - Priority secondaryOffloadMinPriority = 30); + Priority coldTierOffloadMinPriority = 30); ~StorageManager(); StorageManager(StorageManager const&) = delete; @@ -426,7 +426,7 @@ class StorageManager : public std::enable_shared_from_this LifeCycleRegistry const& mLifeCycles; std::shared_ptr mEventSink; - Priority mSecondaryOffloadMinPriority; + Priority mColdTierOffloadMinPriority; LifeCyclePoolGroupMapping mHotPoolGroupMapping; LifeCyclePoolGroupMapping mColdPoolGroupMapping; // Codec batching representative for each lifecycle. diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 349df3d1c15d..fe1416ac96aa 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1749,7 +1749,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) std::optional> initialPoolRatio, std::optional swaScratchReuse, bool commitMinSnapshot, bool enableStats, bool textOnly, bool enablePartialCommit, kv::PoolRebalanceConfig poolRebalance, - kv::Priority secondaryOffloadMinPriority) + kv::Priority coldTierOffloadMinPriority) { new (cfg) kv::KVCacheManagerConfig(); cfg->tokensPerBlock = tokensPerBlock; @@ -1774,7 +1774,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) cfg->textOnly = textOnly; cfg->enablePartialCommit = enablePartialCommit; cfg->poolRebalance = std::move(poolRebalance); - cfg->secondaryOffloadMinPriority = secondaryOffloadMinPriority; + cfg->coldTierOffloadMinPriority = coldTierOffloadMinPriority; // Mirror Python's __post_init__: validate at construction. Config-integrity // failures raise AssertionError (translated below). cfg->validate(); @@ -1786,7 +1786,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::arg("swa_scratch_reuse").none() = std::nullopt, nb::arg("commit_min_snapshot") = false, nb::arg("enable_stats") = true, nb::arg("text_only") = false, nb::arg("enable_partial_commit") = true, nb::arg("pool_rebalance") = kv::PoolRebalanceConfig{}, - nb::arg("secondary_offload_min_priority") = 30) + nb::arg("cold_tier_offload_min_priority") = 30) .def_rw("tokens_per_block", &kv::KVCacheManagerConfig::tokensPerBlock) .def_rw("cache_tiers", &kv::KVCacheManagerConfig::cacheTiers) .def_rw("layers", &kv::KVCacheManagerConfig::layers) @@ -1794,7 +1794,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def_rw("enable_partial_reuse", &kv::KVCacheManagerConfig::enablePartialReuse) .def_rw("reuse_match_backoff", &kv::KVCacheManagerConfig::reuseMatchBackoff) .def_rw("enable_partial_commit", &kv::KVCacheManagerConfig::enablePartialCommit) - .def_rw("secondary_offload_min_priority", &kv::KVCacheManagerConfig::secondaryOffloadMinPriority) + .def_rw("cold_tier_offload_min_priority", &kv::KVCacheManagerConfig::coldTierOffloadMinPriority) .def_rw("typical_step", &kv::KVCacheManagerConfig::typicalStep) .def_rw("constraints", &kv::KVCacheManagerConfig::constraints) .def_rw("initial_pool_ratio", &kv::KVCacheManagerConfig::initialPoolRatio, diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index e6b6b74df22a..6f8a5f0cb7ca 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/docs/source/developer-guide/telemetry.md @@ -134,6 +134,7 @@ unset or when the safety sanitizer rejects the runtime value. | `kv_cache_config.block_reuse_config.max_num_turns` | `int` | `value` | | | `kv_cache_config.block_reuse_config.policy` | `literal` | `categorical` | `all_reusable`, `per_request`, `per_conversation` | | `kv_cache_config.block_reuse_config.swa_endpoint_rewind_tokens` | `int` | `value` | | +| `kv_cache_config.cold_tier_offload_min_priority` | `int\|none` | `value` | | | `kv_cache_config.copy_on_partial_reuse` | `bool` | `value` | | | `kv_cache_config.cross_kv_cache_fraction` | `float\|none` | `value` | | | `kv_cache_config.disk_cache_size` | `int\|none` | `value` | | @@ -164,7 +165,6 @@ unset or when the safety sanitizer rejects the runtime value. | `kv_cache_config.max_tokens` | `int\|none` | `value` | | | `kv_cache_config.max_util_for_resume` | `float` | `value` | | | `kv_cache_config.pool_ratio` | `list[float]\|none` | `value` | | -| `kv_cache_config.secondary_offload_min_priority` | `int\|none` | `value` | | | `kv_cache_config.sink_token_length` | `int\|none` | `value` | | | `kv_cache_config.tokens_per_block` | `int` | `value` | | | `kv_cache_config.use_kv_cache_manager_v2` | `bool\|literal` | `categorical` | `auto` | diff --git a/docs/source/examples/kvcacheretentionconfig.md b/docs/source/examples/kvcacheretentionconfig.md index 4bde2e29decd..627baad08110 100644 --- a/docs/source/examples/kvcacheretentionconfig.md +++ b/docs/source/examples/kvcacheretentionconfig.md @@ -81,12 +81,18 @@ snapshot can reach this state before its originating request finishes. Prefetch preserves the existing expiry. Expired pages return to priority 35; pages still referenced by a request cannot be permanently evicted. -On GPU eviction, unreferenced pages below `secondary_offload_min_priority` +On GPU eviction, unreferenced pages below `KvCacheConfig.cold_tier_offload_min_priority` (default 30) are discarded. Pages at or above it can be offloaded. Offloading such a page preserves its priority and clears its duration, as in V1; a later claim can supply a new duration. Referenced pages may move between tiers but are never discarded by this threshold. +`secondary_offload_min_priority` is deprecated. Existing configurations using it +are accepted with a warning and mapped to `cold_tier_offload_min_priority` for +both managers. Setting both fields to different values is an error. Serialized +configurations use the new name. The threshold is manager-wide, not part of +the per-request `KvCacheRetentionConfig`. + V2 supports the default `transfer_mode=DRAM`. Its host and disk offload tiers are configured at manager creation. The per-request `GDS` and `POSIX_DEBUG_FALLBACK` transfer modes, whose `directory` selects V1 block files, are rejected by V2. diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index b44fc1386846..4c56ede9eba6 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -251,7 +251,7 @@ V2 derives `mm_keys` directly from the cached token sequence. Each entry identif Before a block is evicted from GPU memory, it can optionally be offloaded to host (CPU) memory. The block remains reusable until it is evicted from host memory. When an offloaded block is reused, it is first copied back into GPU memory. Offloading is controlled with property ```host_cache_size``` which specifies how much host memory (in bytes) should be allocated for offloading. The default is 0. -When offloading is enabled, the client can prevent specific blocks from being offloaded by toggling block priority. Blocks with lower priority than a certain threshold are not offloaded; they are evicted directly from GPU memory to reduce traffic between GPU and host. This priority is set with ```secondary_offload_min_priority```. Default value is 35, meaning any block with lower priority than 35 will not be offloaded. +When offloading is enabled, `KvCacheConfig.cold_tier_offload_min_priority` controls which unreferenced blocks can be offloaded. Blocks below the threshold are discarded from GPU memory; blocks at or above it can move to a cold tier. The default threshold is 30. The old name, `secondary_offload_min_priority`, remains accepted as a deprecated alias. Here is an [example](../../../examples/llm-api/llm_kv_cache_offloading.py) to show how to enable host offloading. diff --git a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py index 397e4a8aa434..8775457f82bb 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py +++ b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py @@ -106,7 +106,7 @@ class KvCacheConfigV2: disk_cache_path: Optional[str] = None onboard_blocks: bool = True cross_kv_cache_fraction: Optional[float] = None - secondary_offload_min_priority: Optional[int] = None + cold_tier_offload_min_priority: Optional[int] = None event_buffer_max_size: int = 0 kv_cache_event_hash_algo: str = "auto" max_gpu_total_bytes: Optional[int] = None diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index eafdc7a70de2..5a7141d7979c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3049,10 +3049,10 @@ def _build_base_config( ratio_tolerance=rebalance_cfg.ratio_tolerance, moving_average_decay=rebalance_cfg.moving_average_decay, ), - secondary_offload_min_priority=( + cold_tier_offload_min_priority=( 30 - if kv_cache_config.secondary_offload_min_priority is None - else kv_cache_config.secondary_offload_min_priority + if kv_cache_config.cold_tier_offload_min_priority is None + else kv_cache_config.cold_tier_offload_min_priority ), ) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 3eb291aac58a..1dad4d4336cb 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -757,6 +757,10 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], self.enable_indexer_k_cache = enable_indexer_k_cache + cold_tier_offload_min_priority = ( + kv_cache_config.secondary_offload_min_priority if isinstance( + kv_cache_config, tensorrt_llm.bindings.executor.KvCacheConfig) + else kv_cache_config.cold_tier_offload_min_priority) kwargs = { 'num_kv_heads_per_layer': self.num_kv_heads_per_layer, 'size_per_head': head_dim, @@ -772,8 +776,7 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], 'chunk_size': min(max_num_tokens, self.max_seq_len), 'enable_block_reuse': kv_cache_config.enable_block_reuse, 'cache_type': kv_cache_type, - 'secondary_offload_min_priority': - kv_cache_config.secondary_offload_min_priority, + 'secondary_offload_min_priority': cold_tier_offload_min_priority, 'enable_partial_reuse': kv_cache_config.enable_partial_reuse, 'copy_on_partial_reuse': kv_cache_config.copy_on_partial_reuse, 'kv_connector_manager': self.kv_connector_manager, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 4d3b3b48de4d..55a23ac1a87b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4876,11 +4876,19 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): description= "The fraction of the KV Cache memory should be reserved for cross attention. If set to p, self attention will use 1-p of KV Cache memory and cross attention will use p of KV Cache memory. Defaults to None (unset); must be set when using an encoder-decoder model and must not be set otherwise." ) - secondary_offload_min_priority: Optional[int] = Field( + cold_tier_offload_min_priority: Optional[int] = Field( default=None, description= - "Only blocks with priority > secondary_offload_min_priority can be offloaded to secondary memory." + "Minimum priority for offloading unreferenced blocks from the hot tier to cold tiers. " + "Blocks with priority at or above this value can be offloaded. Defaults to 30 when unset." ) + secondary_offload_min_priority: Optional[int] = Field( + default=None, + status="deprecated", + deprecated="Use cold_tier_offload_min_priority instead.", + telemetry=False, + exclude=True, + description="Deprecated alias for cold_tier_offload_min_priority.") event_buffer_max_size: int = Field( default=0, description= @@ -5095,7 +5103,7 @@ def _to_pybind(self): free_gpu_memory_fraction=self.free_gpu_memory_fraction, host_cache_size=self.host_cache_size, cross_kv_cache_fraction=self.cross_kv_cache_fraction, - secondary_offload_min_priority=self.secondary_offload_min_priority, + secondary_offload_min_priority=self.cold_tier_offload_min_priority, event_buffer_max_size=self.event_buffer_max_size, enable_partial_reuse=self.enable_partial_reuse, copy_on_partial_reuse=self.copy_on_partial_reuse, @@ -5105,6 +5113,30 @@ def _to_pybind(self): max_gpu_total_bytes=self.max_gpu_total_bytes) return config + @model_validator(mode='after') + def migrate_legacy_offload_min_priority(self) -> 'KvCacheConfig': + legacy = self.__dict__.get('secondary_offload_min_priority') + if legacy is None: + return self + if (self.cold_tier_offload_min_priority is not None + and self.cold_tier_offload_min_priority != legacy): + raise ValueError( + "cold_tier_offload_min_priority and secondary_offload_min_priority " + "must match when both are set.") + logger.warning( + "'kv_cache_config.secondary_offload_min_priority' is deprecated; " + "use 'kv_cache_config.cold_tier_offload_min_priority' instead.") + self.cold_tier_offload_min_priority = legacy + return self + + def __setattr__(self, name: str, value: Any) -> None: + super().__setattr__(name, value) + if name == 'secondary_offload_min_priority': + super().__setattr__('cold_tier_offload_min_priority', value) + elif (name == 'cold_tier_offload_min_priority' + and 'secondary_offload_min_priority' in self.model_fields_set): + self.__dict__['secondary_offload_min_priority'] = value + @field_validator('free_gpu_memory_fraction') @classmethod def validate_free_gpu_memory_fraction(cls, v: float): diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 722ec078d2f2..402f87be5eb1 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -133,7 +133,7 @@ class _KVCacheManagerConfigFieldSpec: text_only: bool = False enable_partial_commit: bool = True pool_rebalance: object = None - secondary_offload_min_priority: int = 30 + cold_tier_offload_min_priority: int = 30 KVCacheManagerConfig.__dataclass_fields__ = _KVCacheManagerConfigFieldSpec.__dataclass_fields__ diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 81517522e770..feebc7c93b2a 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -330,7 +330,7 @@ class KVCacheManagerConfig: text_only: bool = False enable_partial_commit: bool = True pool_rebalance: PoolRebalanceConfig = ... - secondary_offload_min_priority: int = 30 + cold_tier_offload_min_priority: int = 30 @property def enable_swa_scratch_reuse(self) -> bool: ... diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 4f60a9fa0d35..e58f42598f2d 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -621,6 +621,11 @@ "kind": "value", "path": "kv_cache_config.block_reuse_config.swa_endpoint_rewind_tokens" }, + { + "capture_policy": "int|none", + "kind": "value", + "path": "kv_cache_config.cold_tier_offload_min_priority" + }, { "capture_policy": "bool", "kind": "value", @@ -826,11 +831,6 @@ "kind": "value", "path": "kv_cache_config.pool_ratio" }, - { - "capture_policy": "int|none", - "kind": "value", - "path": "kv_cache_config.secondary_offload_min_priority" - }, { "capture_policy": "int|none", "kind": "value", diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 3355ebd7d408..da37a69e3f1b 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -2426,7 +2426,7 @@ def _build_v2_hybrid_with_mamba_layer( mamba_ssm_cache_dtype=torch.float16, kda_replay_num_spec=None, enable_kv_pool_rebalance=False, - secondary_offload_min_priority=None, + cold_tier_offload_min_priority=None, host_cache_size=None, ): """Construct a real MambaHybridCacheManagerV2.""" @@ -2456,7 +2456,7 @@ def _build_v2_hybrid_with_mamba_layer( ), dtype=kv_cache_dtype, enable_kv_pool_rebalance=enable_kv_pool_rebalance, - secondary_offload_min_priority=secondary_offload_min_priority, + cold_tier_offload_min_priority=cold_tier_offload_min_priority, host_cache_size=host_cache_size, ) return MambaHybridCacheManagerV2( @@ -2955,7 +2955,7 @@ def test_v2_hybrid_retains_configured_number_of_conversation_turns(): block_reuse_policy="per_conversation", max_num_turns=2, additional_snapshot_offsets_from_end=[0], - secondary_offload_min_priority=50, + cold_tier_offload_min_priority=50, host_cache_size=0, ) request_a = _make_v2_conversation_request( @@ -2989,7 +2989,7 @@ def test_v2_hybrid_retains_configured_number_of_conversation_turns(): attention_group = mgr.impl.get_layer_group_id(1) priorities = cache_a.get_page_priorities(attention_group) assert priorities and set(priorities) == {80} - assert mgr.impl.init_config.secondary_offload_min_priority == 50 + assert mgr.impl.init_config.cold_tier_offload_min_priority == 50 assert not mgr.can_evict assert mgr.prefetch_for_context_tokens([request_a]) is True request_a_state_index = mgr.get_state_indices([request_a.py_request_id], [False])[0] diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 29df16273a38..6eed8d0620aa 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -370,6 +370,7 @@ def test_context_model_capacity_profile_uses_real_dummy_requests( ) executor, _ = _make_remote_tail_executor_and_request() executor.is_warmup = is_warmup + executor._is_kv_manager_v2 = False if not has_transceiver: executor.kv_cache_transceiver = None executor.max_beam_width = 1 diff --git a/tests/unittest/_torch/executor/test_resource_manager.py b/tests/unittest/_torch/executor/test_resource_manager.py index 60103269ee9e..cceded8d5dec 100644 --- a/tests/unittest/_torch/executor/test_resource_manager.py +++ b/tests/unittest/_torch/executor/test_resource_manager.py @@ -1307,8 +1307,13 @@ class TestKVCacheManagerConfigForwarding(unittest.TestCase): def test_secondary_offload_min_priority_forwarded_to_cpp_manager(self): """Regression guard: preserve explicit priority 0 when building C++ kwargs.""" - kv_cache_config = KvCacheConfig(max_tokens=64, - secondary_offload_min_priority=0) + for config_type in (KvCacheConfig, tllm.KvCacheConfig): + with self.subTest(config_type=config_type): + self._check_offload_priority_forwarding(config_type) + + def _check_offload_priority_forwarding(self, config_type): + kv_cache_config = config_type(max_tokens=64, + secondary_offload_min_priority=0) mock_impl = MagicMock() mock_impl.get_block_pool_pointers.return_value = torch.empty(0) diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index 1aeb0bc55d25..c970628a81af 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -197,7 +197,7 @@ class KvCacheConfigV2: disk_cache_path: Optional[str] = None onboard_blocks: bool = True cross_kv_cache_fraction: Optional[float] = None - secondary_offload_min_priority: Optional[int] = None + cold_tier_offload_min_priority: Optional[int] = None event_buffer_max_size: int = 0 kv_cache_event_hash_algo: str = "auto" max_gpu_total_bytes: Optional[int] = None @@ -2063,7 +2063,7 @@ def _offload_cycle_manager( enable_partial_reuse=enable_block_reuse, copy_on_partial_reuse=False, host_cache_size=host_cache_size, - secondary_offload_min_priority=0, + cold_tier_offload_min_priority=0, ), CacheTypeCpp.SELF, num_layers=NUM_LAYERS, diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 721b1d062537..af59cf40fe54 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -93,7 +93,7 @@ class KvCacheConfigV2: disk_cache_size: Optional[int] = None disk_cache_path: Optional[str] = None cross_kv_cache_fraction: Optional[float] = None - secondary_offload_min_priority: Optional[int] = None + cold_tier_offload_min_priority: Optional[int] = None event_buffer_max_size: int = 0 kv_cache_event_hash_algo: str = "auto" diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index be76b50ef358..d831bf5af169 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -628,10 +628,10 @@ def test_offload_reload_and_error_cleanup(self): self.engine.execute([Step(copy, [], list(range(3)))], self.stream.handle) self.stream.synchronize() - @parameterized.expand([(29, 30), (30, 30), (49, 50), (50, 50)]) + @parameterized.expand([(29, 30), (30, 30), (49, 50), (50, 50), (0, -1), (100, 101)]) def test_released_offload_threshold_and_duration(self, priority, threshold): self.cfg = create_config(4, 2 << 20, 2 << 20, 0, 1, None, 0, 65536) - self.cfg.secondary_offload_min_priority = threshold + self.cfg.cold_tier_offload_min_priority = threshold self.engine = FakeEngine(self.cfg) self.manager = KVCacheManager(self.cfg) tokens = [1, 2, 3, 4] diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_optional_checkpoint.py b/tests/unittest/kv_cache_manager_v2_tests/test_optional_checkpoint.py index d0552c5c5f9a..7705f8fe097c 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_optional_checkpoint.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_optional_checkpoint.py @@ -2,9 +2,13 @@ # SPDX-License-Identifier: Apache-2.0 """Python bindings for native OPTIONAL selection; storage contracts live in C++.""" +import time +from datetime import timedelta + import pytest import torch +from tensorrt_llm.bindings.executor import KvCacheRetentionConfig from tensorrt_llm.runtime import kv_cache_manager_v2 as kv pytestmark = [ @@ -12,7 +16,8 @@ ] -def test_reuse_status_and_optional_selection_bindings(): +@pytest.mark.parametrize("priority", [None, 80]) +def test_reuse_status_and_optional_selection_bindings(priority): layers = [ kv.AttentionLayerConfig( layer_id=i, @@ -32,14 +37,34 @@ def test_reuse_status_and_optional_selection_bindings(): caches = [] tokens = list(range(8)) try: - source = manager.create_kv_cache(None, []) + source = manager.create_kv_cache( + None, + [], + expected_prompt_length=8, + kv_cache_retention_config=KvCacheRetentionConfig( + [ + KvCacheRetentionConfig.TokenRangeRetentionConfig( + 0, None, 60, timedelta(seconds=1) + ) + ] + ), + ) caches.append(source) assert source.resume(stream.cuda_stream) assert source.resize(8) source.commit(tokens) - reader = manager.create_kv_cache(None, tokens) - caches.append(reader) groups = [manager.get_layer_group_id(i) for i in range(3)] + source_indices = [source.get_base_page_indices(group).copy() for group in groups] + source.close() + retention = ( + KvCacheRetentionConfig( + [KvCacheRetentionConfig.TokenRangeRetentionConfig(0, None, priority)] + ) + if priority is not None + else None + ) + reader = manager.create_kv_cache(None, tokens, kv_cache_retention_config=retention) + caches.append(reader) global_group, encoder, peer = groups for layer, group in enumerate(groups): status = reader.reuse_status[group] @@ -52,11 +77,16 @@ def test_reuse_status_and_optional_selection_bindings(): reader.resume(stream.cuda_stream, optional_reuse_groups=invalid) assert not reader.is_active assert reader.resume(stream.cuda_stream, optional_reuse_groups=[encoder]) - assert reader.get_base_page_indices(encoder)[1] == source.get_base_page_indices(encoder)[1] - assert reader.get_base_page_indices(peer)[1] != source.get_base_page_indices(peer)[1] + assert reader.get_base_page_indices(encoder)[1] == source_indices[1][1] + assert reader.get_base_page_indices(peer)[1] != source_indices[2][1] + assert reader.get_page_priorities(encoder)[1] == (60 if priority is None else priority) + assert reader.get_page_priorities(peer)[1] == (35 if priority is None else priority) + time.sleep(1.05) + manager.refresh_retention() + assert reader.get_page_priorities(encoder)[1] == (60 if priority is None else priority) reader.suspend() assert reader.resume(stream.cuda_stream) - assert reader.get_base_page_indices(encoder)[1] == source.get_base_page_indices(encoder)[1] + assert reader.get_base_page_indices(encoder)[1] == source_indices[1][1] finally: for cache in reversed(caches): cache.close() diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 1fbc46e23b19..4c5ca7e8495e 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1208,7 +1208,10 @@ def test_registered_models_keep_v2_on_nixl(self, timeout_ms) -> None: @pytest.mark.cpu_only -def test_KvCacheConfig_declaration(): +@pytest.mark.parametrize( + "offload_field", + ["cold_tier_offload_min_priority", "secondary_offload_min_priority"]) +def test_KvCacheConfig_declaration(offload_field): assert KvCacheConfig().mamba_state_cache_interval is None assert KvCacheConfig().mamba_state_config.periodic_snapshot_interval == 0 assert KvCacheConfig().kv_cache_event_hash_algo == "auto" @@ -1232,7 +1235,7 @@ def test_KvCacheConfig_declaration(): disk_cache_size=2048, disk_cache_path="/tmp", cross_kv_cache_fraction=0.5, - secondary_offload_min_priority=1, + **{offload_field: 1}, event_buffer_max_size=0, kv_cache_event_hash_algo="v2_sha256_64", enable_swa_scratch_reuse=True, @@ -1261,6 +1264,12 @@ def test_KvCacheConfig_declaration(): assert KvCacheConfig().enable_swa_scratch_reuse is False assert pybind_config.cross_kv_cache_fraction == 0.5 assert pybind_config.secondary_offload_min_priority == 1 + assert config.cold_tier_offload_min_priority == 1 + dumped = config.model_dump(exclude_unset=True) + assert dumped["cold_tier_offload_min_priority"] == 1 + assert "secondary_offload_min_priority" not in dumped + assert KvCacheConfig.model_validate_json( + config.model_dump_json()).cold_tier_offload_min_priority == 1 assert pybind_config.event_buffer_max_size == 0 assert config.kv_cache_event_hash_algo == "v2_sha256_64" assert config.pool_ratio == [0.25, 0.75] @@ -1297,6 +1306,36 @@ def test_KvCacheConfig_declaration(): BlockReuseConfig(max_num_turns=0) +@pytest.mark.cpu_only +def test_KvCacheConfig_offload_priority_compatibility(): + assert KvCacheConfig().cold_tier_offload_min_priority is None + for field in ("cold_tier_offload_min_priority", + "secondary_offload_min_priority"): + for threshold in (-1, 101): + config = KvCacheConfig(**{field: threshold}) + assert config.cold_tier_offload_min_priority == threshold + assert config._to_pybind( + ).secondary_offload_min_priority == threshold + with pytest.raises(ValidationError, match="must match"): + KvCacheConfig(cold_tier_offload_min_priority=50, + secondary_offload_min_priority=40) + with patch.object(llm_args_mod.logger, "warning") as warning: + config = KvCacheConfig(cold_tier_offload_min_priority="50", + secondary_offload_min_priority=50) + assert config.cold_tier_offload_min_priority == 50 + warning.assert_called_once() + assert "deprecated" in warning.call_args.args[0] + with pytest.warns(DeprecationWarning, match="cold_tier"): + assert config.secondary_offload_min_priority == 50 + config = KvCacheConfig(secondary_offload_min_priority=0) + config.secondary_offload_min_priority = 50 + assert TorchLlmArgs(model=llama_model_path, kv_cache_config=config + ).kv_cache_config.cold_tier_offload_min_priority == 50 + config.cold_tier_offload_min_priority = 80 + assert KvCacheConfig.model_validate( + config).cold_tier_offload_min_priority == 80 + + @pytest.mark.cpu_only @pytest.mark.parametrize("rewind", [0, 1024, 4096]) def test_BlockReuseConfig_swa_endpoint_rewind_round_trip(rewind: int) -> None: From abc7d4a38ad7892ba3f482802639e2eeec7c5ff7 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sun, 11 Oct 2026 21:56:52 -0700 Subject: [PATCH 13/23] fix: address KV cache retention review feedback Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../tensorrt_llm/batch_manager/common.h | 2 + .../batch_manager/evictionPolicy.cpp | 5 +- .../kv_cache_manager_v2/AGENTS.md | 30 ++ .../kv_cache_manager_v2/config.h | 2 +- .../kv_cache_manager_v2/kvCache.cpp | 6 +- .../kv_cache_manager_v2/storageManager.h | 2 +- .../batch_manager/kvCacheManagerV2.cpp | 4 +- .../source/examples/kvcacheretentionconfig.md | 41 +-- docs/source/features/kv-cache-connector.md | 12 +- .../connectors/kv_cache_connector.py | 11 +- .../kv_cache/kv_cache_manager_v2.py | 3 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 10 +- .../_torch/pyexecutor/py_executor_creator.py | 7 +- .../runtime/kv_cache_manager_v2/__init__.py | 5 +- .../runtime/kv_cache_manager_v2/__init__.pyi | 3 +- .../defs/llmapi/test_llm_api_connector.py | 24 +- .../test_kv_connector_initialization.py | 101 +++++++ .../test_kv_cache_manager_v2.py | 280 +++++++++++++----- 18 files changed, 410 insertions(+), 138 deletions(-) create mode 100644 tests/unittest/_torch/executor/test_kv_connector_initialization.py diff --git a/cpp/include/tensorrt_llm/batch_manager/common.h b/cpp/include/tensorrt_llm/batch_manager/common.h index bd16d2b038ae..f256997c1322 100644 --- a/cpp/include/tensorrt_llm/batch_manager/common.h +++ b/cpp/include/tensorrt_llm/batch_manager/common.h @@ -38,6 +38,8 @@ using RequestIdType = std::uint64_t; using RequestVector = std::vector>; using ReqIdsSet = std::unordered_set; +inline constexpr int kDefaultColdTierOffloadMinPriority = 30; + class ScheduledRequests { public: diff --git a/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp b/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp index 42b5bd9e6879..d4518ffe91d0 100644 --- a/cpp/tensorrt_llm/batch_manager/evictionPolicy.cpp +++ b/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& mAllBlocksById, std::v startIdx += sizes[cacheLevel]; } - mSecondaryOffloadMinPriority = secondaryOffloadMinPriority.value_or(kDefaultSecondaryOffloadMinPriority); + mSecondaryOffloadMinPriority = secondaryOffloadMinPriority.value_or(kDefaultColdTierOffloadMinPriority); } void LRUEvictionPolicy::initializePlaceholders(std::vector& allPlaceholderBlocksById) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md index dd2ae9beade1..4208d1cf85b9 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md +++ b/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 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h index 0ea731080f58..adb134953938 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h @@ -451,7 +451,7 @@ struct KVCacheManagerConfig bool enablePartialCommit = true; // Unreferenced GPU pages below this priority are dropped instead of offloaded. - Priority coldTierOffloadMinPriority = 30; + Priority coldTierOffloadMinPriority = kDefaultColdTierOffloadMinPriority; bool enableSwaScratchReuse() const noexcept { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index 1b81d1a1804c..823948a74622 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -198,9 +198,9 @@ std::vector KvCache::getPagePriorities(LayerGroupId layerGroupId, Beam = std::holds_alternative(mManager->storage().getLifeCycle(LifeCycleId{layerGroupId.value()})); for (auto const& block : mBlocks) { - auto const* page = beamIndex < block.pages.size() - ? blockPageGetPage(block.pages[beamIndex].at(layerGroupId)).get() - : nullptr; + // 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). diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h index 81c24b91f9aa..2b423590aeb7 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h @@ -130,7 +130,7 @@ class StorageManager : public std::enable_shared_from_this std::optional const& typicalBatch = std::nullopt, std::vector const& constraints = {}, std::optional> const& initialPoolRatio = std::nullopt, std::shared_ptr eventSink = nullptr, float maxUtilForResume = 1.0f, - Priority coldTierOffloadMinPriority = 30); + Priority coldTierOffloadMinPriority = kDefaultColdTierOffloadMinPriority); ~StorageManager(); StorageManager(StorageManager const&) = delete; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index fe1416ac96aa..c5e06a4f1c36 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -825,6 +825,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) // Export the C++ debug mode as an immutable Python bool snapshot. m.attr("NDEBUG") = nb::bool_(!kv::gDebug); + m.attr("DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY") + = nb::int_(tensorrt_llm::batch_manager::kDefaultColdTierOffloadMinPriority); // ---- Exceptions -------------------------------------------------------- static nb::object sOutOfMemoryError = nb::exception(m, "OutOfMemoryError"); @@ -1786,7 +1788,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::arg("swa_scratch_reuse").none() = std::nullopt, nb::arg("commit_min_snapshot") = false, nb::arg("enable_stats") = true, nb::arg("text_only") = false, nb::arg("enable_partial_commit") = true, nb::arg("pool_rebalance") = kv::PoolRebalanceConfig{}, - nb::arg("cold_tier_offload_min_priority") = 30) + nb::arg("cold_tier_offload_min_priority") = tensorrt_llm::batch_manager::kDefaultColdTierOffloadMinPriority) .def_rw("tokens_per_block", &kv::KVCacheManagerConfig::tokensPerBlock) .def_rw("cache_tiers", &kv::KVCacheManagerConfig::cacheTiers) .def_rw("layers", &kv::KVCacheManagerConfig::layers) diff --git a/docs/source/examples/kvcacheretentionconfig.md b/docs/source/examples/kvcacheretentionconfig.md index 627baad08110..fb885a0dd08b 100644 --- a/docs/source/examples/kvcacheretentionconfig.md +++ b/docs/source/examples/kvcacheretentionconfig.md @@ -67,34 +67,13 @@ if __name__ == '__main__': This example uses a single ```kv_cache_retention_config``` object for all the prompts. You can also provide a list that must have the same length as the list of prompts. -KV cache manager V2 applies this configuration to native pages, including reused -prefix pages. A token range selects blocks whose first token falls within the -range. The final partial prompt block uses the prompt policy; subsequent blocks -use the decode policy. - -When requests share a page, the latest claim with an explicit priority sets its -priority. Releasing an older request does not restore that request's priority. -A claim without an explicit priority preserves the existing priority, but replaces -the duration, including clearing it when no duration is supplied. The duration -starts when a page has no request references. An independently published -snapshot can reach this state before its originating request finishes. Prefetch -preserves the existing expiry. Expired pages return to priority 35; pages still -referenced by a request cannot be permanently evicted. - -On GPU eviction, unreferenced pages below `KvCacheConfig.cold_tier_offload_min_priority` -(default 30) are discarded. Pages at or above it can be offloaded. Offloading -such a page preserves its priority and clears its duration, as in V1; a later -claim can supply a new duration. Referenced pages may move between tiers but -are never discarded by this threshold. - -`secondary_offload_min_priority` is deprecated. Existing configurations using it -are accepted with a warning and mapped to `cold_tier_offload_min_priority` for -both managers. Setting both fields to different values is an error. Serialized -configurations use the new name. The threshold is manager-wide, not part of -the per-request `KvCacheRetentionConfig`. - -V2 supports the default `transfer_mode=DRAM`. Its host and disk offload tiers are -configured at manager creation. The per-request `GDS` and `POSIX_DEBUG_FALLBACK` -transfer modes, whose `directory` selects V1 block files, are rejected by V2. -They are not aliases for V2's manager-wide disk tier. With `DRAM`, `directory` -does not select storage. +## KV Cache Manager V2 + +KV cache manager V2 accepts the same `KvCacheRetentionConfig`, so the examples +above work unchanged. For `transfer_mode`, V2 supports only the default `DRAM` and rejects +requests that set `GDS` or `POSIX_DEBUG_FALLBACK`. V2's host and disk offload +tiers are configured manager-wide when the manager is created and are +independent of these per-request transfer modes. + +The priority threshold that decides which evicted blocks are offloaded is also +manager-wide; see [Enable Offloading to Host Memory](../features/kvcache.md#enable-offloading-to-host-memory). diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 2d93e25c9194..aa115b1b73a5 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -76,13 +76,18 @@ These methods run on the leader process and drive the connector's behavior. | Set | Models it covers | |---|---| - | per-layer-group | every model, VSWA included | + | per-layer-group | single-layer-group models; multi-group configurations are currently rejected | | flat | non-VSWA, non-hybrid only | - A set is complete when both of its methods are defined. Mixing the two — one method from each — is rejected during executor bring-up, naming the method that is missing. An existing flat connector needs no change on the models it already covers: with a single layer group the base per-layer-group implementation folds back to the flat call. Hybrid / linear-attention models are not enabled for the connector yet; the per-layer-group set is the shape their cache will need. See [Running under VSWA](#running-under-vswa). + A set is complete when both of its methods are defined. An existing flat connector needs no change on the models it already covers: with a single layer group the base per-layer-group implementation folds back to the flat call. Hybrid / linear-attention models are not enabled for the connector yet; the per-layer-group set is the shape their cache will need. See [Running under VSWA](#running-under-vswa). ##### Running under VSWA +Multi-layer-group configurations currently fail during connector initialization because retention +priorities are represented by a single list. This applies to both KV cache managers, even when the +connector implements the grouped APIs below. Single-window SWA remains allowed. The grouped +interfaces and example describe the addressing needed for future VSWA connector support. + Under variable sliding-window attention the KV cache allocates one pool per attention window size, and a page index only means something inside its own layer group. A single tensor and a single flat block list cannot describe that, so three methods have to be implemented together: | Method | Replaces | @@ -93,7 +98,8 @@ Under variable sliding-window attention the KV cache allocates one pool per atte Implementing the per-layer-group form of a pair is enough — the flat method it replaces does not also have to be defined. -All three are checked during executor bring-up, before any request is admitted. `register_kv_cache_layout` refuses there, naming the group and region counts it could not describe, and the two scheduler methods are checked alongside it. Nothing is deferred to the first request, so a partial implementation costs a start-up failure rather than one after the model is loaded. +The current single-group restriction is checked during executor bring-up, before worker registration +or request admission. Implementing all three methods does not bypass that restriction. `examples/llm-api/llm_kv_cache_connector_vswa.py` is a worked connector for this case. A VSWA connector has to do five things: diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 48d555679f57..726418620f43 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -623,13 +623,10 @@ def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManag priorities = None if req.kv_cache_retention_config is not None: if is_v2: - priorities = ( - kv_cache_manager.kv_cache_map[req.py_request_id].get_page_priorities(0)[ - len(self.block_ids) - len(new_block_ids) : len(self.block_ids) - ] - if len(new_block_ids_by_layer_group) == 1 - else [] - ) + kv_cache = kv_cache_manager.kv_cache_map[req.py_request_id] + priorities = kv_cache.get_page_priorities(0)[ + len(self.block_ids) - len(new_block_ids) : len(self.block_ids) + ] else: priorities = [ kv_cache_manager.get_priority_by_block_id(block_id) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 5a7141d7979c..7851bae82da5 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -64,6 +64,7 @@ BAD_PAGE_INDEX, CACHE_LEVEL1, DEFAULT_BEAM_INDEX, + DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY, GPU_LEVEL, AttentionLayerConfig, AttnLifeCycle, @@ -3050,7 +3051,7 @@ def _build_base_config( moving_average_decay=rebalance_cfg.moving_average_decay, ), cold_tier_offload_min_priority=( - 30 + DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY if kv_cache_config.cold_tier_offload_min_priority is None else kv_cache_config.cold_tier_offload_min_priority ), diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4ffb3e1275d4..1792e063fce8 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1301,10 +1301,12 @@ def _maybe_init_kv_connector_manager(self): self._reject_connector_prefix_without_block_reuse( self.kv_cache_manager) layout = build_kv_cache_layout_v2(self.kv_cache_manager) - self.kv_connector_manager.reject_flat_only_scheduler( - len(layout.groups)) - window = (layout.groups[0].window_size - if len(layout.groups) == 1 else None) + if len(layout.groups) != 1: + raise NotImplementedError( + "KV Cache Connector requires a single layer group. " + "Retention priorities for multiple layer groups " + "are not supported.") + window = layout.groups[0].window_size self.kv_connector_manager.warn_flat_scheduler_under_swa(window) self.kv_connector_manager.worker.register_kv_cache_layout( layout) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 7a95ac8c2682..ec205fdd4537 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -1034,15 +1034,14 @@ def allocation_scope(current_stage: ExecutorMemoryType): "KV connector is not supported with speculative decoding. " "Disable speculative decoding to run a connector.") - # VSWA allocates one pool per window size, which only - # `register_kv_cache_layout` can describe. + # V1 registers a single pool. V2 checks its resolved layer groups + # during executor initialization. max_attention_window = kv_cache_config.max_attention_window if (uses_vswa_kv_cache_layout(max_attention_window) and v2_selection is False): raise NotImplementedError( "KV connector is not supported with VSWA (Variable Sliding Window Attention) " - "in this configuration. Set kv_cache_config.use_kv_cache_manager_v2=True." - ) + "in this configuration. Use a single attention window.") if mapping.enable_attention_dp: raise NotImplementedError( diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 402f87be5eb1..3eb903af9973 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -52,6 +52,8 @@ def _load_cpp_module(): _cpp = _load_cpp_module() +DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY = _cpp.DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY + AggregatedPageDesc = _cpp.AggregatedPageDesc AttentionLayerConfig = _cpp.AttentionLayerConfig AttentionReusePolicy = _cpp.AttentionReusePolicy @@ -133,7 +135,7 @@ class _KVCacheManagerConfigFieldSpec: text_only: bool = False enable_partial_commit: bool = True pool_rebalance: object = None - cold_tier_offload_min_priority: int = 30 + cold_tier_offload_min_priority: int = DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY KVCacheManagerConfig.__dataclass_fields__ = _KVCacheManagerConfigFieldSpec.__dataclass_fields__ @@ -227,6 +229,7 @@ def typed_range(*args: int) -> range: "CacheTierConfig", "CudaStream", "DEFAULT_BEAM_INDEX", + "DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY", "DataRole", "DiskCacheTierConfig", "ExpandedBuffer", diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index feebc7c93b2a..6d4ed5f3f88b 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -37,6 +37,7 @@ from tensorrt_llm.bindings.executor import KvCacheRetentionConfig # From _common.py NDEBUG: Final[int] DEFAULT_BEAM_INDEX: Final[BeamIndex] +DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY: Final[int] BAD_PAGE_INDEX: Final[int] GPU_LEVEL: Final[CacheLevel] CACHE_LEVEL1: Final[CacheLevel] @@ -330,7 +331,7 @@ class KVCacheManagerConfig: text_only: bool = False enable_partial_commit: bool = True pool_rebalance: PoolRebalanceConfig = ... - cold_tier_offload_min_priority: int = 30 + cold_tier_offload_min_priority: int = DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY @property def enable_swa_scratch_reuse(self) -> bool: ... diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index dac65f7001ef..07193c05910e 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -722,6 +722,13 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, SWA_MAX_SEQ_LEN = 512 SWA_NUM_INPUT_TOKENS = 256 +MULTI_GROUP_CONNECTOR_UNSUPPORTED = pytest.mark.xfail( + raises=NotImplementedError, + strict=True, + reason= + "Connector initialization rejects multiple layer groups until grouped retention priorities are supported", +) + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_kv_cache_manager_v2", [True], @@ -1029,7 +1036,10 @@ def test_connector_sliding_window_prefix_is_backed_by_real_pages( @pytest.mark.threadleak(enabled=False) -@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], +@pytest.mark.parametrize("use_kv_cache_manager_v2", [ + False, + pytest.param(True, marks=MULTI_GROUP_CONNECTOR_UNSUPPORTED), +], ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], indirect=True) def test_connector_vswa_reports_page_indices_per_layer_group( @@ -1506,6 +1516,14 @@ def test_connector_max_utilization_is_rejected_on_v1_only( "beam", id="beam_search", ), + pytest.param( + dict(kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + max_attention_window=[SWA_WINDOW, SWA_MAX_SEQ_LEN])), + "VSWA", + "single layer group", + id="multiple_layer_groups", + ), pytest.param( dict(enable_attention_dp=True), "attention data parallelism", @@ -2129,6 +2147,7 @@ def test_connector_cancel_holds_pages_until_the_save_reads_them( @pytest.mark.threadleak(enabled=False) +@MULTI_GROUP_CONNECTOR_UNSUPPORTED def test_connector_multi_pool_e2e_persistent_cache(enforce_single_worker, monkeypatch): """The multi-pool data path, end to end, with a connector that moves real bytes. @@ -2330,6 +2349,7 @@ def build(): @pytest.mark.threadleak(enabled=False) +@MULTI_GROUP_CONNECTOR_UNSUPPORTED def test_connector_vswa_e2e_gemma3(enforce_single_worker, monkeypatch): """A KV connector on a model whose architecture is variable-window. @@ -2635,6 +2655,7 @@ def build(): @pytest.mark.threadleak(enabled=False) +@MULTI_GROUP_CONNECTOR_UNSUPPORTED def test_connector_vswa_out_of_window_blocks_reach_the_connector( enforce_single_worker, monkeypatch): """The sliding window passing blocks, delivered per layer group. @@ -2866,6 +2887,7 @@ def recording_register(self, layout): get_sm_version() not in (100, 103), reason="gpt-oss-20b is MXFP4 and needs SM100/SM103")), ]) +@MULTI_GROUP_CONNECTOR_UNSUPPORTED def test_connector_transfers_only_in_window_blocks_to_the_sliding_group( enforce_single_worker, monkeypatch, model_rel, assert_round_trip): """The connector transfers each layer group only what that group can read. diff --git a/tests/unittest/_torch/executor/test_kv_connector_initialization.py b/tests/unittest/_torch/executor/test_kv_connector_initialization.py new file mode 100644 index 000000000000..4fac32b3dc80 --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_connector_initialization.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( + KvCacheConnectorSchedulerOutputRequest, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm.bindings import LlmRequestState +from tensorrt_llm.bindings.executor import KvCacheRetentionConfig +from tensorrt_llm.runtime.kv_cache_manager_v2 import CacheTier + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize("windows", [(None,), (64,), (64, None), (32, 64, None)]) +def test_v2_connector_requires_one_layer_group( + monkeypatch: pytest.MonkeyPatch, windows: tuple[int | None, ...] +) -> None: + executor = object.__new__(PyExecutor) + executor.kv_connector_manager = Mock(capacity_only=False) + executor.kv_cache_transceiver = None + executor.dist = SimpleNamespace(pp_size=1, cp_size=1) + executor.max_beam_width = 1 + executor.max_draft_len = 0 + executor.enable_attention_dp = False + executor.llm_args = SimpleNamespace(kv_cache_config=None, scheduler_config=None) + executor.kv_cache_manager = Mock(spec=KVCacheManagerV2, enable_block_reuse=True) + executor.kv_cache_manager.impl = SimpleNamespace( + init_config=SimpleNamespace(cache_tiers=[SimpleNamespace(tier=CacheTier.GPU_MEM)]) + ) + executor.model_engine = Mock() + executor.model_engine.model.named_modules.return_value = [] + layout = SimpleNamespace(groups=[SimpleNamespace(window_size=window) for window in windows]) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.build_kv_cache_layout_v2", + lambda manager: layout, + ) + + connector = executor.kv_connector_manager + if len(windows) > 1: + with pytest.raises(NotImplementedError, match="requires a single layer group"): + executor._maybe_init_kv_connector_manager() + connector.worker.register_kv_cache_layout.assert_not_called() + connector.wait_for_initialization.assert_not_called() + else: + executor._maybe_init_kv_connector_manager() + connector.worker.register_kv_cache_layout.assert_called_once_with(layout) + connector.warn_flat_scheduler_under_swa.assert_called_once_with(windows[0]) + connector.wait_for_initialization.assert_called_once() + + +def test_single_group_priorities_follow_new_page_ordinals() -> None: + manager = Mock(spec=KVCacheManagerV2) + cache = Mock() + manager.kv_cache_map = {7: cache} + manager.get_page_indices_by_layer_group.return_value = [[12, 4]] + cache.get_page_priorities.return_value = [80, 10] + req = Mock( + request_id=7, + py_request_id=7, + state=LlmRequestState.CONTEXT_INIT, + context_current_position=0, + context_remaining_length=64, + context_chunk_size=64, + kv_cache_retention_config=KvCacheRetentionConfig([]), + cache_salt=None, + ) + req.get_tokens.return_value = list(range(64)) + tracker = KvCacheConnectorSchedulerOutputRequest() + + first = tracker.update_and_build_data(req, manager) + assert first.new_block_ids == [12, 4] + assert first.priorities == [80, 10] + + manager.get_page_indices_by_layer_group.return_value = [[12, 4, 9]] + cache.get_page_priorities.return_value = [80, 10, 20] + second = tracker.update_and_build_data(req, manager) + assert second.new_block_ids == [9] + assert second.priorities == [20] + + unchanged = tracker.update_and_build_data(req, manager) + assert unchanged.new_block_ids == [] + assert unchanged.priorities == [] diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index d831bf5af169..eac26247828c 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -37,6 +37,7 @@ from kv_cache_manager_v2 import ( BAD_PAGE_INDEX, DEFAULT_BEAM_INDEX, + DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY, GPU_LEVEL, AttentionLayerConfig, Batch, @@ -95,6 +96,7 @@ from tensorrt_llm.runtime.kv_cache_manager_v2 import ( BAD_PAGE_INDEX, DEFAULT_BEAM_INDEX, + DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY, GPU_LEVEL, AttentionLayerConfig, Batch, @@ -419,22 +421,53 @@ def prepare( class TestRetention(TestKVCacheManagerV2): def setUp(self) -> None: super().setUp() - self.caches = [] + self.caches: list[_KVCache] = [] self.stream = CachedCudaStream() def tearDown(self) -> None: + self._close_caches() + super().tearDown() + + def _close_caches(self) -> None: for cache in reversed(self.caches): cache.close() self.caches.clear() - super().tearDown() - def _cache(self, tokens, config=None, prompt_length=None): + def _prepare_retention( + self, + *, + host_quota: int = 0, + window_size: SlidingWindowSize = None, + enable_partial_commit: bool = True, + offload_threshold: int = DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY, + ) -> None: + self.cfg = create_config(4, 2 << 20, host_quota, 0, 1, window_size, 0, 65536) + self.cfg.enable_partial_commit = enable_partial_commit + self.cfg.cold_tier_offload_min_priority = offload_threshold + self.engine = FakeEngine(self.cfg) + self.manager = KVCacheManager(self.cfg) + + def _new_cache( + self, + tokens: list[TokenIdExt], + config: KvCacheRetentionConfig | None = None, + prompt_length: int | None = None, + ) -> _KVCache: cache = self.manager.create_kv_cache( input_tokens=tokens, expected_prompt_length=len(tokens) if prompt_length is None else prompt_length, kv_cache_retention_config=config, ) self.caches.append(cache) + return cache + + def _cache( + self, + tokens: list[TokenIdExt], + config: KvCacheRetentionConfig | None = None, + prompt_length: int | None = None, + ) -> _KVCache: + cache = self._new_cache(tokens, config, prompt_length) self.assertTrue(cache.resume(self.stream.handle)) self.assertTrue(cache.resize(len(tokens))) remaining = tokens[cache.num_committed_tokens :] @@ -446,24 +479,113 @@ def _cache(self, tokens, config=None, prompt_length=None): cache.commit(remaining) return cache + def _apply_memory_pressure(self) -> _KVCache: + free = self.manager.get_storage_statistics(GPU_LEVEL)[0].free + pressure = self._cache([]) + self.assertTrue(pressure.resize((free + 1) * self.cfg.tokens_per_block)) + return pressure + @staticmethod - def _config(priority, duration=None): + def _config( + priority: int | None, duration: timedelta | None = None + ) -> KvCacheRetentionConfig | None: if priority is None: return None return KvCacheRetentionConfig( [KvCacheRetentionConfig.TokenRangeRetentionConfig(0, None, priority, duration)] ) + @parameterized.expand([(4, False), (4, True), (5, False), (5, True)]) + def test_beam_priorities_share_prompt_and_keep_private_tail( + self, prompt_length: int, commit: bool + ) -> None: + self._prepare_retention(enable_partial_commit=False) + config = KvCacheRetentionConfig( + [KvCacheRetentionConfig.TokenRangeRetentionConfig(0, None, 80)], + decode_retention_priority=10, + ) + tokens = list(range(prompt_length)) + cache = self._new_cache(tokens, config) + self.assertTrue(cache.resume(self.stream.handle)) + self.assertTrue(cache.resize(12)) + self.engine.execute([Step(cache, tokens, [])], self.stream.handle) + if commit: + cache.commit(tokens) + cache.stop_committing() + cache.beam_width = 3 + + prompt_blocks = div_up(prompt_length, 4) + shared_blocks = prompt_length // 4 + for capacity in (12, 16): + self.assertTrue(cache.resize(capacity)) + indices = [cache.get_aggregated_page_indices(0, beam) for beam in range(3)] + for beam in range(3): + self.assertEqual( + cache.get_page_priorities(0, beam), + [80] * prompt_blocks + [10] * (capacity // 4 - prompt_blocks), + ) + if beam: + self.assertEqual( + indices[beam][:shared_blocks], [BAD_PAGE_INDEX] * shared_blocks + ) + # The partial prompt tail and both preallocated and newly allocated decode pages + # are writable independently by every beam. + for ordinal in range(shared_blocks, capacity // 4): + pages = {row[ordinal] for row in indices} + self.assertNotIn(BAD_PAGE_INDEX, pages) + self.assertEqual(len(pages), 3) + for invalid_beam in (-1, 3): + with self.assertRaises(IndexError): + cache.get_page_priorities(0, invalid_beam) + self.stream.synchronize() + + def test_beam_priorities_read_each_private_page(self) -> None: + self._prepare_retention(enable_partial_commit=False) + priority = 60 + tokens = list(range(5)) + cache = self.manager.create_kv_cache( + input_tokens=tokens, + expected_prompt_length=len(tokens), + custom_priority_callback=lambda ordinal, lifecycle: priority, + ) + self.caches.append(cache) + self.assertTrue(cache.resume(self.stream.handle)) + self.assertTrue(cache.resize(12)) + self.engine.execute([Step(cache, tokens, [])], self.stream.handle) + cache.stop_committing() + # Give each new beam's private pages a distinct allocation-time priority. + for beam_width, priority in ((2, 80), (3, 90)): + cache.beam_width = beam_width + self.assertEqual(cache.get_page_priorities(0, 0), [60, 60, 60]) + self.assertEqual(cache.get_page_priorities(0, 1), [60, 80, 80]) + self.assertEqual(cache.get_page_priorities(0, 2), [60, 90, 90]) + self.stream.synchronize() + + def test_offload_default_and_dataclass_replace(self) -> None: + config = create_config(4, 2 << 20, 0, 0, 1, None, 0, 65536) + self.assertEqual( + config.cold_tier_offload_min_priority, DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY + ) + self.assertEqual( + replace(config).cold_tier_offload_min_priority, DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY + ) + for threshold in (0, 80): + self.assertEqual( + replace( + config, cold_tier_offload_min_priority=threshold + ).cold_tier_offload_min_priority, + threshold, + ) + @parameterized.expand( [(80, 10, False), (80, 10, True), (10, 80, False), (10, 80, True), (80, None, False)] ) - def test_shared_claim_order_controls_eviction(self, first, second, close_newer_first): - self.prepare(2 << 20, 0, 0, 1, None, 0, tokens_per_block=4, kv_buf_size=65536) + def test_shared_claim_order_controls_eviction( + self, first: int, second: int | None, close_newer_first: bool + ) -> None: + self._prepare_retention() prompt = [1, 2, 3, 4] - late = self.manager.create_kv_cache( - input_tokens=prompt, kv_cache_retention_config=self._config(100) - ) - self.caches.append(late) + late = self._new_cache(prompt, self._config(100)) self.assertTrue(late.resume(self.stream.handle)) self.assertTrue(late.resize(len(prompt))) self.engine.execute([Step(late, prompt, [])], self.stream.handle) @@ -482,25 +604,21 @@ def test_shared_claim_order_controls_eviction(self, first, second, close_newer_f competitor = self._cache([5, 6, 7, 8], self._config(50)) competitor.close() # Referenced pages cannot be permanently evicted, irrespective of priority. - free = self.manager.get_storage_statistics()[0].free - pin = self._cache([]) - self.assertTrue(pin.resize((free + 1) * 4)) + pin = self._apply_memory_pressure() self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 4) self.assertEqual(self.manager.probe_reuse(input_tokens=[5, 6, 7, 8]), 0) pin.close() requests[1].close() competitor = self._cache([5, 6, 7, 8], self._config(50)) competitor.close() - free = self.manager.get_storage_statistics()[0].free - pressure = self._cache([]) - self.assertTrue(pressure.resize((free + 1) * 4)) + self._apply_memory_pressure() self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 4 if priority > 50 else 0) self.assertEqual( self.manager.probe_reuse(input_tokens=[5, 6, 7, 8]), 0 if priority > 50 else 4 ) - def test_prompt_decode_duration_and_unconfigured_claim(self): - self.prepare(2 << 20, 0, 0, 1, None, 0, tokens_per_block=4, kv_buf_size=65536) + def test_prompt_and_decode_priorities_expire_after_last_release(self) -> None: + self._prepare_retention() duration = timedelta(milliseconds=20) config = KvCacheRetentionConfig( [KvCacheRetentionConfig.TokenRangeRetentionConfig(1, None, 80, duration)], @@ -515,12 +633,6 @@ def test_prompt_decode_duration_and_unconfigured_claim(self): first.commit(tokens[5:]) second = self._cache(tokens, config, prompt_length=5) self.assertEqual(first.get_page_priorities(0), [35, 80, 10]) - with self.assertRaises(IndexError): - first.get_page_priorities(0, -1) - with self.assertRaises(IndexError): - self.manager.prefetch_reuse(ReuseScope(), tokens, CacheLevel(-1)) - with self.assertRaises(IndexError): - self.manager.prefetch_reuse(ReuseScope(), tokens, CacheLevel(1)) first.close() time.sleep(0.04) self.manager.refresh_retention() @@ -531,7 +643,10 @@ def test_prompt_decode_duration_and_unconfigured_claim(self): self.manager.refresh_retention() reclaimed = self._cache(tokens) self.assertEqual(reclaimed.get_page_priorities(0), [35, 35, 35]) - reclaimed.close() + + def test_unconfigured_claim_preserves_priority_and_clears_expiry(self) -> None: + self._prepare_retention() + tokens = list(range(12)) long_duration = timedelta(seconds=1) config = KvCacheRetentionConfig( [KvCacheRetentionConfig.TokenRangeRetentionConfig(1, None, 80, long_duration)], @@ -546,34 +661,35 @@ def test_prompt_decode_duration_and_unconfigured_claim(self): self.manager.refresh_retention() self.assertEqual(self._cache(tokens).get_page_priorities(0), [35, 80, 10]) - for shared in (False, True): - with self.subTest(shared_swa_page=shared): - for cache in self.caches: - cache.close() - self.caches.clear() - self.manager.shutdown() - self.prepare(2 << 20, 0, 0, 1, 4, 0, tokens_per_block=4, kv_buf_size=65536) - config = self._config(80, duration) - source = self._cache(tokens[:4], config) - sharer = self._cache(tokens[:4], config) if shared else None - self.assertTrue(source.resize(8)) - self.engine.execute([Step(source, tokens[4:8], tokens[:4])], self.stream.handle) - source.commit(tokens[4:8]) - source.close() - if sharer is not None: - sharer.close() - time.sleep(0.04) - self.manager.refresh_retention() - self.assertEqual(self._cache(tokens[:4]).get_page_priorities(0), [80]) - - def test_offload_reload_and_error_cleanup(self): - self.prepare(2 << 20, 2 << 20, 0, 1, None, 0, tokens_per_block=4, kv_buf_size=65536) + @parameterized.expand([(False,), (True,)]) + def test_sliding_window_release_preserves_priority_without_expiry(self, shared: bool) -> None: + self._prepare_retention(window_size=4) + tokens = list(range(8)) + config = self._config(80, timedelta(milliseconds=20)) + source = self._cache(tokens[:4], config) + sharer = self._cache(tokens[:4], config) if shared else None + self.assertTrue(source.resize(8)) + self.engine.execute([Step(source, tokens[4:], tokens[:4])], self.stream.handle) + source.commit(tokens[4:]) + source.close() + if sharer is not None: + sharer.close() + time.sleep(0.04) + self.manager.refresh_retention() + self.assertEqual(self._cache(tokens[:4]).get_page_priorities(0), [80]) + + def test_prefetch_rejects_invalid_cache_levels(self) -> None: + self._prepare_retention() + for level in (-1, 1): + with self.subTest(level=level), self.assertRaises(IndexError): + self.manager.prefetch_reuse(ReuseScope(), [1, 2, 3, 4], CacheLevel(level)) + + def test_offload_reload_preserves_priority_and_partial_reuse(self) -> None: + self._prepare_retention(host_quota=2 << 20) tokens = [1, 2, 3, 4] cached = self._cache(tokens, self._config(80)) cached.suspend() - free = self.manager.get_storage_statistics()[0].free - pressure = self._cache([]) - self.assertTrue(pressure.resize((free + 1) * 4)) + pressure = self._apply_memory_pressure() self.assertEqual(_introspection.active_page_stats(cached)[0], [0, 1]) pressure.close() self.assertTrue(cached.resume(self.stream.handle)) @@ -589,6 +705,17 @@ def test_offload_reload_and_error_cleanup(self): self.assertNotEqual( partial.get_aggregated_page_indices(0), cached.get_aggregated_page_indices(0) ) + self._close_caches() + self.manager.clear_reusable_blocks() + for level in (GPU_LEVEL, CacheLevel(1)): + for stats in self.manager.get_storage_statistics(level): + self.assertEqual(stats.free, stats.total) + self.assertEqual(stats.evictable, 0) + + def test_failed_resize_preserves_capacity_priority_and_data(self) -> None: + self._prepare_retention(host_quota=2 << 20) + tokens = [1, 2, 3, 4] + cached = self._cache(tokens, self._config(80)) capacity = cached.capacity total = self.manager.get_storage_statistics()[0].total self.assertFalse(cached.resize((total + 1) * 4)) @@ -596,32 +723,26 @@ def test_offload_reload_and_error_cleanup(self): self.assertEqual(cached.get_page_priorities(0), [80]) self.engine.execute([Step(cached, [], tokens)], self.stream.handle) self.stream.synchronize() - cached.close() + + def test_unsupported_transfer_modes_do_not_claim_reusable_pages(self) -> None: + self._prepare_retention(host_quota=2 << 20) + tokens = [1, 2, 3, 4] + self._cache(tokens, self._config(80)).close() # Unsupported file-transfer semantics fail before claiming any reusable pages. for mode in (KvCacheTransferMode.GDS, KvCacheTransferMode.POSIX_DEBUG_FALLBACK): config = KvCacheRetentionConfig([], transfer_mode=mode, directory="/tmp") - with self.assertRaisesRegex(ValueError, "per-request GDS/POSIX"): + with ( + self.subTest(mode=mode), + self.assertRaisesRegex(ValueError, "per-request GDS/POSIX"), + ): self.manager.create_kv_cache(input_tokens=tokens, kv_cache_retention_config=config) self.assertEqual(self._cache(tokens).get_page_priorities(0), [80]) - for cache in self.caches: - cache.close() - self.manager.clear_reusable_blocks() - for level in (GPU_LEVEL, CacheLevel(1)): - for stats in self.manager.get_storage_statistics(level): - self.assertEqual(stats.free, stats.total) - self.assertEqual(stats.evictable, 0) + def test_deferred_partial_copies_use_latest_claim_priority(self) -> None: + self._prepare_retention(host_quota=2 << 20) donor = self._cache(list(range(8)), self._config(80)) donor.close() - pending = [] - for priority in (10, 90): - copy = self.manager.create_kv_cache( - input_tokens=list(range(3)), - expected_prompt_length=3, - kv_cache_retention_config=self._config(priority), - ) - self.caches.append(copy) - pending.append(copy) + pending = [self._new_cache(list(range(3)), self._config(priority)) for priority in (10, 90)] for copy in pending: self.assertTrue(copy.resume(self.stream.handle)) self.assertEqual(copy.get_page_priorities(0), [90]) @@ -629,27 +750,32 @@ def test_offload_reload_and_error_cleanup(self): self.stream.synchronize() @parameterized.expand([(29, 30), (30, 30), (49, 50), (50, 50), (0, -1), (100, 101)]) - def test_released_offload_threshold_and_duration(self, priority, threshold): - self.cfg = create_config(4, 2 << 20, 2 << 20, 0, 1, None, 0, 65536) - self.cfg.cold_tier_offload_min_priority = threshold - self.engine = FakeEngine(self.cfg) - self.manager = KVCacheManager(self.cfg) + def test_released_offload_threshold(self, priority: int, threshold: int) -> None: + self._prepare_retention(host_quota=2 << 20, offload_threshold=threshold) tokens = [1, 2, 3, 4] cached = self._cache(tokens, self._config(priority, timedelta(seconds=2))) cached.close() - free = self.manager.get_storage_statistics()[0].free - pressure = self._cache([]) - self.assertTrue(pressure.resize((free + 1) * 4)) + pressure = self._apply_memory_pressure() retained = priority >= threshold self.assertEqual(self.manager.probe_reuse(input_tokens=tokens), 4 if retained else 0) self.assertEqual( self.manager.get_storage_statistics(CacheLevel(1))[0].evictable, 1 if retained else 0 ) pressure.close() + self.assertEqual(self._cache(tokens).get_page_priorities(0), [priority if retained else 35]) + + def test_eviction_offload_preserves_priority_and_clears_expiry(self) -> None: + self._prepare_retention(host_quota=2 << 20, offload_threshold=50) + tokens = [1, 2, 3, 4] + self._cache(tokens, self._config(80, timedelta(seconds=2))).close() + pressure = self._apply_memory_pressure() + self.assertEqual(self.manager.probe_reuse(input_tokens=tokens), 4) + self.assertEqual(self.manager.get_storage_statistics(CacheLevel(1))[0].evictable, 1) + pressure.close() time.sleep(2.05) self.manager.refresh_retention() # GPU eviction-driven offload clears the old duration, retaining the priority. - self.assertEqual(self._cache(tokens).get_page_priorities(0), [priority if retained else 35]) + self.assertEqual(self._cache(tokens).get_page_priorities(0), [80]) class TestStorageStatistics(TestKVCacheManagerV2): From 2d211980086ee96c943a9db61f2d73d855ee5c1a Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:40:04 -0700 Subject: [PATCH 14/23] test: enable V2 in admission fixtures and tidy docs Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- docs/source/examples/kvcacheretentionconfig.md | 11 +++-------- docs/source/features/kv-cache-connector.md | 8 ++------ tests/unittest/_torch/executor/test_py_executor.py | 2 +- tests/unittest/_torch/sampler/test_beam_search.py | 2 +- .../unittest/disaggregated/test_chunked_transfer.py | 12 ++++++------ 5 files changed, 13 insertions(+), 22 deletions(-) diff --git a/docs/source/examples/kvcacheretentionconfig.md b/docs/source/examples/kvcacheretentionconfig.md index fb885a0dd08b..2eaf6aa41c00 100644 --- a/docs/source/examples/kvcacheretentionconfig.md +++ b/docs/source/examples/kvcacheretentionconfig.md @@ -69,11 +69,6 @@ This example uses a single ```kv_cache_retention_config``` object for all the pr ## KV Cache Manager V2 -KV cache manager V2 accepts the same `KvCacheRetentionConfig`, so the examples -above work unchanged. For `transfer_mode`, V2 supports only the default `DRAM` and rejects -requests that set `GDS` or `POSIX_DEBUG_FALLBACK`. V2's host and disk offload -tiers are configured manager-wide when the manager is created and are -independent of these per-request transfer modes. - -The priority threshold that decides which evicted blocks are offloaded is also -manager-wide; see [Enable Offloading to Host Memory](../features/kvcache.md#enable-offloading-to-host-memory). +KV cache manager V2 accepts the same `KvCacheRetentionConfig`, so the examples above work unchanged. For `transfer_mode`, V2 supports only the default `DRAM` and rejects requests that set `GDS` or `POSIX_DEBUG_FALLBACK`. V2's host and disk offload tiers are configured manager-wide when the manager is created and are independent of these per-request transfer modes. + +The priority threshold that decides which evicted blocks are offloaded is also manager-wide; see [Enable Offloading to Host Memory](../features/kvcache.md#enable-offloading-to-host-memory). diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index aa115b1b73a5..e1431841077e 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -83,10 +83,7 @@ These methods run on the leader process and drive the connector's behavior. ##### Running under VSWA -Multi-layer-group configurations currently fail during connector initialization because retention -priorities are represented by a single list. This applies to both KV cache managers, even when the -connector implements the grouped APIs below. Single-window SWA remains allowed. The grouped -interfaces and example describe the addressing needed for future VSWA connector support. +Multi-layer-group configurations currently fail during connector initialization because retention priorities are represented by a single list. This applies to both KV cache managers, even when the connector implements the grouped APIs below. Single-window SWA remains allowed. The grouped interfaces and example describe the addressing needed for future VSWA connector support. Under variable sliding-window attention the KV cache allocates one pool per attention window size, and a page index only means something inside its own layer group. A single tensor and a single flat block list cannot describe that, so three methods have to be implemented together: @@ -98,8 +95,7 @@ Under variable sliding-window attention the KV cache allocates one pool per atte Implementing the per-layer-group form of a pair is enough — the flat method it replaces does not also have to be defined. -The current single-group restriction is checked during executor bring-up, before worker registration -or request admission. Implementing all three methods does not bypass that restriction. +The current single-group restriction is checked during executor bring-up, before worker registration or request admission. Implementing all three methods does not bypass that restriction. `examples/llm-api/llm_kv_cache_connector_vswa.py` is a worked connector for this case. A VSWA connector has to do five things: diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 6eed8d0620aa..cee70b4f612d 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -370,7 +370,7 @@ def test_context_model_capacity_profile_uses_real_dummy_requests( ) executor, _ = _make_remote_tail_executor_and_request() executor.is_warmup = is_warmup - executor._is_kv_manager_v2 = False + executor._is_kv_manager_v2 = True if not has_transceiver: executor.kv_cache_transceiver = None executor.max_beam_width = 1 diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 28e251e1f0e1..43a80f733f87 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -2016,7 +2016,7 @@ def test_vbws_rejects_decreasing_beam_width_array(beam_width_array: list[int], # Exercise the real beam admission checks with unrelated cache, transport, # token-range, sampler, and budget validation disabled in the executor stub. executor = types.SimpleNamespace( - _is_kv_manager_v2=False, + _is_kv_manager_v2=True, max_beam_width=request.py_beam_width, kv_cache_manager=None, kv_cache_transceiver=None, diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index 7d67ecec6fed..28f67dac620f 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -630,7 +630,7 @@ def test_pipelined_transfer_requires_gen_first_flow(): """ValueError when a real request is not using gen-first flow.""" from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - executor = MagicMock(_is_kv_manager_v2=False) + executor = MagicMock(_is_kv_manager_v2=True) executor.is_warmup = False executor.kv_cache_transceiver.pipeline_transfer_enabled = True executor.model_engine.attn_runtime_features.chunked_prefill = True @@ -657,7 +657,7 @@ def test_pipelined_transfer_requires_chunked_prefill(): """Pipelined context transfer requires chunked prefill.""" from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - executor = MagicMock(_is_kv_manager_v2=False) + executor = MagicMock(_is_kv_manager_v2=True) executor.kv_cache_transceiver.pipeline_transfer_enabled = True executor.model_engine = SimpleNamespace( attn_runtime_features=SimpleNamespace(chunked_prefill=False) @@ -680,7 +680,7 @@ def test_pipelined_transfer_rejects_pipeline_parallelism_for_context_request(): """Context workers require all layers on one pipeline rank.""" from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - executor = MagicMock(_is_kv_manager_v2=False) + executor = MagicMock(_is_kv_manager_v2=True) executor.kv_cache_transceiver.pipeline_transfer_enabled = True executor.model_engine.attn_runtime_features.chunked_prefill = True executor.dist.pp_size = 2 @@ -714,7 +714,7 @@ def test_pipelined_transfer_requires_single_beam_for_context_request(): """Context-side pipelining rejects beam search before scheduling.""" from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - executor = MagicMock(_is_kv_manager_v2=False) + executor = MagicMock(_is_kv_manager_v2=True) executor.kv_cache_transceiver.pipeline_transfer_enabled = True executor.model_engine.attn_runtime_features.chunked_prefill = True @@ -734,7 +734,7 @@ def test_pipelined_transfer_allows_non_disaggregated_request(): """Requests without disaggregated parameters do not transfer KV cache.""" from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - executor = MagicMock(_is_kv_manager_v2=False) + executor = MagicMock(_is_kv_manager_v2=True) executor.is_warmup = False executor.max_beam_width = 1 executor.kv_cache_transceiver.pipeline_transfer_enabled = True @@ -757,7 +757,7 @@ def test_pipelined_transfer_allows_generation_only_request(): """Generation workers do not build prefill chunks or enforce sender-only limits.""" from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - executor = MagicMock(_is_kv_manager_v2=False) + executor = MagicMock(_is_kv_manager_v2=True) executor.is_warmup = False executor.max_beam_width = 2 executor.kv_cache_transceiver.pipeline_transfer_enabled = True From b119b99d82236d8e5d34935d65047f6661d79240 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Mon, 12 Oct 2026 07:57:50 +0000 Subject: [PATCH 15/23] fix: validate layer groups in V2 priority queries Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../batch_manager/kv_cache_manager_v2/kvCache.cpp | 6 ++++++ .../test_kv_cache_manager_v2.py | 10 ++++++++++ 2 files changed, 16 insertions(+) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index 823948a74622..a2ec5d33e816 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -191,7 +191,13 @@ std::vector KvCache::getPagePriorities(LayerGroupId layerGroupId, Beam 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 result; result.reserve(mBlocks.stdSize()); bool const isSsm diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index eac26247828c..09443b83a736 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -495,6 +495,16 @@ def _config( [KvCacheRetentionConfig.TokenRangeRetentionConfig(0, None, priority, duration)] ) + def test_page_priorities_reject_invalid_layer_groups(self) -> None: + self._prepare_retention() + for num_tokens in (0, 4): + cache = self._cache(list(range(num_tokens))) + for layer_group in (-1, 1): + with self.subTest(num_tokens=num_tokens, layer_group=layer_group): + with self.assertRaises(IndexError): + cache.get_page_priorities(layer_group) + self.assertEqual(cache.get_page_priorities(0), [35] * (num_tokens // 4)) + @parameterized.expand([(4, False), (4, True), (5, False), (5, True)]) def test_beam_priorities_share_prompt_and_keep_private_tail( self, prompt_length: int, commit: bool From b207266a4b41d044d12d77478b9b31edab6ad3e0 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Mon, 12 Oct 2026 08:06:45 +0000 Subject: [PATCH 16/23] fix: remove obsolete V2 retention warning references Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 1 - tests/integration/defs/llmapi/test_llm_api_connector.py | 2 -- 2 files changed, 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 1792e063fce8..3c9a24c9b317 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -133,7 +133,6 @@ from ..moe.fused_moe.communication.base import CheckpointableCommunication _UNBOUNDED_STATS_MAX_LEN = -1 -_V2_RETENTION_IGNORED_LOG_KEY = "kv_cache_manager_v2_retention_config_ignored" class _ADPForwardIntent(IntEnum): diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index 07193c05910e..eff1a1491648 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -34,8 +34,6 @@ SchedulerOutput) from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import \ KVCacheManagerV2 -from tensorrt_llm._torch.pyexecutor.py_executor import \ - _V2_RETENTION_IGNORED_LOG_KEY from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.bindings.executor import KvCacheTransferMode from tensorrt_llm.llmapi import RequestError From 6f6a51a02980458062a6089b8fcd9b360743a8e4 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Tue, 13 Oct 2026 03:20:54 +0000 Subject: [PATCH 17/23] fix: preserve grouped connectors during retention validation Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- docs/source/features/kv-cache-connector.md | 12 ++-- .../connectors/kv_cache_connector.py | 3 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 16 +++-- .../_torch/pyexecutor/py_executor_creator.py | 4 +- .../defs/llmapi/test_llm_api_connector.py | 60 ++++++++++------- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../test_lists/test-db/l0_b200.yml | 2 + .../integration/test_lists/test-db/l0_cpu.yml | 2 + .../test_kv_connector_initialization.py | 66 +++++++++++++++---- 9 files changed, 118 insertions(+), 48 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index e1431841077e..2a197a02c105 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -76,14 +76,16 @@ These methods run on the leader process and drive the connector's behavior. | Set | Models it covers | |---|---| - | per-layer-group | single-layer-group models; multi-group configurations are currently rejected | + | per-layer-group | single-layer-group models and VSWA with KVCacheManagerV2 | | flat | non-VSWA, non-hybrid only | A set is complete when both of its methods are defined. An existing flat connector needs no change on the models it already covers: with a single layer group the base per-layer-group implementation folds back to the flat call. Hybrid / linear-attention models are not enabled for the connector yet; the per-layer-group set is the shape their cache will need. See [Running under VSWA](#running-under-vswa). ##### Running under VSWA -Multi-layer-group configurations currently fail during connector initialization because retention priorities are represented by a single list. This applies to both KV cache managers, even when the connector implements the grouped APIs below. Single-window SWA remains allowed. The grouped interfaces and example describe the addressing needed for future VSWA connector support. +KVCacheManagerV2 supports multiple layer groups with a connector implementing the grouped APIs below. KVCacheManagerV1 registers a single primary pool and cannot describe VSWA to a connector. Single-window SWA remains allowed with either manager. + +With V2, a connector that transfers KV rejects requests carrying `KvCacheRetentionConfig` when the cache has multiple layer groups: `RequestData.priorities` is a single list and cannot describe priorities per group. This is a request validation error; the service can start and process requests without an explicit retention config. Capacity-only connectors are exempt from this restriction because they do not transfer KV or emit this priority metadata. Native V2 retention without a connector remains supported across layer groups. Under variable sliding-window attention the KV cache allocates one pool per attention window size, and a page index only means something inside its own layer group. A single tensor and a single flat block list cannot describe that, so three methods have to be implemented together: @@ -95,7 +97,7 @@ Under variable sliding-window attention the KV cache allocates one pool per atte Implementing the per-layer-group form of a pair is enough — the flat method it replaces does not also have to be defined. -The current single-group restriction is checked during executor bring-up, before worker registration or request admission. Implementing all three methods does not bypass that restriction. +The executor checks that both per-layer-group scheduler callbacks are implemented during bring-up for a multi-group cache. The worker then receives the layout through `register_kv_cache_layout`; its default implementation rejects multiple groups if only flat registration is implemented. `examples/llm-api/llm_kv_cache_connector_vswa.py` is a worked connector for this case. A VSWA connector has to do five things: @@ -210,9 +212,9 @@ Two `RequestData` fields can be reported empty depending on the cache configurat | Field | When empty | Consequence | |---|---|---| | `block_hashes` | `[]` | No block-hash accessor exists on this path. Nothing in the runtime reads the field, and neither example connector uses it — both hash the token sequence themselves. A connector that keys its external store on `block_hashes` gets no key and therefore no hits and no saves; it does not mis-address a transfer. | -| `priorities` | `None` | `KvCacheRetentionConfig` is not honoured, so every page carries the default priority. A warning is logged the first time a request carrying a retention config is reported. The gap is wider than the connector: the retention config has no effect either way in that configuration. | +| `priorities` | `None` when the request has no explicit `KvCacheRetentionConfig` | Priority metadata is omitted; this does not imply that reused pages have default priority. For a single-group V2 cache, requests with an explicit retention config report native page priorities aligned with `new_block_ids`. Multi-group retention requests are rejected as described above. | -Both are gaps to be closed rather than intended behaviour. +The missing V2 block hashes remain a gap. Omitting priority metadata for requests without an explicit retention config is intentional. ##### Blocks with no page diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 726418620f43..3647d85e8fd0 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -80,7 +80,8 @@ class RequestData: # hash is committed in the same scheduler step. block_hashes: List[int] = field(default_factory=list) # The retention priorities for each new block (same length as new_block_ids). - # Used for priority-based offload filtering. None means use default priority. + # Used for priority-based offload filtering. None omits priority metadata + # when the request has no explicit retention config. priorities: Optional[List[int]] = None # Per-request cache salt that the KV cache manager uses to isolate reuse # between requests carrying different salts. Connectors that key cached diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 3c9a24c9b317..83eefedff435 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1300,12 +1300,10 @@ def _maybe_init_kv_connector_manager(self): self._reject_connector_prefix_without_block_reuse( self.kv_cache_manager) layout = build_kv_cache_layout_v2(self.kv_cache_manager) - if len(layout.groups) != 1: - raise NotImplementedError( - "KV Cache Connector requires a single layer group. " - "Retention priorities for multiple layer groups " - "are not supported.") - window = layout.groups[0].window_size + self.kv_connector_manager.reject_flat_only_scheduler( + len(layout.groups)) + window = (layout.groups[0].window_size + if len(layout.groups) == 1 else None) self.kv_connector_manager.warn_flat_scheduler_under_swa(window) self.kv_connector_manager.worker.register_kv_cache_layout( layout) @@ -6375,6 +6373,12 @@ def _validate_request(self, request: LlmRequest): if not 0 <= retention.decode_retention_priority <= 100: raise ValueError( "Decode retention priority must be between 0 and 100") + if (self._kv_connector_moves_kv + and self.kv_cache_manager.num_pools > 1): + raise ValueError( + "Per-request KV cache retention is not supported with " + "a KV cache connector and multiple layer groups: " + "the connector priority metadata supports only one group.") # Validate context-side pipelined-transfer constraints. disagg_params = request.py_disaggregated_params if (self.kv_cache_transceiver is not None diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index ec205fdd4537..b7c447895378 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -1041,7 +1041,9 @@ def allocation_scope(current_stage: ExecutorMemoryType): and v2_selection is False): raise NotImplementedError( "KV connector is not supported with VSWA (Variable Sliding Window Attention) " - "in this configuration. Use a single attention window.") + "in this configuration. Use KVCacheManagerV2 with a connector " + "that implements the per-layer-group APIs, or use a single attention window." + ) if mapping.enable_attention_dp: raise NotImplementedError( diff --git a/tests/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index eff1a1491648..d87a351ac70f 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -720,13 +720,6 @@ def test_connector_scheduler_output_chunked_context(enforce_single_worker, SWA_MAX_SEQ_LEN = 512 SWA_NUM_INPUT_TOKENS = 256 -MULTI_GROUP_CONNECTOR_UNSUPPORTED = pytest.mark.xfail( - raises=NotImplementedError, - strict=True, - reason= - "Connector initialization rejects multiple layer groups until grouped retention priorities are supported", -) - @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize("use_kv_cache_manager_v2", [True], @@ -1034,10 +1027,7 @@ def test_connector_sliding_window_prefix_is_backed_by_real_pages( @pytest.mark.threadleak(enabled=False) -@pytest.mark.parametrize("use_kv_cache_manager_v2", [ - False, - pytest.param(True, marks=MULTI_GROUP_CONNECTOR_UNSUPPORTED), -], +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["kv_cache_manager_v1", "kv_cache_manager_v2"], indirect=True) def test_connector_vswa_reports_page_indices_per_layer_group( @@ -1082,6 +1072,13 @@ def build(): scheduler.get_num_new_matched_tokens.return_value = 0, False worker.get_finished.return_value = [], [] + with pytest.raises(RequestError, match="retention.*multiple layer groups"): + model.generate([0] * SWA_NUM_INPUT_TOKENS, + sampling_params=SamplingParams(max_tokens=4), + kv_cache_retention_config=KvCacheRetentionConfig([])) + scheduler.update_state_after_alloc_by_layer_group.assert_not_called() + scheduler.request_finished_by_layer_group.assert_not_called() + generate_and_wait(model, scheduler, worker, [0] * SWA_NUM_INPUT_TOKENS, SamplingParams(max_tokens=4, ignore_eos=True)) @@ -1099,6 +1096,7 @@ def build(): req = sched_output.new_requests[0] assert req.new_block_ids == [] + assert req.priorities is None assert len(req.new_block_ids_by_layer_group) == 2 # Ordinals stay positionally aligned across groups: a block with no page in # the sliding group, or one the window has passed, reads back as @@ -1132,6 +1130,34 @@ def build(): ) +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [True], + ids=["kv_cache_manager_v2"], + indirect=True) +def test_capacity_only_connector_accepts_multigroup_retention( + enforce_single_worker, model_with_connector, use_kv_cache_manager_v2): + model_fn, scheduler, worker = model_with_connector + worker.capacity_only = True + model = model_fn(disable_overlap_scheduler=True, + max_seq_len=SWA_MAX_SEQ_LEN, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.1, + max_attention_window=[SWA_WINDOW, SWA_MAX_SEQ_LEN])) + assert_kv_caches_registered(worker, use_kv_cache_manager_v2) + layout = worker.register_kv_cache_layout.call_args.args[0] + assert len(layout.groups) == 2 + + result = model.generate([0] * SWA_NUM_INPUT_TOKENS, + sampling_params=SamplingParams(max_tokens=4, + ignore_eos=True), + kv_cache_retention_config=KvCacheRetentionConfig( + [], decode_retention_priority=80)) + assert len(result.outputs[0].token_ids) == 4 + scheduler.build_connector_meta.assert_not_called() + scheduler.update_state_after_alloc_by_layer_group.assert_not_called() + scheduler.request_finished_by_layer_group.assert_not_called() + + def _disagg_transceiver_config(use_kv_cache_manager_v2): """The transceiver each KV cache manager can actually be driven by. @@ -1514,14 +1540,6 @@ def test_connector_max_utilization_is_rejected_on_v1_only( "beam", id="beam_search", ), - pytest.param( - dict(kv_cache_config=KvCacheConfig( - free_gpu_memory_fraction=0.1, - max_attention_window=[SWA_WINDOW, SWA_MAX_SEQ_LEN])), - "VSWA", - "single layer group", - id="multiple_layer_groups", - ), pytest.param( dict(enable_attention_dp=True), "attention data parallelism", @@ -2145,7 +2163,6 @@ def test_connector_cancel_holds_pages_until_the_save_reads_them( @pytest.mark.threadleak(enabled=False) -@MULTI_GROUP_CONNECTOR_UNSUPPORTED def test_connector_multi_pool_e2e_persistent_cache(enforce_single_worker, monkeypatch): """The multi-pool data path, end to end, with a connector that moves real bytes. @@ -2347,7 +2364,6 @@ def build(): @pytest.mark.threadleak(enabled=False) -@MULTI_GROUP_CONNECTOR_UNSUPPORTED def test_connector_vswa_e2e_gemma3(enforce_single_worker, monkeypatch): """A KV connector on a model whose architecture is variable-window. @@ -2653,7 +2669,6 @@ def build(): @pytest.mark.threadleak(enabled=False) -@MULTI_GROUP_CONNECTOR_UNSUPPORTED def test_connector_vswa_out_of_window_blocks_reach_the_connector( enforce_single_worker, monkeypatch): """The sliding window passing blocks, delivered per layer group. @@ -2885,7 +2900,6 @@ def recording_register(self, layout): get_sm_version() not in (100, 103), reason="gpt-oss-20b is MXFP4 and needs SM100/SM103")), ]) -@MULTI_GROUP_CONNECTOR_UNSUPPORTED def test_connector_transfers_only_in_window_blocks_to_the_sliding_group( enforce_single_worker, monkeypatch, model_rel, assert_round_trip): """The connector transfers each layer group only what that group can read. diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 1fd54e4d1171..de5fcfcaf958 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -238,6 +238,7 @@ l0_a10: - llmapi/test_llm_api_connector.py::test_connector_uniform_sliding_window[kv_cache_manager_v2] - llmapi/test_llm_api_connector.py::test_connector_sliding_window_prefix_is_backed_by_real_pages[kv_cache_manager_v2] - llmapi/test_llm_api_connector.py::test_connector_vswa_reports_page_indices_per_layer_group[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_capacity_only_connector_accepts_multigroup_retention[kv_cache_manager_v2] - llmapi/test_llm_api_connector.py::test_connector_multi_pool_e2e_persistent_cache - llmapi/test_llm_api_connector.py::test_connector_vswa_e2e_gemma3 - llmapi/test_llm_api_connector.py::test_connector_vswa_out_of_window_blocks_reach_the_connector diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 34f3e5b83f0f..ebce3d38635d 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -293,6 +293,8 @@ l0_b200: - unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py::test_v2_hybrid_retains_configured_number_of_conversation_turns - unittest/_torch/executor/kv_cache/test_kvcm2_integration.py::test_helix_retention_uses_rank_local_ranges - llmapi/test_llm_api_connector.py::test_connector_priorities[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_connector_vswa_reports_page_indices_per_layer_group[kv_cache_manager_v2] + - llmapi/test_llm_api_connector.py::test_capacity_only_connector_accepts_multigroup_retention[kv_cache_manager_v2] - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v1] - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v2] - llmapi/test_llm_api_connector.py::test_connector_cancel_holds_pages_until_the_save_reads_them diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index a42393e3d02a..1f239b87ed27 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -49,6 +49,8 @@ l0_cpu: - unittest/_torch/compilation/test_remove_copy_pass.py::test_remove_copy_preserves_minimax_producer_outputs - unittest/_torch/distributed - unittest/_torch/executor + - unittest/_torch/test_connector.py::test_a_partial_grouped_scheduler_is_rejected_at_bring_up + - unittest/_torch/test_connector.py::test_a_grouped_scheduler_and_a_single_group_cache_are_accepted - unittest/test_startup.py - unittest/test_unified_cache.py - unittest/_torch/disaggregation diff --git a/tests/unittest/_torch/executor/test_kv_connector_initialization.py b/tests/unittest/_torch/executor/test_kv_connector_initialization.py index 4fac32b3dc80..659b16566f22 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_initialization.py +++ b/tests/unittest/_torch/executor/test_kv_connector_initialization.py @@ -31,11 +31,12 @@ @pytest.mark.parametrize("windows", [(None,), (64,), (64, None), (32, 64, None)]) -def test_v2_connector_requires_one_layer_group( - monkeypatch: pytest.MonkeyPatch, windows: tuple[int | None, ...] +@pytest.mark.parametrize("capacity_only", [False, True]) +def test_v2_connector_registers_layer_groups( + monkeypatch: pytest.MonkeyPatch, windows: tuple[int | None, ...], capacity_only: bool ) -> None: executor = object.__new__(PyExecutor) - executor.kv_connector_manager = Mock(capacity_only=False) + executor.kv_connector_manager = Mock(capacity_only=capacity_only) executor.kv_cache_transceiver = None executor.dist = SimpleNamespace(pp_size=1, cp_size=1) executor.max_beam_width = 1 @@ -55,16 +56,57 @@ def test_v2_connector_requires_one_layer_group( ) connector = executor.kv_connector_manager - if len(windows) > 1: - with pytest.raises(NotImplementedError, match="requires a single layer group"): - executor._maybe_init_kv_connector_manager() - connector.worker.register_kv_cache_layout.assert_not_called() - connector.wait_for_initialization.assert_not_called() + executor._maybe_init_kv_connector_manager() + connector.reject_flat_only_scheduler.assert_called_once_with(len(windows)) + connector.worker.register_kv_cache_layout.assert_called_once_with(layout) + connector.warn_flat_scheduler_under_swa.assert_called_once_with( + windows[0] if len(windows) == 1 else None + ) + connector.wait_for_initialization.assert_called_once() + + +@pytest.mark.parametrize( + "use_v2,connector_mode,num_groups,with_retention,rejected", + [ + pytest.param(True, "transfer", 2, True, True, id="grouped-retention"), + pytest.param(True, "transfer", 3, True, True, id="three-group-retention"), + pytest.param(True, "transfer", 2, False, False, id="grouped-no-retention"), + pytest.param(True, "capacity", 2, True, False, id="capacity-retention"), + pytest.param(True, None, 2, True, False, id="native-grouped-retention"), + pytest.param(True, "transfer", 1, True, False, id="single-group-retention"), + pytest.param(False, "transfer", 1, True, False, id="v1-retention"), + ], +) +def test_connector_retention_request_validation( + use_v2: bool, + connector_mode: str | None, + num_groups: int, + with_retention: bool, + rejected: bool, +) -> None: + executor = object.__new__(PyExecutor) + executor._is_kv_manager_v2 = use_v2 + executor.kv_connector_manager = ( + None if connector_mode is None else Mock(capacity_only=connector_mode == "capacity") + ) + executor.kv_cache_manager = SimpleNamespace(num_pools=num_groups) + executor.kv_cache_transceiver = None + executor._configure_csa2_remote_tail = Mock() + executor._validate_token_id_range = Mock() + executor._validate_request_budget = Mock() + executor.sampler = Mock() + request = SimpleNamespace( + kv_cache_retention_config=KvCacheRetentionConfig([]) if with_retention else None, + py_disaggregated_params=None, + sampling_config=None, + ) + + if rejected: + with pytest.raises(ValueError, match="retention.*multiple layer groups"): + executor._validate_request(request) else: - executor._maybe_init_kv_connector_manager() - connector.worker.register_kv_cache_layout.assert_called_once_with(layout) - connector.warn_flat_scheduler_under_swa.assert_called_once_with(windows[0]) - connector.wait_for_initialization.assert_called_once() + executor._validate_request(request) + executor.sampler.validate_request.assert_called_once_with(request) def test_single_group_priorities_follow_new_page_ordinals() -> None: From 30f6a53f53d51bf09f652470134d86ee2a3cf52a Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Mon, 12 Oct 2026 22:50:54 -0700 Subject: [PATCH 18/23] refactor: consolidate executor transfer mode import Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 83eefedff435..16622eaa07d5 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -33,13 +33,15 @@ global_mpi_size, is_trace_enabled, mpi_comm, mpi_disabled, nvtx_range, set_thread_local_mpi_comm, trace_func) -from tensorrt_llm.bindings import executor as executor_bindings +# yapf: disable from tensorrt_llm.bindings.executor import (DisServingRequestStats, FinishReason, InflightBatchingStats, IterationStats, KvCacheStats, - RequestStage, RequestStats, - RequestType, SpecDecodingStats, + KvCacheTransferMode, RequestStage, + RequestStats, RequestType, + SpecDecodingStats, StaticBatchingStats) +# yapf: enable from tensorrt_llm.bindings.internal.batch_manager import (LlmRequestType, ReqIdsSet) from tensorrt_llm.executor.request import TruncateKVCacheRequest @@ -6365,7 +6367,7 @@ def _validate_request(self, request: LlmRequest): self._configure_csa2_remote_tail(request) retention = request.kv_cache_retention_config if self._is_kv_manager_v2 and retention is not None: - if retention.transfer_mode != executor_bindings.KvCacheTransferMode.DRAM: + if retention.transfer_mode != KvCacheTransferMode.DRAM: raise ValueError( "KVCacheManagerV2 does not support per-request GDS/POSIX " "transfer_mode or directory; configure the manager's cache tiers instead." From 1949ee4695a1aa7f69e38619103b095952076129 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Mon, 12 Oct 2026 22:51:31 -0700 Subject: [PATCH 19/23] perf: skip connector priority queries without new blocks Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../_torch/pyexecutor/connectors/kv_cache_connector.py | 4 +++- .../_torch/executor/test_kv_connector_initialization.py | 2 ++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 3647d85e8fd0..c5a21aec2b3d 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -623,7 +623,9 @@ def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManag # Report the native pages' current priorities, including shared-page claims. priorities = None if req.kv_cache_retention_config is not None: - if is_v2: + if not new_block_ids: + priorities = [] + elif is_v2: kv_cache = kv_cache_manager.kv_cache_map[req.py_request_id] priorities = kv_cache.get_page_priorities(0)[ len(self.block_ids) - len(new_block_ids) : len(self.block_ids) diff --git a/tests/unittest/_torch/executor/test_kv_connector_initialization.py b/tests/unittest/_torch/executor/test_kv_connector_initialization.py index 659b16566f22..f396881f67cf 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_initialization.py +++ b/tests/unittest/_torch/executor/test_kv_connector_initialization.py @@ -138,6 +138,8 @@ def test_single_group_priorities_follow_new_page_ordinals() -> None: assert second.new_block_ids == [9] assert second.priorities == [20] + cache.get_page_priorities.reset_mock() unchanged = tracker.update_and_build_data(req, manager) assert unchanged.new_block_ids == [] assert unchanged.priorities == [] + cache.get_page_priorities.assert_not_called() From 9bfb1155750ea3b3505bf557b9c72c9b82a0e7a3 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Tue, 13 Oct 2026 02:58:00 -0700 Subject: [PATCH 20/23] fix: select parameterized connector cancellation cases on B200 Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_b200.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index ebce3d38635d..b057b6ab5812 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -297,7 +297,8 @@ l0_b200: - llmapi/test_llm_api_connector.py::test_capacity_only_connector_accepts_multigroup_retention[kv_cache_manager_v2] - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v1] - llmapi/test_llm_api_connector.py::test_connector_priorities_default[kv_cache_manager_v2] - - llmapi/test_llm_api_connector.py::test_connector_cancel_holds_pages_until_the_save_reads_them + - llmapi/test_llm_api_connector.py::test_connector_cancel_holds_pages_until_the_save_reads_them[True] + - llmapi/test_llm_api_connector.py::test_connector_cancel_holds_pages_until_the_save_reads_them[False] # ------------- KV Cache V2 Scheduler IT --------------- - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2Llama::test_v2_vs_v1_basic - kv_cache/test_kv_cache_v2_scheduler.py::TestKVCacheV2Llama::test_token_budget_limited From 5603573b09444b45046eac3e51ba944477629e02 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Tue, 13 Oct 2026 08:08:33 -0700 Subject: [PATCH 21/23] fix: supply retention config on KV cache stats test requests Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py index 5890fc0a6e40..9c49b5a8ce37 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py @@ -57,6 +57,7 @@ class _StatsRequest: is_disagg_generation_transmission_complete: bool = False is_finished_due_to_cancellation: bool = False context_phase_params: None = None + kv_cache_retention_config: None = None return_perf_metrics: bool = False py_draft_tokens: list[int] = field(default_factory=list) draft_tokens: list[int] = field(default_factory=list) From 68b9c2f600517dcf369e26bc25e3ad0e604291fa Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Tue, 13 Oct 2026 08:19:00 -0700 Subject: [PATCH 22/23] fix: include SHA256 headers for KV cache typed index tests Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- cpp/tests/unit_tests/batch_manager/CMakeLists.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index bbe25b6bb339..ff0bd64d1edd 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -49,7 +49,8 @@ add_gtest(kvCacheManagerV2TypedIndexTest "kvCacheManagerV2TypedIndexTest.cpp;kvCacheManagerV2PoisonCheck.cpp") target_include_directories( kvCacheManagerV2TypedIndexTest - PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager) + PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager + ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/sha256) add_gtest(kvCacheManagerV2HostMemTest "kvCacheManagerV2HostMemTest.cpp;kvCacheManagerV2PoisonCheck.cpp") target_include_directories( From 93151008cbdf7da122dcb5e9c5bb304fe25bf696 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Mon, 19 Oct 2026 01:32:53 -0700 Subject: [PATCH 23/23] fix: regenerate serve schemas for cold tier offload priority Update the checked-in trtllm-serve and disagg schemas for the cold_tier_offload_min_priority field and the deprecated secondary_offload_min_priority alias. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> --- .../schemas/trtllm-serve-config.schema.json | 20 ++++++++++++++++++- .../trtllm-serve-disagg-config.schema.json | 20 ++++++++++++++++++- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index ce332dfae73a..ebd67ff54476 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -2720,6 +2720,19 @@ "description": "KV cache manager v2 configuration for block reuse policies.", "status": "prototype" }, + "cold_tier_offload_min_priority": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Minimum priority for offloading unreferenced blocks from the hot tier to cold tiers. Blocks with priority at or above this value can be offloaded. Defaults to 30 when unset.", + "title": "Cold Tier Offload Min Priority" + }, "copy_on_partial_reuse": { "default": true, "description": "Whether partially matched blocks that are in use can be reused after copying them.", @@ -3023,7 +3036,12 @@ } ], "default": null, - "description": "Only blocks with priority > secondary_offload_min_priority can be offloaded to secondary memory.", + "deprecated": true, + "description": "Deprecated alias for cold_tier_offload_min_priority.", + "status": "deprecated", + "telemetry": { + "exclude": true + }, "title": "Secondary Offload Min Priority" }, "sink_token_length": { diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index e34b7f2adaa7..5f3030d6a786 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -4270,6 +4270,19 @@ "description": "KV cache manager v2 configuration for block reuse policies.", "status": "prototype" }, + "cold_tier_offload_min_priority": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Minimum priority for offloading unreferenced blocks from the hot tier to cold tiers. Blocks with priority at or above this value can be offloaded. Defaults to 30 when unset.", + "title": "Cold Tier Offload Min Priority" + }, "copy_on_partial_reuse": { "default": true, "description": "Whether partially matched blocks that are in use can be reused after copying them.", @@ -4573,7 +4586,12 @@ } ], "default": null, - "description": "Only blocks with priority > secondary_offload_min_priority can be offloaded to secondary memory.", + "deprecated": true, + "description": "Deprecated alias for cold_tier_offload_min_priority.", + "status": "deprecated", + "telemetry": { + "exclude": true + }, "title": "Secondary Offload Min Priority" }, "sink_token_length": {