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..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 @@ -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. @@ -93,9 +93,41 @@ 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. 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. +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 +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 +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. +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 +`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..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; } @@ -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..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 @@ -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,15 @@ struct BufferConfig //! densely packed, so a padded page is not expressible. int sparseUnitChunks = 1; + //! 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 + { + 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 +314,6 @@ inline void validateNoDuplicateBufferRoles(std::vector const& buff } // namespace detail -//! Controls whether an attention lifecycle participates in prefix matching and page reuse. -enum class AttentionReusePolicy -{ - //! 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,7 +328,8 @@ struct AttentionLayerConfig std::optional numSinkTokens; // Private working pages never constrain prefix matching or enter the tree. - AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; + // 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 +351,8 @@ 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 { @@ -362,6 +374,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/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index cb791c821abf..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 @@ -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) @@ -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]; @@ -294,13 +293,25 @@ 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)); } } 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(); } @@ -311,16 +322,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 +346,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 +367,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 +379,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 +440,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 +537,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 +546,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,7 +596,7 @@ bool KvCache::resume(std::optional stream, std::optionallifeCycles().isSsm(lc); auto const& lifeCycle = mManager->lifeCycles()[lc]; if (_needsFreshPrefixPages(lc)) { @@ -592,6 +613,7 @@ bool KvCache::resume(std::optional stream, 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 +762,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 +780,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 +822,7 @@ bool KvCache::resume(std::optional stream, std::optionallifeCycles().isSsm(lcIdx)) { targetBp = &mSsmBlocks[beamIdx][lcIdx]; blockOrdinal = kBadBlockOrdinal; @@ -841,6 +862,20 @@ bool KvCache::resume(std::optional stream, std::optionallifeCycles().isSsm(lc)) + { + continue; + } + if (mPendingStats.recordSsmSnapshotLookup(lc, mNumLookupTokens, numCommittedTokens(), mTokensPerBlock)) + { + mManager->markStatsDirty(id); + } + } + } mOptionalCandidates.clear(); // Deferred copies survive a failed decode admission and must not be repeated on retry. mPartialCopyRetentionHolds.clear(); @@ -931,14 +966,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 +1282,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 +1327,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 +1356,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,14 +1477,14 @@ 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. std::vector attachedLcs; for (auto const& [lcIdx, attn] : mManager->lifeCycles().attentionLifeCycles()) { - if (!attn->requiresPrefixCoverage()) + if (!requiresPrefixCoverage(*attn)) { continue; } @@ -1462,8 +1500,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 +1626,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 +1645,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 +1807,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 +2010,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 +2092,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)); } @@ -2105,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) @@ -2149,14 +2190,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 +2270,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 +2304,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 +2401,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 +2410,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 +2464,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 +2499,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 +2511,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 +2660,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 +2680,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 +2698,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 +2813,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 +2853,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 +2964,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(); + mSsmBlocks[kDefaultBeamIndex][lcId] = snapshotPage->hold(); if (claimRetention) { auto const retention = getRetention(snapshotBlock.ordinal()); @@ -2935,9 +2983,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 +3033,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"); } @@ -2997,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)); } @@ -3030,7 +3080,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 +3103,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 @@ -3110,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); @@ -3124,12 +3173,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) @@ -3141,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/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h index b7eef425270c..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 @@ -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. @@ -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 77a7e427b91f..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 @@ -32,6 +32,7 @@ #include #include #include +#include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { @@ -76,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()) { @@ -175,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); @@ -287,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; } @@ -602,11 +611,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); 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( @@ -800,16 +807,46 @@ std::vector KvCacheManager::layerIds() const LayerGroupId KvCacheManager::getLayerGroupId(LayerId layerId) const { - return mStorage->layerToLifeCycleIds().at(layerId); + auto const& lifeCycles = mStorage->layerToLifeCycleIds().at(layerId); + 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(); +} + +CacheGroupId KvCacheManager::getBufferGroupId(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(getLayerGroupId(lid)).push_back(lid); + } + return result; +} + +TypedVec> KvCacheManager::bufferGrouping() const +{ + TypedVec> result(mLifeCycles.size()); + std::unordered_map> buffersByLayer; + for (auto const& buffer : allBufferIds()) { - result.at(lc).push_back(lid); + 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/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..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 @@ -285,15 +285,23 @@ class KvCacheManager : public std::enable_shared_from_this int numLayers() const noexcept; std::vector layerIds() const; + + //! Deprecated: use getBufferGroupId(). Returns the shared group of a uniform layer. + //! Throws common::TllmException if the queried layer spans several lifecycle groups. LayerGroupId getLayerGroupId(LayerId layerId) const; - // Layer grouping: layers with the same lifecycle share pool allocation. - // 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. + //! 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 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 CacheGroupId. 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 f32204aba3e4..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,20 +27,20 @@ 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{}; + return SsmLifeCycle{policy}; } else { - return AttnLifeCycle::make(cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, cfg.reusePolicy); + return AttnLifeCycle::make(cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, policy); } }, layer); @@ -52,9 +52,8 @@ LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock) 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(); @@ -62,7 +61,15 @@ LifeCycleRegistry::LifeCycleRegistry(KVCacheManagerConfig const& config) mLifeCycleList.push_back(lc); mLifeCycleIdMap[lc] = id; if (std::holds_alternative(lc)) - mSsmLifeCycleId = id; + mSsmLifeCycleIds.push_back(id); + } + }; + for (auto const& layer : config.layers) + { + std::visit([](auto const& cfg) { cfg.validate(); }, layer); + for (auto const& buffer : layerBuffers(layer)) + { + 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 55cebefbf206..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 @@ -40,22 +40,7 @@ struct AttnLifeCycle { std::optional windowSize; // nullopt = no sliding window int numSinkBlocks = 0; // divUp(numSinkTokens, tokensPerBlock) - AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; - - bool isPrivate() const noexcept - { - return reusePolicy == AttentionReusePolicy::PRIVATE; - } - - bool requiresPrefixCoverage() const noexcept - { - return reusePolicy == AttentionReusePolicy::REQUIRED; - } - - bool allowsCheckpoint() const noexcept - { - return reusePolicy == AttentionReusePolicy::OPTIONAL; - } + ReusePolicy reusePolicy = ReusePolicy::REQUIRED; HalfOpenRange getStaleRange(int historyLength, int tokensPerBlock) const { @@ -89,7 +74,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); @@ -101,25 +86,27 @@ 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). -// Singleton: all SSM layers share the same lifecycle. +// SSM buffers with the same effective 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 +115,24 @@ 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; +} + +inline bool allowsCheckpoint(LifeCycle const& lc) +{ + return reusePolicy(lc) == ReusePolicy::OPTIONAL; } // Free function: compute stale range via std::visit. @@ -166,13 +161,16 @@ inline HalfOpenRange computeScratchRange( return intersect(capStale, inputRange); } -// Factory: create a LifeCycle from a LayerConfig variant. -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; -// 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; // --------------------------------------------------------------------------- @@ -210,14 +208,31 @@ class LifeCycleRegistry } // SSM helpers. - std::optional ssmLifeCycleId() const noexcept + std::vector const& ssmLifeCycleIds() const noexcept { - return mSsmLifeCycleId; + return mSsmLifeCycleIds; + } + + bool isSsm(LifeCycleId id) const + { + 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 +308,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..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 @@ -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; @@ -474,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/batch_manager/kv_cache_manager_v2/storage/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp index b4dcd77c153b..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; } } @@ -86,63 +92,22 @@ 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); } } - return map; -} - -std::map StorageConfig::layerAttributes() const -{ - std::map ret; - for (auto const& pg : slotDescList) + for (auto& [layerId, ids] : map) { - 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) - { - auto it = ret.find(layerId); - 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; - } - auto& attr = it->second; - TLLM_CHECK_DEBUG(attr.lifeCycleId == lcId); - attr.slotUtil[poolIndex] = count; - Rational frac{count, buffersPerSlot}; - if (frac > attr.slotUtilFracMax) - { - attr.slotUtilFracMax = frac; - } - } - } - } + std::sort(ids.begin(), ids.end()); } - return ret; + return map; } // --------------------------------------------------------------------------- @@ -167,24 +132,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); - - BufferId bid{cfg.layerId, buf.role}; - expansionMap[bid] = exp; - bufferGroups[lcId][{buf.coalescingGroup, expandedSize, buf.isSparse}].push_back(bid); - } - }, - layer); + 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{layerId, buf.role}; + expansionMap[bid] = exp; + bufferGroups[lcId][{buf.coalescingGroup, expandedSize, buf.isSparse}].push_back(bid); + } } // Build one SlotDescVariant per life cycle. @@ -216,7 +175,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 +190,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/storage/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h index 8be1a9eccf6e..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 @@ -28,6 +28,7 @@ #include #include #include +#include #include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 @@ -195,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 { @@ -227,22 +228,6 @@ struct Rational } }; -// --------------------------------------------------------------------------- -// LayerAttr — per-layer storage attributes for scratch slot management. -// 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; -}; - // --------------------------------------------------------------------------- // StorageConfig — complete storage layout for the KV cache. // Mirrors _storage/_config.py::StorageConfig. @@ -264,12 +249,9 @@ struct StorageConfig // Attribute map for each buffer. std::map bufferAttributes() const; - // Map LayerId → LifeCycleId. - std::unordered_map layerToLifeCycleIds() const; - - // Per-layer storage attributes (slot utilization within coalesced slots). - // Mirrors _storage/_config.py::StorageConfig.layer_attributes(). - std::map layerAttributes() const; + // Map LayerId → the lifecycles of its buffers, in ascending order. A layer spans several + // lifecycles when its buffers use different reuse policies. + std::unordered_map> layerToLifeCycleIds() 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 0735340b929f..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& [layerId, 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; + } + } + } } } @@ -432,12 +447,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; @@ -1630,14 +1644,6 @@ TypedVec StorageManager::poolGroupDescs() const return result; } -LayerAttr const& StorageManager::getLayerAttr(LayerId layerId) const -{ - auto it = mLayerAttributes.find(layerId); - if (it == mLayerAttributes.end()) - throw std::out_of_range("Unknown LayerId for LayerAttr"); - return it->second; -} - SlotCount StorageManager::numSlots(PoolGroupIndex pgIdx, CacheLevel level) const { return mLevels.at(level).storage->numSlots(pgIdx); @@ -1885,14 +1891,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 +1958,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 +2057,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/batch_manager/kv_cache_manager_v2/storageManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.h index fd4b8366fd10..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; - // Per-layer storage attributes. - LayerAttr const& getLayerAttr(LayerId layerId) 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 +328,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,14 +474,11 @@ 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; - // 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 f7d9ab23f7e0..4a8eb684fdcb 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==); @@ -1721,21 +1722,23 @@ 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"), - nb::arg("tokens_per_block_override") = std::nullopt, nb::arg("is_sparse") = false, - nb::arg("coalescing_group") = 0, nb::arg("sparse_unit_chunks") = 1) + .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) .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, - 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 +1753,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); @@ -2442,7 +2447,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; @@ -2625,10 +2630,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 @@ -3661,14 +3666,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. 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 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 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")) + nb::arg("layer_id"), + "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"), "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/kvCacheManagerV2ColdPageTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp index 496d2c413acf..04c7e5b590c0 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)); @@ -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) @@ -237,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/kvCacheManagerV2ConfigTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp index 78995508390d..0b3ff6ac90a5 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,76 @@ TEST(KvCacheManagerV2ConfigTest, SparseBuffersCoalesceBySize) EXPECT_EQ(second.expansion, 1); } +TEST(KvCacheManagerV2ConfigTest, AttentionBufferPoliciesOverrideLayerDefault) +{ + 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_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); + } +} + +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]); +} + } // namespace diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index eed2adf2cf9f..eb7c82a71b10 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 @@ -67,7 +71,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 +90,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 +146,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 +166,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 +184,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 +225,528 @@ 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::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); + 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); + 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->commitPendingStats(); + EXPECT_TRUE(manager->getAndResetSsmSnapshotIterationStats().empty()); + reader->suspend(); + 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::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->commitPendingStats(); + EXPECT_TRUE(manager->getAndResetSsmSnapshotIterationStats().empty()); + reader->suspend(); + ASSERT_TRUE(reader->resume(mStream)); + reader->commitPendingStats(); + EXPECT_TRUE(manager->getAndResetSsmSnapshotIterationStats().empty()); + reader->suspend(); + 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(); +} + +// 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.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) + * 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->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])); + EXPECT_TRUE(isPrivateLifeCycle(manager->lifeCycles()[privateGroup])); + for (auto const& buffer : splitSsmBuffers()) + { + 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), tensorrt_llm::common::TllmException); + } + 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})); + EXPECT_THROW(manager->layerGrouping(), tensorrt_llm::common::TllmException); + auto const grouping = manager->bufferGrouping(); + ASSERT_EQ(grouping.size(), LayerGroupId{3}); + size_t count = 0; + for (LayerGroupId group{0}; group < grouping.size(); ++group) + { + for (auto const& buffer : grouping[group]) + { + EXPECT_EQ(manager->getBufferGroupId(buffer.layerId, buffer.role), group); + ++count; + } + } + EXPECT_EQ(count, manager->allBufferIds().size()); + 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->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); + 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->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) + { + 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->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); + 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& [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); + 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(); @@ -252,14 +778,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 +828,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 +907,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 +944,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 +977,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..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 @@ -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: @@ -3906,7 +3906,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") @@ -4044,7 +4048,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..4ca5f0ab7842 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 @@ -183,6 +183,7 @@ class _KVCacheManagerConfigFieldSpec: take_poison = _cpp.take_poison BeamIndex = int +CacheGroupId = int CacheLevel = int CacheTierConfig = Union[GpuCacheTierConfig, HostCacheTierConfig, DiskCacheTierConfig] BatchRow = int @@ -218,7 +219,7 @@ def typed_range(*args: int) -> range: __all__ = [ "AggregatedPageDesc", "AttentionLayerConfig", - "AttentionReusePolicy", + "ReusePolicy", "BAD_PAGE_INDEX", "CACHE_LEVEL1", "BatchDesc", @@ -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 1211ecfd3d89..a7a86cc9d105 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -91,6 +91,7 @@ class PageIndexMode(enum.IntEnum): PER_LAYER = 1 LifeCycleId = NewType("LifeCycleId", int) +CacheGroupId: TypeAlias = LifeCycleId LayerGroupId: TypeAlias = LifeCycleId class AttnLifeCycle: @@ -100,21 +101,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) @@ -274,6 +275,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`). @@ -336,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: @@ -394,7 +401,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 +420,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 +432,16 @@ class AttentionLayerConfig: @dataclass(slots=True) class SsmLayerConfig: + """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. + """ + layer_id: LayerId buffers: list[BufferConfig] + reuse_policy: ReusePolicy = ReusePolicy.REQUIRED LayerConfig = AttentionLayerConfig | SsmLayerConfig @@ -657,7 +672,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: ... @@ -1845,10 +1860,17 @@ class KVCacheManager: def num_layers(self) -> int: ... @property def layer_ids(self) -> Iterator[LayerId]: ... + # 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) -> 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]]: + """Buffer membership indexed by CacheGroupId.""" + ... + @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/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_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index addb326f0a91..83bd77f0c2f4 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,62 @@ 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 + + @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( 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..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 = [ @@ -23,9 +60,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 +126,291 @@ 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]) +@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)]), + *[ + 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, + ) + ], + ] + 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)], + 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_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 + 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 + 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 [] + ) + 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) + ] + 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() + + +@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_buffer_group_id(1, "state") + private = manager.get_buffer_group_id(layer_id=1, role="old_x") + assert required != private + 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(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(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] + 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) + 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("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, + 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), + 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 + 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 [] + ) + reader.suspend() + assert reader.resume(stream.cuda_stream) + finally: + 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()