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 355f6adde0ba..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 @@ -450,6 +450,9 @@ struct KVCacheManagerConfig // Beam search disables this while retaining full-block reuse. bool enablePartialCommit = true; + // Unreferenced GPU pages below this priority are dropped instead of offloaded. + Priority coldTierOffloadMinPriority = kDefaultColdTierOffloadMinPriority; + bool enableSwaScratchReuse() const noexcept { return swaScratchReuse.has_value(); 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..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 @@ -311,6 +311,24 @@ void EventManager::addCacheLevelUpdated( lifeCycle.value()); } +void EventManager::addPriorityUpdated( + Digest const& blockKey, Priority oldPriority, Priority newPriority, LifeCycleId lifeCycle) +{ + if (mMaxKvEventEntries <= 0) + { + return; + } + std::lock_guard lock(mMutex); + auto const state = mStoredBlocks.find(blockKey); + if (state == mStoredBlocks.end() || state->second.lifeCycleIds.count(lifeCycle.value()) == 0) + { + return; + } + addEventUnlocked( + KVCacheUpdatedData{state->second.blockHash, std::nullopt, KVCacheEventDiff{oldPriority, newPriority}}, + lifeCycle.value()); +} + void EventManager::addStoredEventUnlocked(KVCacheStoredData data, EventLayerGroupId layerGroupId) { bool const hasPendingRemovedEvents = !mLatestRemovedBlockHashes.empty(); @@ -547,7 +565,7 @@ std::optional 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..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 @@ -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, bool isPrefetch) : id(mId) , mManager(manager.shared_from_this()) , mReuseScope(std::move(reuseScope)) , mPriorityCb(priorityCb ? std::move(priorityCb) : [](BlockOrdinal, LifeCycleId) { return kPriorityDefault; }) + , mRetentionConfig(std::move(retentionConfig)) , mStatus(Status::SUSPENDED) , mCommitState(CommitState::ALLOWED) , mBeamWidth(BeamIndex{1}) @@ -74,6 +76,26 @@ KvCache::KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::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 @@ -99,7 +121,7 @@ KvCache::KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional(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(); + 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 + = std::holds_alternative(mManager->storage().getLifeCycle(LifeCycleId{layerGroupId.value()})); + for (auto const& block : mBlocks) + { + // Shared prompt blocks keep only beam 0; writable tails have a row per beam. + auto const pageBeamIndex = block.pages.size() == BeamIndex{1} ? kDefaultBeamIndex : beamIndex; + auto const* page = blockPageGetPage(block.pages[pageBeamIndex].at(layerGroupId)).get(); + if (isSsm && block.treeBlock) + page = block.treeBlock->getPage(LifeCycleId{layerGroupId.value()}); + // Preserve ordinals, including absent pages, just like getAggregatedPageIndices(validOnly=false). + result.push_back(page ? page->priority() : kPriorityDefault); + } + return result; +} + StorageManager* KvCache::storageManager() const { return &mManager->storage(); @@ -248,11 +317,18 @@ void KvCache::_selectOptionalGroups(std::vector const& groups) } for (auto lc : groups) { + bool const newlySelected = std::find(mSelectedOptionalGroups.begin(), mSelectedOptionalGroups.end(), lc) + == mSelectedOptionalGroups.end(); auto const stale = _getStaleRange(mHistoryLength, mManager->lifeCycles()[lc]); for (auto const& [ordinal, holder] : mOptionalCandidates[lc]) { if (!stale.contains(ordinal)) { + if (newlySelected) + { + auto const retention = getRetention(ordinal); + holder->page->claimRetention(retention.retentionPriority, retention.durationMs); + } mBlocks[ordinal].pages[kDefaultBeamIndex][lc] = holder; } } @@ -772,6 +848,7 @@ bool KvCache::resume(std::optional 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()) @@ -793,6 +870,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; @@ -851,6 +929,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); @@ -875,6 +959,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) { @@ -1266,6 +1351,7 @@ bool KvCache::_hasReuseSource(BlockPage const& page) void KvCache::_clearBlocks() { + mPartialCopyRetentionHolds.clear(); // Drop last block first (mirrors Python: while self._blocks: self._blocks.pop()). while (!mBlocks.empty()) mBlocks.pop_back(); @@ -1315,6 +1401,15 @@ CommittedPage* KvCache::_copyPageToTreeBlock( newSlot.readyEvent = CachedCudaEvent(reinterpret_cast(stream)); auto committed = makeShared( &storageMgr, treeBlock, lcIdx, lvl, numTokensInBlock, getPriority(treeBlock->ordinal(), lcIdx)); + auto const retention = getRetention(treeBlock->ordinal()); + if (mManager->lifeCycles().ssmLifeCycleId() == lcIdx) + { + committed->claimRetention(retention.retentionPriority, retention.durationMs); + } + else + { + committed->inheritRetention(*srcPage); + } committed->setSlot(newSlot); // Drops the superseded page, deferred until the copy is issued: an // OutOfPagesError above must not destroy a usable shorter snapshot. @@ -1819,6 +1914,11 @@ bool KvCache::resize(std::optional 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(); @@ -2787,7 +2887,7 @@ void KvCache::dropPartialBlockCachedTokenAttribution() // _setupForReuse: find existing blocks in radix tree matching input tokens. // --------------------------------------------------------------------------- -void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match) +void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match, bool claimRetention) { auto const& matched = match.blocks; int const numTokens = match.numTokens; @@ -2866,7 +2966,18 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match) TLLM_CHECK_WITH_INFO(page, "Expected page in non-stale block"); CacheLevel const level = page->cacheLevel; auto& bpSlot = mBlocks[ordinal].pages[beamIdx][lcId]; + bool const partialCopy = hasPartialMatch && ordinal == fullReusedEnd; + bool const borrowPartial = partialCopy && page->hasRetentionConsumers(); bpSlot = page->hold(); + if (claimRetention && !borrowPartial) + { + auto const retention = getRetention(ordinal); + page->claimRetention(retention.retentionPriority, retention.durationMs); + } + if (partialCopy) + { + mPartialCopyRetentionHolds.push_back(page->borrowRetentionForCopy()); + } if (!isAttention) { return; @@ -2937,6 +3048,11 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match) TLLM_CHECK_WITH_INFO(snapshotPage, "Last matched block must have SSM snapshot after truncation"); ssmLevel = snapshotPage->cacheLevel; mSsmBlocks[kDefaultBeamIndex][*ssmLcId] = snapshotPage->hold(); + if (claimRetention) + { + auto const retention = getRetention(snapshotBlock.ordinal()); + snapshotPage->claimRetention(retention.retentionPriority, retention.durationMs); + } } _finalizeCachedTokensByLevel(numTokens, attentionLevels, ssmLevel); 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..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 @@ -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, bool isPrefetch = false); ~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; @@ -615,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. @@ -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; @@ -784,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/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp index f91f62018dee..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.initialPoolRatio, mEventSink, mConfig.maxUtilForResume, mConfig.coldTierOffloadMinPriority); 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,27 @@ 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(); +} + +bool KvCacheManager::prefetchReuse(ReuseScope reuseScope, TokenSpan inputTokens, CacheLevel target, bool knownNoDigest) +{ + 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); + auto close = FuncGuard([&cache]() { cache->close(); }); + return cache->prefetch(target); } 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..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 @@ -147,7 +147,14 @@ 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(); + + // 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). 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..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 @@ -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()) @@ -98,6 +99,35 @@ 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; +} + +void Page::inheritRetention(Page const& source) +{ + claimRetention(source.mPriority, source.mRetentionDuration); +} + +bool Page::hasRetentionConsumers() const +{ + auto const h = holder.lock(); + // A pending partial copy protects the source bytes but does not claim its retention policy. + return h && (!h->uniqLock.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); @@ -165,6 +195,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 +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); + 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); @@ -257,6 +301,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..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,9 +23,12 @@ #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 #include +#include #include #include @@ -38,6 +41,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 +55,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 +66,17 @@ 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); + void inheritRetention(Page const& source); + bool hasRetentionConsumers() const; + FuncGuard> borrowRetentionForCopy(); + PageStatus status() const noexcept; bool scheduledForEviction() const noexcept @@ -82,6 +95,13 @@ 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; + 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 0f19a38841a6..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 @@ -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 coldTierOffloadMinPriority) : mLifeCycles(lifeCycles) , mEventSink(std::move(eventSink)) + , mColdTierOffloadMinPriority(coldTierOffloadMinPriority) , 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(); + // A prefetch may hold storage without claiming or extending its retention interval. + updatePriority(*page, kPriorityDefault); + } + } } void StorageManager::excludeFromEviction(Page& page) @@ -809,6 +879,13 @@ void StorageManager::excludeFromEviction(Page& page) void StorageManager::prepareFreeSlots(CacheLevel level, TypedVec const& requirements, MigrationRecorder const& migrationRecorder, DropRecorder const& dropRecorder) +{ + refreshRetention(); + _prepareFreeSlots(level, requirements, migrationRecorder, dropRecorder); +} + +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) @@ -883,6 +960,37 @@ 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->hasValidSlot() && !page->scheduledForEviction() && isEvictable(*page)) + mLevels.at(page->cacheLevel).controller.scheduleForEviction(*page, /*evictFirst=*/true); + } + } + }); + for (auto& pages : fallenPages) + { + for (auto iter = pages.begin(); iter != pages.end();) + { + auto const& page = *iter; + if (page->cacheLevel == kHotLevel && page->status() == PageStatus::DROPPABLE + && page->priority() < mColdTierOffloadMinPriority) + { + belowOffloadThreshold.at(page->lifeCycle).push_back(std::move(*iter)); + iter = pages.erase(iter); + } + else + { + ++iter; + } + } + } if (TLLM_UNLIKELY(gDebug)) { TLLM_CHECK_WITH_INFO(goals.size() == numCacheLevels(), "goals.rows must equal numCacheLevels"); @@ -976,7 +1084,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 +1096,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,11 +1203,21 @@ void StorageManager::_prepareFreeSlots(TypedVecstatus() == PageStatus::DROPPABLE) + { + // Offloading an unreferenced page ends its released retention interval. + page->claimRetention(std::nullopt, std::nullopt); + } scheduleForEviction(*page); } } } rescheduleEvictedPagesOnFailure.cancel(); + if (dropRecorder) + for (auto const& pages : belowOffloadThreshold) + if (!pages.empty()) + dropRecorder(pages, kHotLevel); + restoreBelowThreshold.cancel(); } // --------------------------------------------------------------------------- @@ -1719,6 +1837,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; @@ -1804,7 +1923,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 d9c9b64abef9..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 @@ -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 coldTierOffloadMinPriority = kDefaultColdTierOffloadMinPriority); ~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; @@ -387,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 = {}); @@ -416,6 +426,7 @@ class StorageManager : public std::enable_shared_from_this LifeCycleRegistry const& mLifeCycles; std::shared_ptr mEventSink; + Priority mColdTierOffloadMinPriority; LifeCyclePoolGroupMapping mHotPoolGroupMapping; LifeCyclePoolGroupMapping mColdPoolGroupMapping; // Codec batching representative for each lifecycle. @@ -462,6 +473,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..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"); @@ -1748,7 +1750,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 coldTierOffloadMinPriority) { new (cfg) kv::KVCacheManagerConfig(); cfg->tokensPerBlock = tokensPerBlock; @@ -1773,6 +1776,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) cfg->textOnly = textOnly; cfg->enablePartialCommit = enablePartialCommit; cfg->poolRebalance = std::move(poolRebalance); + cfg->coldTierOffloadMinPriority = coldTierOffloadMinPriority; // Mirror Python's __post_init__: validate at construction. Config-integrity // failures raise AssertionError (translated below). cfg->validate(); @@ -1783,7 +1787,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("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) @@ -1791,6 +1796,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("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, @@ -1954,6 +1960,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 +2651,20 @@ 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( + "prefetch_reuse", + [](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, kv::CacheLevel{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()) @@ -2647,7 +2672,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 +2681,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 +2698,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/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( 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/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 e490f223798f..2eaf6aa41c00 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,9 @@ 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 + +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..2a197a02c105 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -76,13 +76,17 @@ 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 and VSWA with KVCacheManagerV2 | | 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 +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: | Method | Replaces | @@ -93,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. -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 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: @@ -208,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/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/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 3eb1f191cd3c..c5a21aec2b3d 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 @@ -575,8 +576,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 +620,21 @@ 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 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) + ] + 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..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 @@ -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, @@ -63,6 +64,7 @@ BAD_PAGE_INDEX, CACHE_LEVEL1, DEFAULT_BEAM_INDEX, + DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY, GPU_LEVEL, AttentionLayerConfig, AttnLifeCycle, @@ -3048,6 +3050,11 @@ def _build_base_config( ratio_tolerance=rebalance_cfg.ratio_tolerance, moving_average_decay=rebalance_cfg.moving_average_decay, ), + cold_tier_offload_min_priority=( + 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 + ), ) def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: @@ -3751,6 +3758,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 @@ -3779,6 +3822,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=self._retention_config_for_request(req), is_dummy=req.is_dummy, enable_request_stats=req.return_perf_metrics, expected_prompt_length=( @@ -4245,6 +4289,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 +4664,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=self._retention_config_for_request(req), ) if kv_cache is None: return None @@ -6535,6 +6582,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 +6623,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 @@ -6662,29 +6711,21 @@ 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. + if not self.can_evict: + return True + # 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/_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..16622eaa07d5 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -33,12 +33,15 @@ global_mpi_size, is_trace_enabled, mpi_comm, mpi_disabled, nvtx_range, set_thread_local_mpi_comm, trace_func) +# 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 @@ -132,7 +135,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): @@ -6363,15 +6365,22 @@ 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 != 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") + 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 7a95ac8c2682..b7c447895378 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -1034,14 +1034,15 @@ 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 KVCacheManagerV2 with a connector " + "that implements the per-layer-group APIs, or use a single attention window." ) if mapping.enable_attention_dp: 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 903458676ba8..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,6 +135,7 @@ class _KVCacheManagerConfigFieldSpec: text_only: bool = False enable_partial_commit: bool = True pool_rebalance: object = None + cold_tier_offload_min_priority: int = DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY KVCacheManagerConfig.__dataclass_fields__ = _KVCacheManagerConfigFieldSpec.__dataclass_fields__ @@ -226,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 c1e3ed04e75d..6d4ed5f3f88b 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -32,9 +32,12 @@ from typing import ( Union, ) +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] @@ -328,6 +331,7 @@ class KVCacheManagerConfig: text_only: bool = False enable_partial_commit: bool = True pool_rebalance: PoolRebalanceConfig = ... + cold_tier_offload_min_priority: int = DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY @property def enable_swa_scratch_reuse(self) -> bool: ... @@ -686,6 +690,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 +848,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 +860,11 @@ class KVCacheManager: it also marks generation-phase allocation stats. """ ... + 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/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": { 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/integration/defs/llmapi/test_llm_api_connector.py b/tests/integration/defs/llmapi/test_llm_api_connector.py index b7921433db26..d87a351ac70f 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector.py @@ -34,15 +34,14 @@ 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 tensorrt_llm.llmapi.llm_args import (CacheTransceiverConfig, CapacitySchedulerPolicy, 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 @@ -264,13 +263,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 @@ -1077,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)) @@ -1094,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 @@ -1127,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. @@ -1237,39 +1268,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 +1283,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 = [], [] @@ -1304,6 +1309,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, @@ -1334,83 +1353,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) @@ -1418,11 +1377,13 @@ def test_connector_warns_that_retention_is_ignored_on_v2( 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 = [], [] @@ -2140,7 +2101,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..de5fcfcaf958 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] @@ -241,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 0061bf2389ad..b057b6ab5812 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -289,6 +289,16 @@ 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] + - 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[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 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/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_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 2a1d353bb9b1..685648f1c1e2 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, 66, 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/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 0387894bf82c..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 @@ -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,8 @@ def _build_v2_hybrid_with_mamba_layer( mamba_ssm_cache_dtype=torch.float16, kda_replay_num_spec=None, enable_kv_pool_rebalance=False, + cold_tier_offload_min_priority=None, + host_cache_size=None, ): """Construct a real MambaHybridCacheManagerV2.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -2453,6 +2456,8 @@ def _build_v2_hybrid_with_mamba_layer( ), dtype=kv_cache_dtype, enable_kv_pool_rebalance=enable_kv_pool_rebalance, + cold_tier_offload_min_priority=cold_tier_offload_min_priority, + host_cache_size=host_cache_size, ) return MambaHybridCacheManagerV2( mamba_d_state=8, @@ -2501,6 +2506,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 +2514,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 +2955,15 @@ 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], + cold_tier_offload_min_priority=50, + host_cache_size=0, + ) + 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 +2985,13 @@ 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.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] mgr.free_resources(request_a) _run_v2_hybrid_context(mgr, request_b) 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 = {} 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..f396881f67cf --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_connector_initialization.py @@ -0,0 +1,145 @@ +# 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)]) +@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=capacity_only) + 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 + 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._validate_request(request) + executor.sampler.validate_request.assert_called_once_with(request) + + +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] + + 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() diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 29df16273a38..cee70b4f612d 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 = True 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/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 72fe06068076..43a80f733f87 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=True, max_beam_width=request.py_beam_width, kv_cache_manager=None, kv_cache_transceiver=None, 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_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index d75c51a6ca77..28f67dac620f 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=True) 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=True) 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=True) 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=True) 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=True) 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=True) 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 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 6573a6c5a8c2..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 @@ -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 @@ -36,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, @@ -50,8 +52,10 @@ GpuCacheTierConfig, HostCacheTierConfig, KVCacheDesc, + KVCacheEventManager, KVCacheManager, KVCacheManagerConfig, + KVCacheUpdatedData, LayerGroupId, LayerId, LogicError, @@ -82,7 +86,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, @@ -90,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, @@ -104,8 +111,10 @@ GpuCacheTierConfig, HostCacheTierConfig, KVCacheDesc, + KVCacheEventManager, KVCacheManager, KVCacheManagerConfig, + KVCacheUpdatedData, LayerGroupId, LayerId, LogicError, @@ -409,6 +418,376 @@ def prepare( self.manager = KVCacheManager(self.cfg) +class TestRetention(TestKVCacheManagerV2): + def setUp(self) -> None: + super().setUp() + 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() + + 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 :] + if tokens: + self.engine.execute( + [Step(cache, remaining, tokens[: cache.num_committed_tokens])], self.stream.handle + ) + if remaining: + 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: int | None, duration: timedelta | None = None + ) -> KvCacheRetentionConfig | None: + if priority is None: + return None + return KvCacheRetentionConfig( + [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 + ) -> 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: int, second: int | None, close_newer_first: bool + ) -> None: + self._prepare_retention() + prompt = [1, 2, 3, 4] + 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) + 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] + 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. + 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() + 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_and_decode_priorities_expire_after_last_release(self) -> None: + self._prepare_retention() + 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.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]) + + 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)], + 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]) + + @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() + pressure = self._apply_memory_pressure() + 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) + ) + 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)) + 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() + + 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.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]) + + 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 = [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]) + self.engine.execute([Step(copy, [], list(range(3)))], self.stream.handle) + self.stream.synchronize() + + @parameterized.expand([(29, 30), (30, 30), (49, 50), (50, 50), (0, -1), (100, 101)]) + 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() + 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), [80]) + + 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) @@ -3402,7 +3781,16 @@ 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, timedelta(milliseconds=20) + ) + ] + ), + ) kv_cache.resume(stream) kv_cache.capacity = len(prompt) kv_cache.history_length = len(prompt) @@ -3414,7 +3802,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 @@ -3431,6 +3822,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) + 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.""" @@ -4813,6 +5212,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 @@ -4841,7 +5241,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: @@ -4984,7 +5384,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)] @@ -5003,6 +5404,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() 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) 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: