From fb3ad4b588cdb68c05a51ded0bd564b6523569ed Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Tue, 27 Oct 2026 15:28:50 +0800 Subject: [PATCH 01/19] [None][feat] BREAKING: Support layer-level reuse policies for attention and SSM Signed-off-by: Jiagan Cheng --- .../kv_cache_manager_v2/AGENTS.md | 10 +- .../kv_cache_manager_v2/blockRadixTree.cpp | 2 +- .../kv_cache_manager_v2/config.h | 7 +- .../kv_cache_manager_v2/kvCache.cpp | 210 +++++++++++------- .../kv_cache_manager_v2/kvCache.h | 6 +- .../kv_cache_manager_v2/lifeCycleRegistry.cpp | 4 +- .../kv_cache_manager_v2/lifeCycleRegistry.h | 58 +++-- .../kv_cache_manager_v2/page.cpp | 3 +- .../kv_cache_manager_v2/storage/config.cpp | 5 +- .../kv_cache_manager_v2/storageManager.cpp | 15 +- .../batch_manager/kvCacheManagerV2.cpp | 33 +-- .../kvCacheManagerV2ColdPageTest.cpp | 2 +- .../kvCacheManagerV2StatsTest.cpp | 199 +++++++++++++++-- .../backends/sparse/csa2/cache_manager.py | 18 +- .../kv_cache/kv_cache_manager_v2.py | 6 +- .../runtime/kv_cache_manager_v2/__init__.py | 4 +- .../runtime/kv_cache_manager_v2/__init__.pyi | 19 +- .../sparse/csa2/test_cache_manager.py | 6 +- .../kv_cache/test_kv_cache_estimation.py | 4 +- .../kv_cache/test_mixed_precision_kv_cache.py | 2 +- .../modeling/test_deepseek_v41_snapshots.py | 14 +- .../test_kv_cache_concurrency.py | 4 +- .../test_optional_checkpoint.py | 68 +++++- 23 files changed, 509 insertions(+), 190 deletions(-) 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 32d6371a3666..cc64681c09a4 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 @@ -95,7 +95,15 @@ must eventually be `close()`d. Its normal flow is: `LifeCycleRegistry` maps model layers to lifecycle groups. Attention lifecycles may have sliding-window and sink-token rules; SSM lifecycles represent a -recurrent-state checkpoint. A pool-group index is a storage-layout index and is +recurrent-state checkpoint. Both layer configs expose `reusePolicy`: `REQUIRED` +constrains the prefix match, `OPTIONAL` offers state at the required endpoint for +explicit selection on first resume, and `PRIVATE` never publishes reusable pages. +All buffers within a layer share that policy. SSM lifecycles are deduplicated by +policy, so a manager can have up to three SSM groups; iterate `ssmLifeCycleIds()` +for working state and use `requiredSsmLifeCycleId()` only for prefix pruning. +Optional SSM checkpoints require an exact endpoint match, including partial +blocks. Fresh or unselected recurrent state must be initialized or reconstructed +by the caller before inference. A pool-group index is a storage-layout index and is not interchangeable with a layer ID or lifecycle ID. Sparsity is a property of a pool, not of a lifecycle: a layer's buffers may be diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp index 18ed536c27f0..3b74d32f5629 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp @@ -973,7 +973,7 @@ BlockRadixTree::ReuseMatch BlockRadixTree::match( auto rawMatched = matchTokenPath(reuseScope, tokens, knownNoDigest, enablePartialMatch); // Content-divergence depth is measured before page or recurrent-snapshot pruning. int const numReusableTokensBeforePruning = numMatchedTokens(rawMatched, mTokensPerBlock); - auto const ssmLcId = mLifeCycles.ssmLifeCycleId(); + auto const ssmLcId = mLifeCycles.requiredSsmLifeCycleId(); // Page requirements depend on the final endpoint. Back off before pruning // so SWA coverage and recurrent snapshots are validated at that endpoint. if (backoff > 0) 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 8415192f24bf..869ea3f065b4 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 @@ -292,8 +292,8 @@ inline void validateNoDuplicateBufferRoles(std::vector const& buff } // namespace detail -//! Controls whether an attention lifecycle participates in prefix matching and page reuse. -enum class AttentionReusePolicy +//! Controls whether a layer lifecycle participates in prefix matching and page reuse. +enum class ReusePolicy { //! Commit reusable pages and require live-window coverage when matching a prefix. REQUIRED = 0, @@ -319,7 +319,7 @@ struct AttentionLayerConfig std::optional numSinkTokens; // Private working pages never constrain prefix matching or enter the tree. - AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; + ReusePolicy reusePolicy = ReusePolicy::REQUIRED; //! The sparse fetch path for this layer's sparse buffers, if it has one. std::optional sparseCache; @@ -341,6 +341,7 @@ struct SsmLayerConfig LayerId layerId = 0; std::vector buffers; + ReusePolicy reusePolicy = ReusePolicy::REQUIRED; void validate() const { 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 cb791c821abf..5b5e09b9e68c 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 @@ -236,8 +236,8 @@ void KvCache::_collectReuseCandidates() for (LifeCycleId lc{0}; lc < lifeCycles.size(); ++lc) { auto const* attention = std::get_if(&lifeCycles[lc]); - auto const policy = attention ? attention->reusePolicy : AttentionReusePolicy::REQUIRED; - ReuseGroupStatus status{lc, policy, end, end > 0 && policy == AttentionReusePolicy::REQUIRED, {}}; + auto const policy = reusePolicy(lifeCycles[lc]); + ReuseGroupStatus status{lc, policy, end, end > 0 && policy == ReusePolicy::REQUIRED, {}}; if (status.complete) { if (!attention) @@ -294,6 +294,18 @@ void KvCache::_collectReuseCandidates() status.coverage.clear(); } } + if (!attention && policy == ReusePolicy::OPTIONAL && end > 0) + { + BlockOrdinal const ordinal{(end - 1) / mTokensPerBlock}; + auto const& block = mBlocks[ordinal].treeBlock; + auto* page = block && !block->isOrphan() ? block->getPage(lc) : nullptr; + if (page && page->numTokensInBlock == end - ordinal.value() * mTokensPerBlock) + { + mOptionalCandidates[lc].emplace_back(ordinal, page->hold()); + status.complete = true; + status.coverage.emplace_back(end, end); + } + } mReuseStatus.push_back(std::move(status)); } } @@ -311,16 +323,21 @@ void KvCache::_selectOptionalGroups(std::vector const& groups) for (auto lc : groups) { if (lc < LifeCycleId{0} || lc >= mOptionalCandidates.size() - || mReuseStatus[lc.value()].policy != AttentionReusePolicy::OPTIONAL || !mReuseStatus[lc.value()].complete) + || mReuseStatus[lc.value()].policy != ReusePolicy::OPTIONAL || !mReuseStatus[lc.value()].complete) { throw std::invalid_argument("Selected OPTIONAL group has no reusable state at the claimed prefix"); } } for (auto lc : mSelectedOptionalGroups) { - for (auto const& [ordinal, holder] : mOptionalCandidates[lc]) + if (mManager->lifeCycles().isSsm(lc)) { - mBlocks[ordinal].pages[kDefaultBeamIndex][lc] = std::monostate{}; + mSsmBlocks[kDefaultBeamIndex][lc] = std::monostate{}; + } + else + { + for (auto const& [ordinal, holder] : mOptionalCandidates[lc]) + mBlocks[ordinal].pages[kDefaultBeamIndex][lc] = std::monostate{}; } } for (auto lc : groups) @@ -330,14 +347,17 @@ void KvCache::_selectOptionalGroups(std::vector const& groups) auto const stale = _getStaleRange(mHistoryLength, mManager->lifeCycles()[lc]); for (auto const& [ordinal, holder] : mOptionalCandidates[lc]) { - if (!stale.contains(ordinal)) + bool const isSsm = mManager->lifeCycles().isSsm(lc); + if (isSsm || !stale.contains(ordinal)) { if (newlySelected) { auto const retention = getRetention(ordinal); holder->page->claimRetention(retention.retentionPriority, retention.durationMs); } - mBlocks[ordinal].pages[kDefaultBeamIndex][lc] = holder; + auto& target + = isSsm ? mSsmBlocks[kDefaultBeamIndex][lc] : mBlocks[ordinal].pages[kDefaultBeamIndex][lc]; + target = holder; } } } @@ -348,7 +368,7 @@ void KvCache::_selectOptionalGroups(std::vector const& groups) for (LifeCycleId lc{0}; lc < mOptionalCandidates.size(); ++lc) { if (std::find(groups.begin(), groups.end(), lc) == groups.end() - && mReuseStatus[lc.value()].policy == AttentionReusePolicy::OPTIONAL) + && mReuseStatus[lc.value()].policy == ReusePolicy::OPTIONAL) { mOptionalCandidates[lc].clear(); mReuseStatus[lc.value()].complete = false; @@ -360,18 +380,18 @@ void KvCache::_selectOptionalGroups(std::vector const& groups) std::vector KvCache::_activePages() const { std::vector result; - auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + auto const& lcs = mManager->lifeCycles(); LifeCycleId numLc = mManager->storage().numLifeCycles(); for (LifeCycleId lcId{0}; lcId < numLc; ++lcId) { // SSM lifecycle → yield from mSsmBlocks (check individual entries). - if (ssmLcId.has_value() && lcId == *ssmLcId) + if (mManager->lifeCycles().isSsm(lcId)) { for (BeamIndex bi{0}; bi < mBeamWidth; ++bi) { - if (!blockPageIsNull(mSsmBlocks[bi][*ssmLcId])) + if (!blockPageIsNull(mSsmBlocks[bi][lcId])) result.push_back({kBadBlockOrdinal, bi, lcId}); } continue; @@ -421,8 +441,8 @@ std::vector KvCache::_activePages() const SharedPtr KvCache::_page(BlockOrdinal ordinal, BeamIndex beamIdx, LifeCycleId lcId) const { bool const isSsm = ordinal == kBadBlockOrdinal; - auto const ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); - TLLM_CHECK_DEBUG((ssmLcId.has_value() && lcId == *ssmLcId) == isSsm); + + TLLM_CHECK_DEBUG((mManager->lifeCycles().isSsm(lcId)) == isSsm); auto const& blockPage = isSsm ? mSsmBlocks.at(beamIdx).at(lcId) : mBlocks.at(ordinal).pages.at(beamIdx).at(lcId); return blockPageGetPage(blockPage); } @@ -518,7 +538,7 @@ bool KvCache::resume(std::optional stream, std::optionalstorage(); - auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + LifeCycleId numLc = storageMgr.numLifeCycles(); // Snapshot source levels before allocation or activation can migrate the selected pages. @@ -527,6 +547,8 @@ bool KvCache::resume(std::optional stream, std::optionallifeCycles().isSsm(lc)) + continue; auto& byLevel = optionalReusedByLevel[lc]; byLevel.full = CountsByLevel(storageMgr.numCacheLevels(), 0); byLevel.partial = CountsByLevel(storageMgr.numCacheLevels(), 0); @@ -575,9 +597,13 @@ bool KvCache::resume(std::optional stream, std::optionallifeCycles().isSsm(lc); auto const& lifeCycle = mManager->lifeCycles()[lc]; - if (_needsFreshPrefixPages(lc)) + if (isSsm) + { + numSlotsNeeded[lc] += 1; + } + else if (_needsFreshPrefixPages(lc)) { auto const stale = _getStaleRange(mHistoryLength, lifeCycle); for (BlockOrdinal ord{0}; ord < stale.beg; ++ord) @@ -590,7 +616,7 @@ bool KvCache::resume(std::optional stream, std::optional(privateOrdinals[lc].size()); } - else if (isSsm || hasPartial) + else if (hasPartial) { numSlotsNeeded[lc] += 1; } @@ -633,8 +659,7 @@ bool KvCache::resume(std::optional stream, std::optionallifeCycles().isSsm(lc) || (!isPrivate && hasPartial)); for (size_t i = 0; i < privateOrdinals[lc].size(); ++i) { privateSlots[lc].push_back(std::move(tmpSlots[lc].back())); @@ -741,10 +766,10 @@ bool KvCache::resume(std::optional stream, std::optionallifeCycles().isSsm(lcIdx)) { - if (numCommittedTokens() == 0) - continue; // fresh SSM — no source to copy from + if (blockPageIsNull(mSsmBlocks[beamIdx][lcIdx])) + continue; // Fresh or unselected SSM state is initialized by the caller. sourcePage = &mSsmBlocks[beamIdx][lcIdx]; } else @@ -759,7 +784,7 @@ bool KvCache::resume(std::optional stream, std::optionalpage()->cacheLevel; storageMgr.copySlotData(lcIdx, kHotLevel, sourceLevel, newSlot.slotId(), lock->page()->slotId(), cudaStr); - if ((!ssmLcId.has_value() || lcIdx != *ssmLcId) && (recordManagerStats || recordRequestStats)) + if ((!mManager->lifeCycles().isSsm(lcIdx)) && (recordManagerStats || recordRequestStats)) { bool const changed = mPendingStats.recordAllocationRange(lcIdx, lastOrdinal, lastOrdinal + 1, /*beamWidth=*/1, /*countAsMissed=*/!hasPartialReuseSource, /*countAsGeneration=*/false, @@ -801,7 +826,7 @@ bool KvCache::resume(std::optional stream, std::optionallifeCycles().isSsm(lcIdx)) { targetBp = &mSsmBlocks[beamIdx][lcIdx]; blockOrdinal = kBadBlockOrdinal; @@ -841,6 +866,19 @@ bool KvCache::resume(std::optional stream, std::optionallifeCycles().ssmLifeCycleIds()) + { + if (reusePolicy(mManager->lifeCycles()[lc]) != ReusePolicy::OPTIONAL) + continue; + bool const selected = std::find(mSelectedOptionalGroups.begin(), mSelectedOptionalGroups.end(), lc) + != mSelectedOptionalGroups.end(); + if (mPendingStats.recordSsmSnapshotLookup( + lc, mNumLookupTokens, selected ? numCommittedTokens() : 0, mTokensPerBlock)) + mManager->markStatsDirty(id); + } + } mOptionalCandidates.clear(); // Deferred copies survive a failed decode admission and must not be repeated on retry. mPartialCopyRetentionHolds.clear(); @@ -931,14 +969,13 @@ void KvCache::suspend() // SharedPageLock destructors inside the scope use finishEvent() to synchronize. { auto scope = recordEventScope(); - auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); // Convert SharedPageLocks → PageHolders for active (non-stale) pages only. // Mirrors Python: for ordinal, beam_idx, lc_idx in self._active_pages() for (auto const& ap : _activePages()) { - auto& bp = (ap.lcId != ssmLcId) ? mBlocks[ap.ordinal].pages[ap.beamIdx][ap.lcId] - : mSsmBlocks[ap.beamIdx][ap.lcId]; + auto& bp = (!mManager->lifeCycles().isSsm(ap.lcId)) ? mBlocks[ap.ordinal].pages[ap.beamIdx][ap.lcId] + : mSsmBlocks[ap.beamIdx][ap.lcId]; // expect_type(_SharedPageLock, beam_block[lc_idx]) → std::get raises on wrong type auto& lock = std::get(bp); auto holder = lock.page()->hold(); @@ -1248,11 +1285,11 @@ void KvCache::_clearBlocks() mBlocks.pop_back(); _freeReservedGpuSlots(); // Clear SSM blocks. - auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); - if (ssmLcId.has_value()) + + for (auto lcId : mManager->lifeCycles().ssmLifeCycleIds()) { for (auto& beamBlock : mSsmBlocks) - beamBlock[*ssmLcId] = std::monostate{}; + beamBlock[lcId] = std::monostate{}; } } @@ -1293,7 +1330,7 @@ CommittedPage* KvCache::_copyPageToTreeBlock( auto committed = makeShared( &storageMgr, treeBlock, lcIdx, lvl, numTokensInBlock, getPriority(treeBlock->ordinal(), lcIdx)); auto const retention = getRetention(treeBlock->ordinal()); - if (mManager->lifeCycles().ssmLifeCycleId() == lcIdx) + if (mManager->lifeCycles().isSsm(lcIdx)) { committed->claimRetention(retention.retentionPriority, retention.durationMs); } @@ -1322,6 +1359,10 @@ CommittedPage* KvCache::_copyPageToTreeBlock( void KvCache::_snapshotSsmToTreeBlock(SharedPtr const& treeBlock, LifeCycleId ssmLcId, int numTokens, bool move) { + if (isPrivateLifeCycle(mManager->lifeCycles()[ssmLcId])) + { + return; + } int const numTokensInBlock = numTokens - treeBlock->ordinal().value() * mTokensPerBlock; TLLM_CHECK_DEBUG(0 < numTokensInBlock && numTokensInBlock <= mTokensPerBlock); @@ -1439,7 +1480,7 @@ void KvCache::_snapshotPartialBlockToTree(BlockOrdinal ordinal, bool commitSsm) && std::equal(tokens.begin(), tokens.end(), treeBlock->tokens.begin()))); auto& beamBlock = mBlocks.at(ordinal).pages[kDefaultBeamIndex]; - auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + // `treeBlock` may be a longer existing sibling that covers these tokens. Attaching the // partial attention pages to it is still correct because each page records the token // span it covers, and prefix matching honours that span. @@ -1462,8 +1503,8 @@ void KvCache::_snapshotPartialBlockToTree(BlockOrdinal ordinal, bool commitSsm) } if (commitSsm) { - TLLM_CHECK_DEBUG(ssmLcId.has_value()); - _snapshotSsmToTreeBlock(treeBlock, *ssmLcId, start + numTokens); + for (auto lcId : mManager->lifeCycles().ssmLifeCycleIds()) + _snapshotSsmToTreeBlock(treeBlock, lcId, start + numTokens); } if (treeBlock->eventSink) { @@ -1588,7 +1629,6 @@ bool KvCache::resize(std::optional capacity, std::optional historyLeng _checkPageIndexBufferCapacity(newNumBlocks); - auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); auto backupHolders = _unlockStaleBlocks(newHist); if (newNumBlocks < oldNumBlocks) @@ -1608,7 +1648,7 @@ bool KvCache::resize(std::optional capacity, std::optional historyLeng TypedVec> staleRanges(numLc); for (LifeCycleId lc{0}; lc < numLc; ++lc) { - if (ssmLcId.has_value() && lc == *ssmLcId) + if (mManager->lifeCycles().isSsm(lc)) continue; staleRanges[lc] = _getStaleRange(newHist, lcs.getLifeCycle(lc)); @@ -1770,7 +1810,7 @@ bool KvCache::resize(std::optional capacity, std::optional historyLeng { for (LifeCycleId lc{0}; lc < numLc; ++lc) { - if (ssmLcId.has_value() && lc == *ssmLcId) + if (mManager->lifeCycles().isSsm(lc)) continue; // SSM pages live in mSsmBlocks, not in mBlocks if (enableScratch) { @@ -1973,12 +2013,13 @@ void KvCache::_appendBeams(BeamIndex oldBeamWidth, BeamIndex newBeamWidth) } } - auto const ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); - SharedPtr sourceSsmPage; - if (ssmLcId.has_value()) - sourceSsmPage = blockPageGetPage(mSsmBlocks[kDefaultBeamIndex][*ssmLcId]); - if (sourceSsmPage) - slotCounts[*ssmLcId] += numNewBeams; + TypedVec> sourceSsmPages(numLc); + for (auto lcId : mManager->lifeCycles().ssmLifeCycleIds()) + { + sourceSsmPages[lcId] = blockPageGetPage(mSsmBlocks[kDefaultBeamIndex][lcId]); + if (sourceSsmPages[lcId]) + slotCounts[lcId] += numNewBeams; + } MigrationRecorder const migrationRecorder = [this](std::vector> const& pages, std::vector const& slots, CacheLevel srcLevel, @@ -2054,11 +2095,14 @@ void KvCache::_appendBeams(BeamIndex oldBeamWidth, BeamIndex newBeamWidth) for (BeamIndex beamIdx = oldBeamWidth; beamIdx < newBeamWidth; ++beamIdx) { LifeCycleBlockPages newSsmBlock(numLc); - if (sourceSsmPage) + for (auto lcId : mManager->lifeCycles().ssmLifeCycleIds()) { - auto& slot = newSlots[*ssmLcId].back(); - newSsmBlock[*ssmLcId] = copyPage(sourceSsmPage, slot, kBadBlockOrdinal, *ssmLcId, beamIdx); - newSlots[*ssmLcId].pop_back(); + if (auto const& source = sourceSsmPages[lcId]) + { + auto& slot = newSlots[lcId].back(); + newSsmBlock[lcId] = copyPage(source, slot, kBadBlockOrdinal, lcId, beamIdx); + newSlots[lcId].pop_back(); + } } mSsmBlocks.push_back(std::move(newSsmBlock)); } @@ -2149,14 +2193,13 @@ void KvCache::_lockHeldBlocks(std::vector const& backup) // Mirrors Python's _take_uncommitted_page(). // --------------------------------------------------------------------------- -TypedVec KvCache::_takeUncommittedPage( - SeqBlock& sb, BeamIndex beamIdx, std::optional skipLc) +TypedVec KvCache::_takeUncommittedPage(SeqBlock& sb, BeamIndex beamIdx) { LifeCycleId numLc = mManager->storage().numLifeCycles(); TypedVec result(numLc, TakenPage{nullptr, false}); for (LifeCycleId lc{0}; lc < numLc; ++lc) { - if ((skipLc.has_value() && lc == *skipLc) || isPrivateLifeCycle(mManager->lifeCycles()[lc])) + if (mManager->lifeCycles().isSsm(lc) || isPrivateLifeCycle(mManager->lifeCycles()[lc])) { continue; } @@ -2230,14 +2273,13 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) || (newBlock->tokens.size() >= tokenBlock.size() && std::equal(tokenBlock.begin(), tokenBlock.end(), newBlock->tokens.begin()))); - auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); bool didCommit = false; if (blockIsNew) { // New block: take uncommitted pages, convert to committed. // Mirrors Python's _take_uncommitted_page + convert path. - auto taken = _takeUncommittedPage(sb, kDefaultBeamIndex, ssmLcId); + auto taken = _takeUncommittedPage(sb, kDefaultBeamIndex); for (LifeCycleId lc{0}; lc < numLc; ++lc) { auto& [up, locked] = taken[lc]; @@ -2265,7 +2307,7 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) std::vector reuseTasks; for (LifeCycleId lc{0}; lc < numLc; ++lc) { - if (ssmLcId.has_value() && lc == *ssmLcId) + if (mManager->lifeCycles().isSsm(lc)) continue; if (isPrivateLifeCycle(mManager->lifeCycles()[lc])) { @@ -2362,8 +2404,8 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) if (didCommit && commitSsm) { - TLLM_CHECK_DEBUG(ssmLcId.has_value()); - _snapshotSsmToTreeBlock(newBlock, *ssmLcId, start + static_cast(tokenBlock.size()), moveSsm); + for (auto lcId : mManager->lifeCycles().ssmLifeCycleIds()) + _snapshotSsmToTreeBlock(newBlock, lcId, start + static_cast(tokenBlock.size()), moveSsm); } if (sb.isCommitted()) @@ -2371,7 +2413,7 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) auto const& lifeCycles = mManager->lifeCycles(); for (LifeCycleId lcIdx{0}; lcIdx < numLc; ++lcIdx) { - if (ssmLcId.has_value() && lcIdx == *ssmLcId) + if (mManager->lifeCycles().isSsm(lcIdx)) { continue; } @@ -2425,7 +2467,6 @@ void KvCache::commit(TokenSpan tokens, bool isEnd) "Cannot commit digest tokens to a text-only KV cache"); bool const commitMinSnapshot = mManager->commitMinSnapshot(); - auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); int const numCommitted = static_cast(mCommittedTokens.size()) + static_cast(tokens.size()); if (commitMinSnapshot) { @@ -2461,7 +2502,8 @@ void KvCache::commit(TokenSpan tokens, bool isEnd) auto scope = recordEventScope(); for (int ordinal = numCommittedBlocksBefore; ordinal < newNumFullBlocks; ++ordinal) { - bool const commitSsm = commitMinSnapshot && ssmLcId.has_value() && ordinal == ssmSnapshotOrdinal; + bool const commitSsm + = commitMinSnapshot && mManager->lifeCycles().hasSSM() && ordinal == ssmSnapshotOrdinal; _commitBlock(ordinal, /*isLast=*/false, commitSsm, /*moveSsm=*/isEnd); // _commitBlock transitions to USER_STOP on VIRTUAL_STOP internally. if (mCommitState != CommitState::ALLOWED) @@ -2472,12 +2514,12 @@ void KvCache::commit(TokenSpan tokens, bool isEnd) BlockOrdinal const partialOrdinal{newNumFullBlocks}; if (isEnd && mManager->enablePartialCommit()) { - _commitBlock(newNumFullBlocks, /*isLast=*/true, /*commitSsm=*/ssmLcId.has_value(), - /*moveSsm=*/ssmLcId.has_value()); + _commitBlock(newNumFullBlocks, /*isLast=*/true, /*commitSsm=*/mManager->lifeCycles().hasSSM(), + /*moveSsm=*/mManager->lifeCycles().hasSSM()); } else if (!isEnd) { - _snapshotPartialBlockToTree(partialOrdinal, /*commitSsm=*/ssmLcId.has_value()); + _snapshotPartialBlockToTree(partialOrdinal, /*commitSsm=*/mManager->lifeCycles().hasSSM()); } } } @@ -2621,6 +2663,8 @@ std::unique_ptr KvCache::planCommittedBlockDrop() { LifeCycleId const lcIdx = item.id; LifeCycle const& lc = item.lc; + if (isPrivateLifeCycle(lc)) + continue; BlockOrdinal windowStart; if (auto const* attn = std::get_if(&lc)) { @@ -2639,7 +2683,11 @@ std::unique_ptr KvCache::planCommittedBlockDrop() { CommittedPage* page = match.blocks[ordinal]->getPage(lcIdx); if (page == nullptr) - return nullptr; + { + if (requiresPrefixCoverage(lc)) + return nullptr; + continue; + } pagesToDrop.push_back(page); } } @@ -2653,12 +2701,12 @@ std::unique_ptr KvCache::planCommittedBlockDrop() void KvCache::_onStopCommitting() { - auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + auto const& lcs = mManager->lifeCycles(); for (auto [lcIdx, lc] : lcs) { - if (ssmLcId.has_value() && lcIdx == *ssmLcId) + if (mManager->lifeCycles().isSsm(lcIdx)) continue; // SSM pages live in _ssm_blocks, not in _blocks auto staleRange = _getStaleRange(mHistoryLength, lc); @@ -2768,11 +2816,12 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match, bool claim { auto const& matched = match.blocks; int const numTokens = match.numTokens; + mNumLookupTokens = match.numLookupTokens; auto& lifeCycles = mManager->lifeCycles(); auto const& allLc = lifeCycles.getAll(); LifeCycleId numLc = lifeCycles.size(); - auto ssmLcId = lifeCycles.ssmLifeCycleId(); + BlockOrdinal const fullReusedEnd{numTokens / mTokensPerBlock}; bool const hasPartialMatch = numTokens % mTokensPerBlock != 0; bool const recordManagerStats = _shouldRecordManagerStats(); @@ -2807,7 +2856,7 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match, bool claim for (LifeCycleId lcId{0}; lcId < numLc; ++lcId) { // SSM is handled separately below. - if (ssmLcId.has_value() && lcId == *ssmLcId) + if (mManager->lifeCycles().isSsm(lcId)) continue; if (!requiresPrefixCoverage(allLc[lcId])) @@ -2918,13 +2967,15 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match, bool claim // prefix, so its source tier applies to every logical reused token; older checkpoints are // traversal history, not inputs. std::optional ssmLevel; - if (ssmLcId.has_value() && !matched.empty()) + for (auto lcId : lifeCycles.ssmLifeCycleIds()) { + if (!requiresPrefixCoverage(lifeCycles[lcId]) || matched.empty()) + continue; auto& snapshotBlock = *matched.back(); - auto* snapshotPage = snapshotBlock.storage[*ssmLcId]; + auto* snapshotPage = snapshotBlock.storage[lcId]; TLLM_CHECK_WITH_INFO(snapshotPage, "Last matched block must have SSM snapshot after truncation"); - ssmLevel = snapshotPage->cacheLevel; - mSsmBlocks[kDefaultBeamIndex][*ssmLcId] = snapshotPage->hold(); + ssmLevel = std::max(ssmLevel.value_or(kHotLevel), snapshotPage->cacheLevel); + mSsmBlocks[kDefaultBeamIndex][lcId] = snapshotPage->hold(); if (claimRetention) { auto const retention = getRetention(snapshotBlock.ordinal()); @@ -2935,9 +2986,11 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match, bool claim _finalizeCachedTokensByLevel(numTokens, attentionLevels, ssmLevel); // Record one SSM snapshot lookup for this reuse-match onboarding (a miss when // nothing matched). Mirrors Python's record_ssm_snapshot_lookup call. - if (_shouldRecordManagerStats() && ssmLcId.has_value()) + for (auto lcId : lifeCycles.ssmLifeCycleIds()) { - if (mPendingStats.recordSsmSnapshotLookup(*ssmLcId, match.numLookupTokens, numTokens, mTokensPerBlock)) + if (!_shouldRecordManagerStats() || !requiresPrefixCoverage(lifeCycles[lcId])) + continue; + if (mPendingStats.recordSsmSnapshotLookup(lcId, match.numLookupTokens, numTokens, mTokensPerBlock)) { mManager->markStatsDirty(id); } @@ -2983,11 +3036,11 @@ SharedPtr const& KvCache::_getTreeBlock(BlockOrdinal ordinal) const TLLM_CHECK_DEBUG(ret); if (TLLM_UNLIKELY(gDebug)) { - auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); + auto const& beamBlock = mBlocks[ordinal].pages[kDefaultBeamIndex]; for (LifeCycleId lcId{0}; lcId < beamBlock.size(); ++lcId) { - if (ssmLcId.has_value() && lcId == *ssmLcId) + if (mManager->lifeCycles().isSsm(lcId)) { TLLM_CHECK_WITH_INFO(blockPageIsNull(beamBlock[lcId]), "SSM pages live in mSsmBlocks"); } @@ -3030,7 +3083,6 @@ bool KvCache::_checkSanity() const auto const& lcs = mManager->lifeCycles(); LifeCycleId numLc = mManager->storage().numLifeCycles(); - auto ssmLcId = lcs.ssmLifeCycleId(); // Precompute stale and scratch ranges for each lifecycle. TypedVec> staleRanges(numLc); @@ -3054,7 +3106,7 @@ bool KvCache::_checkSanity() const for (LifeCycleId lc{0}; lc < numLc; ++lc) { auto const& bp = beamBlock[lc]; - if (ssmLcId.has_value() && lc == *ssmLcId) + if (mManager->lifeCycles().isSsm(lc)) { // SSM pages live in mSsmBlocks, not in mBlocks. // When mNeverResumed and SSM snapshot is held, the block is committed @@ -3124,12 +3176,12 @@ bool KvCache::_checkSanity() const } } - // Check SSM blocks (mirrors Python lines 1342-1353). - if (ssmLcId.has_value()) + // Each SSM lifecycle owns a separate writable state per beam. + for (auto lcId : lcs.ssmLifeCycleIds()) { for (BeamIndex bi{0}; bi < mBeamWidth; ++bi) { - auto const& bp = mSsmBlocks[bi][*ssmLcId]; + auto const& bp = mSsmBlocks[bi][lcId]; if (!blockPageIsNull(bp)) { if (mNeverResumed) 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 b7eef425270c..e2a53674c92c 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 @@ -243,7 +243,7 @@ struct ReuseGroupStatus // Lifecycle group returned by getLayerGroupId(layerId), not a physical pool index. LayerGroupId groupId; // REQUIRED constrains the prefix match; OPTIONAL is selected at resume; PRIVATE is never reused. - AttentionReusePolicy policy; + ReusePolicy policy; // Exclusive token endpoint of the prefix claimed by REQUIRED groups, shared by all groups. int endpoint; // All state needed by this group at endpoint is available; does not imply GPU residency or selection. @@ -767,8 +767,7 @@ class KvCache : public std::enable_shared_from_this // Extract uncommitted pages from a SeqBlock, resetting block page entries. // Returns one TakenPage per lifecycle. Mirrors Python's _take_uncommitted_page(). - TypedVec _takeUncommittedPage( - SeqBlock& sb, BeamIndex beamIdx, std::optional skipLc = std::nullopt); + TypedVec _takeUncommittedPage(SeqBlock& sb, BeamIndex beamIdx); // Get and validate the tree block at a committed ordinal. // Mirrors Python's _get_tree_block(). Asserts committed pages reference the correct block. @@ -822,6 +821,7 @@ class KvCache : public std::enable_shared_from_this int mCapacity; int mHistoryLength; std::optional mExpectedPromptLength; + int mNumLookupTokens = 0; bool mGenerationAllocReady = false; // Page index tables: [beamIdx][lcId] → either an internal vector or an external span. diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp index f32204aba3e4..0f6a387e5bd1 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp @@ -36,7 +36,7 @@ LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock) using T = std::decay_t; if constexpr (std::is_same_v) { - return SsmLifeCycle{}; + return SsmLifeCycle{cfg.reusePolicy}; } else { @@ -62,7 +62,7 @@ LifeCycleRegistry::LifeCycleRegistry(KVCacheManagerConfig const& config) mLifeCycleList.push_back(lc); mLifeCycleIdMap[lc] = id; if (std::holds_alternative(lc)) - mSsmLifeCycleId = id; + mSsmLifeCycleIds.push_back(id); } } check(); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h index 55cebefbf206..8987165fa0d4 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h @@ -40,21 +40,21 @@ struct AttnLifeCycle { std::optional windowSize; // nullopt = no sliding window int numSinkBlocks = 0; // divUp(numSinkTokens, tokensPerBlock) - AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; + ReusePolicy reusePolicy = ReusePolicy::REQUIRED; bool isPrivate() const noexcept { - return reusePolicy == AttentionReusePolicy::PRIVATE; + return reusePolicy == ReusePolicy::PRIVATE; } bool requiresPrefixCoverage() const noexcept { - return reusePolicy == AttentionReusePolicy::REQUIRED; + return reusePolicy == ReusePolicy::REQUIRED; } bool allowsCheckpoint() const noexcept { - return reusePolicy == AttentionReusePolicy::OPTIONAL; + return reusePolicy == ReusePolicy::OPTIONAL; } HalfOpenRange getStaleRange(int historyLength, int tokensPerBlock) const @@ -89,7 +89,7 @@ struct AttnLifeCycle } static AttnLifeCycle make(std::optional ws, std::optional numSinkTokens, int tokensPerBlock, - AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED) + ReusePolicy reusePolicy = ReusePolicy::REQUIRED) { TLLM_CHECK_DEBUG(tokensPerBlock > 0); TLLM_CHECK_DEBUG(!ws.has_value() || *ws > 0); @@ -103,23 +103,25 @@ struct AttnLifeCycle // --------------------------------------------------------------------------- // SsmLifeCycle — lifecycle for SSM (State Space Model) layers. // All blocks before the last full block are stale (recurrent state). -// Singleton: all SSM layers share the same lifecycle. +// SSM layers with the same reuse policy share a lifecycle. // --------------------------------------------------------------------------- struct SsmLifeCycle { + ReusePolicy reusePolicy = ReusePolicy::REQUIRED; + HalfOpenRange getStaleRange(int historyLength, int tokensPerBlock) const { return {BlockOrdinal{0}, BlockOrdinal{historyLength / tokensPerBlock}}; } - bool operator==(SsmLifeCycle const&) const noexcept + bool operator==(SsmLifeCycle const& other) const noexcept { - return true; + return reusePolicy == other.reusePolicy; } - bool operator<(SsmLifeCycle const&) const noexcept + bool operator<(SsmLifeCycle const& other) const noexcept { - return false; + return reusePolicy < other.reusePolicy; } }; @@ -128,16 +130,19 @@ struct SsmLifeCycle // --------------------------------------------------------------------------- using LifeCycle = std::variant; +inline ReusePolicy reusePolicy(LifeCycle const& lc) +{ + return std::visit([](auto const& value) { return value.reusePolicy; }, lc); +} + inline bool isPrivateLifeCycle(LifeCycle const& lc) { - auto const* attn = std::get_if(&lc); - return attn && attn->isPrivate(); + return reusePolicy(lc) == ReusePolicy::PRIVATE; } inline bool requiresPrefixCoverage(LifeCycle const& lc) { - auto const* attn = std::get_if(&lc); - return !attn || attn->requiresPrefixCoverage(); + return reusePolicy(lc) == ReusePolicy::REQUIRED; } // Free function: compute stale range via std::visit. @@ -210,14 +215,31 @@ class LifeCycleRegistry } // SSM helpers. - std::optional ssmLifeCycleId() const noexcept + std::vector const& ssmLifeCycleIds() const noexcept + { + return mSsmLifeCycleIds; + } + + bool isSsm(LifeCycleId id) const { - return mSsmLifeCycleId; + return std::holds_alternative((*this)[id]); + } + + std::optional requiredSsmLifeCycleId() const + { + for (auto id : mSsmLifeCycleIds) + { + if (requiresPrefixCoverage((*this)[id])) + { + return id; + } + } + return std::nullopt; } bool hasSSM() const noexcept { - return mSsmLifeCycleId.has_value(); + return !mSsmLifeCycleIds.empty(); } // Return (id, AttnLifeCycle*) pairs for attention lifecycles only. Used by _setupForReuse. @@ -293,7 +315,7 @@ class LifeCycleRegistry TypedVec mLifeCycleList; std::map mLifeCycleIdMap; - std::optional mSsmLifeCycleId; + std::vector mSsmLifeCycleIds; }; } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 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 5a82032341fd..6bef15ba65e7 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 @@ -193,8 +193,7 @@ UncommittedPage::~UncommittedPage() // switching to monostate, so during destruction the slot still references this page. if (TLLM_UNLIKELY(gDebug)) { - auto ssmLcId = manager->lifeCycles().ssmLifeCycleId(); - bool isSsm = ssmLcId.has_value() && lifeCycle == *ssmLcId; + bool const isSsm = manager->lifeCycles().isSsm(lifeCycle); if (!isSsm) { bool blockRemoved = kvCache->blocks().size() <= ordinal; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp index b4dcd77c153b..bf28b4708913 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp @@ -216,7 +216,7 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) // per-pool sparsity too keeps a shared pool uniformly host-backed or device-backed. struct PoolGroupKey { - AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; + ReusePolicy reusePolicy = ReusePolicy::REQUIRED; std::vector slotSizes; std::vector sparsity; @@ -231,8 +231,7 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) for (auto& sg : slotGroups) { auto const sizes = sg.slotSizeList(); - auto const* attn = std::get_if(®istry.getLifeCycle(sg.lifeCycleId)); - auto const policy = attn ? attn->reusePolicy : AttentionReusePolicy::REQUIRED; + auto const policy = reusePolicy(registry.getLifeCycle(sg.lifeCycleId)); TypedVec sparsity; sparsity.reserve(sizes.size()); for (PoolIndex pool{0}; pool < sizes.size(); ++pool) 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 0735340b929f..afae999e6449 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 @@ -432,12 +432,11 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi TypedVec coldGrouping(numLifeCycles()); TypedVec coldSlotDescList; - std::map, PoolGroupIndex> coldGroupByPageBytes; + std::map, PoolGroupIndex> coldGroupByPageBytes; for (LifeCycleId lifeCycle{0}; lifeCycle < numLifeCycles(); ++lifeCycle) { size_t const coldPageBytes = coldPageBytesByLifeCycle[lifeCycle]; - auto const* attn = std::get_if(&mLifeCycles[lifeCycle]); - auto const policy = attn ? attn->reusePolicy : AttentionReusePolicy::REQUIRED; + auto const policy = reusePolicy(mLifeCycles[lifeCycle]); auto [it, inserted] = coldGroupByPageBytes.emplace( std::make_pair(policy, coldPageBytes), PoolGroupIndex{static_cast(coldSlotDescList.size().value())}); PoolGroupIndex const coldPgIdx = it->second; @@ -1885,14 +1884,13 @@ TypedVec StorageManager::ratioFromLength( // same boundary KvCache::_appendBeams() uses. HalfOpenRange const sharedRange{0, std::min(std::max(promptLength, 0) / tokensPerBlock, numBlocks)}; TypedVec numBytes(numLifeCycles(), 0); - auto const ssmLcId = mLifeCycles.ssmLifeCycleId(); auto const& lifeCycles = mLifeCycles.getAll(); for (LifeCycleId lifeCycle{0}; lifeCycle < lifeCycles.size(); ++lifeCycle) { auto const poolSizes = slotSize(level, getPoolGroupIndex(level, lifeCycle)); size_t const slotBytes = std::accumulate(poolSizes.begin(), poolSizes.end(), size_t{0}); int numRequiredBlocks; - if (ssmLcId.has_value() && lifeCycle == *ssmLcId) + if (mLifeCycles.isSsm(lifeCycle)) { // One recurrent-state block per beam. numRequiredBlocks = effBeamWidth; @@ -1953,9 +1951,9 @@ TypedVec StorageManager::ratioFromBatch(BatchDesc const& bat auto numSlots = computeSlotsForBatch(batch, tokensPerBlock, swaScratchReuse); // A soft preference, not a feasibility floor: keep one optional endpoint // per typical request in addition to its private working pages. - for (auto const& [lc, attention] : mLifeCycles.attentionLifeCycles()) + for (auto const& [lc, lifeCycle] : mLifeCycles) { - if (attention->allowsCheckpoint()) + if (reusePolicy(lifeCycle) == ReusePolicy::OPTIONAL) { numSlots[lc] += slotCountValueFromSize(batch.kvCaches.size()); } @@ -2052,12 +2050,11 @@ TypedVec StorageManager::computeSlotsForBatch( BatchDesc const& batch, int tokensPerBlock, std::optional const& swaScratchReuse) const { TypedVec numSlots(numLifeCycles(), 0); - auto ssmLcId = mLifeCycles.ssmLifeCycleId(); int sysBlocks = batch.systemPromptLength / tokensPerBlock; for (auto const& [lcIdx, lc] : mLifeCycles) { - if (ssmLcId.has_value() && lcIdx == *ssmLcId) + if (mLifeCycles.isSsm(lcIdx)) { // SSM: one dedicated block per request per beam, never shared. // _appendBeams() gives every beam its own recurrent-state slot. diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index f7d9ab23f7e0..c09f06d608d1 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1492,18 +1492,18 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .value("SHARED", kv::PageIndexMode::SHARED) .value("PER_LAYER", kv::PageIndexMode::PER_LAYER); - nb::enum_(m, "AttentionReusePolicy") - .value("REQUIRED", kv::AttentionReusePolicy::REQUIRED) - .value("PRIVATE", kv::AttentionReusePolicy::PRIVATE) - .value("OPTIONAL", kv::AttentionReusePolicy::OPTIONAL); + nb::enum_(m, "ReusePolicy") + .value("REQUIRED", kv::ReusePolicy::REQUIRED) + .value("PRIVATE", kv::ReusePolicy::PRIVATE) + .value("OPTIONAL", kv::ReusePolicy::OPTIONAL); nb::class_(m, "AttnLifeCycle") - .def(nb::init, int, kv::AttentionReusePolicy>(), nb::arg("window_size").none(), - nb::arg("num_sink_blocks"), nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED) + .def(nb::init, int, kv::ReusePolicy>(), nb::arg("window_size").none(), + nb::arg("num_sink_blocks"), nb::arg("reuse_policy") = kv::ReusePolicy::REQUIRED) // Sink tokens round up to whole blocks. Bound rather than repeated in Python so the // connector's view of a life cycle is built by the same code as the allocator's. .def_static("make", &kv::AttnLifeCycle::make, nb::arg("window_size").none(), nb::arg("num_sink_tokens").none(), - nb::arg("tokens_per_block"), nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED) + nb::arg("tokens_per_block"), nb::arg("reuse_policy") = kv::ReusePolicy::REQUIRED) .def_prop_ro("window_size", [](kv::AttnLifeCycle const& self) { return self.windowSize; }) .def_ro("num_sink_blocks", &kv::AttnLifeCycle::numSinkBlocks) .def_ro("reuse_policy", &kv::AttnLifeCycle::reusePolicy) @@ -1512,7 +1512,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def("__eq__", &kv::AttnLifeCycle::operator==); nb::class_(m, "SsmLifeCycle") - .def(nb::init<>()) + .def(nb::init(), nb::arg("reuse_policy") = kv::ReusePolicy::REQUIRED) + .def_ro("reuse_policy", &kv::SsmLifeCycle::reusePolicy) .def( "get_stale_range", &kv::SsmLifeCycle::getStaleRange, nb::arg("history_length"), nb::arg("tokens_per_block")) .def("__eq__", &kv::SsmLifeCycle::operator==); @@ -1733,9 +1734,9 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::class_(m, "AttentionLayerConfig") .def(nb::init, std::optional, std::optional, - kv::AttentionReusePolicy, std::optional>(), + kv::ReusePolicy, std::optional>(), nb::arg("layer_id"), nb::arg("buffers"), nb::arg("sliding_window_size") = std::nullopt, - nb::arg("num_sink_tokens") = std::nullopt, nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED, + nb::arg("num_sink_tokens") = std::nullopt, nb::arg("reuse_policy") = kv::ReusePolicy::REQUIRED, nb::arg("sparse_cache") = std::nullopt) .def_rw("sparse_cache", &kv::AttentionLayerConfig::sparseCache) .def_rw("layer_id", &kv::AttentionLayerConfig::layerId) @@ -1750,7 +1751,9 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .value("SSM", kv::LayerType::SSM); nb::class_(m, "SsmLayerConfig") - .def(nb::init>(), nb::arg("layer_id"), nb::arg("buffers")) + .def(nb::init, kv::ReusePolicy>(), nb::arg("layer_id"), + nb::arg("buffers"), nb::arg("reuse_policy") = kv::ReusePolicy::REQUIRED) + .def_rw("reuse_policy", &kv::SsmLayerConfig::reusePolicy) .def_rw("layer_id", &kv::SsmLayerConfig::layerId) .def_rw("buffers", &kv::SsmLayerConfig::buffers) DEF_COPY(kv::SsmLayerConfig); @@ -2625,10 +2628,10 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) "ssm_life_cycle_id", [](kv::KvCacheManager& manager) -> std::optional { - auto id = manager.lifeCycles().ssmLifeCycleId(); - if (id.has_value()) - return id->value(); - return std::nullopt; + auto const& ids = manager.lifeCycles().ssmLifeCycleIds(); + if (ids.size() > 1) + throw std::invalid_argument("Multiple SSM lifecycles; look up the layer group by layer ID"); + return ids.empty() ? std::nullopt : std::optional{ids.front().value()}; }, nb::arg("manager"), nb::call_guard()); // Returns (num_tokens, pages) where pages[i] is None for a block with no page in diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp index 496d2c413acf..8d4b6d61701d 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp @@ -61,7 +61,7 @@ KVCacheManagerConfig makeOptionalDiskConfig(bool hostTier) auto layer = std::get(config.layers.front()); layer.layerId = i; layer.slidingWindowSize = i == 1 ? 4 : 2; - layer.reusePolicy = AttentionReusePolicy::OPTIONAL; + layer.reusePolicy = ReusePolicy::OPTIONAL; // Unequal buffer sizes exercise multiple physical pools per page. layer.buffers.push_back(BufferConfig{"score", 1024, std::nullopt}); config.layers.emplace_back(std::move(layer)); diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index eed2adf2cf9f..65ddd0404b7e 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -67,7 +67,7 @@ class ReplayCacheTest : public ::testing::Test return tokens; } - static void addWindow(KVCacheManagerConfig& config, int id, int window, AttentionReusePolicy policy, int sinks = 0) + static void addWindow(KVCacheManagerConfig& config, int id, int window, ReusePolicy policy, int sinks = 0) { auto layer = std::get(config.layers.front()); layer.layerId = id; @@ -86,7 +86,7 @@ TEST_F(ReplayCacheTest, GlobalReuseKeepsPrivateAttentionUnpublished) { SCOPED_TRACE(length); auto config = makeConfig(); - addWindow(config, 1, 8, AttentionReusePolicy::PRIVATE); + addWindow(config, 1, 8, ReusePolicy::PRIVATE); auto manager = std::make_shared(config); auto const globalLc = manager->getLayerGroupId(LayerId{0}); auto const privateLc = manager->getLayerGroupId(LayerId{1}); @@ -142,11 +142,11 @@ TEST_F(ReplayCacheTest, ConcurrentRebaseDoesNotResurrectRetiredPages) { auto const tokens = makeTokens(24); TokenSpan const tokenSpan{tokens.data(), static_cast(tokens.size())}; - for (auto policy : {AttentionReusePolicy::REQUIRED, AttentionReusePolicy::PRIVATE, AttentionReusePolicy::OPTIONAL}) + for (auto policy : {ReusePolicy::REQUIRED, ReusePolicy::PRIVATE, ReusePolicy::OPTIONAL}) { SCOPED_TRACE(static_cast(policy)); auto config = makeConfig(); - if (policy != AttentionReusePolicy::REQUIRED) + if (policy != ReusePolicy::REQUIRED) { addWindow(config, 1, 8, policy); } @@ -162,11 +162,11 @@ TEST_F(ReplayCacheTest, ConcurrentRebaseDoesNotResurrectRetiredPages) first->commit(tokenSpan); second->commit(tokenSpan); EXPECT_EQ(first->getBasePageIndices(globalLc)[5], second->getBasePageIndices(globalLc)[5]); - if (policy != AttentionReusePolicy::REQUIRED) + if (policy != ReusePolicy::REQUIRED) { auto const privateLc = manager->getLayerGroupId(LayerId{1}); EXPECT_EQ(first->getBasePageIndices(privateLc)[5] == second->getBasePageIndices(privateLc)[5], - policy == AttentionReusePolicy::OPTIONAL); + policy == ReusePolicy::OPTIONAL); } second->close(); first->close(); @@ -180,7 +180,7 @@ TEST_F(ReplayCacheTest, ShrinkReleasesRetiredLocksBeforeVariantChanges) { SCOPED_TRACE(privateState); auto config = makeConfig(); - addWindow(config, 1, 8, privateState ? AttentionReusePolicy::PRIVATE : AttentionReusePolicy::REQUIRED); + addWindow(config, 1, 8, privateState ? ReusePolicy::PRIVATE : ReusePolicy::REQUIRED); auto manager = std::make_shared(config); auto cache = manager->createKvCache({}, {}, 1); ASSERT_TRUE(cache->resume(mStream)); @@ -221,6 +221,177 @@ TEST_F(ReplayCacheTest, SsmSnapshotLocksDoNotUseAttentionPageIndices) manager->shutdown(); } +TEST_F(ReplayCacheTest, MixedSsmPoliciesRestoreStateAndCopyEveryBeam) +{ + for (int length : {7, 8}) + { + SCOPED_TRACE(length); + auto config = makeConfig(); + config.commitMinSnapshot = true; + config.enablePartialCommit = false; + std::array const policies{ + ReusePolicy::REQUIRED, ReusePolicy::OPTIONAL, ReusePolicy::PRIVATE, ReusePolicy::REQUIRED}; + for (int i = 0; i < static_cast(policies.size()); ++i) + { + SsmLayerConfig layer; + layer.layerId = i + 1; + layer.buffers = {BufferConfig{"state", 4096}, BufferConfig{"conv", 4096}}; + layer.reusePolicy = policies[i]; + config.layers.emplace_back(layer); + } + auto manager = std::make_shared(config); + auto const required = manager->getLayerGroupId(1); + auto const optional = manager->getLayerGroupId(2); + auto const privateGroup = manager->getLayerGroupId(3); + EXPECT_EQ(required, manager->getLayerGroupId(4)); + EXPECT_NE(required, optional); + EXPECT_NE(optional, privateGroup); + auto const tokens = makeTokens(length); + TokenSpan const span{tokens.data(), length}; + auto address = [&](KvCache const& cache, int layer, DataRole const& role, int beam = 0) + { + auto const group = manager->getLayerGroupId(layer); + auto const slot = cache.getSsmBlockBaseIndex(group, BeamIndex{beam}); + return manager->getMemPoolBaseAddress(layer, role) + + static_cast(slot) * manager->getPageIndexScale(layer, role) + * manager->getPageStride(layer, role); + }; + auto fill = [&](KvCache const& cache, int layer, DataRole const& role, int pattern) + { + EXPECT_EQ(cudaMemsetAsync(reinterpret_cast(address(cache, layer, role)), pattern, 4096, mStream), + cudaSuccess); + }; + auto check = [&](KvCache const& cache, int layer, DataRole const& role, int beam, uint8_t pattern) + { + EXPECT_EQ(cudaStreamSynchronize(mStream), cudaSuccess); + std::vector data(4096); + EXPECT_EQ(cudaMemcpy(data.data(), reinterpret_cast(address(cache, layer, role, beam)), + data.size(), cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_TRUE(std::all_of(data.begin(), data.end(), [pattern](uint8_t value) { return value == pattern; })); + }; + auto source = manager->createKvCache({}, {}, 1); + ASSERT_TRUE(source->resume(mStream)); + ASSERT_TRUE(source->resize(length)); + for (int layer = 1; layer <= 4; ++layer) + { + fill(*source, layer, "state", 10 + layer); + fill(*source, layer, "conv", 20 + layer); + } + source->commit(span); + source->stopCommitting(); + auto reader = manager->createKvCache({}, span, 2, {}, length); + ASSERT_EQ(reader->numCommittedTokens(), length); + EXPECT_TRUE(reader->reuseStatus()[required.value()].complete); + EXPECT_TRUE(reader->reuseStatus()[optional.value()].complete); + EXPECT_FALSE(reader->reuseStatus()[privateGroup.value()].complete); + EXPECT_EQ( + reader->reuseStatus()[optional.value()].coverage, (std::vector>{{length, length}})); + EXPECT_THROW( + reader->resume(mStream, std::nullopt, std::vector{privateGroup}), std::invalid_argument); + ASSERT_TRUE(reader->resume(mStream, std::nullopt, std::vector{optional})); + for (int layer : {1, 2, 4}) + { + check(*reader, layer, "state", 0, 10 + layer); + check(*reader, layer, "conv", 0, 20 + layer); + } + EXPECT_NE(address(*source, 3, "state"), address(*reader, 3, "state")); + fill(*reader, 3, "state", 13); + fill(*reader, 3, "conv", 23); + ASSERT_TRUE(reader->resize(length + 4)); + reader->setBeamWidth(BeamIndex{3}); + for (int layer = 1; layer <= 4; ++layer) + { + for (int beam = 0; beam < 3; ++beam) + { + check(*reader, layer, "state", beam, 10 + layer); + check(*reader, layer, "conv", beam, 20 + layer); + } + } + reader->suspend(); + ASSERT_TRUE(reader->resume(mStream)); + reader->setBeamWidth(BeamIndex{1}); + reader->close(); + source->close(); + manager->shutdown(); + } +} + +TEST_F(ReplayCacheTest, OptionalSsmRequiresExactEndpointAndDoesNotLimitAttentionReuse) +{ + for (auto policy : {ReusePolicy::OPTIONAL, ReusePolicy::PRIVATE}) + { + auto config = makeConfig(); + config.commitMinSnapshot = true; + config.layers.emplace_back(SsmLayerConfig{1, {BufferConfig{"state", 4096}}, policy}); + auto manager = std::make_shared(config); + auto const group = manager->getLayerGroupId(1); + auto const tokens = makeTokens(8); + TokenSpan const full{tokens.data(), 8}; + auto source = manager->createKvCache({}, {}, 1); + ASSERT_TRUE(source->resume(mStream)); + ASSERT_TRUE(source->resize(8)); + source->commit(full, true); + source->close(); + for (int length : {7, 8}) + { + SCOPED_TRACE(length); + auto reader = manager->createKvCache({}, TokenSpan{tokens.data(), length}, 2); + EXPECT_EQ(reader->numCommittedTokens(), length); + bool const available = policy == ReusePolicy::OPTIONAL && length == 8; + EXPECT_EQ(reader->reuseStatus()[group.value()].complete, available); + if (!available) + { + EXPECT_THROW( + reader->resume(mStream, std::nullopt, std::vector{group}), std::invalid_argument); + } + // Unselected OPTIONAL and PRIVATE state both get writable request-local storage. + ASSERT_TRUE(reader->resume(mStream)); + EXPECT_GE(reader->getSsmBlockBaseIndex(group), 0); + reader->suspend(); + ASSERT_TRUE(reader->resume(mStream)); + reader->close(); + } + manager->shutdown(); + } +} + +TEST_F(ReplayCacheTest, OptionalSsmEvictionPreservesRequiredPrefix) +{ + auto config = makeConfig(); + config.commitMinSnapshot = true; + config.layers.emplace_back(SsmLayerConfig{1, {BufferConfig{"state", 4096}}, ReusePolicy::REQUIRED}); + config.layers.emplace_back(SsmLayerConfig{2, {BufferConfig{"state", 4096}}, ReusePolicy::OPTIONAL}); + auto manager = std::make_shared(config); + auto const optional = manager->getLayerGroupId(2); + auto const tokens = makeTokens(8); + TokenSpan const span{tokens.data(), 8}; + auto source = manager->createKvCache({}, {}, 1); + ASSERT_TRUE(source->resume(mStream)); + ASSERT_TRUE(source->resize(8)); + source->commit(span, true); + source->close(); + auto& storage = manager->storage(); + auto const pool = storage.getPoolGroupIndex(kHotLevel, optional); + auto const stats = storage.getStatistics(kHotLevel, pool); + auto reservations = storage.newSlotsForPoolGroup(kHotLevel, pool, stats.free + stats.evictable); + EXPECT_EQ(manager->probeReuse({}, span), 8); + auto reader = manager->createKvCache({}, span, 2); + EXPECT_EQ(reader->numCommittedTokens(), 8); + EXPECT_FALSE(reader->reuseStatus()[optional.value()].complete); + EXPECT_FALSE(reader->resume(mStream)); + for (auto& slot : reservations) + storage.releaseSlot(optional, kHotLevel, std::move(slot)); + ASSERT_TRUE(reader->resume(mStream)); + reader->stopCommitting(); + auto drop = reader->planCommittedBlockDrop(); + EXPECT_NE(drop, nullptr); + if (drop) + drop->drop(); + reader->close(); + manager->shutdown(); +} + TEST_F(ReplayCacheTest, ClosingScratchCacheReleasesEverySlotExactlyOnce) { auto config = makeConfig(); @@ -252,14 +423,14 @@ TEST_F(ReplayCacheTest, OptionalAndPrivateHaveIndependentGroups) config.cacheTiers.emplace_back(HostCacheTierConfig{4 << 20}); auto& encoder = std::get(config.layers.front()); encoder.slidingWindowSize = 8; - encoder.reusePolicy = AttentionReusePolicy::OPTIONAL; + encoder.reusePolicy = ReusePolicy::OPTIONAL; auto decoder = encoder; decoder.layerId = 1; - decoder.reusePolicy = AttentionReusePolicy::PRIVATE; + decoder.reusePolicy = ReusePolicy::PRIVATE; config.layers.emplace_back(decoder); auto required = decoder; required.layerId = 2; - required.reusePolicy = AttentionReusePolicy::REQUIRED; + required.reusePolicy = ReusePolicy::REQUIRED; config.layers.emplace_back(required); config.initialPoolRatio = std::vector{0.25F, 0.5F, 0.25F}; auto manager = std::make_shared(config); @@ -302,7 +473,7 @@ TEST_F(ReplayCacheTest, OptionalWindowClaimsSurviveIndependentPageRemoval) config.commitMinSnapshot = true; for (int id = 1; id <= 2; ++id) { - addWindow(config, id, id == 1 ? c.window : 3, AttentionReusePolicy::OPTIONAL, id == 1 ? c.sinks : 0); + addWindow(config, id, id == 1 ? c.window : 3, ReusePolicy::OPTIONAL, id == 1 ? c.sinks : 0); } auto manager = std::make_shared(config); auto const encoder = manager->getLayerGroupId(LayerId{1}); @@ -381,7 +552,7 @@ TEST_F(ReplayCacheTest, IndependentOptionalDropPreservesMigrationPeer) auto config = makeConfig(); for (int i = 1; i <= 2; ++i) { - addWindow(config, i, i == 1 ? 4 : 2, AttentionReusePolicy::OPTIONAL); + addWindow(config, i, i == 1 ? 4 : 2, ReusePolicy::OPTIONAL); } auto manager = std::make_shared(config); auto const encoder = manager->getLayerGroupId(LayerId{1}); @@ -418,7 +589,7 @@ TEST_F(ReplayCacheTest, DecliningOptionalClaimReleasesCapacityBeforeAllocation) { auto config = makeConfig(); config.maxUtilForResume = 1.0f; - addWindow(config, 1, 4, AttentionReusePolicy::OPTIONAL); + addWindow(config, 1, 4, ReusePolicy::OPTIONAL); auto manager = std::make_shared(config); auto const encoder = manager->getLayerGroupId(LayerId{1}); auto const tokens = makeTokens(8); @@ -451,7 +622,7 @@ TEST_F(ReplayCacheTest, DecliningOptionalClaimReleasesCapacityBeforeAllocation) TEST_F(ReplayCacheTest, OptionalCommitAcceptsGlobalRebase) { auto config = makeConfig(); - addWindow(config, 1, 4, AttentionReusePolicy::OPTIONAL); + addWindow(config, 1, 4, ReusePolicy::OPTIONAL); auto manager = std::make_shared(config); auto const encoder = manager->getLayerGroupId(LayerId{1}); auto const tokens = makeTokens(8); diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py index 7d87e291667f..7e343063105d 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py @@ -32,11 +32,11 @@ from tensorrt_llm.runtime.kv_cache_manager_v2 import ( BAD_PAGE_INDEX, AttentionLayerConfig, - AttentionReusePolicy, BufferConfig, DataRole, LayerId, PageIndexMode, + ReusePolicy, ) from .....configs.deepseek_v41 import decoder_bounded_replay_enabled, encoder_replay_enabled @@ -228,11 +228,11 @@ def __init__( ) self._encoder_replay_enabled = encoder_replay_enabled() self._encoder_reuse_policy = ( - AttentionReusePolicy.PRIVATE + ReusePolicy.PRIVATE if not kv_cache_config.enable_block_reuse - else AttentionReusePolicy.OPTIONAL + else ReusePolicy.OPTIONAL if self._encoder_replay_enabled - else AttentionReusePolicy.REQUIRED + else ReusePolicy.REQUIRED ) self._decoder_replay_window = _get_decoder_replay_window( self.layout, num_layers, mapping, kv_cache_config, pretrained_config, spec_config @@ -284,10 +284,8 @@ def _log_checkpoint_pools(self): grouping = self.impl.get_life_cycle_pool_group_indices(level) for group, stats in enumerate(self.impl.get_storage_statistics(level)): members = [lc for lc, pg in enumerate(grouping) if pg == group] - optional_pages = sum( - policies[lc] == AttentionReusePolicy.OPTIONAL for lc in members - ) - names = sorted({AttentionReusePolicy(policies[lc]).name for lc in members}) + optional_pages = sum(policies[lc] == ReusePolicy.OPTIONAL for lc in members) + names = sorted({ReusePolicy(policies[lc]).name for lc in members}) slot_bytes = sum( stats.slot_sizes if hasattr(stats, "slot_sizes") else stats.slot_size ) @@ -552,7 +550,7 @@ def add( self.layout.kv_source_layer_ids ) private_config["reuse_policy"] = ( - AttentionReusePolicy.PRIVATE if decoder_swa else self._encoder_reuse_policy + ReusePolicy.PRIVATE if decoder_swa else self._encoder_reuse_policy ) layers.append( AttentionLayerConfig( @@ -643,7 +641,7 @@ def has_private_swa_suffix(self, first_layer: int) -> bool: self.impl.init_config.layers[ self._cache_layer_id(layer, CSA2CacheRole.SWA) ].reuse_policy - == AttentionReusePolicy.PRIVATE + == ReusePolicy.PRIVATE for layer in layers ) 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 fc2164f91cc8..f00f9d97d7d6 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 @@ -72,7 +72,6 @@ DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY, GPU_LEVEL, AttentionLayerConfig, - AttentionReusePolicy, AttnLifeCycle, Batch, BatchDesc, @@ -94,6 +93,7 @@ PlannedDropHandle, PoolGroupPeakBlockStats, PoolRebalanceConfig, + ReusePolicy, ReuseScope, SsmLayerConfig, SwaScratchReuseConfig, @@ -309,7 +309,7 @@ class KVCacheLayerSpec: buffer_roles: tuple[DataRole, ...] sliding_window_size: Optional[int] = None num_sink_tokens: Optional[int] = None - reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED + reuse_policy: ReusePolicy = ReusePolicy.REQUIRED def __post_init__(self) -> None: if self.global_layer_id < 0: @@ -4044,7 +4044,7 @@ def _add_mixed_precision_layers(self, config: KVCacheManagerConfigPy) -> KVCache # Native SWA includes the next query's write position. sliding_window_size=self.last_n + 1, num_sink_tokens=self.first_n, - reuse_policy=AttentionReusePolicy.REQUIRED, + reuse_policy=ReusePolicy.REQUIRED, ) ) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 3c710cb115a2..a02df21c8f35 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -56,7 +56,7 @@ def _load_cpp_module(): AggregatedPageDesc = _cpp.AggregatedPageDesc AttentionLayerConfig = _cpp.AttentionLayerConfig -AttentionReusePolicy = _cpp.AttentionReusePolicy +ReusePolicy = _cpp.ReusePolicy AttnLifeCycle = _cpp.AttnLifeCycle BatchDesc = _cpp.BatchDesc # BatchDesc is also consumed via dataclasses.replace(): MambaCacheManager's @@ -218,7 +218,7 @@ def typed_range(*args: int) -> range: __all__ = [ "AggregatedPageDesc", "AttentionLayerConfig", - "AttentionReusePolicy", + "ReusePolicy", "BAD_PAGE_INDEX", "CACHE_LEVEL1", "BatchDesc", diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 1211ecfd3d89..4cc8ef90ac12 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -100,21 +100,21 @@ class AttnLifeCycle: self, window_size: int | None, num_sink_blocks: int, - reuse_policy: AttentionReusePolicy = ..., + reuse_policy: ReusePolicy = ..., ) -> None: ... @staticmethod def make( window_size: int | None, num_sink_tokens: int | None, tokens_per_block: int, - reuse_policy: AttentionReusePolicy = ..., + reuse_policy: ReusePolicy = ..., ) -> "AttnLifeCycle": ... @property def window_size(self) -> int | None: ... @property def num_sink_blocks(self) -> int: ... @property - def reuse_policy(self) -> AttentionReusePolicy: ... + def reuse_policy(self) -> ReusePolicy: ... def get_stale_range(self, history_length: int, tokens_per_block: int) -> HalfOpenRange: ... CacheLevel = NewType("CacheLevel", int) @@ -394,7 +394,7 @@ class SparseCacheConfig: ``max_cached_units`` and ``unit_tokens``. """ -class AttentionReusePolicy(enum.IntEnum): +class ReusePolicy(enum.IntEnum): REQUIRED = 0 PRIVATE = 1 OPTIONAL = 2 @@ -413,7 +413,7 @@ class AttentionLayerConfig: buffers: list[BufferConfig] sliding_window_size: int | None = None num_sink_tokens: int | None = None - reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED + reuse_policy: ReusePolicy = ReusePolicy.REQUIRED sparse_cache: SparseCacheConfig | None = None """The sparse fetch path for this layer's sparse buffers, if it has one. @@ -425,8 +425,15 @@ class AttentionLayerConfig: @dataclass(slots=True) class SsmLayerConfig: + """Recurrent buffers sharing a layer-level reuse policy. + + OPTIONAL state is restored only when selected on first resume at the exact + matched endpoint. The caller initializes or reconstructs all fresh state. + """ + layer_id: LayerId buffers: list[BufferConfig] + reuse_policy: ReusePolicy = ReusePolicy.REQUIRED LayerConfig = AttentionLayerConfig | SsmLayerConfig @@ -657,7 +664,7 @@ class ReuseGroupStatus: def group_id(self) -> LayerGroupId: ... # REQUIRED constrains the match; OPTIONAL is selected at resume; PRIVATE is never reused. @property - def policy(self) -> AttentionReusePolicy: ... + def policy(self) -> ReusePolicy: ... # Exclusive token endpoint claimed by REQUIRED groups, shared by all groups. @property def endpoint(self) -> int: ... diff --git a/tests/unittest/_torch/attention/sparse/csa2/test_cache_manager.py b/tests/unittest/_torch/attention/sparse/csa2/test_cache_manager.py index 958c004d745b..fe0affaa2ee3 100644 --- a/tests/unittest/_torch/attention/sparse/csa2/test_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/csa2/test_cache_manager.py @@ -339,7 +339,7 @@ def test_cache_manager_preserves_disabled_or_enabled_scratch_config(enable_scrat @pytest.mark.parametrize("fallback", ["scratch", "attention_dp", "disaggregated", "late_engram"]) @pytest.mark.parametrize("partial_reuse", [False, True]) def test_decoder_replay_selects_cache_lifecycles(monkeypatch, fallback, partial_reuse): - from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy + from tensorrt_llm.runtime.kv_cache_manager_v2 import ReusePolicy monkeypatch.setenv("TRTLLM_V41_DECODER_BOUNDED_REPLAY", "1") monkeypatch.setenv("TRTLLM_V41_ENCODER_REPLAY", "0") @@ -378,12 +378,12 @@ def test_decoder_replay_selects_cache_lifecycles(monkeypatch, fallback, partial_ for layer in manager.impl.init_config.layers: decoder = layer.layer_id == manager._layer_roles[1, 1, CSA2CacheRole.SWA] assert layer.reuse_policy == ( - AttentionReusePolicy.PRIVATE if decoder else AttentionReusePolicy.REQUIRED + ReusePolicy.PRIVATE if decoder else ReusePolicy.REQUIRED ) else: assert manager._decoder_replay_window == 0 assert all( - layer.reuse_policy == AttentionReusePolicy.REQUIRED + layer.reuse_policy == ReusePolicy.REQUIRED for layer in manager.impl.init_config.layers ) finally: diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index ef1f2d68b17b..2d605e51b91e 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -51,7 +51,7 @@ TorchLlmArgs, ) from tensorrt_llm.mapping import CpType, Mapping -from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy +from tensorrt_llm.runtime.kv_cache_manager_v2 import ReusePolicy pytestmark = pytest.mark.cpu_only @@ -1262,7 +1262,7 @@ def test_v2_layer_spec_costing_preserves_sink_and_reuse_policy() -> None: ) layers = [ windowed, - replace(windowed, global_layer_id=1, reuse_policy=AttentionReusePolicy.PRIVATE), + replace(windowed, global_layer_id=1, reuse_policy=ReusePolicy.PRIVATE), replace(windowed, global_layer_id=2, num_sink_tokens=None), replace(windowed, global_layer_id=3, sliding_window_size=None, num_sink_tokens=None), ] diff --git a/tests/unittest/_torch/executor/kv_cache/test_mixed_precision_kv_cache.py b/tests/unittest/_torch/executor/kv_cache/test_mixed_precision_kv_cache.py index f1ba0fb27ddc..f7d0b47964b5 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mixed_precision_kv_cache.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mixed_precision_kv_cache.py @@ -227,7 +227,7 @@ def test_automatic_pool_sizing_defaults_to_half_max_seq_len(make_manager): def test_native_common_reuse_and_window_retention(make_manager, policy, reused_tokens): manager = make_manager(config=mixed_config(block_reuse_config={"policy": policy})) config = manager.kv_cache_manager_py_config - assert [layer.reuse_policy for layer in config.layers] == [kv.AttentionReusePolicy.REQUIRED] * 2 + assert [layer.reuse_policy for layer in config.layers] == [kv.ReusePolicy.REQUIRED] * 2 assert (config.layers[1].num_sink_tokens, config.layers[1].sliding_window_size) == (1, 66) assert config.commit_min_snapshot == (policy == "per_request") native = manager.impl diff --git a/tests/unittest/_torch/modeling/test_deepseek_v41_snapshots.py b/tests/unittest/_torch/modeling/test_deepseek_v41_snapshots.py index 9f8c41820854..2ec5b90a88e5 100644 --- a/tests/unittest/_torch/modeling/test_deepseek_v41_snapshots.py +++ b/tests/unittest/_torch/modeling/test_deepseek_v41_snapshots.py @@ -17,7 +17,7 @@ from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import SamplingConfig -from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy, _introspection +from tensorrt_llm.runtime.kv_cache_manager_v2 import ReusePolicy, _introspection pytestmark = skip_pre_blackwell @@ -72,18 +72,18 @@ def test_snapshot_policy_follows_reuse_and_encoder_recovery( with cache_case(config, sparse, dtype, [257], [0], 1, reuse=reuse) as case: manager = case[0].kv_cache_manager expected_encoder = ( - AttentionReusePolicy.PRIVATE + ReusePolicy.PRIVATE if not reuse - else AttentionReusePolicy.OPTIONAL + else ReusePolicy.OPTIONAL if replay - else AttentionReusePolicy.REQUIRED + else ReusePolicy.REQUIRED ) for layer in manager.impl.init_config.layers: owner, _, role = manager._physical_roles[layer.layer_id, layer.buffers[0].role] expected = ( - AttentionReusePolicy.REQUIRED + ReusePolicy.REQUIRED if role == CSA2CacheRole.GLOBAL - else AttentionReusePolicy.PRIVATE + else ReusePolicy.PRIVATE if role == CSA2CacheRole.SWA and owner >= 5 else expected_encoder ) @@ -99,7 +99,7 @@ def test_snapshot_policy_follows_reuse_and_encoder_recovery( expected_groups = { manager.impl.get_layer_group_id(layer.layer_id) for layer in manager.impl.init_config.layers - if layer.reuse_policy == AttentionReusePolicy.OPTIONAL + if layer.reuse_policy == ReusePolicy.OPTIONAL } assert set(manager._encoder_optional_groups) == expected_groups assert ( diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_concurrency.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_concurrency.py index ebd69bbd6934..ded454f5127e 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_concurrency.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_concurrency.py @@ -37,12 +37,12 @@ from tensorrt_llm.runtime.kv_cache_manager_v2 import ( GPU_LEVEL, AttentionLayerConfig, - AttentionReusePolicy, BufferConfig, CudaStream, GpuCacheTierConfig, KVCacheManager, KVCacheManagerConfig, + ReusePolicy, ) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -188,7 +188,7 @@ def test_optional_selection_and_resume_with_concurrent_priority_callbacks() -> N layer_id=1, buffers=[BufferConfig(role="key", size=4096)], sliding_window_size=4, - reuse_policy=AttentionReusePolicy.OPTIONAL, + reuse_policy=ReusePolicy.OPTIONAL, ), ] manager = KVCacheManager(config) 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 7705f8fe097c..83f7317e5169 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 @@ -23,9 +23,7 @@ def test_reuse_status_and_optional_selection_bindings(priority): layer_id=i, buffers=[kv.BufferConfig(role="key", size=4096)], sliding_window_size=None if i == 0 else 4 if i == 1 else 2, - reuse_policy=kv.AttentionReusePolicy.OPTIONAL - if i - else kv.AttentionReusePolicy.REQUIRED, + reuse_policy=kv.ReusePolicy.OPTIONAL if i else kv.ReusePolicy.REQUIRED, ) for i in range(3) ] @@ -91,3 +89,67 @@ def test_reuse_status_and_optional_selection_bindings(priority): for cache in reversed(caches): cache.close() manager.shutdown() + + +@pytest.mark.parametrize("length", [7, 8]) +@pytest.mark.parametrize("select_optional", [False, True]) +def test_ssm_layer_reuse_policies(length, select_optional): + layers = [ + kv.AttentionLayerConfig(layer_id=0, buffers=[kv.BufferConfig(role="key", size=4096)]), + *[ + kv.SsmLayerConfig( + layer_id=i, + buffers=[ + kv.BufferConfig(role="state", size=4096), + kv.BufferConfig(role="conv", size=4096), + ], + reuse_policy=policy, + ) + for i, policy in enumerate( + [kv.ReusePolicy.REQUIRED, kv.ReusePolicy.OPTIONAL, kv.ReusePolicy.PRIVATE], + start=1, + ) + ], + ] + config = kv.KVCacheManagerConfig( + tokens_per_block=4, + cache_tiers=[kv.GpuCacheTierConfig(quota=8 << 20)], + layers=layers, + commit_min_snapshot=True, + ) + manager = kv.KVCacheManager(config) + stream = torch.cuda.Stream() + tokens = list(range(length)) + caches = [] + try: + source = manager.create_kv_cache(input_tokens=[]) + caches.append(source) + assert source.resume(stream.cuda_stream) + assert source.resize(length) + source.commit(tokens, is_end=True) + source.close() + required, optional, private = [manager.get_layer_group_id(i) for i in (1, 2, 3)] + assert len({required, optional, private}) == 3 + reader = manager.create_kv_cache(input_tokens=tokens) + caches.append(reader) + assert reader.num_committed_tokens == length + assert reader.reuse_status[required].policy == kv.ReusePolicy.REQUIRED + assert reader.reuse_status[optional].complete + assert list(reader.reuse_status[optional].coverage) == [(length, length)] + assert not reader.reuse_status[private].complete + assert reader.resume( + stream.cuda_stream, optional_reuse_groups=[optional] if select_optional else [] + ) + indices = [ + reader.get_ssm_block_base_index(group) for group in (required, optional, private) + ] + assert all(index >= 0 for index in indices) + reader.suspend() + assert reader.resume(stream.cuda_stream) + assert indices == [ + reader.get_ssm_block_base_index(group) for group in (required, optional, private) + ] + finally: + for cache in reversed(caches): + cache.close() + manager.shutdown() From e6a1dd7b5581ddaff93373ead89e187ee8a1d99f Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Tue, 27 Oct 2026 16:17:35 +0800 Subject: [PATCH 02/19] [None][fix] Exclude unselected optional SSM groups from reuse statistics Signed-off-by: Jiagan Cheng --- .../kv_cache_manager_v2/kvCache.cpp | 13 +++++++------ .../kvCacheManagerV2StatsTest.cpp | 18 ++++++++++++++++++ 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index 5b5e09b9e68c..dff795291bb5 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 @@ -868,15 +868,16 @@ bool KvCache::resume(std::optional stream, std::optionallifeCycles().ssmLifeCycleIds()) + for (auto lc : mSelectedOptionalGroups) { - if (reusePolicy(mManager->lifeCycles()[lc]) != ReusePolicy::OPTIONAL) + if (!mManager->lifeCycles().isSsm(lc)) + { continue; - bool const selected = std::find(mSelectedOptionalGroups.begin(), mSelectedOptionalGroups.end(), lc) - != mSelectedOptionalGroups.end(); - if (mPendingStats.recordSsmSnapshotLookup( - lc, mNumLookupTokens, selected ? numCommittedTokens() : 0, mTokensPerBlock)) + } + if (mPendingStats.recordSsmSnapshotLookup(lc, mNumLookupTokens, numCommittedTokens(), mTokensPerBlock)) + { mManager->markStatsDirty(id); + } } } mOptionalCandidates.clear(); diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index 65ddd0404b7e..da6ea500e9f2 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -290,6 +290,18 @@ TEST_F(ReplayCacheTest, MixedSsmPoliciesRestoreStateAndCopyEveryBeam) EXPECT_THROW( reader->resume(mStream, std::nullopt, std::vector{privateGroup}), std::invalid_argument); ASSERT_TRUE(reader->resume(mStream, std::nullopt, std::vector{optional})); + reader->commitPendingStats(); + auto const snapshotStats = manager->getAndResetSsmSnapshotIterationStats(); + ASSERT_EQ(snapshotStats.size(), 2); + for (auto group : {required, optional}) + { + ASSERT_EQ(snapshotStats.count(group), 1); + auto const& stats = snapshotStats.at(group); + EXPECT_EQ(stats.iterSnapshotLookups, 1); + EXPECT_EQ(stats.iterSnapshotHits, 1); + EXPECT_EQ(stats.iterSnapshotMisses, 0); + EXPECT_EQ(stats.iterReusedTokens, length); + } for (int layer : {1, 2, 4}) { check(*reader, layer, "state", 0, 10 + layer); @@ -311,6 +323,8 @@ TEST_F(ReplayCacheTest, MixedSsmPoliciesRestoreStateAndCopyEveryBeam) reader->suspend(); ASSERT_TRUE(reader->resume(mStream)); reader->setBeamWidth(BeamIndex{1}); + reader->commitPendingStats(); + EXPECT_TRUE(manager->getAndResetSsmSnapshotIterationStats().empty()); reader->close(); source->close(); manager->shutdown(); @@ -348,8 +362,12 @@ TEST_F(ReplayCacheTest, OptionalSsmRequiresExactEndpointAndDoesNotLimitAttention // Unselected OPTIONAL and PRIVATE state both get writable request-local storage. ASSERT_TRUE(reader->resume(mStream)); EXPECT_GE(reader->getSsmBlockBaseIndex(group), 0); + reader->commitPendingStats(); + EXPECT_TRUE(manager->getAndResetSsmSnapshotIterationStats().empty()); reader->suspend(); ASSERT_TRUE(reader->resume(mStream)); + reader->commitPendingStats(); + EXPECT_TRUE(manager->getAndResetSsmSnapshotIterationStats().empty()); reader->close(); } manager->shutdown(); From 46cd9214d5b77a400c4c39b988e87399adfbc83d Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Tue, 27 Oct 2026 16:26:39 +0800 Subject: [PATCH 03/19] [None][doc] Explain SSM allocation branch ordering Signed-off-by: Jiagan Cheng --- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index dff795291bb5..8cb2973f69d7 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 @@ -599,6 +599,8 @@ bool KvCache::resume(std::optional stream, std::optionallifeCycles().isSsm(lc); auto const& lifeCycle = mManager->lifeCycles()[lc]; + // Every SSM policy needs one writable state slot. Handle SSM first because private and + // unselected optional groups also need fresh state, but must not allocate per-token-block pages. if (isSsm) { numSlotsNeeded[lc] += 1; From ff87edcc25c1050f44bda878346524e91fc25a85 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Tue, 27 Oct 2026 16:42:29 +0800 Subject: [PATCH 04/19] [None][refactor] Simplify required SSM cache level assignment Signed-off-by: Jiagan Cheng --- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 8cb2973f69d7..69acce5b52d0 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 @@ -2977,7 +2977,7 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match, bool claim auto& snapshotBlock = *matched.back(); auto* snapshotPage = snapshotBlock.storage[lcId]; TLLM_CHECK_WITH_INFO(snapshotPage, "Last matched block must have SSM snapshot after truncation"); - ssmLevel = std::max(ssmLevel.value_or(kHotLevel), snapshotPage->cacheLevel); + ssmLevel = snapshotPage->cacheLevel; mSsmBlocks[kDefaultBeamIndex][lcId] = snapshotPage->hold(); if (claimRetention) { From 13a5c2649224b43f76f78f7a7c65bef5c45d037c Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 27 Oct 2026 10:42:30 +0000 Subject: [PATCH 05/19] [None][feat] Support buffer-level reuse policy for SSM layers (By Agent) Add an optional per-buffer reuse policy. A BufferConfig without one uses its layer's policy. An SSM buffer may set a different policy and then belongs to the SSM lifecycle of that policy while keeping its layer ID, so one SSM layer can span several layer groups. Attention buffers must use their layer's policy. Lifecycle registration and the storage layout are built per buffer, and per-layer storage attributes are keyed by (layer, lifecycle), so page index converters use the buffer's own lifecycle. Add getLayerGroupId(layerId, role) and its get_layer_group_id(layer_id, role) binding. The one-argument overload throws when a layer's buffers span several groups; layerGrouping lists such a layer under each group. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../kv_cache_manager_v2/AGENTS.md | 14 +- .../kv_cache_manager_v2/config.h | 43 ++- .../kv_cache_manager_v2/kvCacheManager.cpp | 25 +- .../kv_cache_manager_v2/kvCacheManager.h | 9 +- .../kv_cache_manager_v2/lifeCycleRegistry.cpp | 27 +- .../kv_cache_manager_v2/lifeCycleRegistry.h | 9 +- .../kv_cache_manager_v2/storage/config.cpp | 52 ++- .../kv_cache_manager_v2/storage/config.h | 16 +- .../kv_cache_manager_v2/storageManager.cpp | 8 +- .../kv_cache_manager_v2/storageManager.h | 14 +- .../batch_manager/kvCacheManagerV2.cpp | 15 +- .../kvCacheManagerV2ConfigTest.cpp | 93 +++++ .../kvCacheManagerV2StatsTest.cpp | 328 ++++++++++++++++++ .../runtime/kv_cache_manager_v2/__init__.pyi | 10 +- .../test_optional_checkpoint.py | 87 +++++ 15 files changed, 674 insertions(+), 76 deletions(-) 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 cc64681c09a4..44602a8f2116 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 @@ -93,14 +93,20 @@ must eventually be `close()`d. Its normal flow is: for later requests; `stopCommitting()` finalizes this process. 5. Suspend or close the request, returning pages to holding/eviction ownership. -`LifeCycleRegistry` maps model layers to lifecycle groups. Attention lifecycles +`LifeCycleRegistry` maps layer buffers to lifecycle groups. Attention lifecycles may have sliding-window and sink-token rules; SSM lifecycles represent a recurrent-state checkpoint. Both layer configs expose `reusePolicy`: `REQUIRED` constrains the prefix match, `OPTIONAL` offers state at the required endpoint for explicit selection on first resume, and `PRIVATE` never publishes reusable pages. -All buffers within a layer share that policy. SSM lifecycles are deduplicated by -policy, so a manager can have up to three SSM groups; iterate `ssmLifeCycleIds()` -for working state and use `requiredSsmLifeCycleId()` only for prefix pruning. +Attention buffers use their layer's policy. An SSM buffer may override it through +`BufferConfig::reusePolicy` and then belongs to the SSM lifecycle of its own +policy while keeping its `LayerId`, so one SSM layer can span several layer +groups. Look such a layer up per buffer with `getLayerGroupId(layerId, role)`; +the one-argument overload throws for it, `layerGrouping()` lists it under each +of its groups, and per-layer storage attributes are keyed by (layer, lifecycle). +SSM lifecycles are deduplicated by policy, so a manager can have up to three SSM +groups; iterate `ssmLifeCycleIds()` for working state and use +`requiredSsmLifeCycleId()` only for prefix pruning. Optional SSM checkpoints require an exact endpoint match, including partial blocks. Fresh or unselected recurrent state must be initialized or reconstructed by the caller before inference. A pool-group index is a storage-layout index and is 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 869ea3f065b4..9eff68e19c3c 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 @@ -130,8 +130,21 @@ inline std::optional cacheTierMaxQuota(CacheTierConfig const& cfg) return host != nullptr ? host->maxQuota : std::nullopt; } +//! Controls whether a lifecycle participates in prefix matching and page reuse. +enum class ReusePolicy +{ + //! Commit reusable pages and require live-window coverage when matching a prefix. + REQUIRED = 0, + //! Keep request-local working pages; never publish or reuse them through the radix tree. + PRIVATE = 1, + //! Commit reusable pages without constraining the REQUIRED prefix match. The caller selects + //! available groups at resume; missing or unselected groups use fresh pages for recomputation. + //! These pages may be evicted independently without removing the REQUIRED prefix. + OPTIONAL = 2, +}; + // --------------------------------------------------------------------------- -// Buffer configuration (one KV buffer inside an attention layer). +// Buffer configuration (one buffer inside an attention or SSM layer). // --------------------------------------------------------------------------- // --------------------------------------------------------------------------- @@ -193,6 +206,14 @@ struct BufferConfig //! densely packed, so a padded page is not expressible. int sparseUnitChunks = 1; + //! Reuse policy of this buffer; nullopt uses the layer's policy. + std::optional reusePolicy; + + [[nodiscard]] ReusePolicy effectiveReusePolicy(ReusePolicy layerPolicy) const noexcept + { + return reusePolicy.value_or(layerPolicy); + } + void validate() const; //! Where this buffer's selection units sit inside a page, given how many the layer @@ -292,19 +313,6 @@ inline void validateNoDuplicateBufferRoles(std::vector const& buff } // namespace detail -//! Controls whether a layer lifecycle participates in prefix matching and page reuse. -enum class ReusePolicy -{ - //! Commit reusable pages and require live-window coverage when matching a prefix. - REQUIRED = 0, - //! Keep request-local working pages; never publish or reuse them through the radix tree. - PRIVATE = 1, - //! Commit reusable pages without constraining the REQUIRED prefix match. The caller selects - //! available groups at resume; missing or unselected groups use fresh pages for recomputation. - //! These pages may be evicted independently without removing the REQUIRED prefix. - OPTIONAL = 2, -}; - struct AttentionLayerConfig { static constexpr LayerType type = LayerType::ATTENTION; @@ -319,6 +327,7 @@ struct AttentionLayerConfig std::optional numSinkTokens; // Private working pages never constrain prefix matching or enter the tree. + // Every buffer of an attention layer uses this policy. ReusePolicy reusePolicy = ReusePolicy::REQUIRED; //! The sparse fetch path for this layer's sparse buffers, if it has one. std::optional sparseCache; @@ -341,6 +350,7 @@ struct SsmLayerConfig LayerId layerId = 0; std::vector buffers; + // Default policy of the layer's buffers; BufferConfig::reusePolicy may override it per buffer. ReusePolicy reusePolicy = ReusePolicy::REQUIRED; void validate() const @@ -363,6 +373,11 @@ struct SsmLayerConfig using LayerConfig = std::variant; +inline std::vector const& layerBuffers(LayerConfig const& layer) +{ + return std::visit([](auto const& cfg) -> std::vector const& { return cfg.buffers; }, layer); +} + // --------------------------------------------------------------------------- // KVCacheDesc — describes one KV cache request's capacity and history length. // Mirrors _config.py::KVCacheDesc. 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 77a7e427b91f..e72fc23ff140 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 @@ -32,6 +32,7 @@ #include #include #include +#include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { @@ -602,7 +603,7 @@ int KvCacheManager::getPageIndexScale(LayerId layerId, DataRole role) const PageIndexConverter KvCacheManager::getPageIndexConverter(LayerId layerId, DataRole role) const { auto const& attr = mStorage->getBufferAttr(layerId, role); - auto const& layerAttr = mStorage->getLayerAttr(layerId); + auto const& layerAttr = mStorage->getLayerAttr(layerId, attr.lifeCycleId); int scale = mStorage->mSlotToPageIndices.at(attr.lifeCycleId).at(attr.poolIndex); int offset = exactDiv(static_cast(attr.offset), static_cast(attr.size)); int scratchPages = layerAttr.slotUtil.at(attr.poolIndex); @@ -800,16 +801,32 @@ std::vector KvCacheManager::layerIds() const LayerGroupId KvCacheManager::getLayerGroupId(LayerId layerId) const { - return mStorage->layerToLifeCycleIds().at(layerId); + auto const& lifeCycles = mStorage->layerToLifeCycleIds().at(layerId); + if (lifeCycles.size() != 1) + { + throw std::invalid_argument("Layer " + std::to_string(layerId) + " has buffers in " + + std::to_string(lifeCycles.size()) + + " layer groups because its buffers use different reuse policies; look up its layer group per " + "buffer with getLayerGroupId(layerId, role)"); + } + return lifeCycles.front(); +} + +LayerGroupId KvCacheManager::getLayerGroupId(LayerId layerId, DataRole const& role) const +{ + return mStorage->getBufferAttr(layerId, role).lifeCycleId; } TypedVec> KvCacheManager::layerGrouping() const { LifeCycleId numLc = mLifeCycles.size(); TypedVec> result(numLc); - for (auto const& [lid, lc] : mStorage->layerToLifeCycleIds()) + for (auto const& [lid, lifeCycles] : mStorage->layerToLifeCycleIds()) { - result.at(lc).push_back(lid); + for (auto const lc : lifeCycles) + { + result.at(lc).push_back(lid); + } } return result; } 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 3d2a9ec9f36f..49a1cb56054c 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 @@ -285,9 +285,16 @@ class KvCacheManager : public std::enable_shared_from_this int numLayers() const noexcept; std::vector layerIds() const; + + // Layer group of every buffer of a layer. Throws std::invalid_argument when the layer's SSM + // buffers use different reuse policies and therefore span several layer groups; look up such a + // layer per buffer with getLayerGroupId(layerId, role). LayerGroupId getLayerGroupId(LayerId layerId) const; + // Layer group of one buffer. Throws std::out_of_range for an unknown buffer. + LayerGroupId getLayerGroupId(LayerId layerId, DataRole const& role) const; - // Layer grouping: layers with the same lifecycle share pool allocation. + // Layer grouping: layers with the same lifecycle share pool allocation. A layer whose + // buffers span several layer groups is listed under each of them. // NOTE: the iteration order of the layer lists (and of the groups) is NOT // part of the API contract and may differ across backends/runs. Do not rely // on it to infer buffer/pool memory order — query poolGroupDescs() diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp index 0f6a387e5bd1..9fd6966be474 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp @@ -46,15 +46,25 @@ LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock) layer); } +LifeCycle makeLifeCycle(LayerConfig const& layer, BufferConfig const& buffer, int tokensPerBlock) +{ + LifeCycle lc = makeLifeCycle(layer, tokensPerBlock); + // Attention layers reject buffer policies that differ from the layer's, so only SSM lifecycles change here. + if (auto* ssm = std::get_if(&lc)) + { + ssm->reusePolicy = buffer.effectiveReusePolicy(ssm->reusePolicy); + } + return lc; +} + // --------------------------------------------------------------------------- // LifeCycleRegistry // --------------------------------------------------------------------------- LifeCycleRegistry::LifeCycleRegistry(KVCacheManagerConfig const& config) { - for (auto const& layer : config.layers) + auto const registerLifeCycle = [this](LifeCycle const& lc) { - LifeCycle lc = makeLifeCycle(layer, config.tokensPerBlock); if (mLifeCycleIdMap.find(lc) == mLifeCycleIdMap.end()) { check(); @@ -64,6 +74,19 @@ LifeCycleRegistry::LifeCycleRegistry(KVCacheManagerConfig const& config) if (std::holds_alternative(lc)) mSsmLifeCycleIds.push_back(id); } + }; + for (auto const& layer : config.layers) + { + auto const& buffers = layerBuffers(layer); + // A layer without buffers registers its layer-level lifecycle. + if (buffers.empty()) + { + registerLifeCycle(makeLifeCycle(layer, config.tokensPerBlock)); + } + for (auto const& buffer : buffers) + { + registerLifeCycle(makeLifeCycle(layer, buffer, config.tokensPerBlock)); + } } check(); } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h index 8987165fa0d4..1aeeaad4b1af 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h @@ -101,9 +101,9 @@ struct AttnLifeCycle }; // --------------------------------------------------------------------------- -// SsmLifeCycle — lifecycle for SSM (State Space Model) layers. +// SsmLifeCycle — lifecycle for SSM (State Space Model) buffers. // All blocks before the last full block are stale (recurrent state). -// SSM layers with the same reuse policy share a lifecycle. +// SSM buffers with the same effective reuse policy share a lifecycle. // --------------------------------------------------------------------------- struct SsmLifeCycle { @@ -171,9 +171,12 @@ inline HalfOpenRange computeScratchRange( return intersect(capStale, inputRange); } -// Factory: create a LifeCycle from a LayerConfig variant. +// Factory: create the LifeCycle of a layer, using the layer's reuse policy. LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock); +// Factory: create the LifeCycle that holds one buffer of a layer, using the buffer's effective reuse policy. +LifeCycle makeLifeCycle(LayerConfig const& layer, BufferConfig const& buffer, int tokensPerBlock); + // Integer id assigned to each unique LifeCycle. using LifeCycleId = StrongIndex; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp index bf28b4708913..78791ce224a2 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp @@ -86,23 +86,27 @@ std::map StorageConfig::bufferAttributes() const return ret; } -std::unordered_map StorageConfig::layerToLifeCycleIds() const +std::unordered_map> StorageConfig::layerToLifeCycleIds() const { - std::unordered_map map; + std::unordered_map> map; for (auto const& [bufId, attr] : bufferAttributes()) { - auto [it, inserted] = map.emplace(bufId.layerId, attr.lifeCycleId); - if (!inserted) + auto& ids = map[bufId.layerId]; + if (std::find(ids.begin(), ids.end(), attr.lifeCycleId) == ids.end()) { - TLLM_CHECK_DEBUG(it->second == attr.lifeCycleId); + ids.push_back(attr.lifeCycleId); } } + for (auto& [layerId, ids] : map) + { + std::sort(ids.begin(), ids.end()); + } return map; } -std::map StorageConfig::layerAttributes() const +std::map StorageConfig::layerAttributes() const { - std::map ret; + std::map ret; for (auto const& pg : slotDescList) { for (auto const& variant : pg.variants) @@ -121,17 +125,17 @@ std::map StorageConfig::layerAttributes() const for (auto const& [layerId, count] : slotUtilPerLayer) { - auto it = ret.find(layerId); + LayerLifeCycleKey const key{layerId, lcId}; + auto it = ret.find(key); if (it == ret.end()) { LayerAttr attr; attr.lifeCycleId = lcId; attr.slotUtil.resize(variant.coalescedBuffers.size(), 0); attr.slotUtilFracMax = Rational{0, 1}; - it = ret.emplace(layerId, std::move(attr)).first; + it = ret.emplace(key, std::move(attr)).first; } auto& attr = it->second; - TLLM_CHECK_DEBUG(attr.lifeCycleId == lcId); attr.slotUtil[poolIndex] = count; Rational frac{count, buffersPerSlot}; if (frac > attr.slotUtilFracMax) @@ -167,24 +171,18 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) for (auto const& layer : config.layers) { - LifeCycle lc = makeLifeCycle(layer, tokensPerBlock); - LifeCycleId lcId = registry.getId(lc); - - std::visit( - [&](auto const& cfg) - { - for (auto const& buf : cfg.buffers) - { - int tpbo = buf.tokensPerBlockOverride.value_or(tokensPerBlock); - int exp = exactDiv(tokensPerBlock, tpbo); - size_t expandedSize = buf.size * static_cast(exp); + LayerId const layerId = std::visit([](auto const& cfg) { return cfg.layerId; }, layer); + for (auto const& buf : layerBuffers(layer)) + { + LifeCycleId const lcId = registry.getId(makeLifeCycle(layer, buf, tokensPerBlock)); + int tpbo = buf.tokensPerBlockOverride.value_or(tokensPerBlock); + int exp = exactDiv(tokensPerBlock, tpbo); + size_t expandedSize = buf.size * static_cast(exp); - BufferId bid{cfg.layerId, buf.role}; - expansionMap[bid] = exp; - bufferGroups[lcId][{buf.coalescingGroup, expandedSize, buf.isSparse}].push_back(bid); - } - }, - layer); + BufferId bid{layerId, buf.role}; + expansionMap[bid] = exp; + bufferGroups[lcId][{buf.coalescingGroup, expandedSize, buf.isSparse}].push_back(bid); + } } // Build one SlotDescVariant per life cycle. diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h index 8be1a9eccf6e..0d5695b57950 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h @@ -28,6 +28,7 @@ #include #include #include +#include #include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 @@ -228,7 +229,8 @@ struct Rational }; // --------------------------------------------------------------------------- -// LayerAttr — per-layer storage attributes for scratch slot management. +// LayerAttr — storage attributes of one layer's buffers within one lifecycle, for scratch slot +// management. An SSM layer whose buffers use different reuse policies has one LayerAttr per lifecycle. // Mirrors _storage/_config.py::LayerAttr. // --------------------------------------------------------------------------- struct LayerAttr @@ -243,6 +245,9 @@ struct LayerAttr Rational slotUtilFracMax; }; +// (layer, lifecycle) key of a LayerAttr. +using LayerLifeCycleKey = std::pair; + // --------------------------------------------------------------------------- // StorageConfig — complete storage layout for the KV cache. // Mirrors _storage/_config.py::StorageConfig. @@ -264,12 +269,13 @@ struct StorageConfig // Attribute map for each buffer. std::map bufferAttributes() const; - // Map LayerId → LifeCycleId. - std::unordered_map layerToLifeCycleIds() const; + // Map LayerId → the lifecycles of its buffers, in ascending order. A layer spans several + // lifecycles only when its SSM buffers use different reuse policies. + std::unordered_map> layerToLifeCycleIds() const; - // Per-layer storage attributes (slot utilization within coalesced slots). + // Per-(layer, lifecycle) storage attributes (slot utilization within coalesced slots). // Mirrors _storage/_config.py::StorageConfig.layer_attributes(). - std::map layerAttributes() const; + std::map layerAttributes() const; LifeCycleId numLifeCycles() const; }; 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 afae999e6449..69dc94115906 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 @@ -250,7 +250,7 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi // Compute layer attributes and slot utilization fractions for scratch support. mLayerAttributes = config.layerAttributes(); mSlotUtilFracMax.resize(lifeCycles.size(), Rational{0, 1}); - for (auto const& [layerId, layerAttr] : mLayerAttributes) + for (auto const& [key, layerAttr] : mLayerAttributes) { LifeCycleId const lcIdx = layerAttr.lifeCycleId; if (layerAttr.slotUtilFracMax > mSlotUtilFracMax[lcIdx]) @@ -1629,11 +1629,11 @@ TypedVec StorageManager::poolGroupDescs() const return result; } -LayerAttr const& StorageManager::getLayerAttr(LayerId layerId) const +LayerAttr const& StorageManager::getLayerAttr(LayerId layerId, LifeCycleId lifeCycle) const { - auto it = mLayerAttributes.find(layerId); + auto it = mLayerAttributes.find(LayerLifeCycleKey{layerId, lifeCycle}); if (it == mLayerAttributes.end()) - throw std::out_of_range("Unknown LayerId for LayerAttr"); + throw std::out_of_range("Unknown (LayerId, LifeCycleId) for LayerAttr"); return it->second; } 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 fd4b8366fd10..7f0da2630b3a 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 @@ -299,8 +299,8 @@ class StorageManager : public std::enable_shared_from_this // Pool group base address without per-layer offset. MemAddress getMemPoolBaseAddress(PoolGroupIndex pgIdx, PoolIndex poolIdx) const; - // Per-layer storage attributes. - LayerAttr const& getLayerAttr(LayerId layerId) const; + // Storage attributes of a layer's buffers within one lifecycle. + LayerAttr const& getLayerAttr(LayerId layerId, LifeCycleId lifeCycle) const; // Address of a slot's buffer in a specific pool at a cache level. Address slotAddress(CacheLevel level, PoolGroupIndex pgIdx, SlotId slotId, PoolIndex poolIdx) const; @@ -331,8 +331,8 @@ class StorageManager : public std::enable_shared_from_this // Pool-group slot count (number of pages). [[nodiscard]] SlotCount numSlots(PoolGroupIndex pgIdx, CacheLevel level = kHotLevel) const; - // Layer-to-lifecycle mapping (for KvCacheManager queries). - std::unordered_map const& layerToLifeCycleIds() const noexcept + // Layer-to-lifecycles mapping (for KvCacheManager queries); see StorageConfig::layerToLifeCycleIds(). + std::unordered_map> const& layerToLifeCycleIds() const noexcept { return mLayerToLifeCycleIds; } @@ -477,13 +477,13 @@ class StorageManager : public std::enable_shared_from_this TypedVec mBatchingLayerGroupIds; // Codec-selected PageIndexPair memory location for each lifecycle. TypedVec mPageIndexLocations; - std::unordered_map mLayerToLifeCycleIds; + std::unordered_map> mLayerToLifeCycleIds; // slot-to-page-index scale factors: [lcId][poolIdx] TypedVec> mSlotToPageIndices; - // Per-layer storage attributes for scratch slot management. - std::map mLayerAttributes; + // Per-(layer, lifecycle) storage attributes for scratch slot management. + std::map mLayerAttributes; // Max slot utilization fraction per lifecycle (across all layers in that lifecycle). TypedVec mSlotUtilFracMax; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index c09f06d608d1..a2d9e513aafb 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1722,15 +1722,16 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) "may name a layer that named nothing.") DEF_COPY(kv::SparseCacheConfig); nb::class_(m, "BufferConfig") - .def(nb::init, bool, int, int>(), nb::arg("role"), nb::arg("size"), + .def(nb::init, bool, int, int, std::optional>(), nb::arg("role"), nb::arg("size"), nb::arg("tokens_per_block_override") = std::nullopt, nb::arg("is_sparse") = false, - nb::arg("coalescing_group") = 0, nb::arg("sparse_unit_chunks") = 1) + nb::arg("coalescing_group") = 0, nb::arg("sparse_unit_chunks") = 1, nb::arg("reuse_policy") = std::nullopt) .def_rw("role", &kv::BufferConfig::role) .def_rw("size", &kv::BufferConfig::size) .def_rw("tokens_per_block_override", &kv::BufferConfig::tokensPerBlockOverride) .def_rw("is_sparse", &kv::BufferConfig::isSparse) .def_rw("coalescing_group", &kv::BufferConfig::coalescingGroup) - .def_rw("sparse_unit_chunks", &kv::BufferConfig::sparseUnitChunks) DEF_COPY(kv::BufferConfig); + .def_rw("sparse_unit_chunks", &kv::BufferConfig::sparseUnitChunks) + .def_rw("reuse_policy", &kv::BufferConfig::reusePolicy) DEF_COPY(kv::BufferConfig); nb::class_(m, "AttentionLayerConfig") .def(nb::init, std::optional, std::optional, @@ -3664,7 +3665,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def_prop_ro("layer_ids", &kv::KvCacheManager::layerIds) .def_prop_ro( "layer_grouping", [](kv::KvCacheManager const& self) { return self.layerGrouping().raw(); }, - "Layers grouped by shared lifecycle/pool allocation. The iteration order of the " + "Layers grouped by shared lifecycle/pool allocation. A layer whose buffers span several " + "layer groups is listed under each of them. The iteration order of the " "layer lists (and of the groups) is NOT an API contract and may differ across " "backends/runs; do not rely on it for buffer/pool memory order -- use " "pool_group_descs (PoolGroupDesc.pools[i].base_address + coalesced_buffers) instead.") @@ -3672,6 +3674,11 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) "get_layer_group_id", [](kv::KvCacheManager const& self, kv::LayerId layerId) { return self.getLayerGroupId(layerId).value(); }, nb::arg("layer_id")) + .def( + "get_layer_group_id", + [](kv::KvCacheManager const& self, kv::LayerId layerId, kv::DataRole const& role) + { return self.getLayerGroupId(layerId, role).value(); }, + nb::arg("layer_id"), nb::arg("role"), "Layer group of one buffer of the layer.") .def( "get_sparse_cache_max_units", [](kv::KvCacheManager const& self, kv::LayerId layerId) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp index 78995508390d..b5c6d5534f3b 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp @@ -23,7 +23,9 @@ #include #include +#include #include +#include namespace { @@ -169,4 +171,95 @@ TEST(KvCacheManagerV2ConfigTest, SparseBuffersCoalesceBySize) EXPECT_EQ(second.expansion, 1); } +TEST(KvCacheManagerV2ConfigTest, AttentionBuffersMustUseTheirLayerReusePolicy) +{ + auto config = makeConfig(); + auto& buffer = std::get(config.layers.front()).buffers.front(); + buffer.reusePolicy = ReusePolicy::REQUIRED; + EXPECT_NO_THROW(config.validate()); + EXPECT_NO_THROW(createStorageConfig(config)); + for (auto const policy : {ReusePolicy::PRIVATE, ReusePolicy::OPTIONAL}) + { + buffer.reusePolicy = policy; + EXPECT_THROW(config.validate(), std::invalid_argument); + EXPECT_THROW(createStorageConfig(config), std::invalid_argument); + } +} + +TEST(KvCacheManagerV2ConfigTest, SsmBufferReusePoliciesSplitLayerLifecycles) +{ + auto config = makeTieredConfig(); + config.commitMinSnapshot = true; + for (LayerId layerId = 1; layerId <= 2; ++layerId) + { + config.layers.push_back(SsmLayerConfig{.layerId = layerId, + .buffers = {BufferConfig{.role = "state", .size = 4096}, + BufferConfig{.role = "replay", .size = 2048, .reusePolicy = ReusePolicy::PRIVATE}, + BufferConfig{.role = "draft", .size = 2048, .reusePolicy = ReusePolicy::OPTIONAL}}}); + } + // A layer-level PRIVATE layer shares the lifecycle of PRIVATE buffer overrides. + config.layers.push_back(SsmLayerConfig{ + .layerId = 3, .buffers = {BufferConfig{.role = "replay", .size = 2048}}, .reusePolicy = ReusePolicy::PRIVATE}); + ASSERT_NO_THROW(config.validate()); + + LifeCycleRegistry const registry(config); + auto const storage = createStorageConfig(config); + auto const attributes = storage.bufferAttributes(); + auto lifeCycle = [&](LayerId layerId, DataRole const& role) { + return attributes.at(BufferId{layerId, role}).lifeCycleId; + }; + auto const required = lifeCycle(1, "state"); + auto const privateLc = lifeCycle(1, "replay"); + auto const optional = lifeCycle(1, "draft"); + + // Lifecycles are registered in buffer order, after the attention layer's. + EXPECT_EQ(registry.size(), LifeCycleId{4}); + EXPECT_EQ(storage.numLifeCycles(), LifeCycleId{4}); + EXPECT_EQ(required, LifeCycleId{1}); + EXPECT_EQ(privateLc, LifeCycleId{2}); + EXPECT_EQ(optional, LifeCycleId{3}); + EXPECT_EQ(registry.ssmLifeCycleIds(), (std::vector{required, privateLc, optional})); + EXPECT_EQ(registry.requiredSsmLifeCycleId(), required); + EXPECT_EQ(std::get(registry[privateLc]).reusePolicy, ReusePolicy::PRIVATE); + EXPECT_EQ(std::get(registry[optional]).reusePolicy, ReusePolicy::OPTIONAL); + EXPECT_EQ(lifeCycle(2, "state"), required); + EXPECT_EQ(lifeCycle(2, "replay"), privateLc); + EXPECT_EQ(lifeCycle(2, "draft"), optional); + EXPECT_EQ(lifeCycle(3, "replay"), privateLc); + + auto const layerLifeCycles = storage.layerToLifeCycleIds(); + EXPECT_EQ(layerLifeCycles.size(), 4U); + EXPECT_EQ(layerLifeCycles.at(0), (std::vector{LifeCycleId{0}})); + EXPECT_EQ(layerLifeCycles.at(1), (std::vector{required, privateLc, optional})); + EXPECT_EQ(layerLifeCycles.at(2), (std::vector{required, privateLc, optional})); + EXPECT_EQ(layerLifeCycles.at(3), (std::vector{privateLc})); + + auto const grouping = storage.lifeCycleGrouping(); + EXPECT_NE(grouping[required], grouping[privateLc]); + EXPECT_NE(grouping[required], grouping[optional]); + EXPECT_NE(grouping[privateLc], grouping[optional]); + + // Slot utilization is tracked per (layer, lifecycle): every SSM layer owns one buffer in each + // slot of the lifecycles it uses. + auto const layerAttributes = storage.layerAttributes(); + EXPECT_EQ(layerAttributes.size(), 8U); + EXPECT_EQ(layerAttributes.count({3, required}), 0U); + for (LayerId layerId : {1, 2, 3}) + { + for (auto const lc : {required, privateLc, optional}) + { + if (layerId == 3 && lc != privateLc) + { + continue; + } + SCOPED_TRACE(std::to_string(layerId) + "/" + std::to_string(lc.value())); + auto const& attr = layerAttributes.at({layerId, lc}); + EXPECT_EQ(attr.lifeCycleId, lc); + ASSERT_EQ(attr.slotUtil.size(), PoolIndex{1}); + EXPECT_EQ(attr.slotUtil[PoolIndex{0}], 1); + EXPECT_EQ(attr.slotUtilFracMax, (Rational{1, lc == privateLc ? 3 : 2})); + } + } +} + } // namespace diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index da6ea500e9f2..182aafd5af3c 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -30,9 +30,13 @@ #include #include #include +#include #include #include +#include #include +#include +#include #include namespace @@ -410,6 +414,330 @@ TEST_F(ReplayCacheTest, OptionalSsmEvictionPreservesRequiredPrefix) manager->shutdown(); } +// One buffer of an SSM layer whose recurrent state is REQUIRED while its replay buffers override the +// layer's policy with PRIVATE. +struct SplitSsmBuffer +{ + LayerId layer; + DataRole role; + size_t size; + bool isPrivate; + // Distinct fill byte of this buffer. + uint8_t pattern; +}; + +// Layers 1-3 each keep two REQUIRED state roles and two PRIVATE replay roles. Layer 1 also owns a +// PRIVATE role the other layers lack, like a request-wide counter stored once. +std::vector splitSsmBuffers() +{ + std::vector buffers; + for (LayerId layer = 1; layer <= 3; ++layer) + { + auto const pattern = [layer](int role) { return static_cast(layer * 32 + role * 4); }; + buffers.push_back({layer, "state", 4096, false, pattern(0)}); + buffers.push_back({layer, "conv", 4096, false, pattern(1)}); + buffers.push_back({layer, "old_x", 4096, true, pattern(2)}); + buffers.push_back({layer, "old_dt", 1024, true, pattern(3)}); + if (layer == 1) + { + buffers.push_back({layer, "count", 512, true, pattern(4)}); + } + } + return buffers; +} + +KVCacheManagerConfig makeSplitSsmConfig(bool withHostTier) +{ + auto config = makeConfig(); + config.commitMinSnapshot = true; + config.enablePartialCommit = false; + if (withHostTier) + { + config.cacheTiers.emplace_back(HostCacheTierConfig{4 << 20}); + } + std::map layers; + for (auto const& buffer : splitSsmBuffers()) + { + auto& layer = layers[buffer.layer]; + layer.layerId = buffer.layer; + layer.buffers.push_back(BufferConfig{.role = buffer.role, + .size = buffer.size, + .reusePolicy = buffer.isPrivate ? std::optional{ReusePolicy::PRIVATE} : std::nullopt}); + } + for (auto& [layerId, layer] : layers) + { + config.layers.emplace_back(std::move(layer)); + } + return config; +} + +// Address of a buffer within the SSM page that `beam` of `cache` currently holds. +MemAddress ssmBufferAddress(KvCacheManager const& manager, KvCache const& cache, SplitSsmBuffer const& buffer, int beam) +{ + auto const slot = cache.getSsmBlockBaseIndex(manager.getLayerGroupId(buffer.layer, buffer.role), BeamIndex{beam}); + EXPECT_GE(slot, 0); + return manager.getMemPoolBaseAddress(buffer.layer, buffer.role) + + static_cast(slot) * manager.getPageIndexScale(buffer.layer, buffer.role) + * manager.getPageStride(buffer.layer, buffer.role); +} + +// Fill byte of a buffer, offset by `salt` to tell writers apart. +uint8_t splitSsmPattern(SplitSsmBuffer const& buffer, int salt) +{ + return static_cast(buffer.pattern + salt); +} + +TEST_F(ReplayCacheTest, SsmBufferPolicySplitsLayerAcrossGroups) +{ + auto manager = std::make_shared(makeSplitSsmConfig(/*withHostTier=*/true)); + auto const attention = manager->getLayerGroupId(LayerId{0}); + auto const required = manager->getLayerGroupId(1, "state"); + auto const privateGroup = manager->getLayerGroupId(1, "old_x"); + EXPECT_EQ(manager->getLayerGroupId(0, "key"), attention); + EXPECT_EQ(manager->lifeCycles().size(), LifeCycleId{3}); + EXPECT_NE(required, privateGroup); + EXPECT_TRUE(requiresPrefixCoverage(manager->lifeCycles()[required])); + EXPECT_TRUE(isPrivateLifeCycle(manager->lifeCycles()[privateGroup])); + for (auto const& buffer : splitSsmBuffers()) + { + SCOPED_TRACE(std::to_string(buffer.layer) + "/" + buffer.role); + EXPECT_EQ(manager->getLayerGroupId(buffer.layer, buffer.role), buffer.isPrivate ? privateGroup : required); + EXPECT_THROW(manager->getLayerGroupId(buffer.layer), std::invalid_argument); + } + EXPECT_THROW(manager->getLayerGroupId(1, "missing"), std::out_of_range); + + EXPECT_EQ(manager->numLayers(), 4); + auto layerIds = manager->layerIds(); + std::sort(layerIds.begin(), layerIds.end()); + EXPECT_EQ(layerIds, (std::vector{0, 1, 2, 3})); + auto grouping = manager->layerGrouping(); + ASSERT_EQ(grouping.size(), LayerGroupId{3}); + for (auto& layers : grouping) + { + std::sort(layers.begin(), layers.end()); + } + EXPECT_EQ(grouping[attention], (std::vector{0})); + EXPECT_EQ(grouping[required], (std::vector{1, 2, 3})); + EXPECT_EQ(grouping[privateGroup], (std::vector{1, 2, 3})); + for (auto const level : {kHotLevel, CacheLevel{1}}) + { + EXPECT_NE(manager->storage().getPoolGroupIndex(level, required), + manager->storage().getPoolGroupIndex(level, privateGroup)); + } + + // Each buffer's page index addresses its own lifecycle's pool: state and conv share one REQUIRED + // pool, while the PRIVATE pools hold old_x, old_dt and count, largest first. + std::map const scales{{"state", 6}, {"conv", 6}, {"old_x", 3}, {"old_dt", 3}, {"count", 1}}; + std::map const scratchPages{{"state", 2}, {"conv", 2}, {"old_x", 1}, {"old_dt", 1}, {"count", 1}}; + auto cache = manager->createKvCache({}, {}, 1); + ASSERT_TRUE(cache->resume(mStream)); + ASSERT_TRUE(cache->resize(8)); + for (auto const& buffer : splitSsmBuffers()) + { + SCOPED_TRACE(std::to_string(buffer.layer) + "/" + buffer.role); + auto const converter = manager->getPageIndexConverter(buffer.layer, buffer.role); + EXPECT_EQ(converter.scale, scales.at(buffer.role)); + EXPECT_EQ(converter.scale, manager->getPageIndexScale(buffer.layer, buffer.role)); + EXPECT_EQ(converter.scratchPagesPerBlock, scratchPages.at(buffer.role)); + EXPECT_EQ(manager->getPageStride(buffer.layer, buffer.role), static_cast(buffer.size)); + int const slot + = cache->getSsmBlockBaseIndex(manager->getLayerGroupId(buffer.layer, buffer.role), kDefaultBeamIndex); + ASSERT_GE(slot, 0); + auto const perLayerIndex = converter(std::vector{slot}, PageIndexMode::PER_LAYER).at(0); + auto const sharedIndex = converter(std::vector{slot}, PageIndexMode::SHARED).at(0); + EXPECT_EQ(manager->getMemPoolBaseAddress(buffer.layer, buffer.role, PageIndexMode::PER_LAYER) + + static_cast(perLayerIndex) * buffer.size, + manager->getMemPoolBaseAddress(buffer.layer, buffer.role, PageIndexMode::SHARED) + + static_cast(sharedIndex) * buffer.size); + EXPECT_EQ(ssmBufferAddress(*manager, *cache, buffer, 0), + manager->getMemPoolBaseAddress(buffer.layer, buffer.role, PageIndexMode::SHARED) + + static_cast(sharedIndex) * buffer.size); + } + cache->close(); + manager->shutdown(); +} + +TEST_F(ReplayCacheTest, PrivateSsmBuffersStayOutOfSnapshotsAndFollowBeams) +{ + for (int length : {7, 8}) + { + SCOPED_TRACE(length); + auto manager = std::make_shared(makeSplitSsmConfig(/*withHostTier=*/false)); + auto const required = manager->getLayerGroupId(1, "state"); + auto const privateGroup = manager->getLayerGroupId(1, "old_x"); + auto const tokens = makeTokens(length); + TokenSpan const span{tokens.data(), length}; + auto fill = [&](KvCache const& cache, SplitSsmBuffer const& buffer, int salt) + { + EXPECT_EQ(cudaMemsetAsync(reinterpret_cast(ssmBufferAddress(*manager, cache, buffer, 0)), + splitSsmPattern(buffer, salt), buffer.size, mStream), + cudaSuccess); + }; + auto check = [&](KvCache const& cache, SplitSsmBuffer const& buffer, int beam, int salt) + { + EXPECT_EQ(cudaStreamSynchronize(mStream), cudaSuccess); + std::vector data(buffer.size); + EXPECT_EQ( + cudaMemcpy(data.data(), reinterpret_cast(ssmBufferAddress(*manager, cache, buffer, beam)), + data.size(), cudaMemcpyDeviceToHost), + cudaSuccess); + auto const pattern = splitSsmPattern(buffer, salt); + EXPECT_TRUE(std::all_of(data.begin(), data.end(), [pattern](uint8_t value) { return value == pattern; })) + << buffer.layer << "/" << buffer.role << " beam " << beam; + }; + + auto source = manager->createKvCache({}, {}, 1); + ASSERT_TRUE(source->resume(mStream)); + ASSERT_TRUE(source->resize(length)); + for (auto const& buffer : splitSsmBuffers()) + { + fill(*source, buffer, 0); + } + source->commit(span); + source->stopCommitting(); + { + auto match = manager->radixTree().match({}, span, false, true); + ASSERT_EQ(match.numTokens, length); + ASSERT_FALSE(match.blocks.empty()); + EXPECT_NE(match.blocks.back()->getPage(required), nullptr); + for (auto const* block : match.blocks) + { + EXPECT_EQ(block->getPage(privateGroup), nullptr); + } + } + + // The reader restores the REQUIRED snapshot but starts from a fresh PRIVATE slot. + auto reader = manager->createKvCache({}, span, 2, {}, length); + ASSERT_EQ(reader->numCommittedTokens(), length); + EXPECT_TRUE(reader->reuseStatus()[required.value()].complete); + EXPECT_FALSE(reader->reuseStatus()[privateGroup.value()].complete); + ASSERT_TRUE(reader->resume(mStream)); + reader->commitPendingStats(); + auto const snapshotStats = manager->getAndResetSsmSnapshotIterationStats(); + ASSERT_EQ(snapshotStats.size(), 1); + ASSERT_EQ(snapshotStats.count(required), 1); + EXPECT_EQ(snapshotStats.at(required).iterSnapshotHits, 1); + for (auto const& buffer : splitSsmBuffers()) + { + if (buffer.isPrivate) + { + EXPECT_NE( + ssmBufferAddress(*manager, *source, buffer, 0), ssmBufferAddress(*manager, *reader, buffer, 0)); + fill(*reader, buffer, 1); + } + else + { + check(*reader, buffer, 0, 0); + } + } + + // Beam expansion copies both lifecycles into every beam's own slot. + ASSERT_TRUE(reader->resize(length + 4)); + reader->setBeamWidth(BeamIndex{3}); + for (int beam = 0; beam < 3; ++beam) + { + for (auto const& buffer : splitSsmBuffers()) + { + check(*reader, buffer, beam, buffer.isPrivate ? 1 : 0); + if (beam > 0) + { + EXPECT_NE(ssmBufferAddress(*manager, *reader, buffer, beam), + ssmBufferAddress(*manager, *reader, buffer, 0)); + } + } + } + reader->setBeamWidth(BeamIndex{1}); + reader->close(); + source->close(); + manager->shutdown(); + } +} + +TEST_F(ReplayCacheTest, SuspendedSsmStateSurvivesOffloadAndSlotChange) +{ + auto manager = std::make_shared(makeSplitSsmConfig(/*withHostTier=*/true)); + auto const required = manager->getLayerGroupId(1, "state"); + auto const privateGroup = manager->getLayerGroupId(1, "old_x"); + auto const tokens = makeTokens(8); + TokenSpan const span{tokens.data(), 8}; + auto cache = manager->createKvCache({}, {}, 1); + ASSERT_TRUE(cache->resume(mStream)); + ASSERT_TRUE(cache->resize(10)); + for (auto const& buffer : splitSsmBuffers()) + { + EXPECT_EQ(cudaMemsetAsync(reinterpret_cast(ssmBufferAddress(*manager, *cache, buffer, 0)), + splitSsmPattern(buffer, 0), buffer.size, mStream), + cudaSuccess); + } + cache->commit(span); + std::map oldSlots; + for (auto const group : {required, privateGroup}) + { + oldSlots[group] = cache->getSsmBlockBaseIndex(group); + ASSERT_GE(oldSlots[group], 0); + } + cache->suspend(); + + // Claim every GPU slot of both SSM pool groups. The suspended cache's held state is evicted to + // host memory, and its former GPU slots are overwritten while this test owns them. + auto& storage = manager->storage(); + std::vector> oldSlotReservations; + for (auto const group : {required, privateGroup}) + { + auto const pool = storage.getPoolGroupIndex(kHotLevel, group); + auto const stats = storage.getStatistics(kHotLevel, pool); + EXPECT_GE(stats.evictable, 1); + auto slots = storage.newSlotsForPoolGroup(kHotLevel, pool, stats.free + stats.evictable); + auto const coldPool = storage.getPoolGroupIndex(CacheLevel{1}, group); + auto const coldStats = storage.getStatistics(CacheLevel{1}, coldPool); + EXPECT_GE(coldStats.total - coldStats.free, 1); + for (auto& slot : slots) + { + if (slot.slotId().value() == oldSlots[group]) + { + oldSlotReservations.emplace_back(group, std::move(slot)); + } + else + { + storage.releaseSlot(group, kHotLevel, std::move(slot)); + } + } + } + ASSERT_EQ(oldSlotReservations.size(), 2); + for (auto const& buffer : splitSsmBuffers()) + { + auto const group = manager->getLayerGroupId(buffer.layer, buffer.role); + auto const address = manager->getMemPoolBaseAddress(buffer.layer, buffer.role) + + static_cast(oldSlots[group]) * manager->getPageIndexScale(buffer.layer, buffer.role) + * manager->getPageStride(buffer.layer, buffer.role); + EXPECT_EQ(cudaMemsetAsync(reinterpret_cast(address), 0xEE, buffer.size, mStream), cudaSuccess); + } + EXPECT_EQ(cudaStreamSynchronize(mStream), cudaSuccess); + + ASSERT_TRUE(cache->resume(mStream)); + for (auto const group : {required, privateGroup}) + { + EXPECT_NE(cache->getSsmBlockBaseIndex(group), oldSlots[group]); + } + EXPECT_EQ(cudaStreamSynchronize(mStream), cudaSuccess); + for (auto const& buffer : splitSsmBuffers()) + { + std::vector data(buffer.size); + EXPECT_EQ(cudaMemcpy(data.data(), reinterpret_cast(ssmBufferAddress(*manager, *cache, buffer, 0)), + data.size(), cudaMemcpyDeviceToHost), + cudaSuccess); + auto const pattern = splitSsmPattern(buffer, 0); + EXPECT_TRUE(std::all_of(data.begin(), data.end(), [pattern](uint8_t value) { return value == pattern; })) + << buffer.layer << "/" << buffer.role; + } + for (auto& [group, slot] : oldSlotReservations) + { + storage.releaseSlot(group, kHotLevel, std::move(slot)); + } + cache->close(); + manager->shutdown(); +} + TEST_F(ReplayCacheTest, ClosingScratchCacheReleasesEverySlotExactlyOnce) { auto config = makeConfig(); diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 4cc8ef90ac12..05349c5006ef 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -320,6 +320,8 @@ class BufferConfig: is_sparse: bool = False coalescing_group: int = 0 + reuse_policy: ReusePolicy | None = None + sparse_unit_chunks: int = 1 """Pieces a selection unit is gathered from within *this buffer's* page. @@ -425,8 +427,9 @@ class AttentionLayerConfig: @dataclass(slots=True) class SsmLayerConfig: - """Recurrent buffers sharing a layer-level reuse policy. + """Recurrent buffers with a layer-level default reuse policy. + A buffer's ``reuse_policy`` overrides the default for that buffer. OPTIONAL state is restored only when selected on first resume at the exact matched endpoint. The caller initializes or reconstructs all fresh state. """ @@ -1852,7 +1855,12 @@ class KVCacheManager: def num_layers(self) -> int: ... @property def layer_ids(self) -> Iterator[LayerId]: ... + # Raises ValueError when the layer's buffers span several layer groups. + @overload def get_layer_group_id(self, layer_id: LayerId) -> LayerGroupId: ... + @overload + def get_layer_group_id(self, layer_id: LayerId, role: DataRole) -> LayerGroupId: ... + # A layer whose buffers span several layer groups is listed under each of them. @property def layer_grouping(self) -> Sequence[Sequence[LayerId]]: ... @property 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 83f7317e5169..8cf578846827 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 @@ -153,3 +153,90 @@ def test_ssm_layer_reuse_policies(length, select_optional): for cache in reversed(caches): cache.close() manager.shutdown() + + +@pytest.mark.parametrize("length", [7, 8]) +def test_ssm_buffer_reuse_policies(length): + replay = kv.BufferConfig(role="old_x", size=2048, reuse_policy=kv.ReusePolicy.PRIVATE) + assert replay.reuse_policy == kv.ReusePolicy.PRIVATE + assert kv.BufferConfig(role="state", size=4096).reuse_policy is None + layers = [ + kv.AttentionLayerConfig(layer_id=0, buffers=[kv.BufferConfig(role="key", size=4096)]), + *[ + kv.SsmLayerConfig( + layer_id=i, + buffers=[ + kv.BufferConfig(role="state", size=4096), + kv.BufferConfig(role="conv", size=4096), + kv.BufferConfig(role="old_x", size=2048, reuse_policy=kv.ReusePolicy.PRIVATE), + ], + ) + for i in (1, 2) + ], + ] + config = kv.KVCacheManagerConfig( + tokens_per_block=4, + cache_tiers=[kv.GpuCacheTierConfig(quota=8 << 20)], + layers=layers, + commit_min_snapshot=True, + ) + assert config.layers[1].buffers[2].reuse_policy == kv.ReusePolicy.PRIVATE + manager = kv.KVCacheManager(config) + stream = torch.cuda.Stream() + tokens = list(range(length)) + caches = [] + try: + required = manager.get_layer_group_id(1, "state") + private = manager.get_layer_group_id(layer_id=1, role="old_x") + assert required != private + assert manager.get_layer_group_id(2, "conv") == required + assert manager.get_layer_group_id(2, "old_x") == private + assert manager.get_layer_group_id(0) == manager.get_layer_group_id(0, "key") + for layer_id in (1, 2): + with pytest.raises(ValueError, match="layer groups"): + manager.get_layer_group_id(layer_id) + assert manager.num_layers == 3 + assert sorted(manager.layer_ids) == [0, 1, 2] + assert sorted(manager.layer_grouping[required]) == [1, 2] + assert sorted(manager.layer_grouping[private]) == [1, 2] + + source = manager.create_kv_cache(input_tokens=[]) + caches.append(source) + assert source.resume(stream.cuda_stream) + assert source.resize(length) + source.commit(tokens, is_end=True) + source.close() + reader = manager.create_kv_cache(input_tokens=tokens) + caches.append(reader) + assert reader.num_committed_tokens == length + assert reader.reuse_status[required].complete + assert not reader.reuse_status[private].complete + assert reader.resume(stream.cuda_stream) + indices = [reader.get_ssm_block_base_index(group) for group in (required, private)] + assert all(index >= 0 for index in indices) + reader.suspend() + assert reader.resume(stream.cuda_stream) + assert indices == [reader.get_ssm_block_base_index(group) for group in (required, private)] + finally: + for cache in reversed(caches): + cache.close() + manager.shutdown() + + +@pytest.mark.parametrize("policy", [kv.ReusePolicy.PRIVATE, kv.ReusePolicy.OPTIONAL]) +def test_attention_buffer_reuse_policy_must_match_layer(policy): + def make_config(buffer_policy): + return kv.KVCacheManagerConfig( + tokens_per_block=4, + cache_tiers=[kv.GpuCacheTierConfig(quota=4 << 20)], + layers=[ + kv.AttentionLayerConfig( + layer_id=0, + buffers=[kv.BufferConfig(role="key", size=4096, reuse_policy=buffer_policy)], + ) + ], + ) + + make_config(kv.ReusePolicy.REQUIRED) + with pytest.raises(ValueError, match="only supported for SSM layers"): + make_config(policy) From 54d010c51e8ab45e2b75a1f793154bb79eadf248 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 27 Oct 2026 10:42:48 +0000 Subject: [PATCH 06/19] [None][fix] Keep SSM reuse policies when dropping zero-size buffers (By Agent) _remove_zero_size_buffers rebuilt SSM layer configs without their reuse policy, which turned PRIVATE and OPTIONAL SSM layers into REQUIRED ones whenever any layer had a zero-size buffer. Forward the layer policy; the retained buffers keep their own policies. Describe how the layer-group helpers handle an SSM layer listed under several groups and cover them with a split layer. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../kv_cache/kv_cache_manager_v2.py | 26 ++-- .../kv_cache/test_kvcm2_integration.py | 118 ++++++++++++++++++ 2 files changed, 136 insertions(+), 8 deletions(-) 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 f00f9d97d7d6..562419a1185b 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 @@ -3356,10 +3356,12 @@ def _log_window_crossing( def _get_event_window_sizes_by_layer_group( self, attention_only: bool = False ) -> Dict[int, int]: - # Assumes every layer in a group shares the same sliding_window_size, - # which is how `impl.layer_grouping` partitions layers today. Only the + # Assumes every layer listed in a group shares the same sliding_window_size, + # which is how `impl.layer_grouping` groups layers today. Only the # first layer's window is read; if the grouping policy ever permits - # mixed windows in one group, this needs to fan out per-layer. + # mixed windows in one group, this needs to fan out per-layer. A layer + # listed under several groups is an SSM layer whose buffers use + # different reuse policies, so it describes each of them as SSM. # # `attention_only` is set for the streaming event manager, which tracks # attention prefix reuse only: excluding SSM and other non-attention life cycles @@ -3906,7 +3908,11 @@ def _remove_zero_size_buffers(self, config: KVCacheManagerConfigPy) -> KVCacheMa ) ) else: - layers.append(SsmLayerConfig(layer_id=layer_id, buffers=buffers)) + layers.append( + SsmLayerConfig( + layer_id=layer_id, buffers=buffers, reuse_policy=layer.reuse_policy + ) + ) if not layers: raise ValueError("KVCacheManagerV2 requires at least one non-empty local cache layer") @@ -5516,10 +5522,12 @@ def admit_mirror(self, req: LlmRequest) -> bool: def _life_cycle_by_layer_group(self) -> List[AttnLifeCycle]: """Attention life cycle per layer group. - Every layer in a group shares one life cycle, which is how - ``impl.layer_grouping`` partitions layers, so the first layer describes - the group. A group with no window -- full attention, or a group whose - layers carry no window at all -- yields an empty stale range. + The layers listed in a group share its window and sinks, which is how + ``impl.layer_grouping`` groups layers, so the first layer describes the + group. A layer listed under several groups is an SSM layer whose buffers + use different reuse policies; it has no window in any of them. A group + with no window -- full attention, or a group whose layers carry no + window at all -- yields an empty stale range. """ if self._layer_group_life_cycles is None: layers = self.kv_cache_manager_py_config.layers @@ -6208,6 +6216,8 @@ def _stats_life_cycle_metadata(self) -> dict[int, tuple[int, Optional[int], str] for life_cycle_id, layer_ids in enumerate(self.impl.layer_grouping): if not layer_ids: continue + # A layer listed under several groups is an SSM layer whose buffers use + # different reuse policies, so its kind describes each of those groups. layer = self.kv_cache_manager_py_config.layers[int(layer_ids[0])] is_attention = isinstance(layer, AttentionLayerConfig) metadata[life_cycle_id] = ( 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 addb326f0a91..a5460cbdf92f 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -77,6 +77,7 @@ KVCacheManagerConfig, LayerId, PageIndexMode, + ReusePolicy, ReuseScope, SsmLayerConfig, _introspection, @@ -963,6 +964,123 @@ def test_zero_size_filter_rejects_empty_local_cache() -> None: manager._remove_zero_size_buffers(config) +def test_zero_size_filter_preserves_layer_and_buffer_reuse_policies() -> None: + manager = object.__new__(KVCacheManagerV2) + manager.pp_layers = [0, 1, 2, 3] + manager.num_layers = 4 + manager.num_local_layers = 4 + manager.layer_offsets = {0: 0, 1: 1, 2: 2, 3: 3} + manager.num_kv_heads_per_layer = [2, 0, 0, 0] + manager.head_dim_per_layer = [64, 64, 64, 64] + manager.max_attention_window_vec = [None, None, None, None] + + config = KVCacheManagerConfig( + tokens_per_block=TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[BufferConfig(role=Role.KEY, size=1024)], + reuse_policy=ReusePolicy.PRIVATE, + ), + AttentionLayerConfig( + layer_id=LayerId(1), buffers=[BufferConfig(role=Role.KEY, size=0)] + ), + SsmLayerConfig( + layer_id=LayerId(2), + buffers=[ + BufferConfig(role=DataRole("ssm_state"), size=256), + BufferConfig(role=DataRole("conv_state"), size=0), + BufferConfig(role=DataRole("old_x"), size=64, reuse_policy=ReusePolicy.PRIVATE), + ], + reuse_policy=ReusePolicy.OPTIONAL, + ), + SsmLayerConfig( + layer_id=LayerId(3), + buffers=[BufferConfig(role=DataRole("ssm_state"), size=256)], + reuse_policy=ReusePolicy.PRIVATE, + ), + ], + commit_min_snapshot=True, + ) + + filtered = manager._remove_zero_size_buffers(config) + + assert manager.pp_layers == [0, 2, 3] + assert [layer.layer_id for layer in filtered.layers] == [0, 1, 2] + attention_layer, split_layer, private_layer = filtered.layers + assert attention_layer.reuse_policy == ReusePolicy.PRIVATE + assert isinstance(split_layer, SsmLayerConfig) + assert split_layer.reuse_policy == ReusePolicy.OPTIONAL + assert [(buffer.role, buffer.reuse_policy) for buffer in split_layer.buffers] == [ + (DataRole("ssm_state"), None), + (DataRole("old_x"), ReusePolicy.PRIVATE), + ] + assert isinstance(private_layer, SsmLayerConfig) + assert private_layer.reuse_policy == ReusePolicy.PRIVATE + + +def test_layer_group_helpers_describe_split_ssm_layers() -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + config = KVCacheManagerConfig( + tokens_per_block=TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=4 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[BufferConfig(role=Role.KEY, size=4096)], + sliding_window_size=8, + ), + *[ + SsmLayerConfig( + layer_id=LayerId(layer_id), + buffers=[ + BufferConfig(role=DataRole("ssm_state"), size=4096), + BufferConfig( + role=DataRole("old_x"), size=2048, reuse_policy=ReusePolicy.PRIVATE + ), + ], + ) + for layer_id in (1, 2) + ], + ], + commit_min_snapshot=True, + ) + impl = KVCacheManager(config) + try: + manager = object.__new__(KVCacheManagerV2) + manager.max_seq_len = MAX_SEQ_LEN + manager.tokens_per_block = TOKENS_PER_BLOCK + manager.kv_cache_manager_py_config = config + manager.impl = impl + manager._layer_group_life_cycles = None + + attention = impl.get_layer_group_id(LayerId(0)) + required = impl.get_layer_group_id(LayerId(1), DataRole("ssm_state")) + private = impl.get_layer_group_id(LayerId(1), DataRole("old_x")) + assert len({attention, required, private}) == 3 + assert sorted(impl.layer_grouping[private]) == [1, 2] + + assert manager._get_event_window_sizes_by_layer_group(attention_only=True) == {attention: 8} + assert manager._get_event_window_sizes_by_layer_group() == { + attention: 8, + required: MAX_SEQ_LEN, + private: MAX_SEQ_LEN, + } + life_cycles = manager._life_cycle_by_layer_group() + assert life_cycles[attention].window_size == 8 + assert life_cycles[required].window_size is None + assert life_cycles[private].window_size is None + metadata = manager._stats_life_cycle_metadata() + assert metadata[attention][1:] == (8, "attention") + assert metadata[required][1:] == (None, "ssm") + assert metadata[private][1:] == (None, "ssm") + assert len({metadata[group][0] for group in (attention, required, private)}) == 3 + finally: + impl.shutdown() + + @pytest.mark.parametrize("sliding_window_size", [None, 8]) @pytest.mark.parametrize("attention_first", [False, True]) def test_event_window_sizes_filter_attention_without_backend_internals( From dd58f34f8eb14d6a2dc6b2d11c5ab52e52bb96f1 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Wed, 28 Oct 2026 13:31:43 +0800 Subject: [PATCH 07/19] Support per-buffer lifecycle queries and migrate cache storage metadata Signed-off-by: Jiagan Cheng --- .../kv_cache_manager_v2/AGENTS.md | 26 ++++- .../kv_cache_manager_v2/config.h | 3 +- .../kv_cache_manager_v2/kvCacheManager.cpp | 23 ++-- .../kv_cache_manager_v2/kvCacheManager.h | 25 ++-- .../kv_cache_manager_v2/lifeCycleRegistry.cpp | 6 +- .../kv_cache_manager_v2/storage/config.cpp | 53 ++------- .../kv_cache_manager_v2/storage/config.h | 34 +----- .../kv_cache_manager_v2/storageManager.cpp | 35 +++--- .../kv_cache_manager_v2/storageManager.h | 6 - .../batch_manager/kvCacheManagerV2.cpp | 20 ++-- .../kvCacheManagerV2ConfigTest.cpp | 31 +---- .../kvCacheManagerV2StatsTest.cpp | 40 ++++--- .../backends/fp4_mla/cache_manager.py | 2 +- .../backends/sparse/csa2/cache_manager.py | 7 +- .../sparse/minimax_m3/cache_manager.py | 2 +- .../disaggregation/native/mixers/ssm/peer.py | 13 ++- .../disaggregation/resource/kv_extractor.py | 10 +- .../_torch/disaggregation/resource/utils.py | 2 +- .../_torch/disaggregation/transceiver.py | 2 +- .../pyexecutor/connectors/kv_cache_layout.py | 12 +- .../kv_cache/kv_cache_manager_v2.py | 87 +++++++++----- .../kv_cache/mixed_precision_staging.py | 3 +- .../pyexecutor/kv_cache/sharing/_manager.py | 6 +- .../runtime/kv_cache_manager_v2/__init__.pyi | 11 +- .../kv_cache/sharing/test_layout_identity.py | 5 +- .../kv_cache/sharing/test_public_surface.py | 10 +- .../kv_cache/test_kvcm2_integration.py | 54 ++++++++- .../test_optional_checkpoint.py | 107 ++++++++++++++---- 28 files changed, 362 insertions(+), 273 deletions(-) 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 44602a8f2116..1e7f3ec28bec 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 @@ -98,12 +98,26 @@ may have sliding-window and sink-token rules; SSM lifecycles represent a recurrent-state checkpoint. Both layer configs expose `reusePolicy`: `REQUIRED` constrains the prefix match, `OPTIONAL` offers state at the required endpoint for explicit selection on first resume, and `PRIVATE` never publishes reusable pages. -Attention buffers use their layer's policy. An SSM buffer may override it through -`BufferConfig::reusePolicy` and then belongs to the SSM lifecycle of its own -policy while keeping its `LayerId`, so one SSM layer can span several layer -groups. Look such a layer up per buffer with `getLayerGroupId(layerId, role)`; -the one-argument overload throws for it, `layerGrouping()` lists it under each -of its groups, and per-layer storage attributes are keyed by (layer, lifecycle). +A buffer may override its layer's default through `BufferConfig::reusePolicy`. +This applies to attention and SSM buffers, so one model layer may span several +lifecycle groups. Use `getBufferGroupId(layerId, role)` and `bufferGrouping()` +(`get_buffer_group_id` and `buffer_grouping` in Python) for exact membership. +The legacy `getLayerGroupId(layerId)` and `layerGrouping()` are deprecated in +documentation only: uniform layers retain their results, while a split layer +raises an ambiguity error (for the queried layer or the entire partition, +respectively). No deprecation warnings are emitted. + +Page conversion reads `BufferAttr::scratchPagesPerBlock`, counting only the +layer's buffers in that coalesced slot. Scratch capacity uses the maximum exact +layer share across pools within each lifecycle. Layer counts are temporary +layout calculations, not persistent `(layer, lifecycle)` attributes. + +Execution backends that share a page or state index across coupled buffers must +validate that they use the same policy; independent policies require independent +index tables. In particular, the default attention view couples K/V and scales, +and Mamba kernels couple recurrent, convolution, PLE, and replay state. Native +KVCM supports independent policies for callers that provide those indices. + SSM lifecycles are deduplicated by policy, so a manager can have up to three SSM groups; iterate `ssmLifeCycleIds()` for working state and use `requiredSsmLifeCycleId()` only for prefix pruning. 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 9eff68e19c3c..417c9b7ca130 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 @@ -206,7 +206,8 @@ struct BufferConfig //! densely packed, so a padded page is not expressible. int sparseUnitChunks = 1; - //! Reuse policy of this buffer; nullopt uses the layer's policy. + //! Reuse policy of this buffer; nullopt inherits the layer's policy. Overrides can + //! place buffers of one attention or SSM layer in different lifecycle groups. std::optional reusePolicy; [[nodiscard]] ReusePolicy effectiveReusePolicy(ReusePolicy layerPolicy) const noexcept 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 e72fc23ff140..5ceae5c77b0a 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 @@ -603,11 +603,9 @@ int KvCacheManager::getPageIndexScale(LayerId layerId, DataRole role) const PageIndexConverter KvCacheManager::getPageIndexConverter(LayerId layerId, DataRole role) const { auto const& attr = mStorage->getBufferAttr(layerId, role); - auto const& layerAttr = mStorage->getLayerAttr(layerId, attr.lifeCycleId); int scale = mStorage->mSlotToPageIndices.at(attr.lifeCycleId).at(attr.poolIndex); int offset = exactDiv(static_cast(attr.offset), static_cast(attr.size)); - int scratchPages = layerAttr.slotUtil.at(attr.poolIndex); - return PageIndexConverter{scale, attr.expansion, offset, scratchPages}; + return PageIndexConverter{scale, attr.expansion, offset, attr.scratchPagesPerBlock}; } int KvCacheManager::_scratchSlotsFor( @@ -807,12 +805,12 @@ LayerGroupId KvCacheManager::getLayerGroupId(LayerId layerId) const throw std::invalid_argument("Layer " + std::to_string(layerId) + " has buffers in " + std::to_string(lifeCycles.size()) + " layer groups because its buffers use different reuse policies; look up its layer group per " - "buffer with getLayerGroupId(layerId, role)"); + "buffer with get_buffer_group_id(layer_id, role) / getBufferGroupId(layerId, role)"); } return lifeCycles.front(); } -LayerGroupId KvCacheManager::getLayerGroupId(LayerId layerId, DataRole const& role) const +LayerGroupId KvCacheManager::getBufferGroupId(LayerId layerId, DataRole const& role) const { return mStorage->getBufferAttr(layerId, role).lifeCycleId; } @@ -823,10 +821,17 @@ TypedVec> KvCacheManager::layerGrouping() con TypedVec> result(numLc); for (auto const& [lid, lifeCycles] : mStorage->layerToLifeCycleIds()) { - for (auto const lc : lifeCycles) - { - result.at(lc).push_back(lid); - } + result.at(getLayerGroupId(lid)).push_back(lid); + } + return result; +} + +TypedVec> KvCacheManager::bufferGrouping() const +{ + TypedVec> result(mLifeCycles.size()); + for (auto const& buffer : allBufferIds()) + { + result.at(getBufferGroupId(buffer.layerId, buffer.role)).push_back(buffer); } return result; } 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 49a1cb56054c..8941941d30bd 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 @@ -286,21 +286,22 @@ class KvCacheManager : public std::enable_shared_from_this std::vector layerIds() const; - // Layer group of every buffer of a layer. Throws std::invalid_argument when the layer's SSM - // buffers use different reuse policies and therefore span several layer groups; look up such a - // layer per buffer with getLayerGroupId(layerId, role). + //! Deprecated: use getBufferGroupId(). Returns the shared group of a uniform layer. + //! Throws std::invalid_argument if the queried layer spans several lifecycle groups. LayerGroupId getLayerGroupId(LayerId layerId) const; - // Layer group of one buffer. Throws std::out_of_range for an unknown buffer. - LayerGroupId getLayerGroupId(LayerId layerId, DataRole const& role) const; - - // Layer grouping: layers with the same lifecycle share pool allocation. A layer whose - // buffers span several layer groups is listed under each of them. - // NOTE: the iteration order of the layer lists (and of the groups) is NOT - // part of the API contract and may differ across backends/runs. Do not rely - // on it to infer buffer/pool memory order — query poolGroupDescs() - // (PoolGroupDesc::pools[i].baseAddress + coalescedBuffers) for that. + + //! Group of one buffer; throws std::out_of_range for an unknown buffer. + LayerGroupId getBufferGroupId(LayerId layerId, DataRole const& role) const; + + //! Deprecated: use bufferGrouping(). Partitions uniform layers by lifecycle. + //! Throws std::invalid_argument if any layer spans several groups. Member order + //! is unspecified; use poolGroupDescs() to inspect physical buffer ordering. TypedVec> layerGrouping() const; + //! Buffer membership indexed by lifecycle group ID. Each buffer appears once. + //! Member order is unspecified and does not describe physical memory ordering. + TypedVec> bufferGrouping() const; + // Iterator over all buffer identifiers. Mirrors Python's all_buffer_ids property. std::vector allBufferIds() const; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp index 9fd6966be474..00943894f8f3 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp @@ -49,11 +49,7 @@ LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock) LifeCycle makeLifeCycle(LayerConfig const& layer, BufferConfig const& buffer, int tokensPerBlock) { LifeCycle lc = makeLifeCycle(layer, tokensPerBlock); - // Attention layers reject buffer policies that differ from the layer's, so only SSM lifecycles change here. - if (auto* ssm = std::get_if(&lc)) - { - ssm->reusePolicy = buffer.effectiveReusePolicy(ssm->reusePolicy); - } + std::visit([&buffer](auto& value) { value.reusePolicy = buffer.effectiveReusePolicy(value.reusePolicy); }, lc); return lc; } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp index 78791ce224a2..f918ed184eb2 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp @@ -69,6 +69,11 @@ std::map StorageConfig::bufferAttributes() const for (PoolIndex poolIndex{0}; poolIndex < variant.coalescedBuffers.size(); ++poolIndex) { auto const& cb = variant.coalescedBuffers[poolIndex]; + std::unordered_map buffersPerLayer; + for (auto const& bufId : cb.bufferIds) + { + ++buffersPerLayer[bufId.layerId]; + } size_t offset = 0; for (auto const& bufId : cb.bufferIds) { @@ -77,7 +82,8 @@ std::map StorageConfig::bufferAttributes() const if (it != expansion.end()) exp = it->second; - ret[bufId] = BufferAttr{lcId, poolIndex, offset, cb.singleBufferSize, exp}; + ret[bufId] = BufferAttr{ + lcId, poolIndex, offset, cb.singleBufferSize, exp, buffersPerLayer.at(bufId.layerId)}; offset += cb.singleBufferSize; } } @@ -104,51 +110,6 @@ std::unordered_map> StorageConfig::layerToLife return map; } -std::map StorageConfig::layerAttributes() const -{ - std::map ret; - for (auto const& pg : slotDescList) - { - for (auto const& variant : pg.variants) - { - LifeCycleId lcId = variant.lifeCycleId; - for (PoolIndex poolIndex{0}; poolIndex < variant.coalescedBuffers.size(); ++poolIndex) - { - auto const& cb = variant.coalescedBuffers[poolIndex]; - // Count how many buffers per layer in this coalesced buffer. - std::unordered_map slotUtilPerLayer; - for (auto const& bufId : cb.bufferIds) - { - slotUtilPerLayer[bufId.layerId] += 1; - } - int buffersPerSlot = cb.numBuffers(); - - for (auto const& [layerId, count] : slotUtilPerLayer) - { - LayerLifeCycleKey const key{layerId, lcId}; - auto it = ret.find(key); - if (it == ret.end()) - { - LayerAttr attr; - attr.lifeCycleId = lcId; - attr.slotUtil.resize(variant.coalescedBuffers.size(), 0); - attr.slotUtilFracMax = Rational{0, 1}; - it = ret.emplace(key, std::move(attr)).first; - } - auto& attr = it->second; - attr.slotUtil[poolIndex] = count; - Rational frac{count, buffersPerSlot}; - if (frac > attr.slotUtilFracMax) - { - attr.slotUtilFracMax = frac; - } - } - } - } - } - return ret; -} - // --------------------------------------------------------------------------- // createStorageConfig — factory function. // Mirrors _storage/_config.py::create_storage_config. diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h index 0d5695b57950..58425c36ebf6 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h @@ -196,14 +196,14 @@ struct BufferAttr { LifeCycleId lifeCycleId{0}; PoolIndex poolIndex{0}; - size_t offset = 0; // byte offset within the slot - size_t size = 0; // expanded size of the buffer (after expansion) - int expansion = 1; // expansion factor (tokens_per_block / tokens_per_block_override) + size_t offset = 0; // byte offset within the slot + size_t size = 0; // expanded size of the buffer (after expansion) + int expansion = 1; // expansion factor (tokens_per_block / tokens_per_block_override) + int scratchPagesPerBlock = 1; // buffers of this layer in the same coalesced slot }; // --------------------------------------------------------------------------- // Rational — exact fraction for slot utilization computations. -// Mirrors Python's fractions.Fraction used in _storage/_config.py::LayerAttr. // --------------------------------------------------------------------------- struct Rational { @@ -228,26 +228,6 @@ struct Rational } }; -// --------------------------------------------------------------------------- -// LayerAttr — storage attributes of one layer's buffers within one lifecycle, for scratch slot -// management. An SSM layer whose buffers use different reuse policies has one LayerAttr per lifecycle. -// Mirrors _storage/_config.py::LayerAttr. -// --------------------------------------------------------------------------- -struct LayerAttr -{ - LifeCycleId lifeCycleId{0}; - - // Number of sub-pages within a single coalesced slot that belong to this layer, - // per pool within the pool group. - TypedVec slotUtil; - - // Fraction of slot_util to total number of buffers in the slot, max over all pools. - Rational slotUtilFracMax; -}; - -// (layer, lifecycle) key of a LayerAttr. -using LayerLifeCycleKey = std::pair; - // --------------------------------------------------------------------------- // StorageConfig — complete storage layout for the KV cache. // Mirrors _storage/_config.py::StorageConfig. @@ -270,13 +250,9 @@ struct StorageConfig std::map bufferAttributes() const; // Map LayerId → the lifecycles of its buffers, in ascending order. A layer spans several - // lifecycles only when its SSM buffers use different reuse policies. + // lifecycles when its buffers use different reuse policies. std::unordered_map> layerToLifeCycleIds() const; - // Per-(layer, lifecycle) storage attributes (slot utilization within coalesced slots). - // Mirrors _storage/_config.py::StorageConfig.layer_attributes(). - std::map layerAttributes() const; - LifeCycleId numLifeCycles() const; }; 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 69dc94115906..46026240a9d1 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 @@ -247,15 +247,30 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi // their page sizes. Nothing may read a cold level before then. TLLM_CHECK(appendLevelSlotDescList(config.slotDescList) == kHotLevel); - // Compute layer attributes and slot utilization fractions for scratch support. - mLayerAttributes = config.layerAttributes(); + // A scratch slot serves one layer at a time. Reserve the largest layer share + // across the coalesced pools of each lifecycle, retaining exact fractions. mSlotUtilFracMax.resize(lifeCycles.size(), Rational{0, 1}); - for (auto const& [key, layerAttr] : mLayerAttributes) + for (auto const& slotDesc : config.slotDescList) { - LifeCycleId const lcIdx = layerAttr.lifeCycleId; - if (layerAttr.slotUtilFracMax > mSlotUtilFracMax[lcIdx]) + for (auto const& variant : slotDesc.variants) { - mSlotUtilFracMax[lcIdx] = layerAttr.slotUtilFracMax; + for (auto const& buffers : variant.coalescedBuffers) + { + std::unordered_map counts; + for (auto const& buffer : buffers.bufferIds) + { + ++counts[buffer.layerId]; + } + for (auto const& [layer, count] : counts) + { + Rational const fraction{count, buffers.numBuffers()}; + auto& maximum = mSlotUtilFracMax[variant.lifeCycleId]; + if (fraction > maximum) + { + maximum = fraction; + } + } + } } } @@ -1629,14 +1644,6 @@ TypedVec StorageManager::poolGroupDescs() const return result; } -LayerAttr const& StorageManager::getLayerAttr(LayerId layerId, LifeCycleId lifeCycle) const -{ - auto it = mLayerAttributes.find(LayerLifeCycleKey{layerId, lifeCycle}); - if (it == mLayerAttributes.end()) - throw std::out_of_range("Unknown (LayerId, LifeCycleId) for LayerAttr"); - return it->second; -} - SlotCount StorageManager::numSlots(PoolGroupIndex pgIdx, CacheLevel level) const { return mLevels.at(level).storage->numSlots(pgIdx); 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 7f0da2630b3a..67172c116fc3 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 @@ -299,9 +299,6 @@ class StorageManager : public std::enable_shared_from_this // Pool group base address without per-layer offset. MemAddress getMemPoolBaseAddress(PoolGroupIndex pgIdx, PoolIndex poolIdx) const; - // Storage attributes of a layer's buffers within one lifecycle. - LayerAttr const& getLayerAttr(LayerId layerId, LifeCycleId lifeCycle) const; - // Address of a slot's buffer in a specific pool at a cache level. Address slotAddress(CacheLevel level, PoolGroupIndex pgIdx, SlotId slotId, PoolIndex poolIdx) const; @@ -482,9 +479,6 @@ class StorageManager : public std::enable_shared_from_this // slot-to-page-index scale factors: [lcId][poolIdx] TypedVec> mSlotToPageIndices; - // Per-(layer, lifecycle) storage attributes for scratch slot management. - std::map mLayerAttributes; - // Max slot utilization fraction per lifecycle (across all layers in that lifecycle). TypedVec mSlotUtilFracMax; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index a2d9e513aafb..da7298c11061 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -3665,20 +3665,22 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def_prop_ro("layer_ids", &kv::KvCacheManager::layerIds) .def_prop_ro( "layer_grouping", [](kv::KvCacheManager const& self) { return self.layerGrouping().raw(); }, - "Layers grouped by shared lifecycle/pool allocation. A layer whose buffers span several " - "layer groups is listed under each of them. The iteration order of the " - "layer lists (and of the groups) is NOT an API contract and may differ across " - "backends/runs; do not rely on it for buffer/pool memory order -- use " - "pool_group_descs (PoolGroupDesc.pools[i].base_address + coalesced_buffers) instead.") + "Deprecated: use buffer_grouping. Partitions uniform layers by lifecycle; raises ValueError " + "if any layer spans several groups. Member order does not describe physical memory order.") + .def_prop_ro( + "buffer_grouping", [](kv::KvCacheManager const& self) { return self.bufferGrouping().raw(); }, + "Buffer membership indexed by lifecycle group ID. Each buffer appears once; member order is unspecified.") .def( "get_layer_group_id", [](kv::KvCacheManager const& self, kv::LayerId layerId) { return self.getLayerGroupId(layerId).value(); }, - nb::arg("layer_id")) + nb::arg("layer_id"), + "Deprecated: use get_buffer_group_id. Returns a uniform layer's group; raises ValueError for a split " + "layer.") .def( - "get_layer_group_id", + "get_buffer_group_id", [](kv::KvCacheManager const& self, kv::LayerId layerId, kv::DataRole const& role) - { return self.getLayerGroupId(layerId, role).value(); }, - nb::arg("layer_id"), nb::arg("role"), "Layer group of one buffer of the layer.") + { return self.getBufferGroupId(layerId, role).value(); }, + nb::arg("layer_id"), nb::arg("role"), "Lifecycle group of one buffer of the layer.") .def( "get_sparse_cache_max_units", [](kv::KvCacheManager const& self, kv::LayerId layerId) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp index b5c6d5534f3b..0b3ff6ac90a5 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp @@ -171,7 +171,7 @@ TEST(KvCacheManagerV2ConfigTest, SparseBuffersCoalesceBySize) EXPECT_EQ(second.expansion, 1); } -TEST(KvCacheManagerV2ConfigTest, AttentionBuffersMustUseTheirLayerReusePolicy) +TEST(KvCacheManagerV2ConfigTest, AttentionBufferPoliciesOverrideLayerDefault) { auto config = makeConfig(); auto& buffer = std::get(config.layers.front()).buffers.front(); @@ -181,8 +181,11 @@ TEST(KvCacheManagerV2ConfigTest, AttentionBuffersMustUseTheirLayerReusePolicy) for (auto const policy : {ReusePolicy::PRIVATE, ReusePolicy::OPTIONAL}) { buffer.reusePolicy = policy; - EXPECT_THROW(config.validate(), std::invalid_argument); - EXPECT_THROW(createStorageConfig(config), std::invalid_argument); + EXPECT_NO_THROW(config.validate()); + auto const storage = createStorageConfig(config); + LifeCycleRegistry const registry(config); + auto const attr = storage.bufferAttributes().at(BufferId{0, buffer.role}); + EXPECT_EQ(std::get(registry[attr.lifeCycleId]).reusePolicy, policy); } } @@ -238,28 +241,6 @@ TEST(KvCacheManagerV2ConfigTest, SsmBufferReusePoliciesSplitLayerLifecycles) EXPECT_NE(grouping[required], grouping[privateLc]); EXPECT_NE(grouping[required], grouping[optional]); EXPECT_NE(grouping[privateLc], grouping[optional]); - - // Slot utilization is tracked per (layer, lifecycle): every SSM layer owns one buffer in each - // slot of the lifecycles it uses. - auto const layerAttributes = storage.layerAttributes(); - EXPECT_EQ(layerAttributes.size(), 8U); - EXPECT_EQ(layerAttributes.count({3, required}), 0U); - for (LayerId layerId : {1, 2, 3}) - { - for (auto const lc : {required, privateLc, optional}) - { - if (layerId == 3 && lc != privateLc) - { - continue; - } - SCOPED_TRACE(std::to_string(layerId) + "/" + std::to_string(lc.value())); - auto const& attr = layerAttributes.at({layerId, lc}); - EXPECT_EQ(attr.lifeCycleId, lc); - ASSERT_EQ(attr.slotUtil.size(), PoolIndex{1}); - EXPECT_EQ(attr.slotUtil[PoolIndex{0}], 1); - EXPECT_EQ(attr.slotUtilFracMax, (Rational{1, lc == privateLc ? 3 : 2})); - } - } } } // namespace diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index 182aafd5af3c..948c694b84a9 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -474,7 +474,7 @@ KVCacheManagerConfig makeSplitSsmConfig(bool withHostTier) // Address of a buffer within the SSM page that `beam` of `cache` currently holds. MemAddress ssmBufferAddress(KvCacheManager const& manager, KvCache const& cache, SplitSsmBuffer const& buffer, int beam) { - auto const slot = cache.getSsmBlockBaseIndex(manager.getLayerGroupId(buffer.layer, buffer.role), BeamIndex{beam}); + auto const slot = cache.getSsmBlockBaseIndex(manager.getBufferGroupId(buffer.layer, buffer.role), BeamIndex{beam}); EXPECT_GE(slot, 0); return manager.getMemPoolBaseAddress(buffer.layer, buffer.role) + static_cast(slot) * manager.getPageIndexScale(buffer.layer, buffer.role) @@ -491,9 +491,9 @@ TEST_F(ReplayCacheTest, SsmBufferPolicySplitsLayerAcrossGroups) { auto manager = std::make_shared(makeSplitSsmConfig(/*withHostTier=*/true)); auto const attention = manager->getLayerGroupId(LayerId{0}); - auto const required = manager->getLayerGroupId(1, "state"); - auto const privateGroup = manager->getLayerGroupId(1, "old_x"); - EXPECT_EQ(manager->getLayerGroupId(0, "key"), attention); + auto const required = manager->getBufferGroupId(1, "state"); + auto const privateGroup = manager->getBufferGroupId(1, "old_x"); + EXPECT_EQ(manager->getBufferGroupId(0, "key"), attention); EXPECT_EQ(manager->lifeCycles().size(), LifeCycleId{3}); EXPECT_NE(required, privateGroup); EXPECT_TRUE(requiresPrefixCoverage(manager->lifeCycles()[required])); @@ -501,24 +501,28 @@ TEST_F(ReplayCacheTest, SsmBufferPolicySplitsLayerAcrossGroups) for (auto const& buffer : splitSsmBuffers()) { SCOPED_TRACE(std::to_string(buffer.layer) + "/" + buffer.role); - EXPECT_EQ(manager->getLayerGroupId(buffer.layer, buffer.role), buffer.isPrivate ? privateGroup : required); + EXPECT_EQ(manager->getBufferGroupId(buffer.layer, buffer.role), buffer.isPrivate ? privateGroup : required); EXPECT_THROW(manager->getLayerGroupId(buffer.layer), std::invalid_argument); } - EXPECT_THROW(manager->getLayerGroupId(1, "missing"), std::out_of_range); + EXPECT_THROW(manager->getBufferGroupId(1, "missing"), std::out_of_range); EXPECT_EQ(manager->numLayers(), 4); auto layerIds = manager->layerIds(); std::sort(layerIds.begin(), layerIds.end()); EXPECT_EQ(layerIds, (std::vector{0, 1, 2, 3})); - auto grouping = manager->layerGrouping(); + EXPECT_THROW(manager->layerGrouping(), std::invalid_argument); + auto const grouping = manager->bufferGrouping(); ASSERT_EQ(grouping.size(), LayerGroupId{3}); - for (auto& layers : grouping) + size_t count = 0; + for (LayerGroupId group{0}; group < grouping.size(); ++group) { - std::sort(layers.begin(), layers.end()); + for (auto const& buffer : grouping[group]) + { + EXPECT_EQ(manager->getBufferGroupId(buffer.layerId, buffer.role), group); + ++count; + } } - EXPECT_EQ(grouping[attention], (std::vector{0})); - EXPECT_EQ(grouping[required], (std::vector{1, 2, 3})); - EXPECT_EQ(grouping[privateGroup], (std::vector{1, 2, 3})); + EXPECT_EQ(count, manager->allBufferIds().size()); for (auto const level : {kHotLevel, CacheLevel{1}}) { EXPECT_NE(manager->storage().getPoolGroupIndex(level, required), @@ -541,7 +545,7 @@ TEST_F(ReplayCacheTest, SsmBufferPolicySplitsLayerAcrossGroups) EXPECT_EQ(converter.scratchPagesPerBlock, scratchPages.at(buffer.role)); EXPECT_EQ(manager->getPageStride(buffer.layer, buffer.role), static_cast(buffer.size)); int const slot - = cache->getSsmBlockBaseIndex(manager->getLayerGroupId(buffer.layer, buffer.role), kDefaultBeamIndex); + = cache->getSsmBlockBaseIndex(manager->getBufferGroupId(buffer.layer, buffer.role), kDefaultBeamIndex); ASSERT_GE(slot, 0); auto const perLayerIndex = converter(std::vector{slot}, PageIndexMode::PER_LAYER).at(0); auto const sharedIndex = converter(std::vector{slot}, PageIndexMode::SHARED).at(0); @@ -563,8 +567,8 @@ TEST_F(ReplayCacheTest, PrivateSsmBuffersStayOutOfSnapshotsAndFollowBeams) { SCOPED_TRACE(length); auto manager = std::make_shared(makeSplitSsmConfig(/*withHostTier=*/false)); - auto const required = manager->getLayerGroupId(1, "state"); - auto const privateGroup = manager->getLayerGroupId(1, "old_x"); + auto const required = manager->getBufferGroupId(1, "state"); + auto const privateGroup = manager->getBufferGroupId(1, "old_x"); auto const tokens = makeTokens(length); TokenSpan const span{tokens.data(), length}; auto fill = [&](KvCache const& cache, SplitSsmBuffer const& buffer, int salt) @@ -656,8 +660,8 @@ TEST_F(ReplayCacheTest, PrivateSsmBuffersStayOutOfSnapshotsAndFollowBeams) TEST_F(ReplayCacheTest, SuspendedSsmStateSurvivesOffloadAndSlotChange) { auto manager = std::make_shared(makeSplitSsmConfig(/*withHostTier=*/true)); - auto const required = manager->getLayerGroupId(1, "state"); - auto const privateGroup = manager->getLayerGroupId(1, "old_x"); + auto const required = manager->getBufferGroupId(1, "state"); + auto const privateGroup = manager->getBufferGroupId(1, "old_x"); auto const tokens = makeTokens(8); TokenSpan const span{tokens.data(), 8}; auto cache = manager->createKvCache({}, {}, 1); @@ -706,7 +710,7 @@ TEST_F(ReplayCacheTest, SuspendedSsmStateSurvivesOffloadAndSlotChange) ASSERT_EQ(oldSlotReservations.size(), 2); for (auto const& buffer : splitSsmBuffers()) { - auto const group = manager->getLayerGroupId(buffer.layer, buffer.role); + auto const group = manager->getBufferGroupId(buffer.layer, buffer.role); auto const address = manager->getMemPoolBaseAddress(buffer.layer, buffer.role) + static_cast(oldSlots[group]) * manager->getPageIndexScale(buffer.layer, buffer.role) * manager->getPageStride(buffer.layer, buffer.role); diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py index 9f5e80357b7a..715da9a31172 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py @@ -535,7 +535,7 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: if fp8_pool_id in (cache_pool_id, hp_pool_id): raise RuntimeError("FP4 MLA, FP8 MLA, and HP state must use distinct lifecycles.") - num_pools = len(self.impl.layer_grouping) + num_pools = len(self._layer_ids_by_group) pool_pointers = [[[0, 0], [0, 0]] for _ in range(num_pools)] pool_pointers[cache_pool_id] = [ [ diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py index 7e343063105d..0656fb2bbc95 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py @@ -266,7 +266,7 @@ def __init__( self._encoder_optional_groups = tuple( sorted( { - self.impl.get_layer_group_id(self._layer_roles[key]) + self.impl.get_buffer_group_id(self._layer_roles[key], key[2].role) for key in self._encoder_roles if self._encoder_replay_enabled } @@ -277,8 +277,11 @@ def __init__( def _log_checkpoint_pools(self): """Report physical quota once per group, without counting role views twice.""" policies = { - self.impl.get_layer_group_id(layer.layer_id): layer.reuse_policy + self.impl.get_buffer_group_id(layer.layer_id, buffer.role): layer.reuse_policy + if buffer.reuse_policy is None + else buffer.reuse_policy for layer in self.impl.init_config.layers + for buffer in layer.buffers } for level, tier in enumerate(self.impl.cache_tier_list): grouping = self.impl.get_life_cycle_pool_group_indices(level) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index 7722f76ea664..63bb2b9b2572 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -1033,7 +1033,7 @@ def _kv_pool_mapping_offset(self, layer_id, layer_group_id, key_base_addr) -> in pool-mapping loop meaningful). """ layers_by_addr = sorted( - self.impl.layer_grouping[int(layer_group_id)], + self._layer_ids_by_group[int(layer_group_id)], key=lambda lid: self.impl.get_mem_pool_base_address( lid, Role.KEY, PageIndexMode.SHARED ), diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py index 918252a94578..7208870eaf51 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py @@ -478,10 +478,15 @@ def _find_mamba_layer_group( ) -> Optional[MambaLayerGroup]: if page_table is None: return None - return next( - (lg for lg in page_table.layer_groups if lg.kind == CacheKind.STATE), - None, - ) + groups = [ + lg for lg in page_table.layer_groups if lg.kind == CacheKind.STATE and lg.pool_views + ] + if len(groups) > 1: + raise ValueError( + "Mamba transfer requires one recurrent-state lifecycle; " + "multiple transferable state groups need separate slot indices" + ) + return groups[0] if groups else None @staticmethod def validate_peer_compatible( diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 279916d53001..20eecdc10fa3 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -187,7 +187,9 @@ def _compute_global_layer_ids(manager, lg_idx: int) -> List[int]: (model_layer, attn_type) into synthetic IDs via the _layer_attn_to_layer_id inverse mapping. """ - local_layer_ids = manager.impl.layer_grouping[lg_idx] + local_layer_ids = sorted( + {int(buffer.layer_id) for buffer in manager.impl.buffer_grouping[lg_idx]} + ) get_global_layer_ids = getattr(manager, "get_disagg_global_layer_ids", None) if get_global_layer_ids is not None: @@ -406,7 +408,7 @@ def _window_size_for_layer(internal_layer_id: int): pool_groups: List[PhysicalPoolGroup] = [] storage_pg_to_list_idx: Dict[int, int] = {} - layer_groups_by_id: List[LayerGroup | None] = [None] * len(manager.impl.layer_grouping) + layer_groups_by_id: List[LayerGroup | None] = [None] * len(manager.impl.buffer_grouping) for pg_desc in pool_group_descs: storage_pg_idx = int(pg_desc.pool_group_index) @@ -431,7 +433,9 @@ def _window_size_for_layer(internal_layer_id: int): # memory is reused, not concatenated. for variant in pg_desc.slot_desc.variants: layer_group_id = int(variant.layer_group_id) - all_internal_layer_ids = list(manager.impl.layer_grouping[layer_group_id]) + all_internal_layer_ids = sorted( + {int(buffer.layer_id) for buffer in manager.impl.buffer_grouping[layer_group_id]} + ) all_global_layer_ids = _compute_global_layer_ids(manager, layer_group_id) local_layers = [ LocalLayer(local_layer_id=int(iid), global_layer_id=int(gid)) diff --git a/tensorrt_llm/_torch/disaggregation/resource/utils.py b/tensorrt_llm/_torch/disaggregation/resource/utils.py index 579cc42d1d13..341f6c79d8f8 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/utils.py +++ b/tensorrt_llm/_torch/disaggregation/resource/utils.py @@ -301,7 +301,7 @@ def get_layer_to_layer_group( kind = CacheKind.PAGED out: Dict[int, int] = {} for lg_idx, lg in enumerate(page_table.layer_groups): - if lg.kind == kind: + if lg.kind == kind and lg.pool_views: for ll in lg.local_layers: gid = int(ll.global_layer_id) if gid in out: diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 759d408fa362..0ffb10dcdfb9 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -675,7 +675,7 @@ def _describe_local(self, req: LlmRequest) -> Chunk: kinds = [lg.kind for lg in layer_groups] for idx, lg in enumerate(layer_groups): if lg.kind == CacheKind.STATE: - slot = self._get_mamba_slot_for_request(req) + slot = self._get_mamba_slot_for_request(req) if lg.pool_views else None group = np.array([slot], dtype=np.int64) if slot is not None else empty else: # Block lists carry beam 0 only (beam-search attention reads diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py index b0749a168390..8c627f5a0dec 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py @@ -261,7 +261,7 @@ def build_kv_cache_layout_v2(manager: "KVCacheManagerV2") -> KvCacheLayout: so the addresses are device addresses only because tiers below GPU are rejected while a connector is attached (``PyExecutor._reject_non_gpu_cache_tiers``). Reads only the manager's - public layout API -- ``layer_grouping``, ``all_buffer_ids``, + public layout API -- ``buffer_grouping``, ``all_buffer_ids``, ``get_aggregated_pages``, ``pool_group_descs`` -- so it assumes nothing about dimension order, kv factor, or pool count. """ @@ -276,21 +276,15 @@ def build_kv_cache_layout_v2(manager: "KVCacheManagerV2") -> KvCacheLayout: for variant in pool_group.slot_desc.variants: slots_by_group[int(variant.layer_group_id)] = int(pool_group.num_slots) - buffers_by_layer: Dict[int, List] = {} - for buffer_id in impl.all_buffer_ids: - buffers_by_layer.setdefault(int(buffer_id.layer_id), []).append(buffer_id) - groups: List[KvCacheLayerGroupLayout] = [] - for layer_group_id, local_layer_ids in enumerate(impl.layer_grouping): - local_layer_ids = [int(lid) for lid in local_layer_ids] + for layer_group_id, buffer_ids in enumerate(impl.buffer_grouping): + local_layer_ids = sorted({int(buffer.layer_id) for buffer in buffer_ids}) if not local_layer_ids: continue num_slots = slots_by_group[layer_group_id] global_by_local = dict(zip(local_layer_ids, _global_layer_ids(manager, local_layer_ids))) - buffer_ids = [b for lid in local_layer_ids for b in buffers_by_layer.get(lid, ())] - regions: List[KvCacheRegion] = [] for desc in impl.get_aggregated_pages(buffer_ids): if int(desc.layer_group_id) != layer_group_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 562419a1185b..509fdc973b15 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 @@ -2405,10 +2405,10 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self._get_event_layer_group_ids(), ) - # Both backends build layer_grouping on demand, and the layer order + # Buffer membership has no physical ordering guarantee. The layer order # within a group is not part of its contract. Cache a stable physical- # layout representative for each role from the public pool descriptors. - self.num_pools = len(self.impl.layer_grouping) + self.num_pools = len(self._layer_ids_by_group) self._pool_layer_ids_by_role: Dict[Tuple[int, DataRole], LayerId] = {} for pool_group in self.impl.pool_group_descs: for variant in pool_group.slot_desc.variants: @@ -2418,10 +2418,16 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self._pool_layer_ids_by_role.setdefault( (pool_id, buffer_id.role), buffer_id.layer_id ) - num_layers = len(config.layers) - self.layer_to_pool_mapping_dict: dict[int, int] = { - layer_id: self.impl.get_layer_group_id(layer_id) - for layer_id in typed_range(LayerId(num_layers)) + groups_by_layer: dict[int, set[int]] = defaultdict(set) + self.buffer_to_pool_mapping_dict: dict[tuple[int, DataRole], int] = {} + for group_id, buffers in enumerate(self.impl.buffer_grouping): + for buffer in buffers: + groups_by_layer[int(buffer.layer_id)].add(group_id) + self.buffer_to_pool_mapping_dict[(int(buffer.layer_id), buffer.role)] = group_id + self.layer_to_pool_mapping_dict = { + layer_id: next(iter(groups)) + for layer_id, groups in groups_by_layer.items() + if len(groups) == 1 } self.kv_cache_map: dict[int, _KVCache] = {} @@ -2582,7 +2588,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self._sparse_layer_group_ids = tuple( sorted( { - self.layer_to_pool_mapping_dict[buf.layer_id] + self.impl.get_buffer_group_id(buf.layer_id, buf.role) for buf in self.impl.all_buffer_ids if self.impl.is_sparse(buf.layer_id, buf.role) } @@ -2840,7 +2846,7 @@ def _fill_fresh_kv_pages(self, request_id: int) -> None: fill_stream = torch.cuda.current_stream(device=self._stream.device) fill_stream.wait_stream(self._stream) for layer_idx in self.pp_layers: - if self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] != pool_id: + if self._get_layer_execution_group(self.layer_offsets[layer_idx]) != pool_id: continue buffer = self.get_buffers(layer_idx) if buffer is None: @@ -2899,6 +2905,21 @@ def _get_layer_kv_factor_by_local_id(self, local_layer_id: int) -> int: return self.layer_specs[local_layer_id].kv_factor return self.kv_factor + @property + def _layer_ids_by_group(self) -> list[list[int]]: + # A split layer can occur in several groups. This is membership metadata, + # not the partition promised by the deprecated layer_grouping API. + return [ + sorted({int(buffer.layer_id) for buffer in buffers}) + for buffers in self.impl.buffer_grouping + ] + + def _get_layer_execution_group(self, layer_id: int) -> int: + group = self.layer_to_pool_mapping_dict.get(layer_id) + if group is not None: + return group + return int(self.impl.get_buffer_group_id(layer_id, Role.KEY)) + def _get_layer_page_table_roles( self, local_layer_id: int, pool_id: int ) -> Tuple[DataRole, Optional[DataRole]]: @@ -2937,7 +2958,7 @@ def _build_pool_mapping_tensors(self): ) if self._use_per_layer_page_tables: for layer_id in typed_range(LayerId(self.num_local_layers)): - pool_id = self.impl.get_layer_group_id(layer_id) + pool_id = self._get_layer_execution_group(layer_id) role_a, _ = self._get_layer_page_table_roles(layer_id, pool_id) kv_cache_pool_pointers_list.append( [ @@ -2999,7 +3020,7 @@ def _build_pool_mapping_tensors(self): block_scale_pool_pointers_list.append([scale_base_addr, 0]) for layer_id in typed_range(LayerId(self.num_local_layers)): - layer_group_id = self.impl.get_layer_group_id(layer_id) + layer_group_id = self._get_layer_execution_group(layer_id) role_a, role_b = self._get_pool_roles(layer_group_id) index_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] if role_a == Role.KEY: @@ -3071,6 +3092,19 @@ def _build_pool_mapping_tensors(self): return kv_cache_pool_pointers, kv_cache_pool_mapping def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: + # The default attention view shares indices between K/V and their scales. + # Custom backends with separate tables override this initializer. + for layer in self.kv_cache_manager_py_config.layers: + if isinstance(layer, AttentionLayerConfig): + groups = { + self.impl.get_buffer_group_id(layer.layer_id, buffer.role) + for buffer in layer.buffers + } + if len(groups) > 1: + raise ValueError( + f"Attention layer {layer.layer_id} has different buffer reuse policies; " + "the default attention backend requires one shared page-index table" + ) # Keep role-dependent checks in this default page-table initializer: # custom layouts such as DeepSeek V4 override it with their own roles. # A lifecycle group can contain different page sizes in separate @@ -3086,7 +3120,7 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: } ) > 1 - for pool_id, layer_ids in enumerate(self.impl.layer_grouping) + for pool_id, layer_ids in enumerate(self._layer_ids_by_group) ) if self.layer_specs is not None: storage_layouts = {(layer.dtype, layer.kv_factor) for layer in self.layer_specs} @@ -3289,7 +3323,7 @@ def _get_event_num_blocks_per_cache_level( return [int(tier.quota // bytes_per_block) for tier in cache_tiers] def _get_event_layer_group_ids(self) -> List[int]: - return [int(layer_group_id) for layer_group_id in range(len(self.impl.layer_grouping))] + return [int(layer_group_id) for layer_group_id in range(len(self._layer_ids_by_group))] def _get_windowed_layer_group_sizes(self) -> Dict[int, int]: """Layer-group id -> window size, for groups whose window is smaller than max_seq_len. @@ -3360,8 +3394,7 @@ def _get_event_window_sizes_by_layer_group( # which is how `impl.layer_grouping` groups layers today. Only the # first layer's window is read; if the grouping policy ever permits # mixed windows in one group, this needs to fan out per-layer. A layer - # listed under several groups is an SSM layer whose buffers use - # different reuse policies, so it describes each of them as SSM. + # can belong to several groups through buffer policy overrides. # # `attention_only` is set for the streaming event manager, which tracks # attention prefix reuse only: excluding SSM and other non-attention life cycles @@ -3374,7 +3407,7 @@ def get_event_window_size(layer_config: object) -> int: return self.max_seq_len if window_size is None else int(window_size) window_sizes: Dict[int, int] = {} - for layer_group_id, layer_ids in enumerate(self.impl.layer_grouping): + for layer_group_id, layer_ids in enumerate(self._layer_ids_by_group): # Native bindings expose grouping, not the private Python lifecycle registry. layer_config = self.kv_cache_manager_py_config.layers[int(layer_ids[0])] if attention_only and not isinstance(layer_config, AttentionLayerConfig): @@ -3429,7 +3462,7 @@ def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: scratch_pages = torch.empty_like(scales) for local_layer_idx in range(self.num_local_layers): layer_id = LayerId(local_layer_idx) - pool_id = self.layer_to_pool_mapping_dict[layer_id] + pool_id = self._get_layer_execution_group(layer_id) role_a, role_b = self._get_layer_page_table_roles(layer_id, pool_id) roles = [ role_a, @@ -4097,7 +4130,7 @@ def fill_mixed_precision_page_tables( """ local = self.layer_offsets[layer_idx] for role, layer in enumerate((local, self.high_layers[local])): - pool_id = int(self.impl.get_layer_group_id(layer)) + pool_id = int(self.impl.get_buffer_group_id(layer, Role.KEY)) scale = int(self.impl.get_page_index_converter(layer, Role.KEY).scale) for row, req_id in enumerate(request_ids): kv_cache = self.kv_cache_map[req_id] @@ -5524,8 +5557,7 @@ def _life_cycle_by_layer_group(self) -> List[AttnLifeCycle]: The layers listed in a group share its window and sinks, which is how ``impl.layer_grouping`` groups layers, so the first layer describes the - group. A layer listed under several groups is an SSM layer whose buffers - use different reuse policies; it has no window in any of them. A group + group, including when buffer policies split a layer. A group with no window -- full attention, or a group whose layers carry no window at all -- yields an empty stale range. """ @@ -5537,7 +5569,7 @@ def _life_cycle_by_layer_group(self) -> List[AttnLifeCycle]: getattr(layer, "num_sink_tokens", None), self.tokens_per_block, ) - for layer in (layers[int(layer_ids[0])] for layer_ids in self.impl.layer_grouping) + for layer in (layers[int(layer_ids[0])] for layer_ids in self._layer_ids_by_group) ] return self._layer_group_life_cycles @@ -5581,10 +5613,10 @@ def get_page_indices_by_layer_group(self, request: LlmRequest) -> List[List[int] """ kv_cache = self.kv_cache_map.get(request.py_request_id) if kv_cache is None: - return [[] for _ in range(len(self.impl.layer_grouping))] + return [[] for _ in range(len(self._layer_ids_by_group))] history_length = kv_cache.history_length by_group: List[List[int]] = [] - for layer_group_id in range(len(self.impl.layer_grouping)): + for layer_group_id in range(len(self._layer_ids_by_group)): indices = list(kv_cache.get_aggregated_page_indices(layer_group_id, valid_only=False)) stale_start, stale_end = self._stale_block_range(layer_group_id, history_length) for ordinal in range(stale_start, min(stale_end, len(indices))): @@ -6213,11 +6245,10 @@ def _stats_life_cycle_metadata(self) -> dict[int, tuple[int, Optional[int], str] } metadata: dict[int, tuple[int, Optional[int], str]] = {} - for life_cycle_id, layer_ids in enumerate(self.impl.layer_grouping): + for life_cycle_id, layer_ids in enumerate(self._layer_ids_by_group): if not layer_ids: continue - # A layer listed under several groups is an SSM layer whose buffers use - # different reuse policies, so its kind describes each of those groups. + # Buffer policies can split a layer without changing its kind or window. layer = self.kv_cache_manager_py_config.layers[int(layer_ids[0])] is_attention = isinstance(layer, AttentionLayerConfig) metadata[life_cycle_id] = ( @@ -7203,7 +7234,7 @@ def get_batch_cache_indices( else self.kv_factor ) else: - pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] + pool_id = self._get_layer_execution_group(self.layer_offsets[layer_idx]) index_scale = 1 if raw_indices else self.get_layer_page_index_scale(layer_idx) kv_factor = self.get_layer_kv_factor(layer_idx) return self._get_batch_cache_indices_by_pool_id( @@ -7283,7 +7314,7 @@ def get_batch_cache_indices_flat( else self.kv_factor ) else: - pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] + pool_id = self._get_layer_execution_group(self.layer_offsets[layer_idx]) # Layers sharing a pool can still require different page-index # scales (e.g. Gemma4 sliding/global head_dim). Use the per-layer # scale so this flat block table matches get_batch_cache_indices() @@ -7415,7 +7446,7 @@ def _iter_cache_buffers_for_invalid_check(self) -> Iterable[Tuple[int, torch.Ten # layouts that coalesce extra buffers alongside K/V. yield layer_id, buffer continue - pool_id = self.layer_to_pool_mapping_dict[layer_offset] + pool_id = self._get_layer_execution_group(layer_offset) if pool_id in pool_handled: continue yield layer_id, buffer diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mixed_precision_staging.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mixed_precision_staging.py index 72033a3b4a15..16431da0799d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mixed_precision_staging.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mixed_precision_staging.py @@ -10,6 +10,7 @@ from tensorrt_llm._utils import prefer_pinned from ...memory_buffer_utils import Buffers +from .kv_cache_manager_v2 import Role class MixedPrecisionKVStaging: @@ -71,7 +72,7 @@ def __init__( for layer_idx in self.pools: local = manager.layer_offsets[layer_idx] key = tuple( - int(manager.impl.get_layer_group_id(layer)) + int(manager.impl.get_buffer_group_id(layer, Role.KEY)) for layer in (local, manager.high_layers[local]) ) self.table_owner[layer_idx] = owners.setdefault(key, layer_idx) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py index aa1f46cd2813..46dcdfaa31e0 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py @@ -166,7 +166,9 @@ def layer_configs(manager: KVCacheManagerV2) -> dict[int, object]: def layer_group_layer_ids(manager: KVCacheManagerV2) -> list[list[int]]: """Internal layer ids per runtime layer group, in layer-group order.""" - return [[int(layer_id) for layer_id in group] for group in manager.impl.layer_grouping] + return [ + sorted({int(buffer.layer_id) for buffer in group}) for group in manager.impl.buffer_grouping + ] def pool_group_descriptors(manager: KVCacheManagerV2) -> Sequence[PoolGroupDesc]: @@ -180,7 +182,7 @@ def declared_global_layer_ids(manager: KVCacheManagerV2) -> dict[int, int] | Non if getter is None: return None table: dict[int, int] = {} - for layer_group, layer_ids in enumerate(manager.impl.layer_grouping): + for layer_group, layer_ids in enumerate(layer_group_layer_ids(manager)): ids = [int(value) for value in getter(layer_group)] if len(ids) != len(layer_ids): raise ValueError( diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 05349c5006ef..523db9d37e2a 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -30,7 +30,6 @@ from typing import ( Type, TypeAlias, Union, - overload, ) from tensorrt_llm.bindings.executor import KvCacheRetentionConfig @@ -1855,15 +1854,15 @@ class KVCacheManager: def num_layers(self) -> int: ... @property def layer_ids(self) -> Iterator[LayerId]: ... - # Raises ValueError when the layer's buffers span several layer groups. - @overload + # Deprecated: use get_buffer_group_id. Raises ValueError for a split layer. def get_layer_group_id(self, layer_id: LayerId) -> LayerGroupId: ... - @overload - def get_layer_group_id(self, layer_id: LayerId, role: DataRole) -> LayerGroupId: ... - # A layer whose buffers span several layer groups is listed under each of them. + def get_buffer_group_id(self, layer_id: LayerId, role: DataRole) -> LayerGroupId: ... + # Deprecated: use buffer_grouping. Raises ValueError if any layer is split. @property def layer_grouping(self) -> Sequence[Sequence[LayerId]]: ... @property + def buffer_grouping(self) -> Sequence[Sequence[BufferId]]: ... + @property def all_buffer_ids(self) -> Iterator[BufferId]: ... def get_aggregated_pages(self, buffers: Iterable[BufferId]) -> Iterator[AggregatedPageDesc]: ... @property diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py index c4442ac4ed9e..4b20395472fe 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py @@ -404,7 +404,10 @@ def __init__( slot_desc=SimpleNamespace(variants=[variant]), ) self.impl = SimpleNamespace( - layer_grouping=[list(lids) for _, lids, _ in groups], + buffer_grouping=[ + [SimpleNamespace(layer_id=lid, role=role) for lid in lids for role, _, _ in roles] + for _, lids, roles in groups + ], pool_group_descs=list(descs.values()), get_life_cycle_pool_group_indices=lambda: list(pool_group_of), init_config=SimpleNamespace(layers=layers), diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py index d08fa1ea6cd5..a65252473745 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py @@ -68,7 +68,7 @@ INTERNALS = { "impl", "init_config", - "layer_grouping", + "buffer_grouping", "pool_group_descs", "_has_lossy_cold_page_codec", "get_disagg_global_layer_ids", @@ -591,7 +591,13 @@ def test_manager_facade_preserves_layer_group_order_and_normalizes_ids(): from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager grouping = [[np.int64(5), "7"], [2]] - manager = SimpleNamespace(impl=SimpleNamespace(layer_grouping=grouping)) + manager = SimpleNamespace( + impl=SimpleNamespace( + buffer_grouping=[ + [SimpleNamespace(layer_id=layer) for layer in group] for group in grouping + ] + ) + ) actual = _manager.layer_group_layer_ids(manager) assert actual == [[5, 7], [2]] assert all(type(layer_id) is int for group in actual for layer_id in group) 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 a5460cbdf92f..b34f9a298cc6 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -421,7 +421,7 @@ def build_cache_config( None, ) if fake_impl is not None: - fake_impl.layer_grouping = [[0]] + fake_impl.buffer_grouping = [[BufferId(LayerId(0), Role.KEY)]] fake_impl.pool_group_descs = [] fake_impl.get_layer_group_id.side_effect = lambda _: 0 fake_impl.all_buffer_ids = [BufferId(LayerId(0), Role.KEY)] @@ -1057,10 +1057,10 @@ def test_layer_group_helpers_describe_split_ssm_layers() -> None: manager._layer_group_life_cycles = None attention = impl.get_layer_group_id(LayerId(0)) - required = impl.get_layer_group_id(LayerId(1), DataRole("ssm_state")) - private = impl.get_layer_group_id(LayerId(1), DataRole("old_x")) + required = impl.get_buffer_group_id(LayerId(1), DataRole("ssm_state")) + private = impl.get_buffer_group_id(LayerId(1), DataRole("old_x")) assert len({attention, required, private}) == 3 - assert sorted(impl.layer_grouping[private]) == [1, 2] + assert sorted({buffer.layer_id for buffer in impl.buffer_grouping[private]}) == [1, 2] assert manager._get_event_window_sizes_by_layer_group(attention_only=True) == {attention: 8} assert manager._get_event_window_sizes_by_layer_group() == { @@ -1103,7 +1103,11 @@ def test_event_window_sizes_filter_attention_without_backend_internals( ) # Match the C++ binding surface: layer_grouping is public, while the # Python implementation's private _life_cycles registry is absent. - manager.impl = SimpleNamespace(layer_grouping=[[0], [1]] if attention_first else [[1], [0]]) + manager.impl = SimpleNamespace( + buffer_grouping=[ + [BufferId(LayerId(i), Role.KEY)] for i in ([0, 1] if attention_first else [1, 0]) + ] + ) attention_group_id = 0 if attention_first else 1 window_size = MAX_SEQ_LEN if sliding_window_size is None else sliding_window_size @@ -3245,7 +3249,7 @@ def _index_mapper_capacity_for( """ module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" fake_impl = Mock() - fake_impl.layer_grouping = [[0]] + fake_impl.buffer_grouping = [[BufferId(LayerId(0), Role.KEY)]] fake_impl.pool_group_descs = [] fake_impl.get_layer_group_id.side_effect = lambda _: 0 fake_impl.all_buffer_ids = [BufferId(LayerId(0), Role.KEY)] @@ -4081,3 +4085,41 @@ def test_try_commit_generation_blocks_honors_guards( manager.try_commit_generation_blocks(request) assert kv_cache.committed_tokens is None + + +@pytest.mark.parametrize("role", [Role.VALUE, Role.KEY_BLOCK_SCALE]) +def test_default_attention_rejects_split_coupled_buffers(monkeypatch, role): + original = KVCacheManagerV2._build_cache_config + + def build_config(self, config): + config = original(self, config) + layers = list(config.layers) + layer = layers[0] + buffers = [ + BufferConfig( + role=buffer.role, + size=buffer.size, + reuse_policy=ReusePolicy.PRIVATE if buffer.role == role else None, + ) + for buffer in layer.buffers + ] + if role == Role.KEY_BLOCK_SCALE: + buffers.append(BufferConfig(role=role, size=128, reuse_policy=ReusePolicy.PRIVATE)) + layers[0] = AttentionLayerConfig(layer_id=layer.layer_id, buffers=buffers) + return replace(config, layers=layers) + + monkeypatch.setattr(KVCacheManagerV2, "_build_cache_config", build_config) + with pytest.raises(ValueError, match="requires one shared page-index table"): + KVCacheManagerV2( + KvCacheConfig(max_gpu_total_bytes=4 << 20, host_cache_size=0), + CacheType.SELF, + num_layers=1, + num_kv_heads=1, + head_dim=64, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=1, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.HALF, + vocab_size=16, + ) 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 8cf578846827..581ea0b2bbc7 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 @@ -93,7 +93,8 @@ def test_reuse_status_and_optional_selection_bindings(priority): @pytest.mark.parametrize("length", [7, 8]) @pytest.mark.parametrize("select_optional", [False, True]) -def test_ssm_layer_reuse_policies(length, select_optional): +@pytest.mark.parametrize("split_layer", [False, True]) +def test_ssm_layer_reuse_policies(length, select_optional, split_layer): layers = [ kv.AttentionLayerConfig(layer_id=0, buffers=[kv.BufferConfig(role="key", size=4096)]), *[ @@ -111,6 +112,22 @@ def test_ssm_layer_reuse_policies(length, select_optional): ) ], ] + if split_layer: + layers = [ + layers[0], + kv.SsmLayerConfig( + layer_id=1, + reuse_policy=kv.ReusePolicy.PRIVATE, + buffers=[ + kv.BufferConfig(role=role, size=4096, reuse_policy=policy) + for role, policy in ( + ("state", kv.ReusePolicy.REQUIRED), + ("checkpoint", kv.ReusePolicy.OPTIONAL), + ("replay", None), + ) + ], + ), + ] config = kv.KVCacheManagerConfig( tokens_per_block=4, cache_tiers=[kv.GpuCacheTierConfig(quota=8 << 20)], @@ -128,7 +145,11 @@ def test_ssm_layer_reuse_policies(length, select_optional): assert source.resize(length) source.commit(tokens, is_end=True) source.close() - required, optional, private = [manager.get_layer_group_id(i) for i in (1, 2, 3)] + required, optional, private = ( + [manager.get_buffer_group_id(1, role) for role in ("state", "checkpoint", "replay")] + if split_layer + else [manager.get_layer_group_id(i) for i in (1, 2, 3)] + ) assert len({required, optional, private}) == 3 reader = manager.create_kv_cache(input_tokens=tokens) caches.append(reader) @@ -186,19 +207,25 @@ def test_ssm_buffer_reuse_policies(length): tokens = list(range(length)) caches = [] try: - required = manager.get_layer_group_id(1, "state") - private = manager.get_layer_group_id(layer_id=1, role="old_x") + required = manager.get_buffer_group_id(1, "state") + private = manager.get_buffer_group_id(layer_id=1, role="old_x") assert required != private - assert manager.get_layer_group_id(2, "conv") == required - assert manager.get_layer_group_id(2, "old_x") == private - assert manager.get_layer_group_id(0) == manager.get_layer_group_id(0, "key") + assert manager.get_buffer_group_id(2, "conv") == required + assert manager.get_buffer_group_id(2, "old_x") == private + assert manager.get_layer_group_id(0) == manager.get_buffer_group_id(0, "key") for layer_id in (1, 2): with pytest.raises(ValueError, match="layer groups"): manager.get_layer_group_id(layer_id) assert manager.num_layers == 3 assert sorted(manager.layer_ids) == [0, 1, 2] - assert sorted(manager.layer_grouping[required]) == [1, 2] - assert sorted(manager.layer_grouping[private]) == [1, 2] + with pytest.raises(ValueError, match="layer groups"): + _ = manager.layer_grouping + for group in (required, private): + assert sorted({buffer.layer_id for buffer in manager.buffer_grouping[group]}) == [1, 2] + grouped = [ + (buffer.layer_id, buffer.role) for group in manager.buffer_grouping for buffer in group + ] + assert len(grouped) == len(set(grouped)) == len(list(manager.all_buffer_ids)) source = manager.create_kv_cache(input_tokens=[]) caches.append(source) @@ -223,20 +250,50 @@ def test_ssm_buffer_reuse_policies(length): manager.shutdown() -@pytest.mark.parametrize("policy", [kv.ReusePolicy.PRIVATE, kv.ReusePolicy.OPTIONAL]) -def test_attention_buffer_reuse_policy_must_match_layer(policy): - def make_config(buffer_policy): - return kv.KVCacheManagerConfig( - tokens_per_block=4, - cache_tiers=[kv.GpuCacheTierConfig(quota=4 << 20)], - layers=[ - kv.AttentionLayerConfig( - layer_id=0, - buffers=[kv.BufferConfig(role="key", size=4096, reuse_policy=buffer_policy)], - ) - ], +@pytest.mark.parametrize("length", [7, 8]) +@pytest.mark.parametrize("select_optional", [False, True]) +def test_attention_buffer_reuse_policies(length, select_optional): + config = kv.KVCacheManagerConfig( + tokens_per_block=4, + cache_tiers=[kv.GpuCacheTierConfig(quota=4 << 20)], + layers=[ + kv.AttentionLayerConfig( + layer_id=0, + reuse_policy=kv.ReusePolicy.PRIVATE, + buffers=[ + kv.BufferConfig(role="key", size=4096, reuse_policy=kv.ReusePolicy.REQUIRED), + kv.BufferConfig(role="index", size=2048, reuse_policy=kv.ReusePolicy.OPTIONAL), + kv.BufferConfig(role="scratch", size=1024), + ], + ) + ], + ) + manager = kv.KVCacheManager(config) + stream = torch.cuda.Stream() + caches = [] + try: + required, optional, private = [ + manager.get_buffer_group_id(0, role) for role in ("key", "index", "scratch") + ] + assert len({required, optional, private}) == 3 + source = manager.create_kv_cache(input_tokens=[]) + caches.append(source) + assert source.resume(stream.cuda_stream) + assert source.resize(length) + source.commit(list(range(length)), is_end=True) + source.close() + reader = manager.create_kv_cache(input_tokens=list(range(length))) + caches.append(reader) + assert reader.num_committed_tokens == length + assert reader.reuse_status[required].complete + assert reader.reuse_status[optional].complete + assert not reader.reuse_status[private].complete + assert reader.resume( + stream.cuda_stream, optional_reuse_groups=[optional] if select_optional else [] ) - - make_config(kv.ReusePolicy.REQUIRED) - with pytest.raises(ValueError, match="only supported for SSM layers"): - make_config(policy) + reader.suspend() + assert reader.resume(stream.cuda_stream) + finally: + for cache in reversed(caches): + cache.close() + manager.shutdown() From 892c6c8addaee5ac7d2a4dcc7a21f6d7ed10d55b Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Wed, 28 Oct 2026 13:36:32 +0800 Subject: [PATCH 08/19] Cover split-buffer scratch allocation and transfer membership Signed-off-by: Jiagan Cheng --- docs/source/torch/kv_cache_manager.md | 51 +++++++++++++ .../disaggregation/resource/kv_extractor.py | 4 + .../_torch/disaggregation/resource/utils.py | 2 +- .../kv_cache/kv_cache_manager_v2.py | 4 +- .../kv_cache/test_dual_pool_kv_cache.py | 4 +- .../kv_cache/test_mamba_cache_manager.py | 2 +- .../unittest/disaggregated/test_extractor.py | 14 +++- .../test_minimax_m3_kv_transfer.py | 5 +- .../test_optional_checkpoint.py | 74 +++++++++++++++++++ 9 files changed, 152 insertions(+), 8 deletions(-) diff --git a/docs/source/torch/kv_cache_manager.md b/docs/source/torch/kv_cache_manager.md index f6928377ed3a..bbf4f1e13431 100644 --- a/docs/source/torch/kv_cache_manager.md +++ b/docs/source/torch/kv_cache_manager.md @@ -71,3 +71,54 @@ In `KvCacheCreator._create_kv_cache_manager`, the `KVCacheManager` is instantiat For local testing or proof of concept, update these lines to use your implementation. Then, test it to ensure the `PyExecutor` runs with your customized `KVCacheManager`. + + +## Native V2 buffer reuse policies + +Custom cache managers can override a layer's reuse policy per buffer through +`tensorrt_llm.runtime.kv_cache_manager_v2`. Both attention and SSM layers default +to `ReusePolicy.REQUIRED`; a buffer's `reuse_policy=None` inherits its layer's +policy. For example, a custom recurrent consumer with independent state indices +can configure: + +```python +from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, ReusePolicy, SsmLayerConfig + +layer = SsmLayerConfig( + layer_id=0, + reuse_policy=ReusePolicy.REQUIRED, + buffers=[ + BufferConfig(role="state", size=4096), + BufferConfig(role="checkpoint", size=2048, reuse_policy=ReusePolicy.OPTIONAL), + BufferConfig(role="scratch", size=1024, reuse_policy=ReusePolicy.PRIVATE), + ], +) +``` + +REQUIRED buffers constrain the reusable prefix. OPTIONAL buffers are candidates +at that endpoint; select their group IDs through `resume(optional_reuse_groups=...)` +on the first resume. An SSM snapshot must match the endpoint exactly. PRIVATE +buffers remain request-local and are never published as reusable checkpoints. +The consumer must initialize or reconstruct private and unselected optional state +before inference. Unselected optional snapshots count as neither hits nor misses. + +Use `manager.get_buffer_group_id(layer_id, role)` to resolve a buffer and +`manager.buffer_grouping[group_id]` to enumerate membership. Buffers with the same +lifecycle properties share a group, including across layers; selecting a group +selects all of its buffers. Policy overrides add at most three policy variants +per attention window/sink/sparsity combination and at most three SSM groups. +Coalescing remains within lifecycle groups and compatible buffer sizes/coalescing +groups. + +`get_layer_group_id(layer_id)` and `layer_grouping` are deprecated in documentation +only and produce no deprecation warnings. The former returns a uniform layer's +group and raises `ValueError` for a split layer. The latter preserves its +partition-of-layers result and raises if any layer is split. `LayerGroupId` retains +its name; it identifies a lifecycle group, not a physical storage pool. + +Independent policies require independent indices in the consumer. The default +attention view requires compatible policies for K/V and scales because they share +one page table. Current Mamba kernels similarly couple recurrent, convolution, +PLE and replay buffers through one state-slot index. These consumers reject +incompatible policies explicitly; this native configuration does not enable +separate replay indices in those kernels or add an LLM API configuration field. diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 20eecdc10fa3..63a68ffcea1c 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -473,6 +473,10 @@ def _window_size_for_layer(internal_layer_id: int): ignored_roles=ignored_roles, view_partition=view_partition, ) + if not pool_views: + # Ignored-only groups keep their lifecycle index but have no + # transferable layers or regions to match against a peer. + local_layers = [] for pool_view in pool_views: if pool_view.mapper_kind not in _MAPPER_KINDS_BY_CACHE_KIND[cache_kind]: raise ValueError( diff --git a/tensorrt_llm/_torch/disaggregation/resource/utils.py b/tensorrt_llm/_torch/disaggregation/resource/utils.py index 341f6c79d8f8..579cc42d1d13 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/utils.py +++ b/tensorrt_llm/_torch/disaggregation/resource/utils.py @@ -301,7 +301,7 @@ def get_layer_to_layer_group( kind = CacheKind.PAGED out: Dict[int, int] = {} for lg_idx, lg in enumerate(page_table.layer_groups): - if lg.kind == kind and lg.pool_views: + if lg.kind == kind: for ll in lg.local_layers: gid = int(ll.global_layer_id) if gid in out: 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 509fdc973b15..f9cbf1aa7a23 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 @@ -3391,7 +3391,7 @@ def _get_event_window_sizes_by_layer_group( self, attention_only: bool = False ) -> Dict[int, int]: # Assumes every layer listed in a group shares the same sliding_window_size, - # which is how `impl.layer_grouping` groups layers today. Only the + # which is part of the lifecycle identity. Only the # first layer's window is read; if the grouping policy ever permits # mixed windows in one group, this needs to fan out per-layer. A layer # can belong to several groups through buffer policy overrides. @@ -5556,7 +5556,7 @@ def _life_cycle_by_layer_group(self) -> List[AttnLifeCycle]: """Attention life cycle per layer group. The layers listed in a group share its window and sinks, which is how - ``impl.layer_grouping`` groups layers, so the first layer describes the + the lifecycle identity groups buffers, so the first layer describes the group, including when buffer policies split a layer. A group with no window -- full attention, or a group whose layers carry no window at all -- yields an empty stale range. diff --git a/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py b/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py index 621d27f1d83b..82500f3d79b8 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py @@ -449,7 +449,9 @@ def test_create_cross_kv_cache_manager_uses_tp_sharded_encoder_heads(self): distributed = Mock() distributed.allreduce.side_effect = lambda value, op: value impl = Mock( - layer_grouping=[list(range(num_decoder_layers))], + buffer_grouping=[ + [SimpleNamespace(layer_id=layer, role="key") for layer in range(num_decoder_layers)] + ], pool_group_descs=[], all_buffer_ids=[], ) 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 945df03801e8..b78a4b1f8dc9 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 @@ -1391,7 +1391,7 @@ def test_v2_disagg_slice_reads_state_index_without_refreshing_batch_mask(): side_effect=AssertionError("state-index lookup must not refresh the dummy mask") ) # Provide a mamba layer group so _create_chunk places the slot ID - mamba_lg = SimpleNamespace(kind=CacheKind.STATE) + mamba_lg = SimpleNamespace(kind=CacheKind.STATE, pool_views=[object()]) transceiver = object.__new__(KvCacheTransceiverV2) transceiver._kv_cache_manager = manager transceiver._reuse_adapter = SimpleNamespace(tokens_per_block=32) diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index ec6ef42ad0f7..d4154e8a5e74 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -300,7 +300,9 @@ def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_l ) impl = SimpleNamespace( - layer_grouping=((0, 1),), + buffer_grouping=( + [buffer for coalesced in coalesced_buffers for buffer in coalesced.buffer_ids], + ), init_config=SimpleNamespace( tokens_per_block=16, layers=[SimpleNamespace(window_size=None), SimpleNamespace(window_size=None)], @@ -597,7 +599,15 @@ def buf(layer_id, role): manager = object.__new__(Qwen4ExpHybridCacheManager) manager.impl = SimpleNamespace( - layer_grouping=((0, 1), (2,)), + buffer_grouping=tuple( + [ + buffer + for variant in pg.slot_desc.variants + for coalesced in variant.coalesced_buffers + for buffer in coalesced.buffer_ids + ] + for pg in (recurrent_pg, attention_pg) + ), init_config=SimpleNamespace( tokens_per_block=16, layers=[SimpleNamespace(window_size=None) for _ in range(3)], diff --git a/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py index 559500719639..4372c4813011 100644 --- a/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py +++ b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py @@ -177,7 +177,10 @@ class _StubGroupingImpl: """Stub ``manager.impl`` exposing only what ``_kv_pool_mapping_offset`` reads.""" def __init__(self, grouping, key_addrs): - self.layer_grouping = grouping + self.buffer_grouping = [ + [SimpleNamespace(layer_id=layer, role=Role.KEY) for layer in group] + for group in grouping + ] self._key_addrs = key_addrs def get_mem_pool_base_address(self, layer_id, role, index_mode=None): 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 581ea0b2bbc7..86c95d75deb9 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 @@ -151,6 +151,7 @@ def test_ssm_layer_reuse_policies(length, select_optional, split_layer): else [manager.get_layer_group_id(i) for i in (1, 2, 3)] ) assert len({required, optional, private}) == 3 + manager.get_and_reset_ssm_snapshot_iteration_stats() reader = manager.create_kv_cache(input_tokens=tokens) caches.append(reader) assert reader.num_committed_tokens == length @@ -161,6 +162,14 @@ def test_ssm_layer_reuse_policies(length, select_optional, split_layer): assert reader.resume( stream.cuda_stream, optional_reuse_groups=[optional] if select_optional else [] ) + reader.commit_pending_stats() + stats = manager.get_and_reset_ssm_snapshot_iteration_stats() + assert stats[required].iter_snapshot_hits == 1 + if select_optional: + assert stats[optional].iter_snapshot_lookups == stats[optional].iter_snapshot_hits == 1 + elif optional in stats: + assert stats[optional].iter_snapshot_lookups == 0 + assert stats[optional].iter_snapshot_hits == stats[optional].iter_snapshot_misses == 0 indices = [ reader.get_ssm_block_base_index(group) for group in (required, optional, private) ] @@ -297,3 +306,68 @@ def test_attention_buffer_reuse_policies(length, select_optional): for cache in reversed(caches): cache.close() manager.shutdown() + + +@pytest.mark.parametrize("split", [False, True]) +def test_scratch_addresses_with_unequal_buffer_membership(split): + layers = [] + for layer in range(3): + buffers = [kv.BufferConfig(role="key", size=2048, tokens_per_block_override=2)] + if layer < 2: + buffers.extend( + [ + kv.BufferConfig(role="value", size=4096), + kv.BufferConfig(role="side", size=512, coalescing_group=7), + ] + ) + buffers.append( + kv.BufferConfig( + role="private", size=1024, reuse_policy=kv.ReusePolicy.PRIVATE if split else None + ) + ) + layers.append( + kv.AttentionLayerConfig(layer_id=layer, buffers=buffers, sliding_window_size=8) + ) + manager = kv.KVCacheManager( + kv.KVCacheManagerConfig( + tokens_per_block=4, + layers=layers, + cache_tiers=[kv.GpuCacheTierConfig(quota=4 << 20)], + swa_scratch_reuse=kv.SwaScratchReuseConfig(), + ) + ) + stream = torch.cuda.Stream() + cache = manager.create_kv_cache(input_tokens=list(range(64))) + try: + assert cache.resume(stream.cuda_stream) + assert cache.resize(64) + required = manager.get_buffer_group_id(0, "key") + scratch = cache.get_scratch_desc(required) + assert scratch is not None + blocks = scratch.range.end - scratch.range.beg + # The side pool has one buffer for each of two layers and needs half + # a slot per logical block, more than K/V's two-fifths share. + assert len(scratch.slot_ids) == (blocks + 1) // 2 + for layer in layers: + spans = [] + for buffer in layer.buffers: + group = manager.get_buffer_group_id(layer.layer_id, buffer.role) + converter = manager.get_page_index_converter(layer.layer_id, buffer.role) + indices = converter( + cache.get_base_page_indices(group), + kv.PageIndexMode.PER_LAYER, + cache.get_scratch_desc(group), + ) + base = manager.get_mem_pool_base_address( + layer.layer_id, buffer.role, kv.PageIndexMode.PER_LAYER + ) + stride = manager.get_page_stride(layer.layer_id, buffer.role) + assert len(indices) == 16 * converter.expansion + for index in indices: + assert index >= 0 + spans.append((base + index * stride, base + (index + 1) * stride)) + spans.sort() + assert all(end <= following for (_, end), (following, _) in zip(spans, spans[1:])) + finally: + cache.close() + manager.shutdown() From 7ef993161d1a7a7c609536fedd1f96c8c43eb799 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Wed, 28 Oct 2026 13:39:47 +0800 Subject: [PATCH 09/19] Reject ambiguous connector layer lookup and cover split transfer layout Signed-off-by: Jiagan Cheng --- .../_torch/pyexecutor/connectors/kv_cache_layout.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py index 8c627f5a0dec..910c01dc8d00 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py @@ -188,10 +188,15 @@ def group(self, layer_group_id: int) -> KvCacheLayerGroupLayout: raise KeyError(f"no layer group {layer_group_id} in layout") def group_of_layer(self, layer_id: int) -> KvCacheLayerGroupLayout: - """The layer group owning a global model layer index.""" - for group in self.groups: - if layer_id in group.layer_ids: - return group + """Deprecated: resolve a buffer group in KVCM and call group(). + + Returns a uniform layer's group; raises ValueError for a split layer. + """ + matches = [group for group in self.groups if layer_id in group.layer_ids] + if len(matches) > 1: + raise ValueError(f"layer {layer_id} spans several groups; resolve a buffer group first") + if matches: + return matches[0] raise KeyError(f"layer {layer_id} is not covered by this layout") def as_single_pool_tensor(self) -> Optional[torch.Tensor]: From 0e8a9cc4f007320462694fdc0e6fcb56db7f193e Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Wed, 28 Oct 2026 14:10:31 +0800 Subject: [PATCH 10/19] Preserve group membership order and pack unequal scratch shares safely Signed-off-by: Jiagan Cheng --- .../kv_cache_manager_v2/kvCacheManager.cpp | 22 ++++++++++-- docs/source/torch/kv_cache_manager.md | 9 ++++- .../disaggregation/resource/kv_extractor.py | 10 +++--- .../kv_cache/kv_cache_manager_v2.py | 16 ++++++++- .../pyexecutor/kv_cache/sharing/_manager.py | 3 +- .../test_kv_cache_v2_extra_buffers.py | 7 ++-- .../kv_cache/test_kvcm2_integration.py | 36 +++++++++++++++++++ .../test_optional_checkpoint.py | 8 ++++- 8 files changed, 98 insertions(+), 13 deletions(-) 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 5ceae5c77b0a..6034a50bfbd7 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 @@ -77,12 +77,20 @@ std::vector PageIndexConverter::operator()(std::vector const& baseIndi // base_indices. int blockPos = ordinal - scratchRange.beg; int totalOffset = blockPos * scratchPages; + bool const crossesSlots = scale % scratchPages != 0; + if (crossesSlots) + { + // Unequal layer shares can put a block across two slots. Pack + // the layer's consecutive roles before resolving the slot ID, + // so a trailing role cannot wrap over an earlier block's data. + totalOffset += appliedLayerOffset % scratchPages; + } int slotIdx = totalOffset / scale; TLLM_CHECK_DEBUG_WITH_INFO(slotIdx < scratchSlots.size(), "Scratch slot %d past the %d the range describes", slotIdx, scratchSlots.size()); int slotId = scratchSlots[slotIdx]; int offset = totalOffset % scale; - index = slotId * scale + (offset + appliedLayerOffset) % scale; + index = slotId * scale + (crossesSlots ? offset : (offset + appliedLayerOffset) % scale); } else if (baseIndices[toSizeT(ordinal)] == kBadPageIndex.value()) { @@ -829,9 +837,19 @@ TypedVec> KvCacheManager::layerGrouping() con TypedVec> KvCacheManager::bufferGrouping() const { TypedVec> result(mLifeCycles.size()); + std::unordered_map> buffersByLayer; for (auto const& buffer : allBufferIds()) { - result.at(getBufferGroupId(buffer.layerId, buffer.role)).push_back(buffer); + buffersByLayer[buffer.layerId].push_back(buffer); + } + // Keep layer order consistent with layerGrouping() for consumers that pair + // group membership with a declared list of global model layer IDs. + for (auto const& [layer, lifeCycles] : mStorage->layerToLifeCycleIds()) + { + for (auto const& buffer : buffersByLayer.at(layer)) + { + result.at(getBufferGroupId(buffer.layerId, buffer.role)).push_back(buffer); + } } return result; } diff --git a/docs/source/torch/kv_cache_manager.md b/docs/source/torch/kv_cache_manager.md index bbf4f1e13431..1dfd710fe3b1 100644 --- a/docs/source/torch/kv_cache_manager.md +++ b/docs/source/torch/kv_cache_manager.md @@ -97,7 +97,8 @@ layer = SsmLayerConfig( REQUIRED buffers constrain the reusable prefix. OPTIONAL buffers are candidates at that endpoint; select their group IDs through `resume(optional_reuse_groups=...)` -on the first resume. An SSM snapshot must match the endpoint exactly. PRIVATE +on the first resume. Attention candidates retain the block-aligned sliding-window +checkpoint contract; an SSM snapshot must match the endpoint exactly. PRIVATE buffers remain request-local and are never published as reusable checkpoints. The consumer must initialize or reconstruct private and unselected optional state before inference. Unselected optional snapshots count as neither hits nor misses. @@ -122,3 +123,9 @@ one page table. Current Mamba kernels similarly couple recurrent, convolution, PLE and replay buffers through one state-slot index. These consumers reject incompatible policies explicitly; this native configuration does not enable separate replay indices in those kernels or add an LLM API configuration field. + +The native page-index converter also supports unequal layer shares in scratch +slots, including a layer's block crossing two physical slots. Existing attention +scratch-table kernels require whole layer blocks per slot. The wrapper rejects +unequal shares that require crossing slots; use separate `coalescing_group` values +or disable SWA scratch reuse for those layouts. diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 63a68ffcea1c..d613a19f68e2 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -187,8 +187,8 @@ def _compute_global_layer_ids(manager, lg_idx: int) -> List[int]: (model_layer, attn_type) into synthetic IDs via the _layer_attn_to_layer_id inverse mapping. """ - local_layer_ids = sorted( - {int(buffer.layer_id) for buffer in manager.impl.buffer_grouping[lg_idx]} + local_layer_ids = list( + dict.fromkeys(int(buffer.layer_id) for buffer in manager.impl.buffer_grouping[lg_idx]) ) get_global_layer_ids = getattr(manager, "get_disagg_global_layer_ids", None) @@ -433,8 +433,10 @@ def _window_size_for_layer(internal_layer_id: int): # memory is reused, not concatenated. for variant in pg_desc.slot_desc.variants: layer_group_id = int(variant.layer_group_id) - all_internal_layer_ids = sorted( - {int(buffer.layer_id) for buffer in manager.impl.buffer_grouping[layer_group_id]} + all_internal_layer_ids = list( + dict.fromkeys( + int(buffer.layer_id) for buffer in manager.impl.buffer_grouping[layer_group_id] + ) ) all_global_layer_ids = _compute_global_layer_ids(manager, layer_group_id) local_layers = [ 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 f9cbf1aa7a23..4c99240fb9dc 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 @@ -2584,6 +2584,20 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: # release half of `preempt_request` has to wait. self._pending_preemption: Dict[int, LlmRequest] = {} self._use_per_layer_page_tables = False + if self.enable_swa_scratch_reuse: + for layer in config.layers: + if not isinstance(layer, AttentionLayerConfig) or layer.sliding_window_size is None: + continue + for buffer in layer.buffers: + if buffer.is_sparse: + continue + converter = self.impl.get_page_index_converter(layer.layer_id, buffer.role) + if converter.scale % converter.scratch_pages_per_block: + raise ValueError( + "The attention scratch-table kernels require whole layer blocks per slot; " + f"layer {layer.layer_id}, role {buffer.role} has an unequal coalesced share. " + "Use separate coalescing groups or disable SWA scratch reuse." + ) self._prepare_page_table_tensor(index_mapper_capacity) self._sparse_layer_group_ids = tuple( sorted( @@ -2910,7 +2924,7 @@ def _layer_ids_by_group(self) -> list[list[int]]: # A split layer can occur in several groups. This is membership metadata, # not the partition promised by the deprecated layer_grouping API. return [ - sorted({int(buffer.layer_id) for buffer in buffers}) + list(dict.fromkeys(int(buffer.layer_id) for buffer in buffers)) for buffers in self.impl.buffer_grouping ] diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py index 46dcdfaa31e0..8e911bbcd265 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py @@ -167,7 +167,8 @@ def layer_configs(manager: KVCacheManagerV2) -> dict[int, object]: def layer_group_layer_ids(manager: KVCacheManagerV2) -> list[list[int]]: """Internal layer ids per runtime layer group, in layer-group order.""" return [ - sorted({int(buffer.layer_id) for buffer in group}) for group in manager.impl.buffer_grouping + list(dict.fromkeys(int(buffer.layer_id) for buffer in group)) + for group in manager.impl.buffer_grouping ] diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py index 8a8f0b47570f..150f201b84f3 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py @@ -168,9 +168,10 @@ def test_page_table_uses_physical_pool_representative(self): layers = real_impl.layer_grouping[pool_id] other_layer = next(layer for layer in layers if layer != physical_layer) impl_proxy = Mock(wraps=real_impl) - impl_proxy.layer_grouping = ( - (other_layer, *(layer for layer in layers if layer != other_layer)), - ) + impl_proxy.buffer_grouping = [ + sorted(group, key=lambda buffer: buffer.layer_id != other_layer) + for group in real_impl.buffer_grouping + ] mgr.impl = impl_proxy mgr._prepare_page_table_tensor(index_mapper_capacity=1) 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 b34f9a298cc6..e0e4e321a9b0 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -4123,3 +4123,39 @@ def build_config(self, config): dtype=DataType.HALF, vocab_size=16, ) + + +def test_attention_scratch_rejects_layout_requiring_split_slot_indices(monkeypatch): + original = KVCacheManagerV2._build_cache_config + + def build_config(self, config): + config = original(self, config) + layers = list(config.layers) + layer = layers[0] + layers[0] = AttentionLayerConfig( + layer_id=layer.layer_id, + sliding_window_size=layer.sliding_window_size, + buffers=[*layer.buffers, BufferConfig(role="side", size=layer.buffers[0].size)], + ) + return replace(config, layers=layers) + + monkeypatch.setattr(KVCacheManagerV2, "_build_cache_config", build_config) + with pytest.raises(ValueError, match="require whole layer blocks per slot"): + KVCacheManagerV2( + KvCacheConfig( + max_gpu_total_bytes=4 << 20, + host_cache_size=0, + max_attention_window=[8, 8], + enable_swa_scratch_reuse=True, + ), + CacheType.SELF, + num_layers=2, + num_kv_heads=1, + head_dim=64, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=1, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.HALF, + vocab_size=16, + ) 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 86c95d75deb9..fc361e9fcd9f 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 @@ -269,6 +269,7 @@ def test_attention_buffer_reuse_policies(length, select_optional): kv.AttentionLayerConfig( layer_id=0, reuse_policy=kv.ReusePolicy.PRIVATE, + sliding_window_size=4, buffers=[ kv.BufferConfig(role="key", size=4096, reuse_policy=kv.ReusePolicy.REQUIRED), kv.BufferConfig(role="index", size=2048, reuse_policy=kv.ReusePolicy.OPTIONAL), @@ -295,8 +296,13 @@ def test_attention_buffer_reuse_policies(length, select_optional): caches.append(reader) assert reader.num_committed_tokens == length assert reader.reuse_status[required].complete - assert reader.reuse_status[optional].complete + optional_complete = length % config.tokens_per_block == 0 + assert reader.reuse_status[optional].complete == optional_complete assert not reader.reuse_status[private].complete + if select_optional and not optional_complete: + with pytest.raises(ValueError): + reader.resume(stream.cuda_stream, optional_reuse_groups=[optional]) + select_optional = False assert reader.resume( stream.cuda_stream, optional_reuse_groups=[optional] if select_optional else [] ) From a4b257eb3e32d3fe4ee2007d3a276cb7bc63c6fd Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Wed, 28 Oct 2026 16:24:30 +0800 Subject: [PATCH 11/19] Register lifecycles only for configured buffers Signed-off-by: Jiagan Cheng --- .../kv_cache_manager_v2/lifeCycleRegistry.cpp | 24 +++++-------------- .../kv_cache_manager_v2/lifeCycleRegistry.h | 3 --- 2 files changed, 6 insertions(+), 21 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp index 00943894f8f3..72899e1da5b6 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp @@ -27,32 +27,25 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 // makeLifeCycle — factory dispatching on LayerConfig variant. // --------------------------------------------------------------------------- -LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock) +LifeCycle makeLifeCycle(LayerConfig const& layer, BufferConfig const& buffer, int tokensPerBlock) { return std::visit( [&](auto const& cfg) -> LifeCycle { - cfg.validate(); + auto const policy = buffer.effectiveReusePolicy(cfg.reusePolicy); using T = std::decay_t; if constexpr (std::is_same_v) { - return SsmLifeCycle{cfg.reusePolicy}; + return SsmLifeCycle{policy}; } else { - return AttnLifeCycle::make(cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, cfg.reusePolicy); + return AttnLifeCycle::make(cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, policy); } }, layer); } -LifeCycle makeLifeCycle(LayerConfig const& layer, BufferConfig const& buffer, int tokensPerBlock) -{ - LifeCycle lc = makeLifeCycle(layer, tokensPerBlock); - std::visit([&buffer](auto& value) { value.reusePolicy = buffer.effectiveReusePolicy(value.reusePolicy); }, lc); - return lc; -} - // --------------------------------------------------------------------------- // LifeCycleRegistry // --------------------------------------------------------------------------- @@ -73,13 +66,8 @@ LifeCycleRegistry::LifeCycleRegistry(KVCacheManagerConfig const& config) }; for (auto const& layer : config.layers) { - auto const& buffers = layerBuffers(layer); - // A layer without buffers registers its layer-level lifecycle. - if (buffers.empty()) - { - registerLifeCycle(makeLifeCycle(layer, config.tokensPerBlock)); - } - for (auto const& buffer : buffers) + std::visit([](auto const& cfg) { cfg.validate(); }, layer); + for (auto const& buffer : layerBuffers(layer)) { registerLifeCycle(makeLifeCycle(layer, buffer, config.tokensPerBlock)); } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h index 1aeeaad4b1af..47137af632f3 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h @@ -171,9 +171,6 @@ inline HalfOpenRange computeScratchRange( return intersect(capStale, inputRange); } -// Factory: create the LifeCycle of a layer, using the layer's reuse policy. -LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock); - // Factory: create the LifeCycle that holds one buffer of a layer, using the buffer's effective reuse policy. LifeCycle makeLifeCycle(LayerConfig const& layer, BufferConfig const& buffer, int tokensPerBlock); From 16773f0d08dec01a08d04c9a94d2945750f1adb0 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Wed, 28 Oct 2026 16:29:13 +0800 Subject: [PATCH 12/19] Defer Python consumer adoption of per-buffer reuse policies Signed-off-by: Jiagan Cheng --- docs/source/torch/kv_cache_manager.md | 19 +-- .../backends/fp4_mla/cache_manager.py | 2 +- .../backends/sparse/csa2/cache_manager.py | 7 +- .../sparse/minimax_m3/cache_manager.py | 2 +- .../disaggregation/native/mixers/ssm/peer.py | 13 +- .../disaggregation/resource/kv_extractor.py | 16 +- .../_torch/disaggregation/transceiver.py | 2 +- .../pyexecutor/connectors/kv_cache_layout.py | 25 +-- .../kv_cache/kv_cache_manager_v2.py | 105 ++++--------- .../kv_cache/mixed_precision_staging.py | 3 +- .../pyexecutor/kv_cache/sharing/_manager.py | 7 +- .../kv_cache/sharing/test_layout_identity.py | 5 +- .../kv_cache/sharing/test_public_surface.py | 10 +- .../kv_cache/test_dual_pool_kv_cache.py | 4 +- .../test_kv_cache_v2_extra_buffers.py | 7 +- .../kv_cache/test_kvcm2_integration.py | 145 +----------------- .../kv_cache/test_mamba_cache_manager.py | 2 +- .../unittest/disaggregated/test_extractor.py | 14 +- .../test_minimax_m3_kv_transfer.py | 5 +- 19 files changed, 82 insertions(+), 311 deletions(-) diff --git a/docs/source/torch/kv_cache_manager.md b/docs/source/torch/kv_cache_manager.md index 1dfd710fe3b1..2a359a4bff4c 100644 --- a/docs/source/torch/kv_cache_manager.md +++ b/docs/source/torch/kv_cache_manager.md @@ -117,15 +117,16 @@ group and raises `ValueError` for a split layer. The latter preserves its partition-of-layers result and raises if any layer is split. `LayerGroupId` retains its name; it identifies a lifecycle group, not a physical storage pool. -Independent policies require independent indices in the consumer. The default -attention view requires compatible policies for K/V and scales because they share -one page table. Current Mamba kernels similarly couple recurrent, convolution, -PLE and replay buffers through one state-slot index. These consumers reject -incompatible policies explicitly; this native configuration does not enable -separate replay indices in those kernels or add an LLM API configuration field. +Existing Python cache-manager consumers continue using the legacy layer APIs +and require one lifecycle per layer. They do not adopt per-buffer policy splits; +the legacy API rejects ambiguous layer mappings. Consumers can migrate to the +buffer APIs when they add support for independent indices. Current Mamba kernels +couple recurrent, convolution, PLE and replay buffers through one state-slot +index; separate replay indices are a follow-up. This native configuration does +not add an LLM API configuration field. The native page-index converter also supports unequal layer shares in scratch slots, including a layer's block crossing two physical slots. Existing attention -scratch-table kernels require whole layer blocks per slot. The wrapper rejects -unequal shares that require crossing slots; use separate `coalescing_group` values -or disable SWA scratch reuse for those layouts. +scratch-table kernels require whole layer blocks per slot. Custom consumers of +unequal shares that cross slots must use the native converter or supply compatible +index generation; existing Python consumers have not adopted that layout. diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py index 715da9a31172..9f5e80357b7a 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py @@ -535,7 +535,7 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: if fp8_pool_id in (cache_pool_id, hp_pool_id): raise RuntimeError("FP4 MLA, FP8 MLA, and HP state must use distinct lifecycles.") - num_pools = len(self._layer_ids_by_group) + num_pools = len(self.impl.layer_grouping) pool_pointers = [[[0, 0], [0, 0]] for _ in range(num_pools)] pool_pointers[cache_pool_id] = [ [ diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py index 0656fb2bbc95..7e343063105d 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py @@ -266,7 +266,7 @@ def __init__( self._encoder_optional_groups = tuple( sorted( { - self.impl.get_buffer_group_id(self._layer_roles[key], key[2].role) + self.impl.get_layer_group_id(self._layer_roles[key]) for key in self._encoder_roles if self._encoder_replay_enabled } @@ -277,11 +277,8 @@ def __init__( def _log_checkpoint_pools(self): """Report physical quota once per group, without counting role views twice.""" policies = { - self.impl.get_buffer_group_id(layer.layer_id, buffer.role): layer.reuse_policy - if buffer.reuse_policy is None - else buffer.reuse_policy + self.impl.get_layer_group_id(layer.layer_id): layer.reuse_policy for layer in self.impl.init_config.layers - for buffer in layer.buffers } for level, tier in enumerate(self.impl.cache_tier_list): grouping = self.impl.get_life_cycle_pool_group_indices(level) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index 63bb2b9b2572..7722f76ea664 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -1033,7 +1033,7 @@ def _kv_pool_mapping_offset(self, layer_id, layer_group_id, key_base_addr) -> in pool-mapping loop meaningful). """ layers_by_addr = sorted( - self._layer_ids_by_group[int(layer_group_id)], + self.impl.layer_grouping[int(layer_group_id)], key=lambda lid: self.impl.get_mem_pool_base_address( lid, Role.KEY, PageIndexMode.SHARED ), diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py index 7208870eaf51..918252a94578 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py @@ -478,15 +478,10 @@ def _find_mamba_layer_group( ) -> Optional[MambaLayerGroup]: if page_table is None: return None - groups = [ - lg for lg in page_table.layer_groups if lg.kind == CacheKind.STATE and lg.pool_views - ] - if len(groups) > 1: - raise ValueError( - "Mamba transfer requires one recurrent-state lifecycle; " - "multiple transferable state groups need separate slot indices" - ) - return groups[0] if groups else None + return next( + (lg for lg in page_table.layer_groups if lg.kind == CacheKind.STATE), + None, + ) @staticmethod def validate_peer_compatible( diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index d613a19f68e2..279916d53001 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -187,9 +187,7 @@ def _compute_global_layer_ids(manager, lg_idx: int) -> List[int]: (model_layer, attn_type) into synthetic IDs via the _layer_attn_to_layer_id inverse mapping. """ - local_layer_ids = list( - dict.fromkeys(int(buffer.layer_id) for buffer in manager.impl.buffer_grouping[lg_idx]) - ) + local_layer_ids = manager.impl.layer_grouping[lg_idx] get_global_layer_ids = getattr(manager, "get_disagg_global_layer_ids", None) if get_global_layer_ids is not None: @@ -408,7 +406,7 @@ def _window_size_for_layer(internal_layer_id: int): pool_groups: List[PhysicalPoolGroup] = [] storage_pg_to_list_idx: Dict[int, int] = {} - layer_groups_by_id: List[LayerGroup | None] = [None] * len(manager.impl.buffer_grouping) + layer_groups_by_id: List[LayerGroup | None] = [None] * len(manager.impl.layer_grouping) for pg_desc in pool_group_descs: storage_pg_idx = int(pg_desc.pool_group_index) @@ -433,11 +431,7 @@ def _window_size_for_layer(internal_layer_id: int): # memory is reused, not concatenated. for variant in pg_desc.slot_desc.variants: layer_group_id = int(variant.layer_group_id) - all_internal_layer_ids = list( - dict.fromkeys( - int(buffer.layer_id) for buffer in manager.impl.buffer_grouping[layer_group_id] - ) - ) + all_internal_layer_ids = list(manager.impl.layer_grouping[layer_group_id]) all_global_layer_ids = _compute_global_layer_ids(manager, layer_group_id) local_layers = [ LocalLayer(local_layer_id=int(iid), global_layer_id=int(gid)) @@ -475,10 +469,6 @@ def _window_size_for_layer(internal_layer_id: int): ignored_roles=ignored_roles, view_partition=view_partition, ) - if not pool_views: - # Ignored-only groups keep their lifecycle index but have no - # transferable layers or regions to match against a peer. - local_layers = [] for pool_view in pool_views: if pool_view.mapper_kind not in _MAPPER_KINDS_BY_CACHE_KIND[cache_kind]: raise ValueError( diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 0ffb10dcdfb9..759d408fa362 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -675,7 +675,7 @@ def _describe_local(self, req: LlmRequest) -> Chunk: kinds = [lg.kind for lg in layer_groups] for idx, lg in enumerate(layer_groups): if lg.kind == CacheKind.STATE: - slot = self._get_mamba_slot_for_request(req) if lg.pool_views else None + slot = self._get_mamba_slot_for_request(req) group = np.array([slot], dtype=np.int64) if slot is not None else empty else: # Block lists carry beam 0 only (beam-search attention reads diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py index 910c01dc8d00..b0749a168390 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_layout.py @@ -188,15 +188,10 @@ def group(self, layer_group_id: int) -> KvCacheLayerGroupLayout: raise KeyError(f"no layer group {layer_group_id} in layout") def group_of_layer(self, layer_id: int) -> KvCacheLayerGroupLayout: - """Deprecated: resolve a buffer group in KVCM and call group(). - - Returns a uniform layer's group; raises ValueError for a split layer. - """ - matches = [group for group in self.groups if layer_id in group.layer_ids] - if len(matches) > 1: - raise ValueError(f"layer {layer_id} spans several groups; resolve a buffer group first") - if matches: - return matches[0] + """The layer group owning a global model layer index.""" + for group in self.groups: + if layer_id in group.layer_ids: + return group raise KeyError(f"layer {layer_id} is not covered by this layout") def as_single_pool_tensor(self) -> Optional[torch.Tensor]: @@ -266,7 +261,7 @@ def build_kv_cache_layout_v2(manager: "KVCacheManagerV2") -> KvCacheLayout: so the addresses are device addresses only because tiers below GPU are rejected while a connector is attached (``PyExecutor._reject_non_gpu_cache_tiers``). Reads only the manager's - public layout API -- ``buffer_grouping``, ``all_buffer_ids``, + public layout API -- ``layer_grouping``, ``all_buffer_ids``, ``get_aggregated_pages``, ``pool_group_descs`` -- so it assumes nothing about dimension order, kv factor, or pool count. """ @@ -281,15 +276,21 @@ def build_kv_cache_layout_v2(manager: "KVCacheManagerV2") -> KvCacheLayout: for variant in pool_group.slot_desc.variants: slots_by_group[int(variant.layer_group_id)] = int(pool_group.num_slots) + buffers_by_layer: Dict[int, List] = {} + for buffer_id in impl.all_buffer_ids: + buffers_by_layer.setdefault(int(buffer_id.layer_id), []).append(buffer_id) + groups: List[KvCacheLayerGroupLayout] = [] - for layer_group_id, buffer_ids in enumerate(impl.buffer_grouping): - local_layer_ids = sorted({int(buffer.layer_id) for buffer in buffer_ids}) + for layer_group_id, local_layer_ids in enumerate(impl.layer_grouping): + local_layer_ids = [int(lid) for lid in local_layer_ids] if not local_layer_ids: continue num_slots = slots_by_group[layer_group_id] global_by_local = dict(zip(local_layer_ids, _global_layer_ids(manager, local_layer_ids))) + buffer_ids = [b for lid in local_layer_ids for b in buffers_by_layer.get(lid, ())] + regions: List[KvCacheRegion] = [] for desc in impl.get_aggregated_pages(buffer_ids): if int(desc.layer_group_id) != layer_group_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 4c99240fb9dc..562419a1185b 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 @@ -2405,10 +2405,10 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self._get_event_layer_group_ids(), ) - # Buffer membership has no physical ordering guarantee. The layer order + # Both backends build layer_grouping on demand, and the layer order # within a group is not part of its contract. Cache a stable physical- # layout representative for each role from the public pool descriptors. - self.num_pools = len(self._layer_ids_by_group) + self.num_pools = len(self.impl.layer_grouping) self._pool_layer_ids_by_role: Dict[Tuple[int, DataRole], LayerId] = {} for pool_group in self.impl.pool_group_descs: for variant in pool_group.slot_desc.variants: @@ -2418,16 +2418,10 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self._pool_layer_ids_by_role.setdefault( (pool_id, buffer_id.role), buffer_id.layer_id ) - groups_by_layer: dict[int, set[int]] = defaultdict(set) - self.buffer_to_pool_mapping_dict: dict[tuple[int, DataRole], int] = {} - for group_id, buffers in enumerate(self.impl.buffer_grouping): - for buffer in buffers: - groups_by_layer[int(buffer.layer_id)].add(group_id) - self.buffer_to_pool_mapping_dict[(int(buffer.layer_id), buffer.role)] = group_id - self.layer_to_pool_mapping_dict = { - layer_id: next(iter(groups)) - for layer_id, groups in groups_by_layer.items() - if len(groups) == 1 + num_layers = len(config.layers) + self.layer_to_pool_mapping_dict: dict[int, int] = { + layer_id: self.impl.get_layer_group_id(layer_id) + for layer_id in typed_range(LayerId(num_layers)) } self.kv_cache_map: dict[int, _KVCache] = {} @@ -2584,25 +2578,11 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: # release half of `preempt_request` has to wait. self._pending_preemption: Dict[int, LlmRequest] = {} self._use_per_layer_page_tables = False - if self.enable_swa_scratch_reuse: - for layer in config.layers: - if not isinstance(layer, AttentionLayerConfig) or layer.sliding_window_size is None: - continue - for buffer in layer.buffers: - if buffer.is_sparse: - continue - converter = self.impl.get_page_index_converter(layer.layer_id, buffer.role) - if converter.scale % converter.scratch_pages_per_block: - raise ValueError( - "The attention scratch-table kernels require whole layer blocks per slot; " - f"layer {layer.layer_id}, role {buffer.role} has an unequal coalesced share. " - "Use separate coalescing groups or disable SWA scratch reuse." - ) self._prepare_page_table_tensor(index_mapper_capacity) self._sparse_layer_group_ids = tuple( sorted( { - self.impl.get_buffer_group_id(buf.layer_id, buf.role) + self.layer_to_pool_mapping_dict[buf.layer_id] for buf in self.impl.all_buffer_ids if self.impl.is_sparse(buf.layer_id, buf.role) } @@ -2860,7 +2840,7 @@ def _fill_fresh_kv_pages(self, request_id: int) -> None: fill_stream = torch.cuda.current_stream(device=self._stream.device) fill_stream.wait_stream(self._stream) for layer_idx in self.pp_layers: - if self._get_layer_execution_group(self.layer_offsets[layer_idx]) != pool_id: + if self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] != pool_id: continue buffer = self.get_buffers(layer_idx) if buffer is None: @@ -2919,21 +2899,6 @@ def _get_layer_kv_factor_by_local_id(self, local_layer_id: int) -> int: return self.layer_specs[local_layer_id].kv_factor return self.kv_factor - @property - def _layer_ids_by_group(self) -> list[list[int]]: - # A split layer can occur in several groups. This is membership metadata, - # not the partition promised by the deprecated layer_grouping API. - return [ - list(dict.fromkeys(int(buffer.layer_id) for buffer in buffers)) - for buffers in self.impl.buffer_grouping - ] - - def _get_layer_execution_group(self, layer_id: int) -> int: - group = self.layer_to_pool_mapping_dict.get(layer_id) - if group is not None: - return group - return int(self.impl.get_buffer_group_id(layer_id, Role.KEY)) - def _get_layer_page_table_roles( self, local_layer_id: int, pool_id: int ) -> Tuple[DataRole, Optional[DataRole]]: @@ -2972,7 +2937,7 @@ def _build_pool_mapping_tensors(self): ) if self._use_per_layer_page_tables: for layer_id in typed_range(LayerId(self.num_local_layers)): - pool_id = self._get_layer_execution_group(layer_id) + pool_id = self.impl.get_layer_group_id(layer_id) role_a, _ = self._get_layer_page_table_roles(layer_id, pool_id) kv_cache_pool_pointers_list.append( [ @@ -3034,7 +2999,7 @@ def _build_pool_mapping_tensors(self): block_scale_pool_pointers_list.append([scale_base_addr, 0]) for layer_id in typed_range(LayerId(self.num_local_layers)): - layer_group_id = self._get_layer_execution_group(layer_id) + layer_group_id = self.impl.get_layer_group_id(layer_id) role_a, role_b = self._get_pool_roles(layer_group_id) index_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] if role_a == Role.KEY: @@ -3106,19 +3071,6 @@ def _build_pool_mapping_tensors(self): return kv_cache_pool_pointers, kv_cache_pool_mapping def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: - # The default attention view shares indices between K/V and their scales. - # Custom backends with separate tables override this initializer. - for layer in self.kv_cache_manager_py_config.layers: - if isinstance(layer, AttentionLayerConfig): - groups = { - self.impl.get_buffer_group_id(layer.layer_id, buffer.role) - for buffer in layer.buffers - } - if len(groups) > 1: - raise ValueError( - f"Attention layer {layer.layer_id} has different buffer reuse policies; " - "the default attention backend requires one shared page-index table" - ) # Keep role-dependent checks in this default page-table initializer: # custom layouts such as DeepSeek V4 override it with their own roles. # A lifecycle group can contain different page sizes in separate @@ -3134,7 +3086,7 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: } ) > 1 - for pool_id, layer_ids in enumerate(self._layer_ids_by_group) + for pool_id, layer_ids in enumerate(self.impl.layer_grouping) ) if self.layer_specs is not None: storage_layouts = {(layer.dtype, layer.kv_factor) for layer in self.layer_specs} @@ -3337,7 +3289,7 @@ def _get_event_num_blocks_per_cache_level( return [int(tier.quota // bytes_per_block) for tier in cache_tiers] def _get_event_layer_group_ids(self) -> List[int]: - return [int(layer_group_id) for layer_group_id in range(len(self._layer_ids_by_group))] + return [int(layer_group_id) for layer_group_id in range(len(self.impl.layer_grouping))] def _get_windowed_layer_group_sizes(self) -> Dict[int, int]: """Layer-group id -> window size, for groups whose window is smaller than max_seq_len. @@ -3405,10 +3357,11 @@ def _get_event_window_sizes_by_layer_group( self, attention_only: bool = False ) -> Dict[int, int]: # Assumes every layer listed in a group shares the same sliding_window_size, - # which is part of the lifecycle identity. Only the + # which is how `impl.layer_grouping` groups layers today. Only the # first layer's window is read; if the grouping policy ever permits # mixed windows in one group, this needs to fan out per-layer. A layer - # can belong to several groups through buffer policy overrides. + # listed under several groups is an SSM layer whose buffers use + # different reuse policies, so it describes each of them as SSM. # # `attention_only` is set for the streaming event manager, which tracks # attention prefix reuse only: excluding SSM and other non-attention life cycles @@ -3421,7 +3374,7 @@ def get_event_window_size(layer_config: object) -> int: return self.max_seq_len if window_size is None else int(window_size) window_sizes: Dict[int, int] = {} - for layer_group_id, layer_ids in enumerate(self._layer_ids_by_group): + for layer_group_id, layer_ids in enumerate(self.impl.layer_grouping): # Native bindings expose grouping, not the private Python lifecycle registry. layer_config = self.kv_cache_manager_py_config.layers[int(layer_ids[0])] if attention_only and not isinstance(layer_config, AttentionLayerConfig): @@ -3476,7 +3429,7 @@ def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: scratch_pages = torch.empty_like(scales) for local_layer_idx in range(self.num_local_layers): layer_id = LayerId(local_layer_idx) - pool_id = self._get_layer_execution_group(layer_id) + pool_id = self.layer_to_pool_mapping_dict[layer_id] role_a, role_b = self._get_layer_page_table_roles(layer_id, pool_id) roles = [ role_a, @@ -4144,7 +4097,7 @@ def fill_mixed_precision_page_tables( """ local = self.layer_offsets[layer_idx] for role, layer in enumerate((local, self.high_layers[local])): - pool_id = int(self.impl.get_buffer_group_id(layer, Role.KEY)) + pool_id = int(self.impl.get_layer_group_id(layer)) scale = int(self.impl.get_page_index_converter(layer, Role.KEY).scale) for row, req_id in enumerate(request_ids): kv_cache = self.kv_cache_map[req_id] @@ -5570,8 +5523,9 @@ def _life_cycle_by_layer_group(self) -> List[AttnLifeCycle]: """Attention life cycle per layer group. The layers listed in a group share its window and sinks, which is how - the lifecycle identity groups buffers, so the first layer describes the - group, including when buffer policies split a layer. A group + ``impl.layer_grouping`` groups layers, so the first layer describes the + group. A layer listed under several groups is an SSM layer whose buffers + use different reuse policies; it has no window in any of them. A group with no window -- full attention, or a group whose layers carry no window at all -- yields an empty stale range. """ @@ -5583,7 +5537,7 @@ def _life_cycle_by_layer_group(self) -> List[AttnLifeCycle]: getattr(layer, "num_sink_tokens", None), self.tokens_per_block, ) - for layer in (layers[int(layer_ids[0])] for layer_ids in self._layer_ids_by_group) + for layer in (layers[int(layer_ids[0])] for layer_ids in self.impl.layer_grouping) ] return self._layer_group_life_cycles @@ -5627,10 +5581,10 @@ def get_page_indices_by_layer_group(self, request: LlmRequest) -> List[List[int] """ kv_cache = self.kv_cache_map.get(request.py_request_id) if kv_cache is None: - return [[] for _ in range(len(self._layer_ids_by_group))] + return [[] for _ in range(len(self.impl.layer_grouping))] history_length = kv_cache.history_length by_group: List[List[int]] = [] - for layer_group_id in range(len(self._layer_ids_by_group)): + for layer_group_id in range(len(self.impl.layer_grouping)): indices = list(kv_cache.get_aggregated_page_indices(layer_group_id, valid_only=False)) stale_start, stale_end = self._stale_block_range(layer_group_id, history_length) for ordinal in range(stale_start, min(stale_end, len(indices))): @@ -6259,10 +6213,11 @@ def _stats_life_cycle_metadata(self) -> dict[int, tuple[int, Optional[int], str] } metadata: dict[int, tuple[int, Optional[int], str]] = {} - for life_cycle_id, layer_ids in enumerate(self._layer_ids_by_group): + for life_cycle_id, layer_ids in enumerate(self.impl.layer_grouping): if not layer_ids: continue - # Buffer policies can split a layer without changing its kind or window. + # A layer listed under several groups is an SSM layer whose buffers use + # different reuse policies, so its kind describes each of those groups. layer = self.kv_cache_manager_py_config.layers[int(layer_ids[0])] is_attention = isinstance(layer, AttentionLayerConfig) metadata[life_cycle_id] = ( @@ -7248,7 +7203,7 @@ def get_batch_cache_indices( else self.kv_factor ) else: - pool_id = self._get_layer_execution_group(self.layer_offsets[layer_idx]) + pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] index_scale = 1 if raw_indices else self.get_layer_page_index_scale(layer_idx) kv_factor = self.get_layer_kv_factor(layer_idx) return self._get_batch_cache_indices_by_pool_id( @@ -7328,7 +7283,7 @@ def get_batch_cache_indices_flat( else self.kv_factor ) else: - pool_id = self._get_layer_execution_group(self.layer_offsets[layer_idx]) + pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] # Layers sharing a pool can still require different page-index # scales (e.g. Gemma4 sliding/global head_dim). Use the per-layer # scale so this flat block table matches get_batch_cache_indices() @@ -7460,7 +7415,7 @@ def _iter_cache_buffers_for_invalid_check(self) -> Iterable[Tuple[int, torch.Ten # layouts that coalesce extra buffers alongside K/V. yield layer_id, buffer continue - pool_id = self._get_layer_execution_group(layer_offset) + pool_id = self.layer_to_pool_mapping_dict[layer_offset] if pool_id in pool_handled: continue yield layer_id, buffer diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mixed_precision_staging.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mixed_precision_staging.py index 16431da0799d..72033a3b4a15 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mixed_precision_staging.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mixed_precision_staging.py @@ -10,7 +10,6 @@ from tensorrt_llm._utils import prefer_pinned from ...memory_buffer_utils import Buffers -from .kv_cache_manager_v2 import Role class MixedPrecisionKVStaging: @@ -72,7 +71,7 @@ def __init__( for layer_idx in self.pools: local = manager.layer_offsets[layer_idx] key = tuple( - int(manager.impl.get_buffer_group_id(layer, Role.KEY)) + int(manager.impl.get_layer_group_id(layer)) for layer in (local, manager.high_layers[local]) ) self.table_owner[layer_idx] = owners.setdefault(key, layer_idx) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py index 8e911bbcd265..aa1f46cd2813 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/sharing/_manager.py @@ -166,10 +166,7 @@ def layer_configs(manager: KVCacheManagerV2) -> dict[int, object]: def layer_group_layer_ids(manager: KVCacheManagerV2) -> list[list[int]]: """Internal layer ids per runtime layer group, in layer-group order.""" - return [ - list(dict.fromkeys(int(buffer.layer_id) for buffer in group)) - for group in manager.impl.buffer_grouping - ] + return [[int(layer_id) for layer_id in group] for group in manager.impl.layer_grouping] def pool_group_descriptors(manager: KVCacheManagerV2) -> Sequence[PoolGroupDesc]: @@ -183,7 +180,7 @@ def declared_global_layer_ids(manager: KVCacheManagerV2) -> dict[int, int] | Non if getter is None: return None table: dict[int, int] = {} - for layer_group, layer_ids in enumerate(layer_group_layer_ids(manager)): + for layer_group, layer_ids in enumerate(manager.impl.layer_grouping): ids = [int(value) for value in getter(layer_group)] if len(ids) != len(layer_ids): raise ValueError( diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py index 4b20395472fe..c4442ac4ed9e 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_layout_identity.py @@ -404,10 +404,7 @@ def __init__( slot_desc=SimpleNamespace(variants=[variant]), ) self.impl = SimpleNamespace( - buffer_grouping=[ - [SimpleNamespace(layer_id=lid, role=role) for lid in lids for role, _, _ in roles] - for _, lids, roles in groups - ], + layer_grouping=[list(lids) for _, lids, _ in groups], pool_group_descs=list(descs.values()), get_life_cycle_pool_group_indices=lambda: list(pool_group_of), init_config=SimpleNamespace(layers=layers), diff --git a/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py b/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py index a65252473745..d08fa1ea6cd5 100644 --- a/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py +++ b/tests/unittest/_torch/executor/kv_cache/sharing/test_public_surface.py @@ -68,7 +68,7 @@ INTERNALS = { "impl", "init_config", - "buffer_grouping", + "layer_grouping", "pool_group_descs", "_has_lossy_cold_page_codec", "get_disagg_global_layer_ids", @@ -591,13 +591,7 @@ def test_manager_facade_preserves_layer_group_order_and_normalizes_ids(): from tensorrt_llm._torch.pyexecutor.kv_cache.sharing import _manager grouping = [[np.int64(5), "7"], [2]] - manager = SimpleNamespace( - impl=SimpleNamespace( - buffer_grouping=[ - [SimpleNamespace(layer_id=layer) for layer in group] for group in grouping - ] - ) - ) + manager = SimpleNamespace(impl=SimpleNamespace(layer_grouping=grouping)) actual = _manager.layer_group_layer_ids(manager) assert actual == [[5, 7], [2]] assert all(type(layer_id) is int for group in actual for layer_id in group) diff --git a/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py b/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py index 82500f3d79b8..621d27f1d83b 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py +++ b/tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py @@ -449,9 +449,7 @@ def test_create_cross_kv_cache_manager_uses_tp_sharded_encoder_heads(self): distributed = Mock() distributed.allreduce.side_effect = lambda value, op: value impl = Mock( - buffer_grouping=[ - [SimpleNamespace(layer_id=layer, role="key") for layer in range(num_decoder_layers)] - ], + layer_grouping=[list(range(num_decoder_layers))], pool_group_descs=[], all_buffer_ids=[], ) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py index 150f201b84f3..8a8f0b47570f 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_extra_buffers.py @@ -168,10 +168,9 @@ def test_page_table_uses_physical_pool_representative(self): layers = real_impl.layer_grouping[pool_id] other_layer = next(layer for layer in layers if layer != physical_layer) impl_proxy = Mock(wraps=real_impl) - impl_proxy.buffer_grouping = [ - sorted(group, key=lambda buffer: buffer.layer_id != other_layer) - for group in real_impl.buffer_grouping - ] + impl_proxy.layer_grouping = ( + (other_layer, *(layer for layer in layers if layer != other_layer)), + ) mgr.impl = impl_proxy mgr._prepare_page_table_tensor(index_mapper_capacity=1) 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 e0e4e321a9b0..83bd77f0c2f4 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -421,7 +421,7 @@ def build_cache_config( None, ) if fake_impl is not None: - fake_impl.buffer_grouping = [[BufferId(LayerId(0), Role.KEY)]] + fake_impl.layer_grouping = [[0]] fake_impl.pool_group_descs = [] fake_impl.get_layer_group_id.side_effect = lambda _: 0 fake_impl.all_buffer_ids = [BufferId(LayerId(0), Role.KEY)] @@ -1020,67 +1020,6 @@ def test_zero_size_filter_preserves_layer_and_buffer_reuse_policies() -> None: assert private_layer.reuse_policy == ReusePolicy.PRIVATE -def test_layer_group_helpers_describe_split_ssm_layers() -> None: - if not torch.cuda.is_available(): - pytest.skip("requires CUDA") - config = KVCacheManagerConfig( - tokens_per_block=TOKENS_PER_BLOCK, - cache_tiers=[GpuCacheTierConfig(quota=4 << 20)], - layers=[ - AttentionLayerConfig( - layer_id=LayerId(0), - buffers=[BufferConfig(role=Role.KEY, size=4096)], - sliding_window_size=8, - ), - *[ - SsmLayerConfig( - layer_id=LayerId(layer_id), - buffers=[ - BufferConfig(role=DataRole("ssm_state"), size=4096), - BufferConfig( - role=DataRole("old_x"), size=2048, reuse_policy=ReusePolicy.PRIVATE - ), - ], - ) - for layer_id in (1, 2) - ], - ], - commit_min_snapshot=True, - ) - impl = KVCacheManager(config) - try: - manager = object.__new__(KVCacheManagerV2) - manager.max_seq_len = MAX_SEQ_LEN - manager.tokens_per_block = TOKENS_PER_BLOCK - manager.kv_cache_manager_py_config = config - manager.impl = impl - manager._layer_group_life_cycles = None - - attention = impl.get_layer_group_id(LayerId(0)) - required = impl.get_buffer_group_id(LayerId(1), DataRole("ssm_state")) - private = impl.get_buffer_group_id(LayerId(1), DataRole("old_x")) - assert len({attention, required, private}) == 3 - assert sorted({buffer.layer_id for buffer in impl.buffer_grouping[private]}) == [1, 2] - - assert manager._get_event_window_sizes_by_layer_group(attention_only=True) == {attention: 8} - assert manager._get_event_window_sizes_by_layer_group() == { - attention: 8, - required: MAX_SEQ_LEN, - private: MAX_SEQ_LEN, - } - life_cycles = manager._life_cycle_by_layer_group() - assert life_cycles[attention].window_size == 8 - assert life_cycles[required].window_size is None - assert life_cycles[private].window_size is None - metadata = manager._stats_life_cycle_metadata() - assert metadata[attention][1:] == (8, "attention") - assert metadata[required][1:] == (None, "ssm") - assert metadata[private][1:] == (None, "ssm") - assert len({metadata[group][0] for group in (attention, required, private)}) == 3 - finally: - impl.shutdown() - - @pytest.mark.parametrize("sliding_window_size", [None, 8]) @pytest.mark.parametrize("attention_first", [False, True]) def test_event_window_sizes_filter_attention_without_backend_internals( @@ -1103,11 +1042,7 @@ def test_event_window_sizes_filter_attention_without_backend_internals( ) # Match the C++ binding surface: layer_grouping is public, while the # Python implementation's private _life_cycles registry is absent. - manager.impl = SimpleNamespace( - buffer_grouping=[ - [BufferId(LayerId(i), Role.KEY)] for i in ([0, 1] if attention_first else [1, 0]) - ] - ) + manager.impl = SimpleNamespace(layer_grouping=[[0], [1]] if attention_first else [[1], [0]]) attention_group_id = 0 if attention_first else 1 window_size = MAX_SEQ_LEN if sliding_window_size is None else sliding_window_size @@ -3249,7 +3184,7 @@ def _index_mapper_capacity_for( """ module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" fake_impl = Mock() - fake_impl.buffer_grouping = [[BufferId(LayerId(0), Role.KEY)]] + fake_impl.layer_grouping = [[0]] fake_impl.pool_group_descs = [] fake_impl.get_layer_group_id.side_effect = lambda _: 0 fake_impl.all_buffer_ids = [BufferId(LayerId(0), Role.KEY)] @@ -4085,77 +4020,3 @@ def test_try_commit_generation_blocks_honors_guards( manager.try_commit_generation_blocks(request) assert kv_cache.committed_tokens is None - - -@pytest.mark.parametrize("role", [Role.VALUE, Role.KEY_BLOCK_SCALE]) -def test_default_attention_rejects_split_coupled_buffers(monkeypatch, role): - original = KVCacheManagerV2._build_cache_config - - def build_config(self, config): - config = original(self, config) - layers = list(config.layers) - layer = layers[0] - buffers = [ - BufferConfig( - role=buffer.role, - size=buffer.size, - reuse_policy=ReusePolicy.PRIVATE if buffer.role == role else None, - ) - for buffer in layer.buffers - ] - if role == Role.KEY_BLOCK_SCALE: - buffers.append(BufferConfig(role=role, size=128, reuse_policy=ReusePolicy.PRIVATE)) - layers[0] = AttentionLayerConfig(layer_id=layer.layer_id, buffers=buffers) - return replace(config, layers=layers) - - monkeypatch.setattr(KVCacheManagerV2, "_build_cache_config", build_config) - with pytest.raises(ValueError, match="requires one shared page-index table"): - KVCacheManagerV2( - KvCacheConfig(max_gpu_total_bytes=4 << 20, host_cache_size=0), - CacheType.SELF, - num_layers=1, - num_kv_heads=1, - head_dim=64, - tokens_per_block=TOKENS_PER_BLOCK, - max_seq_len=MAX_SEQ_LEN, - max_batch_size=1, - mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), - dtype=DataType.HALF, - vocab_size=16, - ) - - -def test_attention_scratch_rejects_layout_requiring_split_slot_indices(monkeypatch): - original = KVCacheManagerV2._build_cache_config - - def build_config(self, config): - config = original(self, config) - layers = list(config.layers) - layer = layers[0] - layers[0] = AttentionLayerConfig( - layer_id=layer.layer_id, - sliding_window_size=layer.sliding_window_size, - buffers=[*layer.buffers, BufferConfig(role="side", size=layer.buffers[0].size)], - ) - return replace(config, layers=layers) - - monkeypatch.setattr(KVCacheManagerV2, "_build_cache_config", build_config) - with pytest.raises(ValueError, match="require whole layer blocks per slot"): - KVCacheManagerV2( - KvCacheConfig( - max_gpu_total_bytes=4 << 20, - host_cache_size=0, - max_attention_window=[8, 8], - enable_swa_scratch_reuse=True, - ), - CacheType.SELF, - num_layers=2, - num_kv_heads=1, - head_dim=64, - tokens_per_block=TOKENS_PER_BLOCK, - max_seq_len=MAX_SEQ_LEN, - max_batch_size=1, - mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), - dtype=DataType.HALF, - vocab_size=16, - ) 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 b78a4b1f8dc9..945df03801e8 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 @@ -1391,7 +1391,7 @@ def test_v2_disagg_slice_reads_state_index_without_refreshing_batch_mask(): side_effect=AssertionError("state-index lookup must not refresh the dummy mask") ) # Provide a mamba layer group so _create_chunk places the slot ID - mamba_lg = SimpleNamespace(kind=CacheKind.STATE, pool_views=[object()]) + mamba_lg = SimpleNamespace(kind=CacheKind.STATE) transceiver = object.__new__(KvCacheTransceiverV2) transceiver._kv_cache_manager = manager transceiver._reuse_adapter = SimpleNamespace(tokens_per_block=32) diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index d4154e8a5e74..ec6ef42ad0f7 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -300,9 +300,7 @@ def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_l ) impl = SimpleNamespace( - buffer_grouping=( - [buffer for coalesced in coalesced_buffers for buffer in coalesced.buffer_ids], - ), + layer_grouping=((0, 1),), init_config=SimpleNamespace( tokens_per_block=16, layers=[SimpleNamespace(window_size=None), SimpleNamespace(window_size=None)], @@ -599,15 +597,7 @@ def buf(layer_id, role): manager = object.__new__(Qwen4ExpHybridCacheManager) manager.impl = SimpleNamespace( - buffer_grouping=tuple( - [ - buffer - for variant in pg.slot_desc.variants - for coalesced in variant.coalesced_buffers - for buffer in coalesced.buffer_ids - ] - for pg in (recurrent_pg, attention_pg) - ), + layer_grouping=((0, 1), (2,)), init_config=SimpleNamespace( tokens_per_block=16, layers=[SimpleNamespace(window_size=None) for _ in range(3)], diff --git a/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py index 4372c4813011..559500719639 100644 --- a/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py +++ b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py @@ -177,10 +177,7 @@ class _StubGroupingImpl: """Stub ``manager.impl`` exposing only what ``_kv_pool_mapping_offset`` reads.""" def __init__(self, grouping, key_addrs): - self.buffer_grouping = [ - [SimpleNamespace(layer_id=layer, role=Role.KEY) for layer in group] - for group in grouping - ] + self.layer_grouping = grouping self._key_addrs = key_addrs def get_mem_pool_base_address(self, layer_id, role, index_mode=None): From 3e3ce5d09400258c092a126aab1a5d97a142c8e2 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Wed, 28 Oct 2026 16:39:42 +0800 Subject: [PATCH 13/19] Remove deferred buffer adoption comments from Python consumers Signed-off-by: Jiagan Cheng --- .../kv_cache/kv_cache_manager_v2.py | 20 +++++++------------ 1 file changed, 7 insertions(+), 13 deletions(-) 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 562419a1185b..3bad6a286b76 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 @@ -3356,12 +3356,10 @@ def _log_window_crossing( def _get_event_window_sizes_by_layer_group( self, attention_only: bool = False ) -> Dict[int, int]: - # Assumes every layer listed in a group shares the same sliding_window_size, - # which is how `impl.layer_grouping` groups layers today. Only the + # Assumes every layer in a group shares the same sliding_window_size, + # which is how `impl.layer_grouping` partitions layers today. Only the # first layer's window is read; if the grouping policy ever permits - # mixed windows in one group, this needs to fan out per-layer. A layer - # listed under several groups is an SSM layer whose buffers use - # different reuse policies, so it describes each of them as SSM. + # mixed windows in one group, this needs to fan out per-layer. # # `attention_only` is set for the streaming event manager, which tracks # attention prefix reuse only: excluding SSM and other non-attention life cycles @@ -5522,12 +5520,10 @@ def admit_mirror(self, req: LlmRequest) -> bool: def _life_cycle_by_layer_group(self) -> List[AttnLifeCycle]: """Attention life cycle per layer group. - The layers listed in a group share its window and sinks, which is how - ``impl.layer_grouping`` groups layers, so the first layer describes the - group. A layer listed under several groups is an SSM layer whose buffers - use different reuse policies; it has no window in any of them. A group - with no window -- full attention, or a group whose layers carry no - window at all -- yields an empty stale range. + Every layer in a group shares one life cycle, which is how + ``impl.layer_grouping`` partitions layers, so the first layer describes + the group. A group with no window -- full attention, or a group whose + layers carry no window at all -- yields an empty stale range. """ if self._layer_group_life_cycles is None: layers = self.kv_cache_manager_py_config.layers @@ -6216,8 +6212,6 @@ def _stats_life_cycle_metadata(self) -> dict[int, tuple[int, Optional[int], str] for life_cycle_id, layer_ids in enumerate(self.impl.layer_grouping): if not layer_ids: continue - # A layer listed under several groups is an SSM layer whose buffers use - # different reuse policies, so its kind describes each of those groups. layer = self.kv_cache_manager_py_config.layers[int(layer_ids[0])] is_attention = isinstance(layer, AttentionLayerConfig) metadata[life_cycle_id] = ( From 74e17aa3add25f16696a8acc51413d5fd6374a25 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Wed, 28 Oct 2026 17:00:10 +0800 Subject: [PATCH 14/19] Remove V2 reuse policy section from V1 cache manager guide Signed-off-by: Jiagan Cheng --- docs/source/torch/kv_cache_manager.md | 59 --------------------------- 1 file changed, 59 deletions(-) diff --git a/docs/source/torch/kv_cache_manager.md b/docs/source/torch/kv_cache_manager.md index 2a359a4bff4c..f6928377ed3a 100644 --- a/docs/source/torch/kv_cache_manager.md +++ b/docs/source/torch/kv_cache_manager.md @@ -71,62 +71,3 @@ In `KvCacheCreator._create_kv_cache_manager`, the `KVCacheManager` is instantiat For local testing or proof of concept, update these lines to use your implementation. Then, test it to ensure the `PyExecutor` runs with your customized `KVCacheManager`. - - -## Native V2 buffer reuse policies - -Custom cache managers can override a layer's reuse policy per buffer through -`tensorrt_llm.runtime.kv_cache_manager_v2`. Both attention and SSM layers default -to `ReusePolicy.REQUIRED`; a buffer's `reuse_policy=None` inherits its layer's -policy. For example, a custom recurrent consumer with independent state indices -can configure: - -```python -from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, ReusePolicy, SsmLayerConfig - -layer = SsmLayerConfig( - layer_id=0, - reuse_policy=ReusePolicy.REQUIRED, - buffers=[ - BufferConfig(role="state", size=4096), - BufferConfig(role="checkpoint", size=2048, reuse_policy=ReusePolicy.OPTIONAL), - BufferConfig(role="scratch", size=1024, reuse_policy=ReusePolicy.PRIVATE), - ], -) -``` - -REQUIRED buffers constrain the reusable prefix. OPTIONAL buffers are candidates -at that endpoint; select their group IDs through `resume(optional_reuse_groups=...)` -on the first resume. Attention candidates retain the block-aligned sliding-window -checkpoint contract; an SSM snapshot must match the endpoint exactly. PRIVATE -buffers remain request-local and are never published as reusable checkpoints. -The consumer must initialize or reconstruct private and unselected optional state -before inference. Unselected optional snapshots count as neither hits nor misses. - -Use `manager.get_buffer_group_id(layer_id, role)` to resolve a buffer and -`manager.buffer_grouping[group_id]` to enumerate membership. Buffers with the same -lifecycle properties share a group, including across layers; selecting a group -selects all of its buffers. Policy overrides add at most three policy variants -per attention window/sink/sparsity combination and at most three SSM groups. -Coalescing remains within lifecycle groups and compatible buffer sizes/coalescing -groups. - -`get_layer_group_id(layer_id)` and `layer_grouping` are deprecated in documentation -only and produce no deprecation warnings. The former returns a uniform layer's -group and raises `ValueError` for a split layer. The latter preserves its -partition-of-layers result and raises if any layer is split. `LayerGroupId` retains -its name; it identifies a lifecycle group, not a physical storage pool. - -Existing Python cache-manager consumers continue using the legacy layer APIs -and require one lifecycle per layer. They do not adopt per-buffer policy splits; -the legacy API rejects ambiguous layer mappings. Consumers can migrate to the -buffer APIs when they add support for independent indices. Current Mamba kernels -couple recurrent, convolution, PLE and replay buffers through one state-slot -index; separate replay indices are a follow-up. This native configuration does -not add an LLM API configuration field. - -The native page-index converter also supports unequal layer shares in scratch -slots, including a layer's block crossing two physical slots. Existing attention -scratch-table kernels require whole layer blocks per slot. Custom consumers of -unequal shares that cross slots must use the native converter or supply compatible -index generation; existing Python consumers have not adopted that layout. From 360a2f29e46c3ead36d511edacd40af74d1dc829 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Wed, 28 Oct 2026 17:22:55 +0800 Subject: [PATCH 15/19] Wait for offload before overwriting recycled slots in tests Signed-off-by: Jiagan Cheng --- .../batch_manager/kvCacheManagerV2ColdPageTest.cpp | 7 +++++++ .../unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp | 5 +++++ 2 files changed, 12 insertions(+) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp index 8d4b6d61701d..a3cb5cdebc8f 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp @@ -227,6 +227,13 @@ TEST_P(OptionalDiskCheckpointTest, PayloadRestoreRetryAndIndependentEviction) auto reader = manager->createKvCache({}, span, 2); ASSERT_TRUE(reader->reuseStatus()[encoder.value()].complete); + auto reservations + = storage.newSlotsForPoolGroup(kHotLevel, hotGroup, storage.getStatistics(kHotLevel, hotGroup).free); + for (auto const& slot : reservations) + { + // Offload can still be reading the recycled GPU slots on another stream. + slot.readyEvent.waitInStream(reinterpret_cast(stream)); + } // Poison the entire hot Encoder group: a false restore cannot pass using // stale bytes from the original slot after it was offloaded to disk. for (int p = 0; p < sizes.size().value(); ++p) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index 948c694b84a9..d2245c7b76e6 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -708,6 +708,11 @@ TEST_F(ReplayCacheTest, SuspendedSsmStateSurvivesOffloadAndSlotChange) } } ASSERT_EQ(oldSlotReservations.size(), 2); + for (auto const& [group, slot] : oldSlotReservations) + { + // Offload can still be reading the recycled GPU slot on another stream. + slot.readyEvent.waitInStream(reinterpret_cast(mStream)); + } for (auto const& buffer : splitSsmBuffers()) { auto const group = manager->getBufferGroupId(buffer.layer, buffer.role); From 4f26fedc514e2dd8265265288aef517981e7221d Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Wed, 28 Oct 2026 19:44:38 +0800 Subject: [PATCH 16/19] Document attention and SSM buffer reuse policy overrides Signed-off-by: Jiagan Cheng --- tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 523db9d37e2a..a4c2a39cf8c0 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -273,6 +273,11 @@ class DiskCacheTierConfig: class BufferConfig: """One buffer of a layer's KV cache. + ``reuse_policy`` of ``None`` uses the layer's policy. Attention and SSM + buffers can override the layer's policy; the buffer then belongs to the + layer group for its effective policy and lifecycle properties while keeping + its layer ID, so one layer can span several layer groups. + ``is_sparse`` marks a buffer whose history is too large for device memory and is read through a per-request device cache rather than directly (see :meth:`KVCacheManager.is_sparse`). From 13ff2398eda4bb9377032d0d062f3d36f07d3580 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Thu, 29 Oct 2026 13:16:37 +0800 Subject: [PATCH 17/19] Clarify lifecycle policy helpers and cache group API contracts Signed-off-by: Jiagan Cheng --- .../kv_cache_manager_v2/AGENTS.md | 8 ++++-- .../kv_cache_manager_v2/blockRadixTree.cpp | 4 +-- .../kv_cache_manager_v2/kvCache.cpp | 28 ++++++++----------- .../kv_cache_manager_v2/kvCache.h | 2 +- .../kv_cache_manager_v2/kvCacheManager.cpp | 17 +++++------ .../kv_cache_manager_v2/kvCacheManager.h | 12 ++++---- .../kv_cache_manager_v2/lifeCycleRegistry.h | 25 ++++++----------- .../kv_cache_manager_v2/page.cpp | 2 +- .../batch_manager/kvCacheManagerV2.cpp | 10 +++---- .../kvCacheManagerV2StatsTest.cpp | 4 +-- .../runtime/kv_cache_manager_v2/__init__.py | 2 ++ .../runtime/kv_cache_manager_v2/__init__.pyi | 11 +++++--- .../test_optional_checkpoint.py | 4 +-- 13 files changed, 60 insertions(+), 69 deletions(-) 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 1e7f3ec28bec..21989ea2bb25 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 @@ -49,7 +49,7 @@ them before changing `batch.*`, `sparseCache.*` or the index conversion. - Prefer composition when it expresses ownership or containment. Use inheritance only for a real IS-A relationship or required virtual dispatch, such as `CommittedPage : Page` and `EventManager : EventSink`. -- Preserve strongly typed indices. `LayerGroupId` is a public lifecycle +- Preserve strongly typed indices. `CacheGroupId` (and legacy `LayerGroupId`) is a public lifecycle semantic, `PoolGroupIndex` is a storage-layout index, `BlockOrdinal` is a sequence position, and `SlotId`/page indices address storage. Do not convert or compare them implicitly. @@ -102,10 +102,12 @@ A buffer may override its layer's default through `BufferConfig::reusePolicy`. This applies to attention and SSM buffers, so one model layer may span several lifecycle groups. Use `getBufferGroupId(layerId, role)` and `bufferGrouping()` (`get_buffer_group_id` and `buffer_grouping` in Python) for exact membership. +These APIs return or index by `CacheGroupId`, a public alias of `LifeCycleId`. +`LayerGroupId` remains a lifecycle alias for legacy APIs. The legacy `getLayerGroupId(layerId)` and `layerGrouping()` are deprecated in documentation only: uniform layers retain their results, while a split layer -raises an ambiguity error (for the queried layer or the entire partition, -respectively). No deprecation warnings are emitted. +fails a `TLLM_CHECK_WITH_INFO` (for the queried layer or the entire partition, +respectively), exposed as `RuntimeError` in Python. No deprecation warnings are emitted. Page conversion reads `BufferAttr::scratchPagesPerBlock`, counting only the layer's buffers in that coalesced slot. Scratch capacity uses the maximum exact diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp index 3b74d32f5629..62d7b6d50507 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp @@ -490,7 +490,7 @@ std::vector> Block::clearStaleBlocksAfterPageUnlink( // If this is a full-attention block or a sink block: evict subtree. // Mirrors Python: pages = remove_subtree(self) - if (alc && alc->requiresPrefixCoverage() + if (alc && requiresPrefixCoverage(*alc) && (!alc->windowSize.has_value() || block.ordinal() < BlockOrdinal{alc->numSinkBlocks})) { pruneStart = block.prev; @@ -899,7 +899,7 @@ std::vector BlockRadixTree::pruneMatch( bool shortened = false; for (auto [lcId, attn] : attnLcs) { - if (!attn->requiresPrefixCoverage()) + if (!requiresPrefixCoverage(*attn)) { continue; } 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 69acce5b52d0..7e7b9e4c1ebd 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 @@ -259,8 +259,7 @@ void KvCache::_collectReuseCandidates() } // OPTIONAL checks coverage at the REQUIRED endpoint without shortening it. // Use the ordinary SWA live ranges, including sinks and multi-page windows. - if (attention && attention->allowsCheckpoint() && end > 0 && end % mTokensPerBlock == 0 - && attention->windowSize) + if (attention && allowsCheckpoint(*attention) && end > 0 && end % mTokensPerBlock == 0 && attention->windowSize) { auto const stale = getStaleRange(lifeCycles[lc], end, mTokensPerBlock); auto& candidates = mOptionalCandidates[lc]; @@ -312,7 +311,7 @@ void KvCache::_collectReuseCandidates() bool KvCache::_needsFreshPrefixPages(LifeCycleId lc) const { - return !requiresPrefixCoverage(mManager->lifeCycles()[lc]) + return !mManager->lifeCycles().isSsm(lc) && !requiresPrefixCoverage(mManager->lifeCycles()[lc]) && std::find(mSelectedOptionalGroups.begin(), mSelectedOptionalGroups.end(), lc) == mSelectedOptionalGroups.end(); } @@ -599,13 +598,7 @@ bool KvCache::resume(std::optional stream, std::optionallifeCycles().isSsm(lc); auto const& lifeCycle = mManager->lifeCycles()[lc]; - // Every SSM policy needs one writable state slot. Handle SSM first because private and - // unselected optional groups also need fresh state, but must not allocate per-token-block pages. - if (isSsm) - { - numSlotsNeeded[lc] += 1; - } - else if (_needsFreshPrefixPages(lc)) + if (_needsFreshPrefixPages(lc)) { auto const stale = _getStaleRange(mHistoryLength, lifeCycle); for (BlockOrdinal ord{0}; ord < stale.beg; ++ord) @@ -618,8 +611,9 @@ bool KvCache::resume(std::optional stream, std::optional(privateOrdinals[lc].size()); } - else if (hasPartial) + else if (isSsm || hasPartial) { + // Every SSM policy needs one writable state slot, independent of prefix length. numSlotsNeeded[lc] += 1; } } @@ -1490,7 +1484,7 @@ void KvCache::_snapshotPartialBlockToTree(BlockOrdinal ordinal, bool commitSsm) std::vector attachedLcs; for (auto const& [lcIdx, attn] : mManager->lifeCycles().attentionLifeCycles()) { - if (!attn->requiresPrefixCoverage()) + if (!requiresPrefixCoverage(*attn)) { continue; } @@ -2152,9 +2146,9 @@ std::vector KvCache::_unlockStaleBlocks(int newHistoryLeng // A partial tail can expire pages still needed at the preceding // aligned checkpoint. Keep its entire live window until publication. int const alignedEnd = newHistoryLength / mTokensPerBlock * mTokensPerBlock; - bool const keepAlignedCheckpoint = alc.allowsCheckpoint() - && ord < BlockOrdinal{alignedEnd / mTokensPerBlock} && !_getStaleRange(alignedEnd, lc).contains(ord); - bool holdForCommit = !alc.isPrivate() && (!mManager->commitMinSnapshot() || keepAlignedCheckpoint) + bool const keepAlignedCheckpoint = allowsCheckpoint(alc) && ord < BlockOrdinal{alignedEnd / mTokensPerBlock} + && !_getStaleRange(alignedEnd, lc).contains(ord); + bool holdForCommit = !isPrivateLifeCycle(alc) && (!mManager->commitMinSnapshot() || keepAlignedCheckpoint) && !isCommitted && (mCommitState == CommitState::ALLOWED); for (BeamIndex bi{0}; bi < sb.pages.size(); ++bi) @@ -3053,7 +3047,7 @@ SharedPtr const& KvCache::_getTreeBlock(BlockOrdinal ordinal) const // committed->block may differ from `ret`: a page can be moved to a longer // sibling block, or replaced there by one with larger token coverage. auto const* attention = std::get_if(&mManager->lifeCycles()[lcId]); - if (attention && attention->allowsCheckpoint()) + if (attention && allowsCheckpoint(*attention)) { TLLM_CHECK(dynamicPointerCast(page) || dynamicPointerCast(page)); } @@ -3165,7 +3159,7 @@ bool KvCache::_checkSanity() const auto page = blockPageGetPage(bp); auto const* attention = std::get_if(&lcs[lc]); auto committed = dynamicPointerCast(page); - if (attention && attention->allowsCheckpoint() && committed) + if (attention && allowsCheckpoint(*attention) && committed) { TLLM_CHECK_DEBUG(isCommitted && committed->numTokensInBlock > 0 && committed->numTokensInBlock <= mTokensPerBlock); 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 e2a53674c92c..5c1850018ef2 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 @@ -865,7 +865,7 @@ class KvCache : public std::enable_shared_from_this std::vector mSelectedOptionalGroups; void _collectReuseCandidates(); void _selectOptionalGroups(std::vector const& groups); - // PRIVATE and unselected OPTIONAL groups need fresh pages for the claimed prefix. + // PRIVATE and unselected OPTIONAL attention groups need fresh pages for the claimed prefix. bool _needsFreshPrefixPages(LifeCycleId lc) const; // SWA scratch slot support. 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 6034a50bfbd7..0a6e5cb398d0 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 @@ -808,17 +808,14 @@ std::vector KvCacheManager::layerIds() const LayerGroupId KvCacheManager::getLayerGroupId(LayerId layerId) const { auto const& lifeCycles = mStorage->layerToLifeCycleIds().at(layerId); - if (lifeCycles.size() != 1) - { - throw std::invalid_argument("Layer " + std::to_string(layerId) + " has buffers in " - + std::to_string(lifeCycles.size()) - + " layer groups because its buffers use different reuse policies; look up its layer group per " - "buffer with get_buffer_group_id(layer_id, role) / getBufferGroupId(layerId, role)"); - } + TLLM_CHECK_WITH_INFO(lifeCycles.size() == 1, + "Layer %d has buffers in %zu layer groups because its buffers use different reuse policies; " + "look up its layer group per buffer with get_buffer_group_id(layer_id, role) / getBufferGroupId(layerId, role)", + layerId, lifeCycles.size()); return lifeCycles.front(); } -LayerGroupId KvCacheManager::getBufferGroupId(LayerId layerId, DataRole const& role) const +CacheGroupId KvCacheManager::getBufferGroupId(LayerId layerId, DataRole const& role) const { return mStorage->getBufferAttr(layerId, role).lifeCycleId; } @@ -834,9 +831,9 @@ TypedVec> KvCacheManager::layerGrouping() con return result; } -TypedVec> KvCacheManager::bufferGrouping() const +TypedVec> KvCacheManager::bufferGrouping() const { - TypedVec> result(mLifeCycles.size()); + TypedVec> result(mLifeCycles.size()); std::unordered_map> buffersByLayer; for (auto const& buffer : allBufferIds()) { 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 8941941d30bd..3066d804ab34 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 @@ -287,20 +287,20 @@ class KvCacheManager : public std::enable_shared_from_this std::vector layerIds() const; //! Deprecated: use getBufferGroupId(). Returns the shared group of a uniform layer. - //! Throws std::invalid_argument if the queried layer spans several lifecycle groups. + //! Throws common::TllmException if the queried layer spans several lifecycle groups. LayerGroupId getLayerGroupId(LayerId layerId) const; - //! Group of one buffer; throws std::out_of_range for an unknown buffer. - LayerGroupId getBufferGroupId(LayerId layerId, DataRole const& role) const; + //! Cache group of one buffer; throws std::out_of_range for an unknown buffer. + CacheGroupId getBufferGroupId(LayerId layerId, DataRole const& role) const; //! Deprecated: use bufferGrouping(). Partitions uniform layers by lifecycle. - //! Throws std::invalid_argument if any layer spans several groups. Member order + //! Throws common::TllmException if any layer spans several groups. Member order //! is unspecified; use poolGroupDescs() to inspect physical buffer ordering. TypedVec> layerGrouping() const; - //! Buffer membership indexed by lifecycle group ID. Each buffer appears once. + //! Buffer membership indexed by CacheGroupId. Each buffer appears once. //! Member order is unspecified and does not describe physical memory ordering. - TypedVec> bufferGrouping() const; + TypedVec> bufferGrouping() const; // Iterator over all buffer identifiers. Mirrors Python's all_buffer_ids property. std::vector allBufferIds() const; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h index 47137af632f3..f578e4d2a8ab 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h @@ -42,21 +42,6 @@ struct AttnLifeCycle int numSinkBlocks = 0; // divUp(numSinkTokens, tokensPerBlock) ReusePolicy reusePolicy = ReusePolicy::REQUIRED; - bool isPrivate() const noexcept - { - return reusePolicy == ReusePolicy::PRIVATE; - } - - bool requiresPrefixCoverage() const noexcept - { - return reusePolicy == ReusePolicy::REQUIRED; - } - - bool allowsCheckpoint() const noexcept - { - return reusePolicy == ReusePolicy::OPTIONAL; - } - HalfOpenRange getStaleRange(int historyLength, int tokensPerBlock) const { int numBlocks = divUp(historyLength, tokensPerBlock); @@ -145,6 +130,11 @@ inline bool requiresPrefixCoverage(LifeCycle const& lc) return reusePolicy(lc) == ReusePolicy::REQUIRED; } +inline bool allowsCheckpoint(LifeCycle const& lc) +{ + return reusePolicy(lc) == ReusePolicy::OPTIONAL; +} + // Free function: compute stale range via std::visit. inline HalfOpenRange getStaleRange(LifeCycle const& lc, int historyLength, int tokensPerBlock) { @@ -177,7 +167,10 @@ LifeCycle makeLifeCycle(LayerConfig const& layer, BufferConfig const& buffer, in // Integer id assigned to each unique LifeCycle. using LifeCycleId = StrongIndex; -// Alias for public exposure (same meaning as LifeCycleId). +// Public cache groups correspond one-to-one with lifecycles. +using CacheGroupId = LifeCycleId; + +// Legacy public name for a lifecycle group. using LayerGroupId = LifeCycleId; // --------------------------------------------------------------------------- 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 6bef15ba65e7..3d41f92e6e09 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 @@ -473,7 +473,7 @@ void SharedPageLock::releasePageIndex() // entry. updateBasePageIndex deliberately returns BAD for this ordinal. auto const& lc = page()->manager->lifeCycles().getLifeCycle(mUser.lifeCycle); auto const* attention = std::get_if(&lc); - TLLM_CHECK_DEBUG(std::holds_alternative(lc) || (attention && attention->allowsCheckpoint())); + TLLM_CHECK_DEBUG(std::holds_alternative(lc) || (attention && allowsCheckpoint(*attention))); return; } int oldBaseIndex diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index da7298c11061..f9e9aca76ec0 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -2446,7 +2446,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) { for (auto const& [lc, attention] : cache.manager().lifeCycles().attentionLifeCycles()) { - if (!attention->allowsCheckpoint() + if (!kv::allowsCheckpoint(*attention) || (groups && std::find(groups->begin(), groups->end(), lc.value()) == groups->end())) { continue; @@ -3665,22 +3665,22 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def_prop_ro("layer_ids", &kv::KvCacheManager::layerIds) .def_prop_ro( "layer_grouping", [](kv::KvCacheManager const& self) { return self.layerGrouping().raw(); }, - "Deprecated: use buffer_grouping. Partitions uniform layers by lifecycle; raises ValueError " + "Deprecated: use buffer_grouping. Partitions uniform layers by lifecycle; raises RuntimeError " "if any layer spans several groups. Member order does not describe physical memory order.") .def_prop_ro( "buffer_grouping", [](kv::KvCacheManager const& self) { return self.bufferGrouping().raw(); }, - "Buffer membership indexed by lifecycle group ID. Each buffer appears once; member order is unspecified.") + "Buffer membership indexed by CacheGroupId. Each buffer appears once; member order is unspecified.") .def( "get_layer_group_id", [](kv::KvCacheManager const& self, kv::LayerId layerId) { return self.getLayerGroupId(layerId).value(); }, nb::arg("layer_id"), - "Deprecated: use get_buffer_group_id. Returns a uniform layer's group; raises ValueError for a split " + "Deprecated: use get_buffer_group_id. Returns a uniform layer's group; raises RuntimeError for a split " "layer.") .def( "get_buffer_group_id", [](kv::KvCacheManager const& self, kv::LayerId layerId, kv::DataRole const& role) { return self.getBufferGroupId(layerId, role).value(); }, - nb::arg("layer_id"), nb::arg("role"), "Lifecycle group of one buffer of the layer.") + nb::arg("layer_id"), nb::arg("role"), "CacheGroupId of one buffer of the layer.") .def( "get_sparse_cache_max_units", [](kv::KvCacheManager const& self, kv::LayerId layerId) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index d2245c7b76e6..97bfa0521b3c 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -502,7 +502,7 @@ TEST_F(ReplayCacheTest, SsmBufferPolicySplitsLayerAcrossGroups) { SCOPED_TRACE(std::to_string(buffer.layer) + "/" + buffer.role); EXPECT_EQ(manager->getBufferGroupId(buffer.layer, buffer.role), buffer.isPrivate ? privateGroup : required); - EXPECT_THROW(manager->getLayerGroupId(buffer.layer), std::invalid_argument); + EXPECT_THROW(manager->getLayerGroupId(buffer.layer), tensorrt_llm::common::TllmException); } EXPECT_THROW(manager->getBufferGroupId(1, "missing"), std::out_of_range); @@ -510,7 +510,7 @@ TEST_F(ReplayCacheTest, SsmBufferPolicySplitsLayerAcrossGroups) auto layerIds = manager->layerIds(); std::sort(layerIds.begin(), layerIds.end()); EXPECT_EQ(layerIds, (std::vector{0, 1, 2, 3})); - EXPECT_THROW(manager->layerGrouping(), std::invalid_argument); + EXPECT_THROW(manager->layerGrouping(), tensorrt_llm::common::TllmException); auto const grouping = manager->bufferGrouping(); ASSERT_EQ(grouping.size(), LayerGroupId{3}); size_t count = 0; diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index a02df21c8f35..4ca5f0ab7842 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -183,6 +183,7 @@ class _KVCacheManagerConfigFieldSpec: take_poison = _cpp.take_poison BeamIndex = int +CacheGroupId = int CacheLevel = int CacheTierConfig = Union[GpuCacheTierConfig, HostCacheTierConfig, DiskCacheTierConfig] BatchRow = int @@ -227,6 +228,7 @@ def typed_range(*args: int) -> range: "SparseCacheConfig", "BufferId", "CoalescedBuffer", + "CacheGroupId", "CacheLevel", "CacheTier", "CacheTierConfig", diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index a4c2a39cf8c0..c8140627b2e5 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -90,6 +90,7 @@ class PageIndexMode(enum.IntEnum): PER_LAYER = 1 LifeCycleId = NewType("LifeCycleId", int) +CacheGroupId: TypeAlias = LifeCycleId LayerGroupId: TypeAlias = LifeCycleId class AttnLifeCycle: @@ -1859,14 +1860,16 @@ class KVCacheManager: def num_layers(self) -> int: ... @property def layer_ids(self) -> Iterator[LayerId]: ... - # Deprecated: use get_buffer_group_id. Raises ValueError for a split layer. + # Deprecated: use get_buffer_group_id. Raises RuntimeError for a split layer. def get_layer_group_id(self, layer_id: LayerId) -> LayerGroupId: ... - def get_buffer_group_id(self, layer_id: LayerId, role: DataRole) -> LayerGroupId: ... - # Deprecated: use buffer_grouping. Raises ValueError if any layer is split. + def get_buffer_group_id(self, layer_id: LayerId, role: DataRole) -> CacheGroupId: ... + # Deprecated: use buffer_grouping. Raises RuntimeError if any layer is split. @property def layer_grouping(self) -> Sequence[Sequence[LayerId]]: ... @property - def buffer_grouping(self) -> Sequence[Sequence[BufferId]]: ... + def buffer_grouping(self) -> Sequence[Sequence[BufferId]]: + """Buffer membership indexed by CacheGroupId.""" + ... @property def all_buffer_ids(self) -> Iterator[BufferId]: ... def get_aggregated_pages(self, buffers: Iterable[BufferId]) -> Iterator[AggregatedPageDesc]: ... 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 fc361e9fcd9f..5523b05c4f34 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 @@ -223,11 +223,11 @@ def test_ssm_buffer_reuse_policies(length): assert manager.get_buffer_group_id(2, "old_x") == private assert manager.get_layer_group_id(0) == manager.get_buffer_group_id(0, "key") for layer_id in (1, 2): - with pytest.raises(ValueError, match="layer groups"): + with pytest.raises(RuntimeError, match="layer groups"): manager.get_layer_group_id(layer_id) assert manager.num_layers == 3 assert sorted(manager.layer_ids) == [0, 1, 2] - with pytest.raises(ValueError, match="layer groups"): + with pytest.raises(RuntimeError, match="layer groups"): _ = manager.layer_grouping for group in (required, private): assert sorted({buffer.layer_id for buffer in manager.buffer_grouping[group]}) == [1, 2] From 2295af02be34cf7160d331436f05e75729d60861 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Fri, 30 Oct 2026 10:50:30 +0800 Subject: [PATCH 18/19] Adapt buffer reuse policies to upstream sparse cache APIs Signed-off-by: Jiagan Cheng --- .../kv_cache_manager_v2/AGENTS.md | 2 + .../kv_cache_manager_v2/kvCacheManager.cpp | 4 +- .../batch_manager/kvCacheManagerV2.cpp | 7 ++-- .../kvCacheManagerV2ColdPageTest.cpp | 2 - .../kvCacheManagerV2StatsTest.cpp | 8 ++-- .../runtime/kv_cache_manager_v2/__init__.pyi | 4 +- .../test_optional_checkpoint.py | 37 +++++++++++++++++++ 7 files changed, 50 insertions(+), 14 deletions(-) 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 21989ea2bb25..821fe3abea6b 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 @@ -119,6 +119,8 @@ validate that they use the same policy; independent policies require independent index tables. In particular, the default attention view couples K/V and scales, and Mamba kernels couple recurrent, convolution, PLE, and replay state. Native KVCM supports independent policies for callers that provide those indices. +The sparse fetch path configured by `sparseCache` still uses one group per layer; +its setup uses the checked layer lookup and rejects split layers. SSM lifecycles are deduplicated by policy, so a manager can have up to three SSM groups; iterate `ssmLifeCycleIds()` for working state and use 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 0a6e5cb398d0..ba0f88a964fe 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 @@ -184,7 +184,7 @@ void KvCacheManager::_computeSparseCacheShape() { continue; } - LifeCycleId const lcId = mLifeCycles.getId(makeLifeCycle(layer, mConfig.tokensPerBlock)); + LifeCycleId const lcId = getLayerGroupId(attn->layerId); auto const& cache = *attn->sparseCache; int const unitsPerPage = cache.unitsPerPage(mConfig.tokensPerBlock); @@ -296,7 +296,7 @@ std::vector KvCacheManager::_sparseBufferPlan(LifeCycleId life for (auto const& layer : mConfig.layers) { auto const* attn = std::get_if(&layer); - if (attn == nullptr || mLifeCycles.getId(makeLifeCycle(layer, mConfig.tokensPerBlock)) != lifeCycle) + if (attn == nullptr || !attn->sparseCache.has_value() || getLayerGroupId(attn->layerId) != lifeCycle) { continue; } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index f9e9aca76ec0..4a8eb684fdcb 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1722,9 +1722,10 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) "may name a layer that named nothing.") DEF_COPY(kv::SparseCacheConfig); nb::class_(m, "BufferConfig") - .def(nb::init, bool, int, int, std::optional>(), nb::arg("role"), nb::arg("size"), - nb::arg("tokens_per_block_override") = std::nullopt, nb::arg("is_sparse") = false, - nb::arg("coalescing_group") = 0, nb::arg("sparse_unit_chunks") = 1, nb::arg("reuse_policy") = std::nullopt) + .def(nb::init, bool, int, int, std::optional>(), + nb::arg("role"), nb::arg("size"), nb::arg("tokens_per_block_override") = std::nullopt, + nb::arg("is_sparse") = false, nb::arg("coalescing_group") = 0, nb::arg("sparse_unit_chunks") = 1, + nb::arg("reuse_policy") = std::nullopt) .def_rw("role", &kv::BufferConfig::role) .def_rw("size", &kv::BufferConfig::size) .def_rw("tokens_per_block_override", &kv::BufferConfig::tokensPerBlockOverride) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp index a3cb5cdebc8f..04c7e5b590c0 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp @@ -244,8 +244,6 @@ TEST_P(OptionalDiskCheckpointTest, PayloadRestoreRetryAndIndependentEviction) cudaSuccess); } ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); - auto reservations - = storage.newSlotsForPoolGroup(kHotLevel, hotGroup, storage.getStatistics(kHotLevel, hotGroup).free); EXPECT_FALSE(reader->resume(stream, std::vector{encoder, peer})); EXPECT_FALSE(reader->isActive()); EXPECT_EQ(reader->numCommittedTokens(), 8); diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index 97bfa0521b3c..17786a60a338 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -291,9 +291,8 @@ TEST_F(ReplayCacheTest, MixedSsmPoliciesRestoreStateAndCopyEveryBeam) EXPECT_FALSE(reader->reuseStatus()[privateGroup.value()].complete); EXPECT_EQ( reader->reuseStatus()[optional.value()].coverage, (std::vector>{{length, length}})); - EXPECT_THROW( - reader->resume(mStream, std::nullopt, std::vector{privateGroup}), std::invalid_argument); - ASSERT_TRUE(reader->resume(mStream, std::nullopt, std::vector{optional})); + EXPECT_THROW(reader->resume(mStream, std::vector{privateGroup}), std::invalid_argument); + ASSERT_TRUE(reader->resume(mStream, std::vector{optional})); reader->commitPendingStats(); auto const snapshotStats = manager->getAndResetSsmSnapshotIterationStats(); ASSERT_EQ(snapshotStats.size(), 2); @@ -360,8 +359,7 @@ TEST_F(ReplayCacheTest, OptionalSsmRequiresExactEndpointAndDoesNotLimitAttention EXPECT_EQ(reader->reuseStatus()[group.value()].complete, available); if (!available) { - EXPECT_THROW( - reader->resume(mStream, std::nullopt, std::vector{group}), std::invalid_argument); + EXPECT_THROW(reader->resume(mStream, std::vector{group}), std::invalid_argument); } // Unselected OPTIONAL and PRIVATE state both get writable request-local storage. ASSERT_TRUE(reader->resume(mStream)); diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index c8140627b2e5..a7a86cc9d105 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -30,6 +30,7 @@ from typing import ( Type, TypeAlias, Union, + overload, ) from tensorrt_llm.bindings.executor import KvCacheRetentionConfig @@ -325,8 +326,6 @@ class BufferConfig: is_sparse: bool = False coalescing_group: int = 0 - reuse_policy: ReusePolicy | None = None - sparse_unit_chunks: int = 1 """Pieces a selection unit is gathered from within *this buffer's* page. @@ -343,6 +342,7 @@ class BufferConfig: ``sparse_cache.unit_tokens`` and the page being densely packed, so a padded page is not expressible. """ + reuse_policy: ReusePolicy | None = None @dataclass(slots=True) class SparseCacheConfig: 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 5523b05c4f34..d11bf8fb9eae 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 @@ -16,6 +16,43 @@ ] +@pytest.mark.parametrize("split", [False, True]) +def test_sparse_fetch_buffer_policy_overrides(split): + buffers = [ + kv.BufferConfig( + role=role, size=4096, is_sparse=is_sparse, reuse_policy=kv.ReusePolicy.PRIVATE + ) + for role, is_sparse in (("key", False), ("value", True)) + ] + if split: + buffers.append(kv.BufferConfig(role="index", size=2048)) + config = kv.KVCacheManagerConfig( + tokens_per_block=4, + cache_tiers=[kv.GpuCacheTierConfig(quota=4 << 20)], + layers=[ + kv.AttentionLayerConfig( + layer_id=0, + buffers=buffers, + sparse_cache=kv.SparseCacheConfig( + num_selected_blocks=2, max_cached_units=4, unit_tokens=4 + ), + ) + ], + ) + if split: + with pytest.raises(RuntimeError, match="layer groups"): + kv.KVCacheManager(config) + return + manager = kv.KVCacheManager(config) + try: + assert manager.get_buffer_group_id(0, "key") == manager.get_buffer_group_id(0, "value") + assert manager.get_sparse_cache_max_units(0) == 4 + assert manager.get_sparse_cache_units_per_page(0) == 1 + assert manager.get_sparse_cache_buffer_count(0) == 1 + finally: + manager.shutdown() + + @pytest.mark.parametrize("priority", [None, 80]) def test_reuse_status_and_optional_selection_bindings(priority): layers = [ From d0911a1c9fcd967db113ae87e6dad68a35343eaf Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Fri, 30 Oct 2026 11:52:42 +0800 Subject: [PATCH 19/19] Validate suspended SSM working state with page holders Signed-off-by: Jiagan Cheng --- .../batch_manager/kv_cache_manager_v2/kvCache.cpp | 4 +++- .../unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp | 2 ++ 2 files changed, 5 insertions(+), 1 deletion(-) 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 7e7b9e4c1ebd..912e4c1f828b 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 @@ -3190,7 +3190,9 @@ bool KvCache::_checkSanity() const } else { - TLLM_CHECK_DEBUG(std::holds_alternative(bp)); + // Working state is locked while active and held while suspended. + TLLM_CHECK_DEBUG((mStatus == Status::ACTIVE && std::holds_alternative(bp)) + || (mStatus == Status::SUSPENDED && std::holds_alternative>(bp))); auto page = blockPageGetPage(bp); TLLM_CHECK_DEBUG(dynamicPointerCast(page) != nullptr); } diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index 17786a60a338..eb7c82a71b10 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -328,6 +328,7 @@ TEST_F(ReplayCacheTest, MixedSsmPoliciesRestoreStateAndCopyEveryBeam) reader->setBeamWidth(BeamIndex{1}); reader->commitPendingStats(); EXPECT_TRUE(manager->getAndResetSsmSnapshotIterationStats().empty()); + reader->suspend(); reader->close(); source->close(); manager->shutdown(); @@ -370,6 +371,7 @@ TEST_F(ReplayCacheTest, OptionalSsmRequiresExactEndpointAndDoesNotLimitAttention ASSERT_TRUE(reader->resume(mStream)); reader->commitPendingStats(); EXPECT_TRUE(manager->getAndResetSsmSnapshotIterationStats().empty()); + reader->suspend(); reader->close(); } manager->shutdown();