diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 4200e690755e..514e6f6e9155 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -115,6 +115,23 @@ std::list> chopVectorIntoBlocks( return blockedVectors; } +//! \brief Configuration of a single KV pool. +//! +//! A pool is uniquely described by its attention @c windowSize, the per-layer +//! @c sizePerHead of the layers it serves, and the KV-cache element @c dtype. +//! A KVCacheManager / BlockManager is constructed from a +//! @c std::vector -- one entry per pool the manager hosts. +//! Multiple entries with the same @c windowSize are legal and reserved for +//! future multi-pool-per-window cases (e.g. mixed head_dim within a single +//! window); call sites that key by window today must be revisited when that +//! lands. +struct PoolConfiguration +{ + SizeType32 windowSize; + SizeType32 sizePerHead; + nvinfer1::DataType dtype; +}; + struct LinearAttentionMetadata { enum LinearCacheType : WindowSizeType @@ -554,7 +571,6 @@ class GenerationRequest , mNumTokens(numTokens) , mBeamWidth(beamWidth) , mKvCacheRetentionConfig(std::move(kvCacheRetentionConfig)) - , mNumFrontBlocksRemoved(0) , mCurrentPrepopulatedPromptLen(std::numeric_limits::max()) { auto const numWindowSizes = windowSizeToMetadata.size(); @@ -563,6 +579,8 @@ class GenerationRequest for (auto const [windowSize, metadata] : windowSizeToMetadata) { mCacheBlockIds[windowSize] = std::vector>(beamWidth); + // Per-window front-eviction counter; SWA windows evict, full-attention windows stay at 0. + mNumFrontBlocksRemovedPerWindow[windowSize] = 0; auto const numPools = metadata.numPools; auto const maxBlocks = metadata.maxBlocksPerSeq; mCacheBlockIndices[windowSize] @@ -598,9 +616,14 @@ class GenerationRequest return mNumTokens; } - [[nodiscard]] SizeType32 getNumFrontBlocksRemoved() const + //! \brief Per-window front-eviction count; full-attention windows always return 0. + [[nodiscard]] SizeType32 getNumFrontBlocksRemoved(SizeType32 windowSize) const { - return mNumFrontBlocksRemoved; + auto it = mNumFrontBlocksRemovedPerWindow.find(windowSize); + TLLM_CHECK_WITH_INFO(it != mNumFrontBlocksRemovedPerWindow.end(), + "GenerationRequest::getNumFrontBlocksRemoved: windowSize=%d not registered for request %lu", windowSize, + static_cast(mRequestId)); + return it->second; } [[nodiscard]] SizeType32 getBeamWidth() const @@ -645,12 +668,18 @@ class GenerationRequest { beamBlockIds.clear(); } - mNumFrontBlocksRemoved = 0; + // Reset only this window's counter, not the others — clearing one pool's blocks + // must not invalidate the eviction state of co-existing pools. + auto it = mNumFrontBlocksRemovedPerWindow.find(windowSize); + if (it != mNumFrontBlocksRemovedPerWindow.end()) + { + it->second = 0; + } } void removeFrontBlock(SizeType32 windowSize) { - ++mNumFrontBlocksRemoved; + ++mNumFrontBlocksRemovedPerWindow.at(windowSize); } void removeLastBlock(SizeType32 windowSize) @@ -708,8 +737,10 @@ class GenerationRequest std::unordered_map mCacheBlockIndices; // The retention priority to assign to decode blocks executor::KvCacheRetentionConfig mKvCacheRetentionConfig; - // Number of front blocks removed from the sequence - SizeType32 mNumFrontBlocksRemoved; + // Number of front blocks evicted from the sequence, tracked per window size. + // Each WindowBlockManager increments only its own entry on detachFrontBlock; full-attn + // windows always stay at 0 since they never trigger SWA eviction. + std::map mNumFrontBlocksRemovedPerWindow; // Set of used blocks by the sequence std::set mUsedBlocks; // Current prepopulated prompt length @@ -968,6 +999,31 @@ class WindowBlockManager return mWindowSize; } + //! \brief Return this manager's KV-cache element data type. + //! \details Each WindowBlockManager carries its own dtype, which lets BlockManager + //! host pools with mixed precisions when constructed with a per-window + //! dtype map. Empty pools or NVFP4-scale pools are routed through the + //! per-pool tensor metadata instead. + [[nodiscard]] nvinfer1::DataType getDataType() const noexcept + { + return mDataType; + } + + //! \brief Return the head_dim shared by all KV pools in this window manager. + //! \details All pools constructed under a WindowBlockManager share a single + //! head_dim (the constructor's @c sizePerHead), so we read it from the + //! first pool. Returns 0 if no pools are present, or if the pool was + //! built with @c sizePerHead<=0 (recurrent-state pools use -1 as a + //! sentinel and do not have a meaningful head dimension). + [[nodiscard]] SizeType32 getSizePerHead() const noexcept + { + if (mPools.empty() || mPools.front().sizePerHead <= 0) + { + return 0; + } + return mPools.front().sizePerHead; + } + [[nodiscard]] std::string const& getLogPrefix() const noexcept { return mLogPrefix; @@ -1377,6 +1433,16 @@ class BlockManager using SizeType32 = tensorrt_llm::runtime::SizeType32; using BaseEvictionPolicy = tensorrt_llm::batch_manager::eviction_policy::BaseEvictionPolicy; + //! \brief Construct a BlockManager. + //! \param sizePerHead Default head size (used for pools whose window_size has no entry in + //! @p poolConfigurations). + //! \param dtype Default KV-cache data type (used for pools whose window_size has no entry in + //! @p poolConfigurations). + //! \param poolConfigurations One PoolConfiguration per pool the manager will host. Each + //! entry pins (windowSize, sizePerHead, dtype) for one pool, letting a single + //! BlockManager host pools with mixed shapes (e.g. Gemma4 SWA head_dim=256 alongside + //! full-attention head_dim=512). Empty vector = uniform @p sizePerHead / @p dtype + //! across all windows. explicit BlockManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, @@ -1388,8 +1454,8 @@ class BlockManager std::shared_ptr kvCacheConnectorManager = nullptr, std::optional agentConfig = std::nullopt, bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, - bool indexerKCacheUseFp4 = false, - std::optional linearAttentionMetadata = std::nullopt); + bool indexerKCacheUseFp4 = false, std::optional linearAttentionMetadata = std::nullopt, + std::vector const& poolConfigurations = {}); [[nodiscard]] bool isEnableIndexerKCache() const { @@ -1502,6 +1568,44 @@ class BlockManager return numFreeBlocksPerWindowSize; } + //! \brief Per-pool configuration view across all WindowBlockManagers. + //! \details One PoolConfiguration per WindowBlockManager, exposing (windowSize, + //! sizePerHead, dtype) for the pool. Lets callers (e.g. Python-side tensor + //! reshaping or kernel dispatch) discover per-pool shape information without + //! reaching into KVCacheBlockPool. + [[nodiscard]] std::vector getPoolConfigurations() const + { + std::vector result; + result.reserve(mWindowBlockManagers.size()); + for (auto const& [windowSize, manager] : mWindowBlockManagers) + { + result.push_back(PoolConfiguration{windowSize, manager.getSizePerHead(), manager.getDataType()}); + } + return result; + } + + //! \brief Convenience: window_size -> dataType, derived from getPoolConfigurations(). + //! For one-pool-per-window managers only; multi-pool-per-window will collide. + [[nodiscard]] std::map getDataTypePerWindow() const + { + std::map result; + for (auto const& [windowSize, manager] : mWindowBlockManagers) + { + result[windowSize] = manager.getDataType(); + } + return result; + } + + [[nodiscard]] SizeType32 getSizePerHeadForWindow(SizeType32 windowSize) const + { + return mWindowBlockManagers.at(windowSize).getSizePerHead(); + } + + [[nodiscard]] nvinfer1::DataType getDataTypeForWindow(SizeType32 windowSize) const + { + return mWindowBlockManagers.at(windowSize).getDataType(); + } + [[nodiscard]] SizeType32 getNumFreeBlocks() const { return sumWindows([](WindowBlockManager const& manager) { return manager.getNumFreeBlocks(); }); @@ -2032,9 +2136,9 @@ class BaseKVCacheManager /// of memory requirements. The weighting considers both the window size and the number of /// layers using each window size, as well as the sum of cache sizes per token for each window. /// @param config KV cache configuration parameters - /// @param dtype Data type used for KV cache values + /// @param dtype Default KV-cache data type (used for windows without a matching pool entry) /// @param numKvHeadsPerLayer Number of KV heads for each local layer (caller selects self/cross attention heads) - /// @param sizePerHead Size of each attention head + /// @param sizePerHead Default head size (used for windows without a matching pool entry) /// @param tokensPerBlock Number of tokens per KV cache block /// @param worldConfig World configuration for multi-GPU setups /// @param windowSizeToLayers Map from attention window size to vector of layer indices using that window size @@ -2043,13 +2147,19 @@ class BaseKVCacheManager /// @param extraCostMemory Additional memory cost to account for CacheTransBufferManager::preAllocBufferSize /// @param kvFactor Factor for KV cache size calculation (typically 2 for key+value) /// @param maxBatchSize Maximum batch size + /// @param linearAttentionMetadata Optional metadata for linear / recurrent state pools. + /// @param poolConfigurations One PoolConfiguration per pool the manager will host. Each entry + /// pins (windowSize, sizePerHead, dtype) for one pool, letting a single BlockManager + /// budget pools with mixed shapes (e.g. Gemma4 SWA head_dim=256 + full-attention + /// head_dim=512). Empty vector = uniform @p sizePerHead / @p dtype across all windows. /// @return Map from window size to tuple of (primary blocks, secondary blocks) [[nodiscard]] static BlocksPerWindow calculateMaxNumBlocks(executor::KvCacheConfig const& config, nvinfer1::DataType dtype, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, tensorrt_llm::runtime::WorldConfig const& worldConfig, std::map> const& windowSizeToLayers, uint64_t allottedPrimaryMemBytes, uint64_t allottedSecondaryMemBytes, size_t extraCostMemory, SizeType32 kvFactor, SizeType32 maxBatchSize, - std::optional const& linearAttentionMetadata = std::nullopt); + std::optional const& linearAttentionMetadata = std::nullopt, + std::vector const& poolConfigurations = {}); /// @brief Calculates the maximum batch size that can fit the kv-cache, given that all sequences in the batch have /// the provided input and output length. @@ -2097,6 +2207,13 @@ class KVCacheManager : public BaseKVCacheManager using CudaStreamPtr = std::shared_ptr; using CacheType = tensorrt_llm::batch_manager::kv_cache_manager::CacheType; + //! \brief Construct a KVCacheManager. + //! \param sizePerHead Default head size used for windows without a matching pool entry. + //! \param dtype Default KV-cache data type used for windows without a matching pool entry. + //! \param poolConfigurations One PoolConfiguration per pool the manager will host. + //! Lets a single manager host mixed-shape pools (e.g. Gemma4 SWA head_dim=256 + //! + full head_dim=512) so the existing block reuse, scheduling, MPI reduction, + //! and disagg transfer machinery applies natively. Empty vector = uniform. KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, @@ -2108,7 +2225,8 @@ class KVCacheManager : public BaseKVCacheManager std::shared_ptr kvCacheConnectorManager = nullptr, bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, - std::optional linearAttentionMetadata = std::nullopt); + std::optional linearAttentionMetadata = std::nullopt, + std::vector const& poolConfigurations = {}); KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, @@ -2121,7 +2239,8 @@ class KVCacheManager : public BaseKVCacheManager std::shared_ptr kvCacheConnectorManager = nullptr, bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, - std::optional linearAttentionMetadata = std::nullopt); + std::optional linearAttentionMetadata = std::nullopt, + std::vector const& poolConfigurations = {}); KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, @@ -2134,7 +2253,8 @@ class KVCacheManager : public BaseKVCacheManager std::shared_ptr kvCacheConnectorManager = nullptr, bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, - std::optional linearAttentionMetadata = std::nullopt); + std::optional linearAttentionMetadata = std::nullopt, + std::vector const& poolConfigurations = {}); KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, @@ -2143,7 +2263,8 @@ class KVCacheManager : public BaseKVCacheManager CacheType cacheType = CacheType::kSELF, bool enablePartialReuse = true, bool copyOnpartialReuse = true, bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false, - std::optional linearAttentionMetadata = std::nullopt); + std::optional linearAttentionMetadata = std::nullopt, + std::vector const& poolConfigurations = {}); ~KVCacheManager() override = default; diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index 3daf2735db96..17dd557be1a3 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -696,6 +696,11 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess { NVTX3_SCOPED_RANGE(formatInputRecvBuffer); + // TODO(disagg-multi-dtype): pool 0's dtype is treated as canonical for the wire + // transport here. Pools with differing dtypes are rejected up-front in + // CacheTransBufferManager's constructor (see cacheTransBuffer.cpp). When + // per-pool dtype dispatch lands, this single dataType variable must be replaced + // with a per-pool lookup keyed by the source pool of each block. auto dataType = mCacheManager->getPrimaryPool(0)->getDataType(); bool layerWise = common::getEnvDisaggLayerwise() && numKvPools == 1; if (layerWise) diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp index 875e3c7e3bee..b6f1e4e8df15 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp @@ -253,6 +253,28 @@ CacheTransBufferManager::CacheTransBufferManager( { // TODO: FP4 dataSize TLLM_CHECK(mCacheManager); + // TODO(disagg-multi-dtype): Per-pool dtype dispatch in formatter / transfer buffer + // not yet implemented. Disagg currently picks pool 0's dtype as the canonical + // transport type (above), so any KV pool with a different dtype would be silently + // miscoerced on the wire. Fail loudly until per-pool dispatch lands. We restrict + // the comparison to KV pools (getNumPools(false, false)) since block-scale and + // indexer-K pools legitimately have their own dtypes and travel through their own + // code paths. + if (!transferIndexerKCache) + { + auto const numKvPools = mCacheManager->getBlockManager().getNumPools( + /*includeBlockScalePools=*/false, /*includeIndexerKCachePools=*/false); + auto const dtype0 = mCacheManager->getPrimaryPool(0)->getDataType(); + for (SizeType32 i = 1; i < numKvPools; ++i) + { + auto const dtypeI = mCacheManager->getPrimaryPool(i)->getDataType(); + TLLM_CHECK_WITH_INFO(dtypeI == dtype0, + "Disaggregated KV cache transfer does not yet support pools with differing dtypes " + "(pool 0 dtype=%d, pool %d dtype=%d). TODO(disagg-multi-dtype): per-pool dtype " + "dispatch in formatter.", + static_cast(dtype0), i, static_cast(dtypeI)); + } + } TLLM_LOG_INFO("CacheTransBufferManager created for KV cache"); } diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index c129a3646214..03a4908e7c27 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -582,7 +582,8 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si std::shared_ptr kvCacheConnectorManager, std::optional agentConfig, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, - std::optional linearAttentionMetadata) + std::optional linearAttentionMetadata, + std::vector const& poolConfigurations) : mNumLayers{static_cast(numKvHeadsPerLayer.size())} , mTokensPerBlock{tokensPerBlock} , mEventManager{std::move(eventManager)} @@ -624,6 +625,18 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si mIsVariableWindow = numUniqueWindowSizes > 1; mIsVariableGQA = std::unordered_set(numKvHeadsPerLayer.begin(), numKvHeadsPerLayer.end()).size() > 1; + // Build a window -> PoolConfiguration index for fast lookup. Today, one pool per window + // is the only case in use; multi-pool-per-window would replace this with an explicit + // layer->pool mapping from the caller. + std::map poolByWindow; + for (auto const& pc : poolConfigurations) + { + TLLM_CHECK_WITH_INFO(poolByWindow.emplace(pc.windowSize, &pc).second, + "Multiple PoolConfigurations share windowSize=%d. Multi-pool-per-window requires an " + "explicit layer->pool mapping, which is not yet wired through BlockManager.", + pc.windowSize); + } + mLayerToWindowSize.resize(mNumLayers); for (auto const& [windowSize, layersWithWindowSize] : uniqueWindowSizeToLayers) { @@ -658,8 +671,15 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si } } - mWindowBlockManagers.try_emplace(SizeType32(windowSize), dtype, windowSize, layersWithWindowSize, - numKvHeadsPerLayer, sizePerHead, tokensPerBlock, + // Pull per-pool head_dim and dtype from poolConfigurations when available; fall back to + // the manager-level scalars otherwise. Lets a single BlockManager host pools with + // distinct shapes (Gemma4-style mixed attention types) without giving up the existing + // radix-tree / scheduler / disagg machinery that already iterates pools internally. + auto const poolIt = poolByWindow.find(windowSize); + auto const windowSizePerHead = (poolIt != poolByWindow.end()) ? poolIt->second->sizePerHead : sizePerHead; + auto const windowDtype = (poolIt != poolByWindow.end()) ? poolIt->second->dtype : dtype; + mWindowBlockManagers.try_emplace(SizeType32(windowSize), windowDtype, windowSize, layersWithWindowSize, + numKvHeadsPerLayer, windowSizePerHead, tokensPerBlock, /*isSWA=*/(windowSize < maxSequenceLength) && (windowSize >= 0), allottedPrimaryBlocks, allottedSecondaryBlocks, maxNumSequences, stream, cacheType, secondaryOffloadMinPriority, mEventManager, enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, mLookupTree, mLoopbackAgent, @@ -2173,7 +2193,7 @@ void WindowBlockManager::adjustBlocksIfNeeded(GenerationRequest& sequence) { auto const minTokensForBlockDetach = mWindowSize + mTokensPerBlock; while (mIsSWA && // A block only go out-of-window in SWA - (sequence.getNumTokens() - sequence.getNumFrontBlocksRemoved() * getTokensPerBlock() + (sequence.getNumTokens() - sequence.getNumFrontBlocksRemoved(mWindowSize) * getTokensPerBlock() >= minTokensForBlockDetach)) { // Detaching block for SWA is non-trivial due to the radix tree structure. @@ -2874,11 +2894,18 @@ void BlockManager::pinBlocks(GenerationRequest& sequence) void BlockManager::unpinBlocksById(std::vector const& blockIds) { - // Use the first window size if (mWindowBlockManagers.empty()) { return; } + // Block IDs are window-scoped (each WindowBlockManager owns its own + // mAllBlocksById vector), so silently dispatching to the first window + // would corrupt refcounts whenever the caller's IDs belong to a different + // window. Multi-window block-ID-keyed retention pinning needs + // window-qualified handles; until that lands, refuse the multi-window case. + TLLM_CHECK_WITH_INFO(mWindowBlockManagers.size() == 1, + "BlockManager::unpinBlocksById currently only supports single-window managers; " + "multi-window block-ID-keyed retention pinning requires window-qualified handles (TODO)."); auto& firstManager = mWindowBlockManagers.begin()->second; firstManager.unpinBlocksById(blockIds); } @@ -3132,13 +3159,14 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size SizeType32 sinkTokenLength, int64_t stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, bool enablePartialReuse, bool copyOnPartialReuse, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, - bool indexerKCacheUseFp4, std::optional linearAttentionMetadata) + bool indexerKCacheUseFp4, std::optional linearAttentionMetadata, + std::vector const& poolConfigurations) : KVCacheManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, dtype, sinkTokenLength, std::make_shared(reinterpret_cast(stream)), maxSequenceLength, chunkSize, enableBlockReuse, cacheType, std::nullopt, nullptr, enablePartialReuse, copyOnPartialReuse, nullptr, enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, - linearAttentionMetadata) + linearAttentionMetadata, poolConfigurations) { } @@ -3150,13 +3178,14 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, - std::optional linearAttentionMetadata) + std::optional linearAttentionMetadata, + std::vector const& poolConfigurations) : KVCacheManager(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, dtype, sinkTokenLength, std::make_shared(reinterpret_cast(stream)), maxSequenceLength, chunkSize, enableBlockReuse, cacheType, secondaryOffloadMinPriority, eventManager, enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, enableIndexerKCache, indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, - indexerKCacheUseFp4, linearAttentionMetadata) + indexerKCacheUseFp4, linearAttentionMetadata, poolConfigurations) { } @@ -3168,7 +3197,8 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, - std::optional linearAttentionMetadata) + std::optional linearAttentionMetadata, + std::vector const& poolConfigurations) : mMaxBeamWidth(maxBeamWidth) , mDataType(dtype) , mMaxAttentionWindow(*std::max_element(maxAttentionWindowVec.begin(), maxAttentionWindowVec.end())) @@ -3180,7 +3210,8 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer std::move(stream), maxSequenceLength, maxBeamWidth, maxAttentionWindowVec, dtype, mSinkBubbleLength, mChunkSize, cacheType, secondaryOffloadMinPriority, std::move(eventManager), enablePartialReuse, copyOnPartialReuse, std::move(kvCacheConnectorManager), std::nullopt, enableIndexerKCache, - indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, linearAttentionMetadata) + indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, linearAttentionMetadata, + poolConfigurations) // disable block reuse for sink bubble since chopVectorIntoBlocks does not match KV cache blocks in this case , mEnableBlockReuse{mSinkBubbleLength > 0 ? false : enableBlockReuse} { @@ -3208,12 +3239,14 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, bool indexerKCacheUseFp4, - std::optional linearAttentionMetadata) + std::optional linearAttentionMetadata, + std::vector const& poolConfigurations) : KVCacheManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, dtype, sinkTokenLength, std::move(stream), maxSequenceLength, chunkSize, enableBlockReuse, cacheType, secondaryOffloadMinPriority, std::move(eventManager), enablePartialReuse, copyOnPartialReuse, std::move(kvCacheConnectorManager), enableIndexerKCache, - indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, linearAttentionMetadata) + indexerKCacheQuantBlockSize, indexerKCacheIndexHeadDim, indexerKCacheUseFp4, linearAttentionMetadata, + poolConfigurations) { } @@ -3225,16 +3258,21 @@ void KVCacheManager::allocatePools(bool useUvm) uint64_t cacheSizeBytes = 0; for (SizeType32 poolIdx = 0; poolIdx < numPools; poolIdx++) { - auto const cacheShape = mBlockManager.getPrimaryPool(poolIdx)->getShape(); + auto const& primaryPool = mBlockManager.getPrimaryPool(poolIdx); + auto const cacheShape = primaryPool->getShape(); auto const cacheVolume = ITensor::volume(cacheShape); + // Query the per-pool dtype from the tensor itself rather than the manager-level + // mDataType. Each pool may carry its own dtype: NVFP4 element / scale pools, or + // a future per-window override map can mix precisions inside a single manager. + auto const poolDataType = primaryPool->getDataType(); #ifdef ENABLE_FP4 - auto const isFp4 = mDataType == nvinfer1::DataType::kFP4; + auto const isFp4 = poolDataType == nvinfer1::DataType::kFP4; #else auto const isFp4 = false; #endif if (!isFp4) { - cacheSizeBytes += cacheVolume * BufferDataType(mDataType).getSize(); + cacheSizeBytes += cacheVolume * BufferDataType(poolDataType).getSize(); } else { @@ -3442,7 +3480,8 @@ SizeType32 KVCacheManager::getRemainingBlocksToCompletion( auto const& seq = seqIt->second; // Subtract detached front blocks (from SWA sliding) which are still // in the cache block ID list but no longer held by the sequence. - numAllocBlocksPerBeam = seq.getCacheBlockIds(windowSize).at(0).size() - seq.getNumFrontBlocksRemoved(); + numAllocBlocksPerBeam + = seq.getCacheBlockIds(windowSize).at(0).size() - seq.getNumFrontBlocksRemoved(windowSize); } } @@ -3599,7 +3638,7 @@ void WindowBlockManager::detachFrontBlock(GenerationRequest& sequence) auto const requestId = sequence.getRequestId(); auto const beamWidth = sequence.getBeamWidth(); auto& allocatedBlocks = mAllocatedBlocksPerSeq.at(requestId); - SizeType32 outOfWindowBlockIdx = sequence.getNumFrontBlocksRemoved(); + SizeType32 outOfWindowBlockIdx = sequence.getNumFrontBlocksRemoved(mWindowSize); for (auto beamIdx = 0; beamIdx < beamWidth; ++beamIdx) { @@ -3642,6 +3681,12 @@ void WindowBlockManager::detachFrontBlock(GenerationRequest& sequence) // Disconnect first block from sequence and remove it from allocated blocks sequence.removeFrontBlock(mWindowSize); + + // Surface the per-window eviction count post-increment so logs show the + // running counter for *this* window (the existing block-id log above only + // identifies which physical block was detached). + TLLM_LOG_DEBUG("%s::detachFrontBlock window=%d count=%d", mLogPrefix.c_str(), mWindowSize, + sequence.getNumFrontBlocksRemoved(mWindowSize)); } PrefixReuseSummary KVCacheManager::analyzePrefixReuse( @@ -4139,7 +4184,8 @@ BlocksPerWindow BaseKVCacheManager::calculateMaxNumBlocks(executor::KvCacheConfi SizeType32 tokensPerBlock, WorldConfig const& worldConfig, std::map> const& windowSizeToLayers, uint64_t allottedPrimaryMemBytes, uint64_t allottedSecondaryMemBytes, size_t extraCostMemory, SizeType32 kvFactor, SizeType32 maxBatchSize, - std::optional const& linearAttentionMetadata) + std::optional const& linearAttentionMetadata, + std::vector const& poolConfigurations) { TLLM_LOG_DEBUG("Calculating max num blocks: {.allottedPrimaryMemBytes=%" PRIu64 ", .allottedSecondaryMemBytes=%" PRIu64 "}", @@ -4154,15 +4200,34 @@ BlocksPerWindow BaseKVCacheManager::calculateMaxNumBlocks(executor::KvCacheConfi "configuration you meant to use."); } + // Build a windowSize -> PoolConfiguration lookup so the per-pool loop below stays terse. + // Multi-pool-per-window would replace this with an explicit layer->pool mapping; today + // the 1:1 mapping is the only case in use. + std::map poolByWindow; + for (auto const& pc : poolConfigurations) + { + TLLM_CHECK_WITH_INFO(poolByWindow.emplace(pc.windowSize, &pc).second, + "Multiple PoolConfigurations share windowSize=%d. Multi-pool-per-window requires an " + "explicit layer->pool mapping, which is not yet wired through calculateMaxNumBlocks.", + pc.windowSize); + } + std::map cacheSizeBytesPerTokenPerWindow; for (auto const& [windowSize, managedLayers] : windowSizeToLayers) { - auto const cacheSizePerToken = LinearAttentionMetadata::hasLinearCache(windowSize) - ? 1 - : BaseKVCacheManager::calculateCacheSizePerTokenForSingleWindowSize( - numKvHeadsPerLayer, sizePerHead, managedLayers, kvFactor); - auto const cacheSizeBytesPerToken = cacheSizePerToken * BufferDataType(dtype).getSize(); - cacheSizeBytesPerTokenPerWindow[windowSize] = cacheSizeBytesPerToken; + if (LinearAttentionMetadata::hasLinearCache(windowSize)) + { + cacheSizeBytesPerTokenPerWindow[windowSize] = 1; + continue; + } + // Resolve the pool serving this window: use its (sizePerHead, dtype) when the caller + // supplied a PoolConfiguration; otherwise fall back to the manager-level scalars. + auto const poolIt = poolByWindow.find(windowSize); + auto const windowSizePerHead = (poolIt != poolByWindow.end()) ? poolIt->second->sizePerHead : sizePerHead; + auto const windowDtype = (poolIt != poolByWindow.end()) ? poolIt->second->dtype : dtype; + auto const cacheSizePerToken = BaseKVCacheManager::calculateCacheSizePerTokenForSingleWindowSize( + numKvHeadsPerLayer, windowSizePerHead, managedLayers, kvFactor); + cacheSizeBytesPerTokenPerWindow[windowSize] = cacheSizePerToken * BufferDataType(windowDtype).getSize(); } // When the model is mamba hybrid (i.e. has only 2 windows sizes: max_seq_len and diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 325fbbe17a0c..63c07b4f2de2 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -346,6 +346,14 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::enum_(m, "LinearCacheType") .value("RECURRENT_STATES", tbk::LinearAttentionMetadata::LinearCacheType::kRecurrentStates); + nb::class_(m, "PoolConfiguration") + .def(nb::init<>()) + .def(nb::init(), nb::arg("window_size"), nb::arg("size_per_head"), + nb::arg("dtype")) + .def_rw("window_size", &tbk::PoolConfiguration::windowSize) + .def_rw("size_per_head", &tbk::PoolConfiguration::sizePerHead) + .def_rw("dtype", &tbk::PoolConfiguration::dtype); + nb::class_(m, "PrefixReuseSummary") .def(nb::init<>()) .def_ro("reusable_blocks_allocated", &tbk::PrefixReuseSummary::reusableBlocksAllocated) @@ -415,6 +423,7 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::arg("world_config"), nb::arg("window_size_to_layers"), nb::arg("allotted_primary_mem_bytes"), nb::arg("allotted_secondary_mem_bytes"), nb::arg("extra_cost_memory"), nb::arg("kv_factor"), nb::arg("max_batch_size"), nb::arg("linear_attention_metadata") = std::nullopt, + nb::arg("pool_configurations") = std::vector{}, nb::call_guard()) .def("allocate_pools", &BaseKVCacheManager::allocatePools, nb::call_guard()) .def("release_pools", &BaseKVCacheManager::releasePools, nb::call_guard()) @@ -597,6 +606,18 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) .def("analyze_prefix_reuse", &BaseKVCacheManager::analyzePrefixReuse, nb::arg("unique_tokens"), nb::arg("llm_request"), nb::call_guard()) .def("get_cache_block_ids", &BaseKVCacheManager::getCacheBlockIds, nb::call_guard()) + .def( + "get_num_front_blocks_removed", + [](BaseKVCacheManager const& self, tb::LlmRequest::RequestIdType requestId, SizeType32 windowSize) + { + auto const& seq = self.getSequence(requestId); + // Per-window query. windowSize is required (no aggregation): the Python + // wrapper in resource_manager.py provides a "default to first window" + // convenience layer; the C++/nanobind boundary stays explicit so that + // every caller is forced to think about which pool's counter it wants. + return seq.getNumFrontBlocksRemoved(windowSize); + }, + nb::arg("request_id"), nb::arg("window_size"), nb::call_guard()) .def("get_batch_cache_block_ids", &BaseKVCacheManager::getBatchCacheBlockIds, nb::call_guard()) .def("flush_iteration_events", &BaseKVCacheManager::flushIterationEvents, @@ -631,7 +652,8 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) std::vector const&, nvinfer1::DataType, SizeType32, int64_t, SizeType32, SizeType32, bool, tbk::CacheType, std::optional, std::shared_ptr, bool, bool, std::shared_ptr, - bool, SizeType32, SizeType32, bool, std::optional>(), + bool, SizeType32, SizeType32, bool, std::optional, + std::vector const&>(), nb::arg("num_kv_heads_per_layer"), nb::arg("size_per_head"), nb::arg("tokens_per_block"), nb::arg("blocks_per_window"), nb::arg("max_num_sequences"), nb::arg("max_beam_width"), nb::arg("max_attention_window_vec"), nb::arg("dtype"), nb::arg("sink_token_length"), nb::arg("stream"), @@ -641,7 +663,9 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::arg("copy_on_partial_reuse") = true, nb::arg("kv_connector_manager") = nullptr, nb::arg("enable_indexer_k_cache") = false, nb::arg("indexer_k_cache_quant_block_size") = 128, nb::arg("indexer_k_cache_index_head_dim") = 0, nb::arg("indexer_k_cache_use_fp4") = false, - nb::arg("linear_attention_metadata").none() = std::nullopt, nb::call_guard()) + nb::arg("linear_attention_metadata").none() = std::nullopt, + nb::arg("pool_configurations") = std::vector{}, + nb::call_guard()) .def( "scheduling_has_free_blocks", [](tbk::KVCacheManager& self, SizeType32 numRequired, SizeType32 windowSize) @@ -649,6 +673,11 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::arg("num_required"), nb::arg("window_size"), nb::call_guard()) .def_prop_ro( "is_variable_window", [](tbk::KVCacheManager& self) { return self.getBlockManager().isVariableWindow(); }) + // Per-pool introspection: lets Python discover (windowSize, sizePerHead, dtype) per + // hosted pool so a single KVCacheManager can host mixed-shape pools without a + // Python-side wrapper duplicating the layer->pool routing. + .def_prop_ro("pool_configurations", + [](tbk::KVCacheManager& self) { return self.getBlockManager().getPoolConfigurations(); }) .def("copy_linear_attention_block", &tbk::KVCacheManager::copyLinearAttentionBlock, nb::arg("llm_request"), nb::call_guard()) .def("copy_linear_attention_block_batch", &tbk::KVCacheManager::copyLinearAttentionBlockBatch, diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index fd6011e807ec..8b15651e2c97 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -16,6 +16,7 @@ */ #include "tensorrt_llm/batch_manager/kvCacheManager.h" +#include "tensorrt_llm/batch_manager/cacheTransBuffer.h" #include "tensorrt_llm/batch_manager/common.h" #include "tensorrt_llm/batch_manager/kvCacheConnector.h" #include "tensorrt_llm/batch_manager/kvCacheEventManager.h" @@ -9798,3 +9799,407 @@ TEST_F(KVCacheManagerTest, KvCacheConnector_DecodeBlockBoundary_ParityWithBaseli << ", withConnector=" << withConnector[i] << "); see issue #13320"; } } + +// ============================================================================= +// Per-window head_dim / dtype overrides on BlockManager and KVCacheManager, +// plus per-window front-eviction bookkeeping on GenerationRequest. +// ============================================================================= + +TEST_F(KVCacheManagerTest, BlockManagerTestPerWindowFallback) +{ + auto constexpr numLayers = 2; + auto constexpr numKvHeads = 2; + auto constexpr scalarSizePerHead = 64; + auto constexpr tokensPerBlock = 4; + auto constexpr blocksInPrimary = 4; + auto constexpr blocksInSecondary = 0; + auto constexpr maxNumSequences = 4; + auto constexpr maxBeamWidth = 1; + auto constexpr smallWindow = 1024; + auto constexpr largeWindow = 4096; + auto constexpr scalarDtype = nvinfer1::DataType::kHALF; + auto const stream = std::make_shared(); + auto const maxAttentionWindowVec = std::vector{smallWindow, largeWindow}; + auto const blocksPerWindow = BlocksPerWindow{ + {smallWindow, {blocksInPrimary, blocksInSecondary}}, {largeWindow, {blocksInPrimary, blocksInSecondary}}}; + + // No per-pool configurations — expect both windows to fall back to scalars. + BlockManager blockManager(std::vector(numLayers, numKvHeads), scalarSizePerHead, + tokensPerBlock, blocksPerWindow, maxNumSequences, stream, /*maxSequenceLength=*/largeWindow, maxBeamWidth, + maxAttentionWindowVec, scalarDtype, /*sinkBubbleLength=*/0, /*chunkSize=*/0); + blockManager.allocatePools(/*useUvm=*/false); + + auto const pools = blockManager.getPoolConfigurations(); + ASSERT_EQ(pools.size(), 2u); + std::map poolByWindow; + for (auto const& pc : pools) + { + poolByWindow.emplace(pc.windowSize, pc); + } + EXPECT_EQ(poolByWindow.at(smallWindow).sizePerHead, scalarSizePerHead); + EXPECT_EQ(poolByWindow.at(largeWindow).sizePerHead, scalarSizePerHead); + EXPECT_EQ(poolByWindow.at(smallWindow).dtype, scalarDtype); + EXPECT_EQ(poolByWindow.at(largeWindow).dtype, scalarDtype); + EXPECT_EQ(blockManager.getSizePerHeadForWindow(smallWindow), blockManager.getSizePerHeadForWindow(largeWindow)); + EXPECT_EQ(blockManager.getDataTypeForWindow(smallWindow), blockManager.getDataTypeForWindow(largeWindow)); +} + +// Static; does not require a GPU device. +TEST(BaseKVCacheManagerCalculateMaxNumBlocks, PerWindowOverrideDivergesByteBudget) +{ + using SizeType32 = tensorrt_llm::runtime::SizeType32; + auto constexpr numLayers = 2; + auto constexpr numKvHeads = 2; + auto constexpr scalarSizePerHead = 64; + auto constexpr largeSizePerHead = 128; + auto constexpr tokensPerBlock = 64; + auto constexpr smallWindow = 1024; + auto constexpr largeWindow = 4096; + auto constexpr kvFactor = 2; + auto constexpr maxBatchSize = 1; + + auto const numKvHeadsPerLayer = std::vector(numLayers, numKvHeads); + // Layer 0 → smallWindow, layer 1 → largeWindow. + auto const windowSizeToLayers = std::map>{ + {smallWindow, std::vector{0}}, {largeWindow, std::vector{1}}}; + uint64_t const allottedPrimaryMemBytes = static_cast(1) << 30; // 1 GiB + uint64_t const allottedSecondaryMemBytes = static_cast(1) << 30; + size_t const extraCostMemory = 0; + auto const dtype = nvinfer1::DataType::kHALF; + tensorrt_llm::executor::KvCacheConfig const config{}; + tensorrt_llm::runtime::WorldConfig const worldConfig{}; + + auto const baseline = BaseKVCacheManager::calculateMaxNumBlocks(config, dtype, numKvHeadsPerLayer, + scalarSizePerHead, tokensPerBlock, worldConfig, windowSizeToLayers, allottedPrimaryMemBytes, + allottedSecondaryMemBytes, extraCostMemory, kvFactor, maxBatchSize); + + auto const poolConfigurations = std::vector{ + {smallWindow, scalarSizePerHead, dtype}, {largeWindow, largeSizePerHead, dtype}}; + auto const overridden + = BaseKVCacheManager::calculateMaxNumBlocks(config, dtype, numKvHeadsPerLayer, scalarSizePerHead, + tokensPerBlock, worldConfig, windowSizeToLayers, allottedPrimaryMemBytes, allottedSecondaryMemBytes, + extraCostMemory, kvFactor, maxBatchSize, /*linearAttentionMetadata=*/std::nullopt, poolConfigurations); + + ASSERT_NE(baseline.find(smallWindow), baseline.end()); + ASSERT_NE(baseline.find(largeWindow), baseline.end()); + ASSERT_NE(overridden.find(smallWindow), overridden.end()); + ASSERT_NE(overridden.find(largeWindow), overridden.end()); + + auto const& [baselineSmallPrim, baselineSmallSec] = baseline.at(smallWindow); + auto const& [baselineLargePrim, baselineLargeSec] = baseline.at(largeWindow); + auto const& [overSmallPrim, overSmallSec] = overridden.at(smallWindow); + auto const& [overLargePrim, overLargeSec] = overridden.at(largeWindow); + + // Sanity: every count is positive. + EXPECT_GT(baselineSmallPrim, 0); + EXPECT_GT(baselineLargePrim, 0); + EXPECT_GT(overSmallPrim, 0); + EXPECT_GT(overLargePrim, 0); + + // Default share allocation gives equal share (1/N) to each window, so baseline gives the + // same number of blocks per window when both windows have identical bytes/token. + EXPECT_EQ(baselineSmallPrim, baselineLargePrim); + + // Override only doubles bytes/token for the *large* window. Equal share is unchanged, + // so the small window's block count stays the same and the large window's halves. + EXPECT_EQ(overSmallPrim, baselineSmallPrim); + EXPECT_EQ(overSmallSec, baselineSmallSec); + EXPECT_LT(overLargePrim, baselineLargePrim); + EXPECT_LT(overLargeSec, baselineLargeSec); + // Large window pays double bytes/token under override → roughly half the blocks. + EXPECT_NEAR(static_cast(overLargePrim) / baselineLargePrim, 0.5, 0.01); +} + +TEST_F(KVCacheManagerTest, KVCacheManagerSWAEvictionCountPerWindow) +{ + auto constexpr numLayers = 2; + auto constexpr numKvHeads = 2; + auto constexpr sizePerHead = 64; + auto constexpr tokensPerBlock = 4; + auto constexpr maxNumSequences = 4; + auto constexpr maxBeamWidth = 1; + auto constexpr sinkTokenLength = 0; + auto constexpr dtype = nvinfer1::DataType::kHALF; + auto const stream = std::make_shared(); + auto constexpr maxSequenceLength = 128; + auto constexpr maxNewTokens = 40; + + // Layer 0 lives in the SWA pool; layer 1 lives in the full-attention pool. + auto constexpr swaWindow = 8; + auto constexpr fullWindow = 128; + auto const maxAttentionWindowVec = std::vector{swaWindow, fullWindow}; + + auto constexpr blocksInPrimary = 8; + auto constexpr blocksInSecondary = 0; + auto const blocksPerWindow = BlocksPerWindow{ + {swaWindow, {blocksInPrimary, blocksInSecondary}}, {fullWindow, {blocksInPrimary, blocksInSecondary}}}; + + KVCacheManager kvCacheManager(numLayers, numKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, + maxBeamWidth, maxAttentionWindowVec, dtype, sinkTokenLength, stream, maxSequenceLength, /*chunkSize=*/0, + /*enableBlockReuse=*/false); + kvCacheManager.allocatePools(/*useUvm=*/false); + + auto constexpr beamWidth = maxBeamWidth; + auto constexpr beamIdx = 0; + tr::SamplingConfig const samplingConfig{beamWidth}; + bool constexpr isStreaming{false}; + TokenIdType constexpr firstToken = 1000; + + int inputLength = 11; // > swaWindow (8); < fullWindow (128) + auto inputTokens = std::make_shared(inputLength); + std::iota(inputTokens->begin(), inputTokens->end(), firstToken); + auto llmRequest + = std::make_shared(/*requestId=*/0, maxNewTokens, inputTokens, samplingConfig, isStreaming); + + kvCacheManager.addSequenceBatch({{{/*requestId=*/0, inputLength, beamWidth}}}, {std::ref(*llmRequest)}); + GenerationRequest const& seq = kvCacheManager.getSequence(/*requestId=*/0); + + // No tokens added past the SWA window yet — both counters are zero. + EXPECT_EQ(seq.getNumFrontBlocksRemoved(swaWindow), 0); + EXPECT_EQ(seq.getNumFrontBlocksRemoved(fullWindow), 0); + EXPECT_EQ(seq.getNumFrontBlocksRemoved(swaWindow) + seq.getNumFrontBlocksRemoved(fullWindow), 0); + + // Add one token past the window. SWA pool detaches block 0; full pool does not. + llmRequest->addNewToken(firstToken + inputLength, beamIdx); + kvCacheManager.addToken(/*requestId=*/0); + + EXPECT_EQ(seq.getNumFrontBlocksRemoved(swaWindow), 1) << "SWA pool should have evicted exactly one front block"; + EXPECT_EQ(seq.getNumFrontBlocksRemoved(fullWindow), 0) << "Full-attention pool must never front-evict"; + EXPECT_EQ(seq.getNumFrontBlocksRemoved(swaWindow) + seq.getNumFrontBlocksRemoved(fullWindow), 1) + << "Aggregate sum must equal the SWA-only count"; + + // Drive a second SWA eviction; full-attn count stays at zero. + llmRequest->addNewToken(firstToken + inputLength + 1, beamIdx); + kvCacheManager.addToken(/*requestId=*/0); + llmRequest->addNewToken(firstToken + inputLength + 2, beamIdx); + kvCacheManager.addToken(/*requestId=*/0); + llmRequest->addNewToken(firstToken + inputLength + 3, beamIdx); + kvCacheManager.addToken(/*requestId=*/0); + llmRequest->addNewToken(firstToken + inputLength + 4, beamIdx); + kvCacheManager.addToken(/*requestId=*/0); + + auto const swaEvictedAfter = seq.getNumFrontBlocksRemoved(swaWindow); + EXPECT_GT(swaEvictedAfter, 1) << "Continued generation should keep evicting SWA blocks"; + EXPECT_EQ(seq.getNumFrontBlocksRemoved(fullWindow), 0); + EXPECT_EQ(seq.getNumFrontBlocksRemoved(swaWindow) + seq.getNumFrontBlocksRemoved(fullWindow), swaEvictedAfter); + + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*llmRequest); + EXPECT_NO_THROW(static_cast(kvCacheManager.removeSequence(/*requestId=*/0, llmRequest))); +} + +TEST_F(KVCacheManagerTest, GenerationRequestClearCacheBlocksPerWindowResetsOnlyThatWindow) +{ + // Use a real BlockManager solely to obtain a valid windowSizeToMetadata map; the + // bookkeeping under test (removeFrontBlock / clearCacheBlocks / getNumFrontBlocksRemoved) + // lives entirely on GenerationRequest. + auto constexpr numLayers = 2; + auto constexpr numKvHeads = 2; + auto constexpr sizePerHead = 64; + auto constexpr tokensPerBlock = 4; + auto constexpr blocksInPrimary = 4; + auto constexpr blocksInSecondary = 0; + auto constexpr maxNumSequences = 4; + auto constexpr maxBeamWidth = 1; + auto constexpr swaWindow = 8; + auto constexpr fullWindow = 128; + auto const stream = std::make_shared(); + auto const maxAttentionWindowVec = std::vector{swaWindow, fullWindow}; + auto const blocksPerWindow = BlocksPerWindow{ + {swaWindow, {blocksInPrimary, blocksInSecondary}}, {fullWindow, {blocksInPrimary, blocksInSecondary}}}; + + BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, + blocksPerWindow, maxNumSequences, stream, /*maxSequenceLength=*/fullWindow, maxBeamWidth, maxAttentionWindowVec, + nvinfer1::DataType::kHALF, /*sinkBubbleLength=*/0, /*chunkSize=*/0); + blockManager.allocatePools(/*useUvm=*/false); + + auto constexpr requestId = 7; + auto constexpr numTokens = 1; + GenerationRequest seq{requestId, numTokens, maxBeamWidth, blockManager.getWindowSizesMetadata()}; + + // Drive each window's counter to a distinct non-zero value. + seq.removeFrontBlock(swaWindow); + seq.removeFrontBlock(swaWindow); + seq.removeFrontBlock(fullWindow); + + EXPECT_EQ(seq.getNumFrontBlocksRemoved(swaWindow), 2); + EXPECT_EQ(seq.getNumFrontBlocksRemoved(fullWindow), 1); + EXPECT_EQ(seq.getNumFrontBlocksRemoved(swaWindow) + seq.getNumFrontBlocksRemoved(fullWindow), 3); + + // Clearing one window's blocks must not touch the other window's counter. + seq.clearCacheBlocks(swaWindow); + EXPECT_EQ(seq.getNumFrontBlocksRemoved(swaWindow), 0); + EXPECT_EQ(seq.getNumFrontBlocksRemoved(fullWindow), 1); + EXPECT_EQ(seq.getNumFrontBlocksRemoved(swaWindow) + seq.getNumFrontBlocksRemoved(fullWindow), 1); + + // Symmetrically: clearing the other window resets only that one. + seq.clearCacheBlocks(fullWindow); + EXPECT_EQ(seq.getNumFrontBlocksRemoved(swaWindow), 0); + EXPECT_EQ(seq.getNumFrontBlocksRemoved(fullWindow), 0); + EXPECT_EQ(seq.getNumFrontBlocksRemoved(swaWindow) + seq.getNumFrontBlocksRemoved(fullWindow), 0); +} + +// A5: Smoke-test VSWA + radix prefix reuse with mixed head_dim. +// +// Coverage: constructs a KVCacheManager whose two windows have different head_dim +// (poolConfigurations supplying smallWindow head_dim=256, largeWindow head_dim=512) with reuse + +// partial reuse enabled, runs addSequenceBatch / storeContextBlocks / removeSequence +// for two sequences sharing a prefix, and asserts pool routing is not cross- +// contaminated and refcounts are balanced. This is the smoke variant called for in +// the spec: a complete prefix-reuse hit-rate scenario across mixed head_dim is more +// brittle to set up and is left as a follow-up. +TEST_F(KVCacheManagerTest, VswaMixedHeadDimReuseSmoke) +{ + auto constexpr numLayers = 2; + auto constexpr numKvHeads = 2; + auto constexpr scalarSizePerHead = 64; + auto constexpr tokensPerBlock = 4; + auto constexpr blocksInPrimary = 16; + auto constexpr blocksInSecondary = 0; + auto constexpr maxNumSequences = 4; + auto constexpr maxBeamWidth = 1; + auto constexpr smallWindow = 32; + auto constexpr largeWindow = 64; + auto constexpr smallSizePerHead = 256; + auto constexpr largeSizePerHead = 512; + auto constexpr sinkTokenLength = 0; + auto constexpr dtype = nvinfer1::DataType::kHALF; + auto constexpr maxSequenceLength = 64; + auto constexpr maxNewTokens = 0; + auto const stream = std::make_shared(); + auto const maxAttentionWindowVec = std::vector{smallWindow, largeWindow}; + auto const blocksPerWindow = BlocksPerWindow{ + {smallWindow, {blocksInPrimary, blocksInSecondary}}, {largeWindow, {blocksInPrimary, blocksInSecondary}}}; + auto const poolConfigurations = std::vector{ + {smallWindow, smallSizePerHead, dtype}, {largeWindow, largeSizePerHead, dtype}}; + + KVCacheManager kvCacheManager(numLayers, numKvHeads, scalarSizePerHead, tokensPerBlock, blocksPerWindow, + maxNumSequences, maxBeamWidth, maxAttentionWindowVec, dtype, sinkTokenLength, stream, maxSequenceLength, + /*chunkSize=*/0, /*enableBlockReuse=*/true, CacheType::kSELF, + /*secondaryOffloadMinPriority=*/std::nullopt, /*eventManager=*/nullptr, + /*enablePartialReuse=*/true, /*copyOnpartialReuse=*/true, + /*kvCacheConnectorManager=*/nullptr, + /*enableIndexerKCache=*/false, /*indexerKCacheQuantBlockSize=*/128, /*indexerKCacheIndexHeadDim=*/0, + /*indexerKCacheUseFp4=*/false, + /*linearAttentionMetadata=*/std::nullopt, poolConfigurations); + kvCacheManager.allocatePools(/*useUvm=*/false); + + // Both windows registered with their distinct head_dims. + EXPECT_EQ(kvCacheManager.getBlockManager().getSizePerHeadForWindow(smallWindow), smallSizePerHead); + EXPECT_EQ(kvCacheManager.getBlockManager().getSizePerHeadForWindow(largeWindow), largeSizePerHead); + + auto const freeBefore = kvCacheManager.getBlockManager().getNumFreeBlocksPerWindowSize(); + + auto constexpr beamWidth = maxBeamWidth; + tr::SamplingConfig const samplingConfig{beamWidth}; + bool constexpr isStreaming{false}; + + // Sequence A: 12-token prompt across two blocks (block-aligned at 8 tokens, partial third). + auto inputTokensA = std::make_shared(VecTokens{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}); + auto const inputLengthA = static_cast(inputTokensA->size()); + auto llmRequestA + = std::make_shared(/*requestId=*/0, maxNewTokens, inputTokensA, samplingConfig, isStreaming); + + EXPECT_NO_THROW( + kvCacheManager.addSequenceBatch({{{/*requestId=*/0, inputLengthA, beamWidth}}}, {std::ref(*llmRequestA)})); + + // Both windows must have allocated some blocks for this sequence — neither pool + // can be empty (i.e. neither pool was skipped due to mixed head_dim routing). + auto const& seqA = kvCacheManager.getSequence(/*requestId=*/0); + EXPECT_FALSE(seqA.getCacheBlockIds(smallWindow).at(0).empty()) << "small-window pool must allocate"; + EXPECT_FALSE(seqA.getCacheBlockIds(largeWindow).at(0).empty()) << "large-window pool must allocate"; + + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*llmRequestA); + EXPECT_NO_THROW(static_cast(kvCacheManager.removeSequence(/*requestId=*/0, llmRequestA))); + + // Sequence B: shares an 8-token prefix with sequence A, then diverges. + auto inputTokensB = std::make_shared(VecTokens{1, 2, 3, 4, 5, 6, 7, 8, 99, 100}); + auto const inputLengthB = static_cast(inputTokensB->size()); + auto llmRequestB + = std::make_shared(/*requestId=*/1, maxNewTokens, inputTokensB, samplingConfig, isStreaming); + + EXPECT_NO_THROW( + kvCacheManager.addSequenceBatch({{{/*requestId=*/1, inputLengthB, beamWidth}}}, {std::ref(*llmRequestB)})); + + auto const& seqB = kvCacheManager.getSequence(/*requestId=*/1); + EXPECT_FALSE(seqB.getCacheBlockIds(smallWindow).at(0).empty()); + EXPECT_FALSE(seqB.getCacheBlockIds(largeWindow).at(0).empty()); + + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*llmRequestB); + EXPECT_NO_THROW(static_cast(kvCacheManager.removeSequence(/*requestId=*/1, llmRequestB))); + + // After both sequences have been removed and their reuse-eligible blocks have + // been returned, the per-window free-block counts must equal what we started + // with (no leaked refcounts in either pool). + auto const freeAfter = kvCacheManager.getBlockManager().getNumFreeBlocksPerWindowSize(); + ASSERT_EQ(freeAfter.size(), freeBefore.size()); + for (auto const& [windowSize, free] : freeBefore) + { + ASSERT_NE(freeAfter.find(windowSize), freeAfter.end()); + EXPECT_EQ(freeAfter.at(windowSize), free) + << "Window " << windowSize << " has leaked or extra blocks after removeSequence."; + } +} + +// A6: VSWA + disagg dtype mismatch must fire the A4 guard. +// +// The constructor of CacheTransBufferManager picks pool 0's dtype as canonical for +// the wire transport. When a KVCacheManager hosts pools with differing dtypes +// (mixed-precision per-window), that silent coercion would corrupt the wire format. +// The guard added in cacheTransBuffer.cpp must throw at construction time. +// +// This test only exercises the helper / construction path that runs the guard; it +// does not stand up a full disaggregated transfer (out of scope at unit-test +// granularity). +TEST_F(KVCacheManagerTest, VswaDisaggDtypeMismatchTriggersGuard) +{ + auto constexpr numLayers = 2; + auto constexpr numKvHeads = 2; + auto constexpr sizePerHead = 64; + auto constexpr tokensPerBlock = 4; + auto constexpr blocksInPrimary = 4; + auto constexpr blocksInSecondary = 0; + auto constexpr maxNumSequences = 4; + auto constexpr maxBeamWidth = 1; + auto constexpr smallWindow = 32; + auto constexpr largeWindow = 64; + auto constexpr sinkTokenLength = 0; + auto constexpr maxSequenceLength = 64; + auto const stream = std::make_shared(); + auto const maxAttentionWindowVec = std::vector{smallWindow, largeWindow}; + auto const blocksPerWindow = BlocksPerWindow{ + {smallWindow, {blocksInPrimary, blocksInSecondary}}, {largeWindow, {blocksInPrimary, blocksInSecondary}}}; + auto const poolConfigurations = std::vector{ + {smallWindow, sizePerHead, nvinfer1::DataType::kHALF}, {largeWindow, sizePerHead, nvinfer1::DataType::kBF16}}; + + auto kvCacheManager = std::make_unique(numLayers, numKvHeads, sizePerHead, tokensPerBlock, + blocksPerWindow, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, /*dtype=*/nvinfer1::DataType::kHALF, + sinkTokenLength, stream, maxSequenceLength, + /*chunkSize=*/0, /*enableBlockReuse=*/false, CacheType::kSELF, + /*secondaryOffloadMinPriority=*/std::nullopt, /*eventManager=*/nullptr, + /*enablePartialReuse=*/true, /*copyOnpartialReuse=*/true, + /*kvCacheConnectorManager=*/nullptr, + /*enableIndexerKCache=*/false, /*indexerKCacheQuantBlockSize=*/128, /*indexerKCacheIndexHeadDim=*/0, + /*indexerKCacheUseFp4=*/false, + /*linearAttentionMetadata=*/std::nullopt, poolConfigurations); + kvCacheManager->allocatePools(/*useUvm=*/false); + + // Sanity: the manager really does host KV pools with two different dtypes. + auto const numKvPools = kvCacheManager->getBlockManager().getNumPools( + /*includeBlockScalePools=*/false, /*includeIndexerKCachePools=*/false); + ASSERT_GE(numKvPools, 2); + auto const dtype0 = kvCacheManager->getPrimaryPool(0)->getDataType(); + bool foundMismatch = false; + for (SizeType32 i = 1; i < numKvPools; ++i) + { + if (kvCacheManager->getPrimaryPool(i)->getDataType() != dtype0) + { + foundMismatch = true; + break; + } + } + ASSERT_TRUE(foundMismatch) << "Test setup must produce at least two pools with different dtypes"; + + // The guard in CacheTransBufferManager's constructor must reject this. + EXPECT_THROW({ CacheTransBufferManager const ctbm(kvCacheManager.get(), /*maxNumTokens=*/std::nullopt); }, + tensorrt_llm::common::TllmException); +} diff --git a/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py b/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py index 71752a2734fe..80226aca6dff 100644 --- a/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py +++ b/tensorrt_llm/_torch/auto_deploy/compile/backends/torch_cudagraph.py @@ -1043,8 +1043,8 @@ def _setup_piecewise_mixed_batch(seq_info: Any, num_tokens: int) -> None: input_ids=input_ids_flat, cu_seqlen=cu_seqlen, input_pos=0, - cache_loc=cache_loc, - cu_num_pages=cu_num_pages, + cache_loc_per_pool=[cache_loc], + cu_num_pages_per_pool=[cu_num_pages], slot_idx=slot_idx, ) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py index 67ec45d04156..529556729c27 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py @@ -147,9 +147,12 @@ def plan_generate_only( cu_num_pages: torch.Tensor, cache_loc: torch.Tensor, last_page_len: torch.Tensor, + pool_window_left: Optional[int] = None, ): for plan_params in self.cached_cuda_graph_decode_wrappers: if plan_params.num_seq == num_seq: + if pool_window_left is not None and plan_params.window_left != pool_window_left: + continue wrapper = self.cached_cuda_graph_decode_wrappers[plan_params] flashinfer.decode.fast_decode_plan( wrapper, @@ -325,6 +328,7 @@ def prepare_flashinfer_metadata_host( cu_num_pages_host: torch.Tensor, cache_loc_host: torch.Tensor, last_page_len_host: torch.Tensor, + pool_window_left: Optional[int] = None, ) -> None: batch_info = BatchInfo(batch_info_host) num_prefill, num_prefill_tokens, num_decode = batch_info.get_absorbed_info() @@ -335,6 +339,7 @@ def prepare_flashinfer_metadata_host( cu_num_pages_host[: num_decode + 1], cache_loc_host, last_page_len_host[:num_decode], + pool_window_left=pool_window_left, ) @@ -577,10 +582,15 @@ def get_prepare_extra_metadata_info( def get_cache_initializers( cls, source_attn_node: Node, cache_config: KvCacheConfig ) -> ResourceHandlerDict: + """Build the per-layer KV handler used by the kvcache transform.""" # source op is [bsnd] layout already k_fake: FakeTensor = source_attn_node.args[1].meta["val"] num_kv_heads = k_fake.shape[2] head_dim = k_fake.shape[3] + # ``sliding_window`` is propagated into the handler so layers + # with different windows land in separate pools. + (sw,) = extract_op_args(source_attn_node, "sliding_window") + sliding_window = sw if isinstance(sw, int) and sw > 0 else 0 return { "kv_cache": KVPagedResourceHandler( @@ -589,6 +599,7 @@ def get_cache_initializers( dtype=cls.resolve_cache_dtype(cache_config.dtype, k_fake.dtype), kv_factor=2, kv_layout=_GlobalFlashInferPlanner.kv_layout, + sliding_window=sliding_window, ) } diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_paged_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_paged_attention.py index 2667f4e18ac5..192a2e659d44 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_paged_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_paged_attention.py @@ -1528,9 +1528,14 @@ def get_prepare_extra_metadata_info( def get_cache_initializers( cls, source_attn_node: Node, cache_config: KvCacheConfig ) -> ResourceHandlerDict: + """Build the per-layer KV handler used by the kvcache transform.""" k_fake: FakeTensor = source_attn_node.args[1].meta["val"] num_kv_heads = k_fake.shape[2] head_dim = k_fake.shape[3] + # ``sliding_window`` is propagated into the handler so layers + # with different windows land in separate pools. + (sw,) = extract_op_args(source_attn_node, "sliding_window") + sliding_window = sw if isinstance(sw, int) and sw > 0 else 0 return { "kv_cache": KVPagedResourceHandler( @@ -1539,6 +1544,7 @@ def get_cache_initializers( dtype=cls.resolve_cache_dtype(cache_config.dtype, k_fake.dtype), kv_factor=2, kv_layout=KV_LAYOUT, + sliding_window=sliding_window, ) } diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py index 875ea2d2e5a5..dfd04e4caefa 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py @@ -852,7 +852,7 @@ def get_standard_metadata_args(cls) -> List[str]: def get_cache_initializers( cls, source_attn_node: Node, cache_config: KvCacheConfig ) -> ResourceHandlerDict: - """Return only KV cache handler (no workspace handler, managed like flashinfer).""" + """Return only the KV cache handler (no workspace handler; managed like flashinfer).""" # In fused QKV mode, K arg is the same as Q (flat fused tensor), so use hints. if source_attn_node.meta.get("_trtllm_fused_qkv"): num_kv_heads = source_attn_node.meta["_trtllm_num_kv_heads"] @@ -876,6 +876,10 @@ def get_cache_initializers( num_kv_heads = k_fake.shape[2] head_dim = k_fake.shape[3] kv_dtype = k_fake.dtype + # ``sliding_window`` is propagated into the handler so layers + # with different windows land in separate pools. + (sw,) = extract_op_args(source_attn_node, "sliding_window") + sliding_window = sw if isinstance(sw, int) and sw > 0 else 0 return { "kv_cache": KVPagedResourceHandler( @@ -884,6 +888,7 @@ def get_cache_initializers( dtype=cls.resolve_cache_dtype(cache_config.dtype, kv_dtype), kv_factor=2, kv_layout="HND", + sliding_window=sliding_window, ) } diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index 50d163209f80..838a1731791a 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -354,6 +354,29 @@ def resize(self, name: str, new_capacity: int) -> None: self._device_views[name] = self._trunc_device_bufs[name].view(dtype) self._host_views[name] = self._trunc_host_bufs[name].view(dtype) + def add_truncatable_tensor(self, name: str, max_numel: int, dtype: torch.dtype) -> None: + """Add a new truncatable tensor after initialization. + + Args: + name: Name of the tensor. + max_numel: Maximum number of elements. + dtype: Data type. + """ + assert name not in self._tensor_specs, f"Tensor '{name}' already registered" + self._tensor_specs[name] = (max_numel, dtype) + self._tensor_order.append(name) + self._truncatable_names.add(name) + self._current_lengths[name] = 0 + + byte_size = max_numel * dtype.itemsize + device = self._device_buffer.device + self._trunc_device_bufs[name] = torch.empty(byte_size, dtype=torch.uint8, device=device) + self._trunc_host_bufs[name] = torch.empty( + byte_size, dtype=torch.uint8, device="cpu", pin_memory=prefer_pinned() + ) + self._device_views[name] = self._trunc_device_bufs[name].view(dtype) + self._host_views[name] = self._trunc_host_bufs[name].view(dtype) + def to(self, *args, **kwargs) -> None: """Move all device buffers to a new device/dtype.""" old_device = self._device_buffer.device @@ -721,6 +744,11 @@ def __init__( self._extra_args: Dict[str, Optional[torch.Tensor]] = {} ############################################################################################ + # VSWA WINDOW GROUPS ####################################################################### + self._window_groups: List[int] = [] + self._window_group_map: Dict[int, int] = {} + ############################################################################################ + # HOST PREPARE FOR ATTENTION FORWARD ####################################################### self._host_prepare_functions: List[Tuple[PrepareMetadataHostCallable, List[str]]] = [] @@ -864,13 +892,20 @@ def estimate_cache_tokens_per_forward(self) -> int: num_blocks_estimate = num_blocks_estimate_per_seq * self.max_batch_size return num_blocks_estimate * self.tokens_per_block - def update_cache_information(self, num_blocks: int, block_offset_multiplier: int = 0) -> None: + def update_cache_information( + self, + num_blocks: int, + block_offset_multiplier: int = 0, + ) -> None: """Update cache information after cache manager creation. Sets num_blocks and block_offset_multiplier, writes max_seq_info into BatchInfo (constant after this call), and resizes cache_loc if needed. + + Args: + num_blocks: Number of blocks in the primary pool. + block_offset_multiplier: Block offset multiplier derived from kv_cache strides. """ - # set num_blocks and block_offset_multiplier self._num_blocks = num_blocks # write max_seq_info once into BatchInfo (constant after this call) @@ -891,6 +926,97 @@ def update_cache_information(self, num_blocks: int, block_offset_multiplier: int if estimated_capacity > cache_loc_capacity: self._input_buffer.resize("cache_loc", estimated_capacity) + # Keep per-group cache_loc buffers in sync + for group_idx in range(1, self.num_window_groups): + self._input_buffer.resize(f"cache_loc_g{group_idx}", estimated_capacity) + + def register_window_groups(self, window_sizes: List[int]) -> None: + """Register KV window groups and create per-group cache tensors. + + Group ordering here matches the KV-cache group/pool ordering produced by + the kvcache transform — ``window_sizes[g]`` belongs to pool ``g``, so a + single ``group_idx`` routes graph metadata, storage pools, and runtime + inputs together. + + Group 0 reuses the existing ``cache_loc`` / ``cu_num_pages`` / + ``last_page_len`` / ``extra_page_per_seq`` tensors. + Groups 1..N-1 get new dedicated tensors (``cache_loc_g{i}``, + ``cu_num_pages_g{i}``, ...). For a single-pool (non-VSWA) model only + group 0 is registered and no new tensors are created. + + Args: + window_sizes: List of per-group window sizes (one entry per group, in + the same order as the KV groups). Full-attention groups use + ``max_seq_len``. + """ + assert len(window_sizes) >= 1, "register_window_groups requires at least 1 window" + self._window_groups = list(window_sizes) + self._window_group_map = {ws: idx for idx, ws in enumerate(window_sizes)} + + cache_loc_cap = self._input_buffer.get_capacity("cache_loc") + for group_idx in range(1, len(window_sizes)): + suffix = f"_g{group_idx}" + self._input_buffer.add_truncatable_tensor( + f"cache_loc{suffix}", cache_loc_cap, torch.int + ) + self._input_buffer.add_truncatable_tensor( + f"cu_num_pages{suffix}", self.max_batch_size + 1, torch.int + ) + self._input_buffer.add_truncatable_tensor( + f"last_page_len{suffix}", self.max_batch_size, torch.int + ) + self._input_buffer.add_truncatable_tensor( + f"extra_page_per_seq{suffix}", self.max_batch_size, torch.int + ) + # Per-group seq_len_with_cache: under SWA front-eviction the window's + # live cache length diverges from the global (input_pos + seq_len), + # so groups 1..N-1 carry their own window-capped values for both the + # kernel's mask math and the prepare-extra-metadata op (which uses + # it to compute write positions for new KV tokens). + self._input_buffer.add_truncatable_tensor( + f"seq_len_with_cache{suffix}", self.max_batch_size, torch.int + ) + # Register as available args (device + host variants) + group_names = [ + f"cache_loc{suffix}", + f"cu_num_pages{suffix}", + f"last_page_len{suffix}", + f"extra_page_per_seq{suffix}", + f"seq_len_with_cache{suffix}", + ] + for base_name in group_names: + self._available_args.add(base_name) + self._available_args.add(base_name + self._host_suffix) + + # Zero-fill all per-group device buffers so shape-checking forward + # passes (resize_kv_cache) that run before nest_sequences don't read + # uninitialized data. With zeros: cache_loc→block 0 (safe), + # cu_num_pages→0 pages. + for base_name in group_names: + self._input_buffer._trunc_device_bufs[base_name].zero_() + + @property + def window_groups(self) -> List[int]: + """Per-pool window sizes (one entry per registered pool, in pool order).""" + return self._window_groups + + @property + def window_group_map(self) -> Dict[int, int]: + """Map from window_size to group index.""" + return self._window_group_map + + @property + def num_window_groups(self) -> int: + """Number of window groups (0 or 1 for non-VSWA, 2+ for VSWA).""" + return len(self._window_groups) + + def get_cache_loc_for_group(self, group_idx: int) -> str: + """Return the cache_loc tensor name for a given window group.""" + return "cache_loc" if group_idx == 0 else f"cache_loc_g{group_idx}" + + def get_cu_num_pages_for_group(self, group_idx: int) -> str: + """Return the cu_num_pages tensor name for a given window group.""" + return "cu_num_pages" if group_idx == 0 else f"cu_num_pages_g{group_idx}" def activate_arg(self, arg_name: str) -> bool: """Activate a desired argument. @@ -946,12 +1072,18 @@ def set_example_sequence( slot_idx = torch.arange(bs) assert len(slot_idx) >= bs + # Replicate the vanilla single-pool warm-up data for every registered + # pool so VSWA models (num_window_groups >= 2) receive one entry per + # pool. Block geometry (tokens_per_block, max_blocks_per_seq) is + # uniform across pools, so the same cache_loc / cu_num_pages tensors + # are valid metadata for every pool. + n_pools = max(1, self.num_window_groups) self.nest_sequences( input_ids.flatten(), cu_seqlen=torch.arange(bs + 1, dtype=torch.int) * seq_len, input_pos=0, # no cache history - cache_loc=cache_loc, # vanilla page assignments - cu_num_pages=cu_num_pages, # vanilla page assignments + cache_loc_per_pool=[cache_loc for _ in range(n_pools)], + cu_num_pages_per_pool=[cu_num_pages for _ in range(n_pools)], slot_idx=slot_idx, # vanilla slot indices **extra_args, ) @@ -1098,11 +1230,18 @@ def nest_sequences( cu_seqlen: Union[Sequence[int], torch.Tensor], input_pos: Union[Sequence[int], int, torch.Tensor], batch_info: Union[Sequence[int], torch.Tensor, None] = None, - cache_loc: Union[Sequence[int], torch.Tensor, None] = None, - cu_num_pages: Union[Sequence[int], torch.Tensor, None] = None, - extra_page_per_seq: Optional[Sequence[int]] = None, slot_idx: Union[Sequence[int], torch.Tensor, None] = None, prompt_lens: Union[Sequence[int], torch.Tensor, None] = None, + ### PER-POOL CACHE DATA (pool 0 is the only pool for non-VSWA configurations) ############# + cache_loc_per_pool: Optional[Sequence[Sequence[int]]] = None, + cu_num_pages_per_pool: Optional[Sequence[Sequence[int]]] = None, + extra_page_per_seq_per_pool: Optional[Sequence[Sequence[int]]] = None, + # Phase 2: per-pool window-capped values used by VSWA front-eviction so + # the live cache length / last-page length can diverge from the global + # (input_pos + seq_len). When None, non-zero pools replicate the global + # value (resp. fall back to ``lpl_host``). + seq_len_with_cache_per_pool: Optional[Sequence[Sequence[int]]] = None, + last_page_len_per_pool: Optional[Sequence[Sequence[int]]] = None, ### RUNTIME ARGUMENTS ###################################################################### gather_context_logits: bool = False, _gather_idx: Union[Sequence[int], torch.Tensor, None] = None, @@ -1124,12 +1263,21 @@ def nest_sequences( heuristic is used to compute it from seq_len. NOTE: the heuristic makes potentially incorrect assumptions about the batch composition. batch_info should be provided explicitly to ensure correctness. - cache_loc: Flat list of page indices for all sequences. Must be provided together with - cu_num_pages. - cu_num_pages: Cumulative number of pages for all sequences. Must be provided together with - cache_loc. - extra_page_per_seq: Extra page per sequence for deferred page insertion. slot_idx: State slot index for each sequence in the batch. + cache_loc_per_pool: Flat list of page indices for all sequences, per pool. Must be + provided together with cu_num_pages_per_pool. Indexed by pool_idx (0-based); + pool 0 is the only pool for non-VSWA configurations. + cu_num_pages_per_pool: Cumulative number of pages for all sequences, per pool. Must be + provided together with cache_loc_per_pool. + extra_page_per_seq_per_pool: Extra page per sequence for deferred page insertion, + per pool. + seq_len_with_cache_per_pool: Per-pool window-capped seq_len_with_cache + used by VSWA front-eviction. Pool 0 is unused (it carries the + unclamped global value). When None, non-zero pools replicate the + global ``input_pos + seq_len``. + last_page_len_per_pool: Per-pool window-capped last_page_len used by + VSWA front-eviction. Pool 0 is unused. When None, non-zero pools + fall back to ``lpl_host`` (the global value). gather_context_logits: If True, keep all context logits (no selective gathering). If False (default), only the last token per context sequence is gathered while all extend/decode tokens are kept. @@ -1182,17 +1330,68 @@ def nest_sequences( ### UPDATE CACHE ASSIGNMENTS (needs to be provided if required!) ########################### # check for updated page assignments - assert (cache_loc is None) == (cu_num_pages is None), "Both must be provided together!" - self._stage_arg("cache_loc", cache_loc) - self._stage_arg("cu_num_pages", cu_num_pages) + assert (cache_loc_per_pool is None) == (cu_num_pages_per_pool is None), ( + "cache_loc_per_pool and cu_num_pages_per_pool must be provided together!" + ) + # Global (unclamped) last_page_len, used as the per-pool fallback below + # when the caller doesn't supply ``last_page_len_per_pool`` (warmup, + # set_example_sequence, CUDA-graph capture). lpl_host = (ip_host + sl_host - 1) % self.tokens_per_block + 1 - self._stage_arg("last_page_len", lpl_host) # check for updated slot_idx self._stage_arg("slot_idx", slot_idx) - # check for updated extra_page_per_seq - self._stage_arg("extra_page_per_seq", extra_page_per_seq) + # cu_num_pages from pool 0 is reused below by the pages_per_seq derivative + # update; capture it before the per-pool staging loop. + cu_num_pages = cu_num_pages_per_pool[0] if cu_num_pages_per_pool is not None else None + + # Guard against partial lists: every registered pool must be present + # in each of the per-pool lists, otherwise omitted pools would silently + # reuse the previous batch's staged tensors. + if cache_loc_per_pool is not None and self.num_window_groups >= 2: + for name, seq in ( + ("cache_loc_per_pool", cache_loc_per_pool), + ("cu_num_pages_per_pool", cu_num_pages_per_pool), + ("extra_page_per_seq_per_pool", extra_page_per_seq_per_pool), + ("seq_len_with_cache_per_pool", seq_len_with_cache_per_pool), + ("last_page_len_per_pool", last_page_len_per_pool), + ): + if seq is None: + continue + assert len(seq) == self.num_window_groups, ( + f"{name} has {len(seq)} entries, expected {self.num_window_groups}" + ) + + # Stage per-pool cache data uniformly: pool 0 uses the unsuffixed + # base names; pools 1..N-1 use the f"_g{pool_idx}" suffix. When the + # caller doesn't provide per-pool data (warmup, set_example_sequence, + # CUDA graph capture), pools 1..N-1 fall back to pool 0's data so every + # kernel receives valid metadata. + if cache_loc_per_pool is not None: + num_pools_provided = len(cache_loc_per_pool) + pool0_cache_loc = cache_loc_per_pool[0] + pool0_cu_num_pages = ( + cu_num_pages_per_pool[0] if cu_num_pages_per_pool is not None else None + ) + pool0_extra_page = ( + extra_page_per_seq_per_pool[0] if extra_page_per_seq_per_pool is not None else None + ) + num_pools_active = max(num_pools_provided, self.num_window_groups) + for pool_idx in range(num_pools_active): + suffix = "" if pool_idx == 0 else f"_g{pool_idx}" + provides_pool = pool_idx < num_pools_provided + pool_cache_loc = cache_loc_per_pool[pool_idx] if provides_pool else pool0_cache_loc + self._stage_arg(f"cache_loc{suffix}", pool_cache_loc) + if cu_num_pages_per_pool is not None: + pool_cu_num_pages = ( + cu_num_pages_per_pool[pool_idx] if provides_pool else pool0_cu_num_pages + ) + self._stage_arg(f"cu_num_pages{suffix}", pool_cu_num_pages) + if extra_page_per_seq_per_pool is not None: + pool_extra_page = ( + extra_page_per_seq_per_pool[pool_idx] if provides_pool else pool0_extra_page + ) + self._stage_arg(f"extra_page_per_seq{suffix}", pool_extra_page) ### UPDATE OPTIONAL DERIVATIVE METADATA #################################################### if self._is_required("position_ids"): @@ -1213,9 +1412,26 @@ def nest_sequences( pages_per_seq = cu_num_pages_host[1:] - cu_num_pages_host[:-1] self._stage_arg("pages_per_seq", pages_per_seq) - # update sequence length with cache + # Stage seq_len_with_cache and last_page_len per pool through a single + # loop. Pool 0 uses the unsuffixed name; pools 1..N-1 use f"_g{i}". + # The unsuffixed name is a wire-level alias for pool 0 -- non-VSWA + # graphs, prepare-extra ops, and host-prep functions still reference + # it -- not a "full-attention pool" privilege; pool 0 is whichever KV + # handler the kvcache transform encountered first, which can be the + # SWA pool (Gemma-3n where layer 0 is sliding, or any pure-SWA model + # such as Mistral / Phi-3). When the caller doesn't supply per-pool + # values (warmup, set_example_sequence, CUDA-graph capture), every + # pool falls back to the global value. Dropping the unsuffixed alias + # entirely is a follow-up refactor that touches the transform, + # backend descriptors, and op signatures. seq_len_with_cache = ip_host + sl_host - self._stage_arg("seq_len_with_cache", seq_len_with_cache) + n_pools = max(1, self.num_window_groups) + per_pool_swc = seq_len_with_cache_per_pool or [seq_len_with_cache] * n_pools + per_pool_lpl = last_page_len_per_pool or [lpl_host] * n_pools + for pool_idx in range(n_pools): + suffix = "" if pool_idx == 0 else f"_g{pool_idx}" + self._stage_arg(f"seq_len_with_cache{suffix}", per_pool_swc[pool_idx]) + self._stage_arg(f"last_page_len{suffix}", per_pool_lpl[pool_idx]) # prompt_lens: original context length per sequence, constant across iterations. # Defaults to seq_len when not provided (correct for prefill). @@ -1404,6 +1620,25 @@ def offset_pos_and_cache_(self, offset: torch.Tensor) -> None: last_page_len %= self.tokens_per_block last_page_len += 1 + # Adjust per-group cache assignments for VSWA groups 1..N-1 + for group_idx in range(1, self.num_window_groups): + suffix = f"_g{group_idx}" + if self._is_active(f"cache_loc{suffix}"): + lpl_g = self.get_arg(f"last_page_len{suffix}", truncate=True) + lpl_g += offset + delta_g = (lpl_g > self.tokens_per_block).int() - (lpl_g <= 0).int() + torch.ops.auto_deploy.adjust_ragged_triton( + cache_loc=self.get_arg(f"cache_loc{suffix}"), + cu_num_blocks=self.get_arg(f"cu_num_pages{suffix}"), + extra_idx=self.get_arg(f"extra_page_per_seq{suffix}"), + delta=delta_g, + num_sequences=num_sequences, + max_blocks_per_seq=self.max_blocks_per_seq, + ) + lpl_g -= 1 + lpl_g %= self.tokens_per_block + lpl_g += 1 + # --- position_ids (device) --- position_ids = self.get_arg("position_ids", truncate=True, unflatten=False) # position_ids is per-token while offset is per-sequence; expand if needed @@ -1426,6 +1661,15 @@ def offset_pos_and_cache_(self, offset: torch.Tensor) -> None: swc = self.get_arg("seq_len_with_cache", truncate=True) swc += offset + # Keep per-group seq_len_with_cache in sync with the global value for + # the overlap scheduler path. This is exact in the non-eviction regime; + # under eviction the next nest_sequences re-caps to the window. + for group_idx in range(1, self.num_window_groups): + suffix = f"_g{group_idx}" + if self._is_active(f"seq_len_with_cache{suffix}"): + swc_g = self.get_arg(f"seq_len_with_cache{suffix}", truncate=True) + swc_g += offset + # --- use_initial_states (device) --- use_initial_states = self.get_arg("use_initial_states", truncate=True) use_initial_states[:] = input_pos > 0 @@ -1597,15 +1841,24 @@ def allocate(self, sequence_info: SequenceInfo) -> torch.Tensor: class KVPagedResourceHandler(ResourceHandler): """Handler for paged KV cache resources. - This handler indicates the resource should be managed by the standard KVCacheManager. + This handler indicates the resource should be managed by the standard + KVCacheManager. Handler equality (``__eq__``) is the single source of + truth for KV-cache grouping: layers whose handlers compare equal share a + storage pool and a metadata set (``group_idx == pool_idx``), and differ + otherwise. Including ``sliding_window`` in the equality means same-head-dim + layers with different windows land in separate pools with their own + ``max_attention_window``. Args: num_kv_heads: Number of key-value heads. head_dim: Dimension of each head. dtype: The dtype of the KV cache. kv_factor: The factor of the KV cache. Default is 2 for combined k/v cache. - kv_layout: Memory layout for the KV cache. Either "HND" (head-num-dim) or "NHD" (num-head-dim). - Default is "HND" which is the standard layout for flashinfer. + kv_layout: Memory layout for the KV cache. Either "HND" (head-num-dim) or + "NHD" (num-head-dim). Default is "HND" which is the standard layout + for flashinfer. + sliding_window: Sliding window size for this layer. ``0`` means full + attention; a positive value puts this layer in its own VSWA group. """ @property @@ -1620,6 +1873,7 @@ def __init__( dtype: torch.dtype, kv_factor: int = 2, kv_layout: Literal["HND", "NHD"] = "HND", + sliding_window: int = 0, ) -> None: """Initialize the KVPagedResourceHandler. @@ -1629,6 +1883,7 @@ def __init__( dtype: The dtype of the KV cache. kv_factor: The factor of the KV cache. Default is 2. kv_layout: Memory layout - "HND" or "NHD". Default is "HND". + sliding_window: Sliding window size for this layer. 0 means full attention. """ self.num_kv_heads = num_kv_heads self.head_dim = head_dim @@ -1636,9 +1891,18 @@ def __init__( self.kv_factor = kv_factor assert kv_factor in [1, 2], f"Invalid kv_factor: {kv_factor}" self.kv_layout = kv_layout + self.sliding_window = ( + sliding_window if isinstance(sliding_window, int) and sliding_window > 0 else 0 + ) def __eq__(self, other: Optional[ResourceHandler]) -> bool: - """Check compatibility for KVCacheManager (head_dim and dtype must match).""" + """Check compatibility — layers in the same group share a KVCacheManager pool. + + Including sliding_window means same-head-dim layers with different windows + get separate pools. This trades slightly higher memory (one pool per unique + window instead of a shared pool with multiple windows) for a simpler design + where group = pool = metadata set, with no cross-referencing needed. + """ if type(other) is not type(self): return False return ( @@ -1646,6 +1910,7 @@ def __eq__(self, other: Optional[ResourceHandler]) -> bool: and self.dtype == other.dtype and self.kv_factor == other.kv_factor and self.kv_layout == other.kv_layout + and self.sliding_window == other.sliding_window ) def _get_bytes_per_token(self) -> int: diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_gemma4.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_gemma4.py index b8e56d42c107..fa83ed0fb399 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_gemma4.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_gemma4.py @@ -2913,16 +2913,21 @@ class Gemma4TextExportInfo(TextModelExportInfo): def post_process(self, sub_mod: nn.Module, sub_gm: GraphModule): super().post_process(sub_mod, sub_gm) + # Gemma4Model.forward calls language_model.get_per_layer_inputs unconditionally, + # so the GraphModule must expose it on every variant. The method itself short-circuits + # to None when embed_tokens_per_layer is None, which is the non-E2B case. embed_tokens_per_layer = getattr(sub_mod, "embed_tokens_per_layer", None) + sub_gm.get_per_layer_inputs = types.MethodType( + sub_mod.get_per_layer_inputs.__func__, sub_gm + ) + if embed_tokens_per_layer is None: + sub_gm.embed_tokens_per_layer = None return sub_gm.config = sub_mod.config sub_gm.hidden_size_per_layer_input = sub_mod.hidden_size_per_layer_input sub_gm.vocab_size_per_layer_input = sub_mod.vocab_size_per_layer_input - sub_gm.get_per_layer_inputs = types.MethodType( - sub_mod.get_per_layer_inputs.__func__, sub_gm - ) for embed_name, subsubmod in sub_mod.named_modules(): if subsubmod is embed_tokens_per_layer: diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index d5da4e26641a..b633d6a3b625 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -14,7 +14,7 @@ from collections import abc, defaultdict from dataclasses import dataclass from types import SimpleNamespace -from typing import Dict, List, Optional, Tuple +from typing import Dict, List, Optional, Sequence, Tuple import torch from strenum import StrEnum @@ -202,6 +202,105 @@ def _call_func(): return wrapper +def _compute_window_local_view( + all_indices: Sequence[int], + front_removed: int, + end_compute_i: int, + group_window: int, + tokens_per_block: int, +) -> Tuple[List[int], int, int, int]: + """Compute the window-coherent metadata view for one (request, window) pair. + + The C++ KVCacheManager allocates blocks linearly during prefill and only + front-evicts during generation (kvCacheManager.cpp::adjustBlocksIfNeeded + is invoked from addToken, not addSequenceBatch). Two regimes follow: + + * Pre-eviction (typically a single long prefill that has not yet been + decoded past the window): ``front_removed == 0`` and the live page + list covers the full sequence. The kernel needs every live page and + the unclamped local cache length; the sliding-window mask is applied + inside the kernel. + + * Post-eviction (decode/extend has advanced past the window): the C++ + side has bumped ``mNumFrontBlocksRemovedPerWindow`` without popping + ``mCacheBlockIds``, so the head of the page list is stale. After + slicing those entries off, the live cache length in window-local + coords is ``end_compute_i - front_removed * tokens_per_block`` — + bounded above by ``group_window + tokens_per_block - 1`` by the + ``while (live >= window + page_size) detachFrontBlock`` loop. + + Both regimes collapse to a single rule: trust the C++ accounting and + derive the local cache length from ``end_compute_i`` and ``front_removed`` + rather than artificially clamping to ``group_window``. Disagreement + between the Python and C++ sides (e.g. a caller passing a stale + ``end_compute_i`` against a fresher ``front_removed``) is asserted on + rather than silently masked. + + Args: + all_indices: Full historical page list from ``get_cache_indices``, + including stale front entries when eviction has fired. + front_removed: Count of stale front entries, from + ``get_num_front_blocks_removed``. + end_compute_i: Global token position after this step's tokens are + processed (``input_pos + seq_len``). + group_window: Sliding-window size for this pool. Currently only used + for the defensive sanity assertion below; the window mask itself + is applied inside the kernel via the ``sliding_window`` constant. + tokens_per_block: Cache page size. + + Returns: + active_indices: the live slice of all_indices to hand the kernel. + extra_page: the next page after the live slice (used by the overlap + scheduler's deferred-page insertion), or -1 when there is no + slot past the live region. + seq_len_with_cache: window-local cache length the kernel should see. + last_page_len: derived from seq_len_with_cache and tokens_per_block. + """ + live_len = len(all_indices) - front_removed + # Window-local cache length: how many tokens the live pages actually hold. + # Pre-eviction: end_compute_i (the full prefill so far). + # Post-eviction: end_compute_i - front_removed * page_size (capped above by + # window + page_size - 1 thanks to the C++ eviction loop). + local_cache_len = end_compute_i - front_removed * tokens_per_block + # Live cache length must fit inside the live pages. The C++ KVCacheManager + # keeps the two in sync (it bumps front_removed exactly as it frees blocks, + # and allocates new ones for incoming tokens), so a violation means the + # Python side and C++ side disagree on either end_compute_i or + # front_removed -- fail loudly rather than silently clamp. + live_capacity = live_len * tokens_per_block + assert local_cache_len <= live_capacity, ( + f"window-local cache length {local_cache_len} exceeds live page capacity " + f"{live_capacity} (live_len={live_len}, end_compute_i={end_compute_i}, " + f"front_removed={front_removed}, tokens_per_block={tokens_per_block}). " + f"C++ KVCacheManager accounting is expected to keep these in sync." + ) + active_token_count = local_cache_len + # Sanity: live cache should never exceed window + one spillover page once + # front-eviction has fired (the C++ ``while (live >= window + page_size) + # detachFrontBlock`` loop guarantees this). Pre-eviction (front_removed=0, + # e.g. a single long prefill) may legitimately exceed the window because + # blocks are allocated linearly for the full prompt and only detached + # during generation -- so we skip the check there. + assert active_token_count <= group_window + tokens_per_block or front_removed == 0, ( + f"window-local cache length {active_token_count} exceeds " + f"window {group_window} + page_size {tokens_per_block} " + f"(end_compute_i={end_compute_i}, front_removed={front_removed})" + ) + # Pages needed to address active_token_count tokens; equals live_len in + # the common case (post-eviction live = window + spillover, pre-eviction + # live = ceil(end_compute_i / page_size)). + pages_for_active = (active_token_count + tokens_per_block - 1) // tokens_per_block + num_active = min(live_len, pages_for_active) + active_indices = list(all_indices[front_removed : front_removed + num_active]) + extra_slot_idx = front_removed + num_active + extra_page = all_indices[extra_slot_idx] if len(all_indices) > extra_slot_idx else -1 + if active_token_count > 0: + last_page_len = (active_token_count - 1) % tokens_per_block + 1 + else: + last_page_len = 0 + return active_indices, extra_page, active_token_count, last_page_len + + class ADEngine(ModelEngine): """The AutoDeploy Engine (ADEngine) is the main engine interface to execute AutoDeploy models. @@ -658,17 +757,36 @@ def _prepare_inputs( if is_overlap: mask_scatter_indices.extend(list(range(cu_seqlen[-2], cu_seqlen[-1]))) - # store cache information for all requests now + # Store cache information for all requests. + # All KV pools are hosted by a single C++ KVCacheManager; pool index is + # the position of the window in self.cache_seq_interface.kv_group_windows + # (the same source the kvcache transform used to register window groups + # on SequenceInfo). Per-window queries on the manager route to the + # correct C++ pool via mLayerToWindowSize. + kv_group_windows = self.cache_seq_interface.kv_group_windows + # Cache hot lookups so the per-request loop avoids repeated C++ + # dispatch / hasattr calls. _tokens_per_block = kv_cache_manager.tokens_per_block _use_mamba = hasattr(kv_cache_manager, "mamba_cache_index") - cache_loc: List[int] = [] - cu_num_pages: List[int] = [0] - extra_page_per_seq: List[int] = [] state_slot_idx: List[int] = [] - batch_cache_indices = kv_cache_manager.get_batch_cache_indices( - [r.py_request_id for r in ordered_requests] - ) + # Uniform per-pool transport: a single-pool deployment is just the + # degenerate case of VSWA with one pool, so all pools (including 0) + # take the same code path below. + num_pools = len(kv_group_windows) + cache_loc_per_pool: List[List[int]] = [[] for _ in range(num_pools)] + cu_num_pages_per_pool: List[List[int]] = [[0] for _ in range(num_pools)] + extra_page_per_seq_per_pool: List[List[int]] = [[] for _ in range(num_pools)] + seq_len_with_cache_per_pool: List[List[int]] = [[] for _ in range(num_pools)] + last_page_len_per_pool: List[List[int]] = [[] for _ in range(num_pools)] + + # Batched per-pool lookup preserves the #13560 host-overhead + # optimization: one C++ call per pool instead of one per (request, pool). + request_ids = [r.py_request_id for r in ordered_requests] + batch_cache_indices_per_pool: List[List[List[int]]] = [ + kv_cache_manager.get_batch_cache_indices(request_ids, window_size=group_window) + for group_window in kv_group_windows + ] for i, request in enumerate(ordered_requests): # store seq slot idx (use mamba_cache_index if available) @@ -682,17 +800,45 @@ def _prepare_inputs( # get some info on the current request seq_len_i = cu_seqlen[i + 1] - cu_seqlen[i] end_compute_i = input_pos[i] + seq_len_i - # Inline get_num_kv_blocks (pure Python, saves function-call overhead per request) - num_active_blocks_i = (end_compute_i + _tokens_per_block - 1) // _tokens_per_block - - # construct cache information for the current request - cache_indices = batch_cache_indices[i] - cache_loc.extend(cache_indices[:num_active_blocks_i]) - cu_num_pages.append(cu_num_pages[i] + num_active_blocks_i) - if len(cache_indices) > num_active_blocks_i: - extra_page_per_seq.append(cache_indices[num_active_blocks_i]) - else: - extra_page_per_seq.append(-1) + + for pool_idx, group_window in enumerate(kv_group_windows): + all_indices = batch_cache_indices_per_pool[pool_idx][i] + # SWA front-eviction: get_batch_cache_indices returns the FULL + # historical page list including front-evicted entries (the + # C++ side bumps a counter rather than popping mCacheBlockIds). + # _compute_window_local_view slices it down to the live window + # in window-local coords. + front_removed = kv_cache_manager.get_num_front_blocks_removed( + request.py_request_id, window_size=group_window + ) + ( + active_indices, + extra_page, + active_token_count, + lpl_i, + ) = _compute_window_local_view( + all_indices, + front_removed=front_removed, + end_compute_i=end_compute_i, + group_window=group_window, + tokens_per_block=_tokens_per_block, + ) + num_active = len(active_indices) + + cache_loc_per_pool[pool_idx].extend(active_indices) + cu_num_pages_per_pool[pool_idx].append( + cu_num_pages_per_pool[pool_idx][i] + num_active + ) + extra_page_per_seq_per_pool[pool_idx].append(extra_page) + # Window-local seq_len_with_cache / last_page_len for every + # pool (including 0). For full-attention pools the helper + # returns the unclamped global value (group_window equals + # max_seq_len, no clamping kicks in), so this is identical to + # the legacy single-pool path for non-SWA models. For SWA + # pools (whether pool 0 or pool 1+), it carries the + # window-local coords the kernel needs under front-eviction. + seq_len_with_cache_per_pool[pool_idx].append(active_token_count) + last_page_len_per_pool[pool_idx].append(lpl_i) # Store batch information based on prefill, decode, and extend requests. num_decode = len(generation_requests) @@ -717,9 +863,11 @@ def _prepare_inputs( cu_seqlen=cu_seqlen, input_pos=input_pos, batch_info=batch_info, - cache_loc=cache_loc, - cu_num_pages=cu_num_pages, - extra_page_per_seq=extra_page_per_seq, + cache_loc_per_pool=cache_loc_per_pool if num_pools > 0 else None, + cu_num_pages_per_pool=cu_num_pages_per_pool if num_pools > 0 else None, + extra_page_per_seq_per_pool=extra_page_per_seq_per_pool if num_pools > 0 else None, + seq_len_with_cache_per_pool=seq_len_with_cache_per_pool if num_pools > 0 else None, + last_page_len_per_pool=last_page_len_per_pool if num_pools > 0 else None, slot_idx=state_slot_idx, prompt_lens=prompt_lens, gather_context_logits=gather_context_logits, diff --git a/tensorrt_llm/_torch/auto_deploy/shim/demollm.py b/tensorrt_llm/_torch/auto_deploy/shim/demollm.py index cd9ed04cb503..4c48239ae98b 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/demollm.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/demollm.py @@ -173,8 +173,8 @@ def generate_tokens_batched( input_ids=input_ids_flat, cu_seqlen=cu_seqlen, input_pos=[0] * len(total_lens), - cache_loc=cache_loc, - cu_num_pages=cu_num_pages, + cache_loc_per_pool=[cache_loc], + cu_num_pages_per_pool=[cu_num_pages], slot_idx=list(range(len(total_lens))), **extra_args, ) @@ -217,8 +217,8 @@ def _generate_single_step(idx: int): input_ids=input_ids_flat, cu_seqlen=cu_seqlen, input_pos=input_pos_next, - cache_loc=cache_loc, - cu_num_pages=cu_num_pages, + cache_loc_per_pool=[cache_loc], + cu_num_pages_per_pool=[cu_num_pages], slot_idx=list(range(batch_size)), prompt_lens=total_lens, ) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index 4bfc0cf00f68..e6714f513184 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -29,7 +29,7 @@ MambaHybridCacheManager, MixedMambaHybridCacheManager, ) - from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager + from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager, PoolConfiguration from tensorrt_llm._utils import torch_dtype_to_binding from tensorrt_llm.mapping import Mapping @@ -40,6 +40,7 @@ # CachedSequenceInterface can still be instantiated and used for transforms, # but initialize_resources() and other cache-manager methods will raise. KVCacheManager = None + PoolConfiguration = None MambaHybridCacheManager = None MixedMambaHybridCacheManager = None CacheTypeCpp = None @@ -135,6 +136,11 @@ def __init__( self._caches: Dict[str, torch.Tensor] = {} # KVCacheManager (or MambaHybridCacheManager) for managed resources self._kv_cache_manager: Optional[Union[KVCacheManager, MambaHybridCacheManager]] = None + # Per-pool sliding window sizes, published by the kvcache transform and + # consumed by the executor. Pool index == position in this list, in the + # same order as the C++ manager's internal pool ordering (i.e. the + # insertion order of the per-window shape map keys). + self._kv_group_windows: List[int] = [] # lookup of unmanaged resources self._unmanaged_resources: List[str] = [] self._spec_config = spec_config @@ -296,37 +302,92 @@ def _get_mamba_state_params( def _identify_managed_kv_resources( self, - ) -> Tuple[Optional[KVPagedResourceHandler], ResourceHandlerDict]: - """Identify KV resources compatible with the reference handler for KVCacheManager. - - The first KVPagedResourceHandler becomes the reference. All handlers matching - the reference (via __eq__) are collected for managed allocation. + ) -> Tuple[ + ResourceHandlerDict, + List[PoolConfiguration], + ]: + """Identify managed KV resources and the per-pool configurations. + + Each ``KVPagedResourceHandler`` contributes to a pool keyed by its + effective sliding window (``sliding_window`` if > 0 else + ``max_seq_len``). Today, layers sharing a window must share the + same ``head_dim`` and ``dtype``; the returned list carries one + ``PoolConfiguration`` per window. The list shape is deliberately + flat -- future multi-pool-per-window cases append additional + entries with the same ``window_size`` rather than nesting shape + info under a window key. Insertion order fixes the pool index + used by the C++ ``mLayerToWindowSize`` routing. Returns: - Tuple of (reference_handler, managed_resources_dict). - reference_handler is None if no KV paged resources exist. + Tuple of: + - kv_managed: ResourceHandlerDict -- every KVPagedResourceHandler + in registration order; passed to the C++ ctor as the layer + list. + - pool_configurations: List[PoolConfiguration] -- one entry + per pool, carrying (window_size, head_dim, dtype). + + Raises: + RuntimeError: if two layers share an effective window but disagree + on head_dim or dtype, or if + ``self._requires_uniform_kv_caches`` is True and more than one + distinct pool is present. """ - kv_ref: Optional[KVPagedResourceHandler] = None kv_managed: ResourceHandlerDict = {} + pool_by_window: Dict[int, PoolConfiguration] = {} + + max_seq_len = self.info.max_seq_len for name, handler in self._resource_lookup.items(): if not isinstance(handler, KVPagedResourceHandler): continue - if kv_ref is None: - kv_ref = handler - if handler == kv_ref: - kv_managed[name] = handler - elif self._requires_uniform_kv_caches: - raise RuntimeError( - f"KV resource {name} is not compatible with the managed KV reference " - f"(reference: head_dim={kv_ref.head_dim}, dtype={kv_ref.dtype}, " - f"kv_factor={kv_ref.kv_factor}, kv_layout={kv_ref.kv_layout}; " - f"candidate: head_dim={handler.head_dim}, dtype={handler.dtype}, " - f"kv_factor={handler.kv_factor}, kv_layout={handler.kv_layout}). " - "This configuration requires all KV caches to be managed." + # Effective window: full-attention layers (sliding_window == 0) use + # max_seq_len so the C++ side gets a single concrete window key. + effective_window = handler.sliding_window if handler.sliding_window > 0 else max_seq_len + + kv_managed[name] = handler + + handler_dtype = torch_dtype_to_binding(handler.dtype) + existing = pool_by_window.get(effective_window) + if existing is not None: + if existing.head_dim != handler.head_dim: + raise RuntimeError( + f"KV layer {name} has head_dim={handler.head_dim} but window " + f"{effective_window} already has head_dim={existing.head_dim}. " + "The C++ KVCacheManager keys pools by window, so within a window all " + "layers must share head_dim. Place mixed-shape layers in " + "distinct windows (e.g. via sliding_window)." + ) + if existing.dtype != handler_dtype: + raise RuntimeError( + f"KV layer {name} has dtype={handler.dtype} but window " + f"{effective_window} already has dtype={existing.dtype}. " + "The C++ KVCacheManager keys pools by window, so within a window all " + "layers must share dtype. Place mixed-dtype layers in distinct " + "windows (e.g. via sliding_window)." + ) + else: + pool_by_window[effective_window] = PoolConfiguration( + window_size=effective_window, + head_dim=handler.head_dim, + dtype=handler_dtype, ) - return kv_ref, kv_managed + pool_configurations: List[PoolConfiguration] = list(pool_by_window.values()) + + # If the runtime requires uniform KV caches (e.g. legacy single-pool + # path), more than one distinct pool is incompatible. + if len(pool_configurations) > 1 and self._requires_uniform_kv_caches: + pools_repr = ", ".join( + f"window={pc.window_size} head_dim={pc.head_dim} dtype={pc.dtype}" + for pc in pool_configurations + ) + raise RuntimeError( + "KV resources are not uniform: " + f"{pools_repr}. " + "This configuration requires all KV caches to share a single pool." + ) + + return kv_managed, pool_configurations def _identify_managed_state_resources( self, @@ -421,8 +482,15 @@ def _prepare_kv_cache_config( ) -> KvCacheConfig: """Prepare and configure KvCacheConfig for cache manager creation. - Handles deep copy, max_tokens synchronization across ranks, block reuse settings, - copy_on_partial_reuse validation, and free_gpu_memory_fraction normalization. + Handles deep copy, max_tokens synchronization across ranks, block reuse + settings, copy_on_partial_reuse validation, and free_gpu_memory_fraction + normalization. + + ``max_attention_window`` is the per-layer window list (one entry per + managed KV layer, in the order ``kv_managed`` enumerates them). The C++ + ctor groups these layers internally by window into pools using its + ``mLayerToWindowSize`` map. Full-attention layers report + ``max_seq_len`` (i.e. their effective window). Args: max_tokens: Maximum tokens to allocate, or None to use config defaults. @@ -434,6 +502,16 @@ def _prepare_kv_cache_config( # Make a deep copy of the kv_cache_config to avoid modifying the original object kv_cache_config = copy.deepcopy(self._kv_cache_config_original) + # Build per-layer max_attention_window in kv_managed insertion order. + # The C++ side groups layers by this value into pools. + if kv_managed: + kv_cache_config.max_attention_window = [ + handler.sliding_window if handler.sliding_window > 0 else self.info.max_seq_len + for handler in kv_managed.values() + ] + else: + kv_cache_config.max_attention_window = None + # Update kv_cache_config based on max_tokens if provided if max_tokens is not None: # sync max_tokens across ranks @@ -474,31 +552,64 @@ def _prepare_kv_cache_config( def _build_kv_cache_kwargs( self, - kv_ref: Optional[KVPagedResourceHandler], kv_managed: ResourceHandlerDict, kv_cache_config: KvCacheConfig, + pool_configurations: Optional[List[PoolConfiguration]] = None, ) -> Dict: """Build common kwargs for KVCacheManager or MambaHybridCacheManager. + With per-pool ``PoolConfiguration`` carried through to the C++ ctor, + a single manager can host pools with mixed shapes. We pass: + + - ``head_dim`` / ``dtype`` as scalar defaults (fall-back values for + managers that run in uniform-shape mode). We pick the maximum + head_dim and the dtype of the first pool so the default is a sane + upper bound on per-pool memory accounting. + - ``pool_configurations`` as the per-pool config list; the C++ ctor + uses it to build pools with the correct shapes. + Args: - kv_ref: Reference KV handler defining head_dim and dtype, or None. kv_managed: Dict of KV resources to be managed. kv_cache_config: Configured KvCacheConfig. + pool_configurations: One PoolConfiguration per pool, in pool-index + order. Empty / None means uniform shape. Returns: Dict of kwargs suitable for both KVCacheManager and MambaHybridCacheManager. """ + pool_configurations = list(pool_configurations) if pool_configurations else [] + # create arguments first that differ whether we have managed kv caches or not - kv_cache_kwargs = {} + kv_cache_kwargs: Dict = {} if kv_managed: - kv_cache_type = CacheTypeCpp.SELFKONLY if kv_ref.kv_factor == 1 else CacheTypeCpp.SELF + # kv_factor is uniform across the managed set: __init__ asserts + # kv_factor in {1, 2}, and AutoDeploy only mixes kv_factor=2 layers + # in the same model today. We take the first handler's value. + ref_handler = next(iter(kv_managed.values())) + kv_cache_type = ( + CacheTypeCpp.SELFKONLY if ref_handler.kv_factor == 1 else CacheTypeCpp.SELF + ) + # Default head_dim: the largest across all pools, so any layer + # that falls back to the scalar gets a safe upper bound. + default_head_dim = ( + max(pc.head_dim for pc in pool_configurations) + if pool_configurations + else ref_handler.head_dim + ) + # Default dtype: dtype of the first pool. + default_dtype = ( + pool_configurations[0].dtype + if pool_configurations + else torch_dtype_to_binding(ref_handler.dtype) + ) kv_cache_kwargs.update( { "kv_cache_type": kv_cache_type, "num_layers": len(kv_managed), "num_kv_heads": [h.num_kv_heads for h in kv_managed.values()], - "head_dim": kv_ref.head_dim, - "dtype": torch_dtype_to_binding(kv_ref.dtype), + "head_dim": default_head_dim, + "dtype": default_dtype, + "pool_configurations": pool_configurations, } ) else: @@ -509,6 +620,7 @@ def _build_kv_cache_kwargs( "num_kv_heads": [1], "head_dim": 1, "dtype": DataType.HALF, + "pool_configurations": [], } ) # remaining arguments are the same for both cases @@ -610,6 +722,12 @@ def _create_and_assign_state_views( def _assign_kv_cache_views(self, kv_managed: Dict[str, KVPagedResourceHandler]) -> int: """Retrieve and assign buffer views for managed KV paged resources. + ``get_buffers`` on ``self._kv_cache_manager`` is per-layer-aware: it + reads the layer's head_dim from the underlying C++ pool configuration + when mixed-shape pools are present (Gemma4-style VSWA), and falls back + to the manager-level scalar otherwise. We can therefore use a single + call regardless of single-pool vs. multi-pool deployments. + Args: kv_managed: Dict of KV resources managed by the cache manager. @@ -675,13 +793,73 @@ def _allocate_unmanaged_state_resource(self, handler: StateResourceHandler) -> t dtype=handler.dtype, ) + def _has_swa_window(self, pool_configurations: List[PoolConfiguration]) -> bool: + """Return True if any pool's window is smaller than max_seq_len (i.e. SWA pool exists).""" + max_seq_len = self.info.max_seq_len + return any(pc.window_size < max_seq_len for pc in pool_configurations) + + def _compute_total_token_budget( + self, + kv_managed: ResourceHandlerDict, + pool_configurations: List[PoolConfiguration], + total_max_tokens: Optional[int], + ) -> Optional[int]: + """Compute the total max_tokens budget for the unified KVCacheManager. + + With a single C++ manager hosting all pools, ``BaseKVCacheManager:: + calculateMaxNumBlocks`` does the per-pool split internally using each + pool's ``PoolConfiguration``. We just need to compute the total token + budget; the C++ side derives N (max concurrent sequences) from the + total byte budget divided by the combined per-sequence cost across + all pools, then gives each pool N × its per-window tokens. + + We retain a Python-side cap to keep the budget feasible during prefill + warmup before the C++ side has a chance to validate against + ``calculateMaxNumBlocks``. + + Args: + kv_managed: All managed KV layers. + pool_configurations: Per-pool configurations, used here purely to + detect whether at least one SWA pool exists. + total_max_tokens: Caller-provided budget, or None to defer to + ``free_gpu_memory_fraction``. + + Returns: + Total max_tokens for the manager, or None to defer to the + ``free_gpu_memory_fraction`` path. + """ + # If the caller didn't pass a budget and there's no SWA pool, defer to + # free_gpu_memory_fraction inside the manager — same as the legacy path. + if total_max_tokens is None and not self._has_swa_window(pool_configurations): + return None + + if total_max_tokens is None: + # If the user already pinned max_tokens via KvCacheConfig, keep it: + # _prepare_kv_cache_config will preserve that value untouched. + if self._kv_cache_config_original.max_tokens is not None: + return None + # SWA present but no total budget — conservative single-sequence estimate. + # The C++ side will still bound this against available memory. + return self.info.max_seq_len + + tpb = self.info.tokens_per_block + # Cap at max_batch_size × max_seq_len (can't need more than one + # full sequence per slot during prefill). + capped = min(total_max_tokens, self.info.max_batch_size * self.info.max_seq_len) + # Floor: at least one block per sequence for warmup feasibility. + min_tokens = self.info.max_batch_size * tpb + return max(capped, min_tokens) + def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: - """Create KVCacheManager or MambaHybridCacheManager with standard layout. + """Create a single KVCacheManager that hosts every KV pool. For paged resources (KVPagedResourceHandler): - - Uses the first KVPagedResourceHandler's head_dim and dtype as reference - - Compatible resources (matching head_dim and dtype) go into KVCacheManager - - Incompatible resources are allocated locally via handler.allocate() + - All managed layers are passed to one ``KVCacheManager`` instance. + - One ``PoolConfiguration`` per pool flows through the + ``pool_configurations`` ctor kwarg; the C++ side routes each layer + to the correct pool via its ``mLayerToWindowSize`` map. + - Cross-pool admission, event flush, disagg transfer, and the shared + radix tree are all handled in C++. For state resources (SSMResourceHandler, CausalConvResourceHandler, StateResourceHandler): - SSMResourceHandler maps to MambaHybridCacheManager's ssm_states buffer @@ -699,20 +877,33 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: 1. the final number of tokens is synced (min) across ranks 2. rounding for getting a multiple of tokens_per_block """ - # 1. Identify managed resources - kv_ref, kv_managed = self._identify_managed_kv_resources() + # 1. Identify managed resources and per-pool configurations + kv_managed, pool_configurations = self._identify_managed_kv_resources() + ssm_ref, ssm_managed, ssm_spec, conv_ref, conv_managed, conv_spec = ( self._identify_managed_state_resources() ) + has_state_resources = ssm_managed or conv_managed - # 2. Prepare configuration - kv_cache_config = self._prepare_kv_cache_config(max_tokens, kv_managed) - kv_cache_kwargs = self._build_kv_cache_kwargs(kv_ref, kv_managed, kv_cache_config) + # 2. Compute the total token budget; the C++ side splits it across pools. + total_max_tokens = self._compute_total_token_budget( + kv_managed, pool_configurations, max_tokens + ) - # 3. Create cache manager (delegate to state helper if state resources exist) - has_state_resources = ssm_managed or conv_managed - if has_state_resources: - # NOTE: +1 for cuda graph padding + kv_cache_config = self._prepare_kv_cache_config(total_max_tokens, kv_managed) + kv_cache_kwargs = self._build_kv_cache_kwargs( + kv_managed, + kv_cache_config, + pool_configurations=pool_configurations, + ) + + # 3. Build the unified manager (single instance covers every pool). + if not kv_managed and not has_state_resources: + # Pure dummy manager: no KV layers, no state — used for state-only or + # cache-less models. We still need a manager so PyExecutor has a + # handle to call. + self._kv_cache_manager = KVCacheManager(**kv_cache_kwargs) + elif has_state_resources: kv_cache_kwargs["max_batch_size"] = self.info.max_num_state_slots self._kv_cache_manager, _ = self._create_and_assign_state_views( kv_cache_kwargs, @@ -724,18 +915,45 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: conv_spec, ) else: - # No typed state resources - use pure KVCacheManager self._kv_cache_manager = KVCacheManager(**kv_cache_kwargs) - # 4. Store tuned config + ad_logger.info( + f"KV manager: {len(kv_managed)} layers across " + f"{len(pool_configurations)} pool(s); " + f"pool_configurations={pool_configurations}, " + f"max_attention_window={kv_cache_config.max_attention_window}, " + f"max_tokens={total_max_tokens}" + ) + + # 4. Store tuned config (mirrors the kv_cache_config used at ctor time). self._kv_cache_config_tuned = kv_cache_config - # 5. Assign KV views (compute block_offset_multiplier from first view's strides) - block_offset_multiplier = self._assign_kv_cache_views(kv_managed) + # 5. Refresh per-pool window list from the finalized manager: KVCacheManager + # may clamp each window to min(window, max_seq_len) during construction, so + # the transform-time _kv_group_windows can become stale. Downstream callers + # (ad_executor.get_cache_indices(window_size=...)) need the post-clamp + # values to hit the correct pool key in C++. + if kv_managed and self._kv_group_windows: + manager_windows = list(dict.fromkeys(self._kv_cache_manager.max_attention_window_vec)) + if manager_windows and manager_windows != self._kv_group_windows: + ad_logger.info( + f"Refreshing kv_group_windows from manager: " + f"{self._kv_group_windows} -> {manager_windows}" + ) + self._kv_group_windows = manager_windows - # 6. Update cache information (resize cache_loc, set max_seq_info with all max sizes) + # 6. Assign KV views and update cache information. + block_offset_multiplier = 0 + if kv_managed: + block_offset_multiplier = self._assign_kv_cache_views(kv_managed) + + num_blocks = getattr( + self._kv_cache_manager, + "blocks_in_primary_pool", + self._kv_cache_manager.get_max_resource_count(), + ) self.info.update_cache_information( - num_blocks=self._kv_cache_manager.blocks_in_primary_pool, + num_blocks=num_blocks, block_offset_multiplier=block_offset_multiplier, ) @@ -748,11 +966,11 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: self._clear_caches, ) - # 9. Compute final token count and cache statistics + # 8. Compute final token count and cache statistics max_resource_count = self._kv_cache_manager.get_max_resource_count() max_tokens_final = max_resource_count * self._kv_cache_manager.tokens_per_block - # 10. Collect statistics of different types of resources + # 9. Collect statistics of different types of resources num_state_total = sum( 1 for h in self._resource_lookup.values() if isinstance(h, StateResourceHandler) ) @@ -871,12 +1089,16 @@ def resize_kv_cache_manager(self, mem_exclude: int = 0) -> None: This implements the two-phase approach: after running a forward pass during estimation to allocate intermediate memory, call this method to recreate the cache manager. - The new manager will compute optimal capacity based on current free GPU memory. + + For multi-pool (dual head_dim): the C++ + ``BaseKVCacheManager::calculateMaxNumBlocks`` distributes the budget + across pools using the per-window head_dim/dtype overrides — the + Python side just provides the total token budget. """ if not self.needs_resize(): return - # Calculate bytes-per-token for paged (resizable) resources + # Calculate bytes-per-token for ALL paged resources (across all groups) paged_bytes_per_token = sum( h.bytes_per_token for h in self._resource_lookup.values() if h.is_paged ) @@ -895,14 +1117,14 @@ def resize_kv_cache_manager(self, mem_exclude: int = 0) -> None: _, free_mem, *_ = get_mem_info(empty_cache=True) # Compute available memory for paged caches - # Reserve space for non-paged caches and mem_exclude, then apply free_gpu_memory_fraction free_gpu_memory_fraction = self._kv_cache_config_original.free_gpu_memory_fraction mem_for_paged_optimal = ( free_mem - non_paged_bytes_total - mem_exclude ) * free_gpu_memory_fraction max_tokens_optimal = int(mem_for_paged_optimal // paged_bytes_per_token) - # Create new cache manager with optimal capacity + # Recreate the unified manager — the C++ side splits the total token + # budget across pools internally based on per-window head_dim/dtype. cache_stats = self._create_kv_cache_manager(max_tokens=max_tokens_optimal) max_tokens_final = cache_stats["max_tokens"] @@ -920,9 +1142,27 @@ def resize_kv_cache_manager(self, mem_exclude: int = 0) -> None: f"total={bytes_to(total_cache_bytes, unit='GB'):.2f}GB" ) + @property + def kv_group_windows(self) -> List[int]: + """Per-pool window sizes, in the order the C++ manager exposes them. + + Pool index == position in this list, matching the order of the + unified manager's ``pool_configurations`` list. + """ + return self._kv_group_windows + + def set_kv_groups(self, group_windows: List[int]) -> None: + """Store per-pool window sizes (called by the kvcache transform). + + ``group_windows[i]`` is the effective window of pool ``i``; the order + must match the unified ``KVCacheManager``'s pool ordering (i.e., the + insertion order of the per-window keys). + """ + self._kv_group_windows = list(group_windows) + @property def kv_cache_manager(self) -> Optional[KVCacheManager]: - """Return the KVCacheManager managing paged resources, or None if not initialized.""" + """Return the unified KVCacheManager, or None if not initialized.""" assert self._kv_cache_manager is not None, "KVCacheManager not initialized." return self._kv_cache_manager diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py index 2b32b1987d0e..79afb6f6ace4 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py @@ -14,7 +14,30 @@ # limitations under the License. -"""Graph transformation to automatically add kv cache into fused MHA op.""" +"""Graph transformation to automatically add kv cache into fused MHA op. + +The transform runs in two passes so that KV grouping is driven by a single +source of truth — ``KVPagedResourceHandler.__eq__``: + + Pass 1: Walk each source attention node, ask the backend descriptor for the + KV handler via ``get_cache_initializers``, and assign a ``group_idx`` by + comparing the handler against previously-seen handlers (`_find_or_add_group` + semantics). Same equivalence class → same group. Because the handler's + equality includes ``sliding_window``, same-head-dim layers with different + windows automatically land in different groups. + + Pass 2: For multi-group models, publish the per-group window sizes to the + interface (``set_kv_groups``) and to SequenceInfo + (``register_window_groups``), allocate per-group metadata placeholders for + groups ``1..N-1``, and insert cached attention ops with their layer's + group's metadata nodes. + +After the transform, ``group_idx == pool_idx`` holds everywhere: the executor +reads per-pool windows from ``cache_seq_interface.kv_group_windows`` and +issues per-window queries (``get_cache_indices(request, window_size=W)``) +against the unified ``KVCacheManager``; the C++ side routes each call to the +correct pool via its ``mLayerToWindowSize`` map. +""" import inspect import operator @@ -29,6 +52,7 @@ AttentionDescriptor, AttentionRegistry, Constant, + KVPagedResourceHandler, PrepareMetadataCallable, ) from ...custom_ops.semantic_mask_registry import SemanticMaskRegistry @@ -183,18 +207,96 @@ def _process_metadata_extra( gm, prep_meta_op, inputs_for_prep_meta, const_args, num_meta_out ) - def _process_metadata_host(self, cm: CachedSequenceInterface): - """Process the host-side prepare metadata function.""" + def _process_metadata_host( + self, + cm: CachedSequenceInterface, + handler_groups: Optional[List[KVPagedResourceHandler]] = None, + ): + """Process the host-side prepare metadata function. + + When ``handler_groups`` has 2+ entries and the backend's host-prep + declares ``pool_window_left``, registers the host-prep once per pool: + pool 0 uses unsuffixed arg names, pools 1..N-1 use ``_g{i}`` / + ``_g{i}_host`` suffixed names plus a closure that remaps those + suffixed kwargs back to the function's canonical parameter names and + binds the pool's ``pool_window_left``. Each invocation only refreshes + cached wrappers belonging to its pool. + """ prep_meta_host_op = self.attn_descriptor.get_host_prepare_metadata_function() if prep_meta_host_op is None: return - # Register the host-side prepare metadata function with SequenceInfo. - # Arg availability is validated by require_copy() inside register_host_prepare. sig = inspect.signature(prep_meta_host_op) - cm.info.register_host_prepare_for_attention_forward( - prep_meta_host_op, list(sig.parameters.keys()) - ) + sig_arg_names = list(sig.parameters.keys()) + + # Backend opts in to per-pool dispatch by declaring `pool_window_left`. + backend_supports_per_pool = "pool_window_left" in sig.parameters + num_groups = len(handler_groups) if handler_groups is not None else 0 + + if not backend_supports_per_pool or num_groups < 2: + # Register the host-side prepare metadata function with SequenceInfo. + # Arg availability is validated by require_copy() inside register_host_prepare. + cm.info.register_host_prepare_for_attention_forward( + prep_meta_host_op, + [name for name in sig_arg_names if name != "pool_window_left"], + ) + return + + # VSWA: register one host-prep per pool with per-pool args + bound + # pool_window_left. Mirrors the per-group prepare_extra_metadata wiring + # below (vswa_swappable_bases) to keep the two routing paths consistent. + # + # The framework's run_host_prepare_for_attention_forward calls each + # registered function as ``host_function(**{arg: get_arg(arg)})`` — i.e. + # the registered arg names become the *kwarg keys* at call time. For + # pool > 0 the source tensors live under suffixed names + # (``cache_loc_g1_host`` etc.) but the underlying host-prep function + # only declares the unsuffixed parameter names. Each per-pool + # registration therefore wraps the host-prep in a closure that remaps + # suffixed kwargs back to the function's canonical parameter names + # before delegating, and binds this pool's ``pool_window_left``. + swappable_bases = { + "cache_loc", + "cu_num_pages", + "last_page_len", + "seq_len_with_cache", + } + host_suffix = "_host" + + def _make_pool_host_prep(canonical_op, pool_window_left, kwarg_remap): + """Closure: rename suffixed kwargs → canonical params, bind pool_window_left.""" + + def _invoke(**kwargs): + canonical = {kwarg_remap.get(k, k): v for k, v in kwargs.items()} + return canonical_op(**canonical, pool_window_left=pool_window_left) + + return _invoke + + for gi, handler in enumerate(handler_groups): + per_group_args: List[str] = [] + kwarg_remap: Dict[str, str] = {} # suffixed → canonical + for arg_name in sig_arg_names: + if arg_name == "pool_window_left": + continue # bound by closure + if gi == 0: + per_group_args.append(arg_name) + continue + base = arg_name.removesuffix(host_suffix) + if base in swappable_bases: + is_host = arg_name.endswith(host_suffix) + suffixed = f"{base}_g{gi}{host_suffix if is_host else ''}" + per_group_args.append(suffixed) + kwarg_remap[suffixed] = arg_name + else: + per_group_args.append(arg_name) + + # FlashInfer's _to_flashinfer_window_left convention: + # window_left = sliding_window - 1 for SWA, -1 for full attention. + sw = handler.sliding_window + pool_window_left = sw - 1 if isinstance(sw, int) and sw > 0 else -1 + + bound = _make_pool_host_prep(prep_meta_host_op, pool_window_left, kwarg_remap) + cm.info.register_host_prepare_for_attention_forward(bound, per_group_args) def _process_cache_node(self, gm: GraphModule, cache_name: str) -> Node: """Process the cache nodes by inserting a cached attention replacement op.""" @@ -256,17 +358,46 @@ def _apply( # insert metadata computation and extract each argument as a node meta_nodes_extra = self._process_metadata_extra(gm, cm, source_attn_nodes[0]) - # Register host-side prepare_metadata function for attention descriptor. - self._process_metadata_host(cm) + # Seed KvCacheConfig.max_attention_window from per-layer sliding_window + # annotations so the rest of the KV-cache config plumbing (e.g. + # downstream C++ validators that inspect the vector) sees the intended + # per-layer windows. The per-group KVCacheManager pools are configured + # separately in _prepare_kv_cache_config, using each group's reference + # handler. + per_layer_sliding_windows = [] + for attn_node in source_attn_nodes: + (sw,) = extract_op_args(attn_node, "sliding_window") + per_layer_sliding_windows.append(sw) + + has_any_sliding_window = any( + isinstance(sw, int) and sw > 0 for sw in per_layer_sliding_windows + ) + if cm.kv_cache_config.max_attention_window is None and has_any_sliding_window: + max_attention_window = [ + sw if isinstance(sw, int) and sw > 0 else cm.info.max_seq_len + for sw in per_layer_sliding_windows + ] + cm.update_kv_cache_config(max_attention_window=max_attention_window) + + # --- Pass 1: register resources and assign per-layer group_idx --- + # Group identity comes from KVPagedResourceHandler.__eq__ (which + # includes sliding_window). A group IS a pool IS a metadata set. + from ...custom_ops.attention_interface import KVPagedResourceHandler + + handler_groups: list[KVPagedResourceHandler] = [] + per_layer_group_idx: list[int] = [] + group_idx_by_layer_idx: dict[int, int] = {} - # replace fused attention node with attention node that has kv cache num_cached_attn_replacements = 0 cache_nodes_by_layer_idx = {} semantic_mask_cache: Dict[Node, Node] = {} + # Collect per-layer info for the second pass (node insertion). + # Tuple layout: + # (attn_node, qkv, cache_in_nodes, constants, group_idx, prepared_mask) + layer_infos: list[tuple] = [] + for attn_node in source_attn_nodes: - # pick out GEMMs qkv = attn_node.args[: attn_descriptor.get_num_qkv_args()] - layer_idx = attn_descriptor.get_layer_idx(attn_node) shared_kv_source_layer_idx = attn_descriptor.get_shared_kv_source_layer_idx(attn_node) @@ -287,23 +418,36 @@ def _apply( f"Missing shared-KV source layer {shared_kv_source_layer_idx}." ) cache_in_nodes = cache_nodes_by_layer_idx[shared_kv_source_layer_idx] + # Shared-KV layers inherit their source layer's group + group_idx = group_idx_by_layer_idx.get(shared_kv_source_layer_idx, 0) else: - # setup + store cache initializers and caches as input nodes if layer_idx is not None and layer_idx in cache_nodes_by_layer_idx: raise RuntimeError( f"Duplicate KV cache owner detected for layer {layer_idx}. " "Each non-shared attention layer must own exactly one cache." ) cache_in_nodes = [] + group_idx = 0 for k, resource_handler in attn_descriptor.get_cache_initializers( attn_node, cm.kv_cache_config ).items(): resource_name = cm.add_resource(k, resource_handler) cache_in_nodes.append(self._process_cache_node(gm, resource_name)) + # Determine group from handler equality + if isinstance(resource_handler, KVPagedResourceHandler): + for gi, ref in enumerate(handler_groups): + if resource_handler == ref: + group_idx = gi + break + else: + group_idx = len(handler_groups) + handler_groups.append(resource_handler) if layer_idx is not None: cache_nodes_by_layer_idx[layer_idx] = cache_in_nodes + group_idx_by_layer_idx[layer_idx] = group_idx + + per_layer_group_idx.append(group_idx) - # allow backend-specific prep before constants are extracted attn_descriptor.prepare_node_for_cache_insertion(gm, attn_node) prepared_mask = self._process_semantic_mask( @@ -316,20 +460,151 @@ def _apply( ) constants = attn_descriptor.get_constants(attn_node) - # insert cached attention replacement op + layer_infos.append( + (attn_node, qkv, cache_in_nodes, constants, group_idx, prepared_mask) + ) + + # --- Group setup: register metadata groups and create graph placeholders --- + # Register every group (including the single-pool case) so downstream + # consumers see a consistent representation: num_pools == num_groups. + # Per-group graph placeholders are still only created for groups 1..N-1 + # because group 0 reuses the legacy unsuffixed tensors. + num_groups = len(handler_groups) + is_multi_group = num_groups >= 2 + vswa_group_nodes: dict[int, dict[str, "Node"]] = {} + # Per-group extra-metadata: group 0 reuses the previously-built nodes; + # groups 1..N-1 get their own prepare-extra call wired to per-group + # swappable inputs. This is required because some backends' prepare-extra + # ops (e.g. Triton's, which consumes seq_len_with_cache) would otherwise + # produce group-0 outputs that silently feed every layer. + meta_nodes_extra_by_group: Dict[int, List[Node]] = {0: meta_nodes_extra} + host_suffix = "_host" + + has_unmanaged_paged = num_groups == 0 and any( + h.is_paged for h in cm._resource_lookup.values() + ) + + if num_groups >= 1: + group_windows = [ + h.sliding_window if h.sliding_window > 0 else cm.info.max_seq_len + for h in handler_groups + ] + cm.info.register_window_groups(group_windows) + cm.set_kv_groups(group_windows) + self._process_metadata_host(cm, handler_groups=handler_groups) + elif has_unmanaged_paged: + # MLA-only (and any future paged-handler family that does not contribute + # to handler_groups): no KVPagedResourceHandler entries, but at least one + # paged cache still consumes cache_loc/cu_num_pages metadata. Register a + # single full-attention pool so ad_executor.py:740 stages those tensors + # instead of passing None to nest_sequences. + # TODO: unify paged-handler grouping (KV + MLA) under a shared abstraction + # so the predicate at line 361 covers all paged handlers naturally. + group_windows = [cm.info.max_seq_len] + cm.info.register_window_groups(group_windows) + cm.set_kv_groups(group_windows) + self._process_metadata_host(cm, handler_groups=[]) + else: + # No paged caches at all (cache-less or state-only models like Mamba + # without attention). nest_sequences will correctly skip cache_loc + # staging because no kernel in the graph consumes it. Intentional no-op. + ad_logger.debug( + "kvcache transform: no paged KV resources registered; " + "cache_loc staging will be skipped. Expected for state-only or " + "cache-less models." + ) + self._process_metadata_host(cm, handler_groups=[]) + + if is_multi_group: + # Names that are routed per-group end-to-end. seq_len_with_cache is + # in this set because under SWA front-eviction the live window's + # length diverges from the global (input_pos + seq_len); the kernel, + # the prepare-extra op, and the host prepare all need the + # window-capped value. + vswa_swappable_bases = { + "cache_loc", + "cu_num_pages", + "last_page_len", + "seq_len_with_cache", + } + + # Create per-group graph placeholders for groups 1..N-1 + std_arg_names = self.attn_descriptor.get_standard_metadata_args() + for gi in range(1, num_groups): + vswa_group_nodes[gi] = {} + for arg_name in std_arg_names: + base = arg_name.removesuffix(host_suffix) + if base in vswa_swappable_bases: + is_host = arg_name.endswith(host_suffix) + group_arg = f"{base}_g{gi}{host_suffix if is_host else ''}" + vswa_group_nodes[gi][arg_name] = self._add_or_retrieve_input( + gm, cm, group_arg + ) + + # Per-group prepare-extra-metadata: re-invoke the op with this + # group's swappable inputs so each non-zero group's downstream + # consumers (e.g. update_paged_kv_cache write positions) reflect + # the window-capped view. Skipped when the backend has no extra-op. + prep_meta_op, num_meta_out, const_args_extra = ( + self.attn_descriptor.get_prepare_extra_metadata_info(source_attn_nodes[0]) + ) + if prep_meta_op is not None and num_meta_out > 0: + op_arg_names = [ + arg.name + for arg in get_op_schema(prep_meta_op).arguments + if arg.name in cm.info.available_args + ] + for gi in range(1, num_groups): + group_inputs: List[Node] = [] + for arg_name in op_arg_names: + base = arg_name.removesuffix(host_suffix) + if base in vswa_swappable_bases: + is_host = arg_name.endswith(host_suffix) + group_arg = f"{base}_g{gi}{host_suffix if is_host else ''}" + group_inputs.append(self._add_or_retrieve_input(gm, cm, group_arg)) + else: + group_inputs.append(self._add_or_retrieve_input(gm, cm, arg_name)) + meta_nodes_extra_by_group[gi] = self._insert_extra_metadata_op( + gm, prep_meta_op, group_inputs, const_args_extra, num_meta_out + ) + else: + # Backend has no extra-op: every group gets an empty list (which + # is what the original meta_nodes_extra would have been anyway). + for gi in range(1, num_groups): + meta_nodes_extra_by_group[gi] = [] + + # --- Pass 2: insert cached attention nodes with correct group metadata --- + for ( + attn_node, + qkv, + cache_in_nodes, + constants, + group_idx, + prepared_mask, + ) in layer_infos: + layer_meta_nodes_std = meta_nodes_std + if is_multi_group and group_idx > 0: + std_arg_names = self.attn_descriptor.get_standard_metadata_args() + layer_meta_nodes_std = list(meta_nodes_std) + for arg_pos, arg_name in enumerate(std_arg_names): + if arg_name in vswa_group_nodes.get(group_idx, {}): + layer_meta_nodes_std[arg_pos] = vswa_group_nodes[group_idx][arg_name] + + layer_meta_nodes_extra = meta_nodes_extra_by_group.get(group_idx, meta_nodes_extra) + self._insert_cached_attn_node( gm, attn_node, attn_descriptor.get_cached_attention_op(), qkv, - meta_nodes_std, - meta_nodes_extra, + layer_meta_nodes_std, + layer_meta_nodes_extra, cache_in_nodes, constants, prepared_mask, ) - num_cached_attn_replacements += 1 + num_cached_attn_replacements = len(layer_infos) info = TransformInfo( skipped=False, diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 8fbf40599d07..e0c7f6a9871d 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -30,7 +30,8 @@ from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import ( - BaseResourceManager, CacheTypeCpp, DataType, KVCacheManager, get_pp_layers) + BaseResourceManager, CacheTypeCpp, DataType, KVCacheManager, + PoolConfiguration, get_pp_layers) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import (nvtx_range, prefer_pinned, torch_dtype_to_binding) @@ -892,6 +893,10 @@ def __init__( model_type: str = "nemotron_hybrid", is_draft: bool = False, use_replay_state_update: bool = False, + # Per-pool configurations forwarded to the C++ KVCacheManager ctor. + # Lets a single manager host pools with mixed shapes (e.g. Gemma4 + # hybrid attention). See KVCacheManager.__init__. + pool_configurations: Optional[List[PoolConfiguration]] = None, ) -> None: # mamba hybrid cache requires block reuse to be disabled in KV cache config @@ -941,6 +946,7 @@ def __init__( is_estimating_kv_cache=is_estimating_kv_cache, execution_stream=execution_stream, is_draft=is_draft, + pool_configurations=pool_configurations, ) def prepare_resources(self, scheduled_batch: ScheduledRequests): @@ -1207,6 +1213,7 @@ def __init__( is_estimating_kv_cache=is_estimating_kv_cache, is_draft=is_draft, linear_attention_metadata=self.linear_attention_metadata, + **kwargs, ) assert self.local_num_mamba_layers > 0, "At least one mamba layer is required" diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 54f23ddda8ac..8f46d88794cd 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -7,6 +7,7 @@ import os from abc import ABC, abstractmethod from collections import OrderedDict, defaultdict, deque +from dataclasses import dataclass from typing import (TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, Sequence, Set, Tuple, Union) @@ -59,6 +60,7 @@ BufferManagerCpp = tensorrt_llm.bindings.internal.runtime.BufferManager KVCacheManagerCpp = tensorrt_llm.bindings.internal.batch_manager.KVCacheManager +PoolConfigurationCpp = tensorrt_llm.bindings.internal.batch_manager.PoolConfiguration CacheTypeCpp = tensorrt_llm.bindings.internal.batch_manager.CacheType ModelConfigCpp = tensorrt_llm.bindings.ModelConfig DataType = tensorrt_llm.bindings.DataType @@ -76,6 +78,22 @@ int]] # window_size -> (blocks_in_primary_pool, blocks_in_secondary_pool) +@dataclass +class PoolConfiguration: + """Configuration of a single KV pool. + + A pool is uniquely described by its attention ``window_size``, the + ``head_dim`` of the layers it serves, and the cache element ``dtype``. + A KVCacheManager is constructed from a ``list[PoolConfiguration]`` -- + one entry per pool the manager hosts. Multiple entries with the same + ``window_size`` are legal and reserved for future multi-pool-per-window + cases (e.g. mixed head_dim within a single window). + """ + window_size: int + head_dim: int + dtype: "DataType" + + class ResourceManagerType(enum.Enum): KV_CACHE_MANAGER = "KV_CACHE_MANAGER" DRAFT_KV_CACHE_MANAGER = "DRAFT_KV_CACHE_MANAGER" @@ -545,6 +563,14 @@ def __init__( is_estimating_kv_cache: bool = False, execution_stream: Optional[torch.cuda.Stream] = None, linear_attention_metadata: Optional[LinearAttentionMetadata] = None, + # Per-pool configuration list forwarded to the C++ ctor. One entry + # per pool the manager will host; each entry pins (window_size, + # head_dim, dtype) for that pool. None / empty = uniform shape + # across all windows (default behavior); a single KVCacheManager can + # host pools with mixed shapes when a model has heterogeneous + # attention types (e.g. Gemma4 SWA head_dim=256 + full-attention + # head_dim=512). + pool_configurations: Optional[List[PoolConfiguration]] = None, **kwargs, ) -> None: self.mapping = mapping @@ -605,6 +631,21 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], self.num_kv_heads = num_kv_heads self.head_dim = head_dim + # Per-pool configuration list -- the source of truth for per-pool + # (window_size, head_dim, dtype). When non-empty, each pool may + # have its own shape that differs from the manager-level scalars + # (e.g. Gemma4 SWA head_dim=256 alongside full-attention head_dim=512). + # Empty list means uniform shape (every window uses self.head_dim / + # self.dtype). Each pool's window_size is remapped after window + # clamping in _validate_and_adjust_attention_windows; the pool + # *indices* stay stable across that rewrite. + self.pool_configurations: List[PoolConfiguration] = ( + list(pool_configurations) if pool_configurations else []) + # Layer -> pool_idx mapping, built once after max_attention_window_vec + # is initialized below. This is the layer-centric replacement for any + # window-keyed shape dict: multi-pool-per-window is safe because pools + # are identified by index, not by window. + self._layer_to_pool_idx: Dict[int, int] = {} self.tokens_per_block = tokens_per_block self.max_seq_len = max_seq_len self.max_batch_size = max_batch_size @@ -639,6 +680,12 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], layer_mask=layer_mask, ) + # Now that max_attention_window_vec is known, build layer -> pool_idx + # from the (pre-clamp) pool_configurations. Stays valid through the + # window clamping below because that only rewrites per-pool + # window_size fields; pool indices don't shift. + self._layer_to_pool_idx = self._build_layer_to_pool_idx() + # Determine if this is VSWA (Variable Sliding Window Attention). # The `w > 0` check excludes LinearCacheType.RECURRENT_STATES sentinel # values (negative) used by hybrid linear attention models. @@ -702,7 +749,6 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], ), "calculate_max_num_blocks_for_vswa only accepts KvCacheConfig" blocks_per_window = self.calculate_max_num_blocks_for_vswa( kv_cache_config=kv_cache_config, - model_config=model_config, extra_cost_memory=0, ) if mapping.world_size > 1: @@ -753,7 +799,7 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], } # Validate and adjust attention windows against their upper bounds if needed - blocks_per_window, self.max_seq_len, self.max_attention_window_vec = self._validate_and_adjust_attention_windows( + blocks_per_window, self.max_seq_len, self.max_attention_window_vec, window_adjustments = self._validate_and_adjust_attention_windows( max_attention_window_vec=self.max_attention_window_vec, blocks_per_window=blocks_per_window, tokens_per_block=tokens_per_block, @@ -761,6 +807,23 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], max_beam_width=max_beam_width, ) + # Rewrite each pool's window_size to match the post-clamp window. + # Without this, a pool pinned to a pre-clamp window (e.g. 32768) + # would be silently dropped when the validator clamps the window + # down (e.g. to 16384), leaving the C++ side to fall back on the + # manager-level scalar -- which is exactly the heterogeneous case + # these per-pool configs exist to handle. Pool *indices* are + # preserved so self._layer_to_pool_idx stays valid. + if window_adjustments and self.pool_configurations: + self.pool_configurations = [ + PoolConfiguration( + window_size=window_adjustments.get(pc.window_size, + pc.window_size), + head_dim=pc.head_dim, + dtype=pc.dtype, + ) for pc in self.pool_configurations + ] + if kv_cache_type != CacheTypeCpp.SELF: assert len( blocks_per_window @@ -780,6 +843,18 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], ) logger.info(f"[KVCacheManager] execution_stream: {self._stream}") logger.info(f"[KVCacheManager] blocks_per_window: {blocks_per_window}") + + # The Python @dataclass PoolConfiguration is a distinct type from the + # nanobind C++ PoolConfiguration (Python uses ``head_dim``; C++ uses + # ``size_per_head``). Translate at the C++ boundary so nanobind can + # dispatch the ctor. + pool_configurations_cpp = [ + PoolConfigurationCpp(window_size=pc.window_size, + size_per_head=pc.head_dim, + dtype=pc.dtype) + for pc in self.pool_configurations + ] + kwargs = { 'num_kv_heads_per_layer': self.num_kv_heads_per_layer, 'size_per_head': head_dim, @@ -804,6 +879,9 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], 'indexer_k_cache_index_head_dim': indexer_k_cache_index_head_dim, 'indexer_k_cache_use_fp4': indexer_k_cache_use_fp4, 'linear_attention_metadata': linear_attention_metadata, + # Forward the (possibly remapped) per-pool configurations. + # window_size values are aligned with the post-clamp sizes. + 'pool_configurations': pool_configurations_cpp, } if self.event_buffer_max_size > 0: @@ -1395,6 +1473,26 @@ def get_cache_indices(self, assert len(result) == 1 return result[0] + def get_num_front_blocks_removed(self, + request_id: int, + window_size: Optional[int] = None) -> int: + """Get the number of front blocks evicted by SWA for a sequence. + + Args: + request_id: The request id. + window_size: Optional window size. When supplied, returns the + per-window eviction count (zero for non-SWA windows). When + omitted, defaults to ``self.max_attention_window_vec[0]`` — + this matches the historical single-pool behavior for + callers that never thought about windows, while keeping the + C++ contract uniformly per-window. VSWA callers should + always pass ``window_size`` explicitly. + """ + if window_size is None: + window_size = self.max_attention_window_vec[0] + return self.impl.get_num_front_blocks_removed(request_id, + window_size=window_size) + def unpin_blocks_by_id(self, kv_cache_block_id: int): self.impl.unpin_blocks_by_id(kv_cache_block_id) @@ -1424,15 +1522,20 @@ def get_batch_cache_indices( self, request_ids: List[int], layer_idx: Optional[int] = None, + window_size: Optional[int] = None, ) -> List[List[int]]: - if layer_idx is None: - if len(self.max_attention_window_vec) > 1: - raise ValueError("layer_idx must be provided for VSWA") - window_size = self.max_attention_window_vec[0] - else: - layer_offset = self.layer_offsets[layer_idx] - window_size = self.max_attention_window_vec[layer_offset % len( - self.max_attention_window_vec)] + if window_size is None: + if layer_idx is None: + if len(self.max_attention_window_vec) > 1: + raise ValueError( + "layer_idx or window_size must be provided for VSWA") + window_size = self.max_attention_window_vec[0] + else: + layer_offset = self.layer_offsets[layer_idx] + # Explicit layer_offset -> window_size mapping (no modulo + # masking length mismatches between pattern and num_local_layers). + window_size = self._get_layer_offset_to_window_size( + )[layer_offset] result = self.impl.get_batch_cache_block_ids(request_ids, window_size) for i in range(len(result)): @@ -1490,10 +1593,19 @@ def get_buffers(self, Note that different attention backend/implementation can have different KV layouts, "kv_layout" should be set accordingly to avoid surprises. + + Per-layer head_dim: when the underlying C++ manager hosts multiple pools with + distinct head_dim (e.g., Gemma4 SWA head_dim=256 alongside full-attention + head_dim=512), this method reads the layer's effective head_dim from the + layer's assigned ``PoolConfiguration`` rather than the manager-level + scalar. Single-pool managers fall back to ``self.head_dim``. ''' layer_offset = self.layer_offsets[layer_idx] result = self.impl.get_primary_pool_data(layer_offset) + pool = self.get_pool_for_layer(layer_offset) + layer_head_dim = pool.head_dim if pool else self.head_dim + assert kv_layout in ["NHD", "HND"], f"Unsupported kv_layout: {kv_layout}" if kv_layout == "NHD": @@ -1502,7 +1614,7 @@ def get_buffers(self, self.kv_factor, self.tokens_per_block, self.num_kv_heads_per_layer[layer_offset], - self.head_dim, + layer_head_dim, ) else: return result.reshape( @@ -1510,7 +1622,7 @@ def get_buffers(self, self.kv_factor, self.num_kv_heads_per_layer[layer_offset], self.tokens_per_block, - self.head_dim, + layer_head_dim, ) def get_indexer_k_cache_pool_data(self, layer_idx: int) -> torch.Tensor: @@ -1595,6 +1707,138 @@ def get_iteration_stats(self): def rewind_kv_cache(self, request: LlmRequest, rewind_len: int): self.impl.rewind_kv_cache(request.py_request_id, rewind_len) + def calculate_cache_size_per_token(self, + layers: Set[int], + window_size: Optional[int] = None + ) -> int: + """Compute the (raw, dtype-agnostic) KV cache size per token for a set of layers. + + head_dim is resolved per-layer via ``get_pool_for_layer``: each + layer's assigned ``PoolConfiguration`` supplies its ``head_dim``. + When the manager runs in uniform-shape mode (no + ``pool_configurations``), every layer uses ``self.head_dim``. + + Args: + layers: Set of layer offsets. + window_size: Accepted for backward compatibility; ignored. + Per-layer pool lookup already covers the homogeneous case. + + Returns: + cache size per token (number of elements, not bytes). + """ + del window_size # kept for compat; resolution is now per-layer + if not self.pool_configurations: + total_kv_heads = sum(self.num_kv_heads_per_layer[i] for i in layers) + return total_kv_heads * self.kv_factor * self.head_dim + + total = 0 + for i in layers: + pool = self.get_pool_for_layer(i) + layer_head_dim = pool.head_dim if pool else self.head_dim + total += self.num_kv_heads_per_layer[i] * layer_head_dim + return total * self.kv_factor + + def _calculate_cache_bytes_per_token_for_layers( + self, + layers: Set[int], + dtype_default: Optional[DataType] = None) -> int: + """Compute KV cache bytes per token for a set of layers. + + Resolves head_dim and dtype per layer through + ``get_pool_for_layer``; layers whose manager has no + ``pool_configurations`` fall back to ``self.head_dim`` and + ``dtype_default`` (or ``self.dtype``). Handles NVFP4 + scaling-factor overhead, computed per dtype. + """ + if dtype_default is None: + dtype_default = self.dtype + + total_bytes = 0 + for i in layers: + pool = self.get_pool_for_layer(i) + layer_head_dim = pool.head_dim if pool else self.head_dim + layer_dtype = pool.dtype if pool else dtype_default + layer_elements = (self.num_kv_heads_per_layer[i] * self.kv_factor * + layer_head_dim) + layer_bytes = get_size_in_bytes(layer_elements, layer_dtype) + if layer_dtype == DataType.NVFP4: + layer_bytes += KVCacheManager.calculate_scaling_factor_size_bytes( + layer_elements, + quant_vector_size=16, + scaling_factor_dtype=DataType.FP8) + total_bytes += layer_bytes + return total_bytes + + def _build_layer_to_pool_idx(self) -> Dict[int, int]: + """Build the layer_offset -> pool_idx mapping for self.pool_configurations. + + Today, pool assignment is implicit via window_size: each pool has a + unique window_size and a layer joins the pool whose window matches + its effective window. Multiple pools sharing a window_size would + require an explicit per-layer pool index from the caller (a future + ``pool_idx_per_layer`` ctor parameter); this helper raises instead + of silently collapsing them. + """ + if not self.pool_configurations: + return {} + window_to_pool_idx: Dict[int, int] = {} + for idx, pc in enumerate(self.pool_configurations): + if pc.window_size in window_to_pool_idx: + raise RuntimeError( + f"Multiple PoolConfigurations share window_size={pc.window_size}. " + "Multi-pool-per-window requires an explicit layer->pool mapping, " + "which is not yet wired through KVCacheManager.__init__.") + window_to_pool_idx[pc.window_size] = idx + layer_offset_to_window_size = self._get_layer_offset_to_window_size() + return { + offset: window_to_pool_idx[w] + for offset, w in layer_offset_to_window_size.items() + } + + def get_pool_configuration(self, pool_idx: int) -> PoolConfiguration: + """Return the PoolConfiguration at ``pool_idx``.""" + return self.pool_configurations[pool_idx] + + def get_pool_for_layer(self, + layer_offset: int) -> Optional[PoolConfiguration]: + """Return the pool serving ``layer_offset``, or None for uniform managers. + + Layer-centric replacement for any window-keyed shape lookup: the + returned pool's ``head_dim`` and ``dtype`` are authoritative for the + given layer, regardless of how many pools share the layer's window. + """ + if not self.pool_configurations: + return None + pool_idx = self._layer_to_pool_idx.get(layer_offset) + if pool_idx is None: + return None + return self.pool_configurations[pool_idx] + + def _get_layer_offset_to_window_size(self) -> Dict[int, int]: + """Inverse of _get_window_size_to_layers: layer_offset -> window_size. + + Asserts every local layer is mapped exactly once. This is the + explicit, length-mismatch-safe replacement for + ``max_attention_window_vec[layer_offset % len(max_attention_window_vec)]`` + — that modulo silently masks length mismatches between the window + pattern and num_local_layers; this helper catches them via the + assert below. + """ + window_size_to_layers = self._get_window_size_to_layers() + layer_offset_to_window_size: Dict[int, int] = {} + for window_size, layer_offsets in window_size_to_layers.items(): + for layer_offset in layer_offsets: + assert layer_offset not in layer_offset_to_window_size, ( + f"layer_offset {layer_offset} mapped to multiple window " + f"sizes ({layer_offset_to_window_size[layer_offset]} and " + f"{window_size}) — window pattern is malformed.") + layer_offset_to_window_size[layer_offset] = window_size + assert len(layer_offset_to_window_size) == self.num_local_layers, ( + f"layer_offset_to_window_size covers " + f"{len(layer_offset_to_window_size)} layers but num_local_layers " + f"is {self.num_local_layers}.") + return layer_offset_to_window_size + def _get_window_size_to_layers(self) -> dict[int, list[int]]: """ Get the window size to layers mapping. @@ -1636,34 +1880,27 @@ def adjust_window_sizes_for_vswa( window_size_to_layers: Dict[int, List[int]], max_attention_window_vec: List[int], kv_cache_config: KvCacheConfig, - model_config: ModelConfigCpp, pool_memory_bytes: int, kv_factor: int, dtype: DataType, is_cross_attention: bool = False, + model_config: Optional[ModelConfigCpp] = None, ) -> Tuple[Dict[int, List[int]], List[int]]: assert is_cross_attention is False, 'Cross attention is not supported' max_tokens_from_config = kv_cache_config.max_tokens - def calculate_cache_size_per_token(layers: Set[int]) -> int: - # Same as BaseKVCacheManager::calculateCacheSizePerTokenForSingleWindowSize - total_kv_heads = sum(self.num_kv_heads_per_layer[i] for i in layers) - return total_kv_heads * kv_factor * model_config.head_size - - # Calculate the required memory bytes per sequence. + # Calculate the required memory bytes per sequence. Each window's + # bytes-per-token is computed with that window's effective head_dim + # / dtype (per-window override map), falling back to the manager + # scalars when no override is registered. required_mem_bytes_per_seq = 0 for window_size in sorted(window_size_to_layers): layers = window_size_to_layers[window_size] - cache_size_per_token = calculate_cache_size_per_token(layers) - cache_size_bytes_per_token = get_size_in_bytes( - cache_size_per_token, dtype) - if dtype == DataType.NVFP4: - cache_size_bytes_per_token += KVCacheManager.calculate_scaling_factor_size_bytes( - cache_size_per_token, - quant_vector_size=16, - scaling_factor_dtype=DataType.FP8) + cache_size_bytes_per_token = ( + self._calculate_cache_bytes_per_token_for_layers( + layers, dtype_default=dtype)) required_mem_bytes_per_seq += window_size * cache_size_bytes_per_token logger.info( f'Required memory per sequence: {required_mem_bytes_per_seq} bytes') @@ -1692,16 +1929,13 @@ def calculate_cache_size_per_token(layers: Set[int]) -> int: for window_size in sorted(window_size_to_layers): layers = window_size_to_layers[window_size] if remaining_mem_bytes > 0 and remaining_layers: - # Calculate cache size per token for remaining layers only - cache_size_per_token = calculate_cache_size_per_token( - remaining_layers) - cache_size_bytes_per_token = get_size_in_bytes( - cache_size_per_token, dtype) - if dtype == DataType.NVFP4: - cache_size_bytes_per_token += KVCacheManager.calculate_scaling_factor_size_bytes( - cache_size_per_token, - quant_vector_size=16, - scaling_factor_dtype=DataType.FP8) + # Calculate cache size per token for remaining layers only. + # ``remaining_layers`` may span multiple windows with + # different head_dim / dtype, so the helper resolves each + # layer's effective shape and dtype individually. + cache_size_bytes_per_token = ( + self._calculate_cache_bytes_per_token_for_layers( + remaining_layers, dtype_default=dtype)) logger.debug( f'Cache size per token for {len(remaining_layers)} layers: ' f'{cache_size_bytes_per_token} bytes') @@ -1843,7 +2077,7 @@ def _calculate_max_num_blocks_for_linear_attention( def calculate_max_num_blocks_for_vswa( self, kv_cache_config: KvCacheConfig, - model_config: Optional[ModelConfigCpp], + model_config: Optional[ModelConfigCpp] = None, extra_cost_memory: int = 0) -> dict[int, tuple[int, int]]: """ Currently, this function is added to support *ONLY* VSWA. @@ -1889,31 +2123,19 @@ def calculate_max_num_blocks_for_vswa( extra_cost_memory=extra_cost_memory, ) - # VSWA case: use C++ implementation for variable window sizes - if model_config is None: - raise ValueError( - "model_config is required for VSWA (Variable Sliding Window Attention)" - ) - assert model_config.layer_types is not None, "layer_types have to be set correctly for VSWA" - if self.is_vswa: - # Adjust the window sizes to fit the memory if even a single sequence - # cannot fit in the memory. - window_size_to_layers, max_attention_window_vec = self.adjust_window_sizes_for_vswa( - window_size_to_layers=window_size_to_layers, - max_attention_window_vec=self.max_attention_window_vec, - model_config=model_config, - kv_cache_config=kv_cache_config, - pool_memory_bytes=self._primary_pool_memory_bytes, - kv_factor=self.kv_factor, - dtype=self.dtype, - is_cross_attention=is_cross_attention, - ) - self.max_attention_window_vec = max_attention_window_vec - - def calculate_cache_size_per_token(layers: Set[int]) -> int: - # Same as BaseKVCacheManager::calculateCacheSizePerTokenForSingleWindowSize - total_kv_heads = sum(self.num_kv_heads_per_layer[i] for i in layers) - return total_kv_heads * self.kv_factor * model_config.head_size + # VSWA case: adjust window sizes via Python helper that derives + # head_dim from self. model_config is no longer required because + # head_size is read from self.head_dim, set during __init__. + window_size_to_layers, max_attention_window_vec = self.adjust_window_sizes_for_vswa( + window_size_to_layers=window_size_to_layers, + max_attention_window_vec=self.max_attention_window_vec, + kv_cache_config=kv_cache_config, + pool_memory_bytes=self._primary_pool_memory_bytes, + kv_factor=self.kv_factor, + dtype=self.dtype, + is_cross_attention=is_cross_attention, + ) + self.max_attention_window_vec = max_attention_window_vec logger.info( f"Primary pool memory bytes: {self._primary_pool_memory_bytes}") @@ -1944,9 +2166,11 @@ def calculate_cache_size_per_token(layers: Set[int]) -> int: blocks_per_window = {} for window_idx, (window_size, layers) in enumerate( sorted(window_size_to_layers.items())): - cache_size_per_token = calculate_cache_size_per_token(layers) - cache_size_bytes_per_token = get_size_in_bytes( - cache_size_per_token, self.dtype) + # Per-window head_dim and dtype (with scalar fallback) — needed + # for heterogeneous-attention models like Gemma4 where SWA and + # full-attention pools have different head_dim. + cache_size_bytes_per_token = ( + self._calculate_cache_bytes_per_token_for_layers(layers)) primary_tokens = self._primary_pool_memory_bytes * window_size_shares[ window_idx] / cache_size_bytes_per_token @@ -1983,7 +2207,7 @@ def _validate_and_adjust_attention_windows( tokens_per_block: int, max_seq_len: int, max_beam_width: int, - ) -> Tuple[BlocksPerWindow, int, List[int]]: + ) -> Tuple[BlocksPerWindow, int, List[int], Dict[int, int]]: """ Validate and adjust attention windows against their upper bounds if needed. If there is no adjustment, the returned max_attention_window_vec will be the same as the input. @@ -1995,7 +2219,12 @@ def _validate_and_adjust_attention_windows( max_seq_len: Maximum sequence length Returns: - Tuple of (adjusted_blocks_per_window, adjusted_max_seq_len, adjusted_max_attention_window_vec) + Tuple of (adjusted_blocks_per_window, adjusted_max_seq_len, + adjusted_max_attention_window_vec, window_adjustments). + window_adjustments maps pre_clamp -> post_clamp window size for + every window that was clamped (empty if nothing was adjusted) so + callers can rewrite their per-pool configurations to match the + post-clamp window keys. """ window_adjustments = {} # Validate each window size in blocks_per_window against its upper bound @@ -2037,9 +2266,9 @@ def _validate_and_adjust_attention_windows( adjusted_max_seq_len = max(adjusted_window_vec) logger.warning(f"Adjusted max_seq_len to {adjusted_max_seq_len}") - return adjusted_blocks_per_window, adjusted_max_seq_len, adjusted_window_vec + return adjusted_blocks_per_window, adjusted_max_seq_len, adjusted_window_vec, window_adjustments else: - return blocks_per_window, max_seq_len, max_attention_window_vec + return blocks_per_window, max_seq_len, max_attention_window_vec, {} def pin_blocks(self, request_id: int): self.impl.pin_blocks(request_id) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_paged_attention_swa.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_paged_attention_swa.py new file mode 100644 index 000000000000..2da593035874 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_paged_attention_swa.py @@ -0,0 +1,360 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SWA correctness tests for Triton paged attention. + +Two regimes that the existing kernel tests don't exercise: + +1. **Long prefill SWA** (no eviction yet): a single prefill longer than the + sliding window. The kernel must apply the SW mask within the chunk itself. + This is the lightweight regression that catches mask-math regressions + without spinning up a real KVCacheManager. + +2. **Front-eviction SWA** (eviction has happened): the shim hands the kernel + a window-coherent local-coord view — sliced cache_loc, window-capped + seq_len_with_cache. The kernel must produce the same output as SDPA on the + live window only. Pre-fix kernels that mixed coordinate spaces silently + dropped boundary pages and corrupted masks; under Option B (everything + local) the kernel needs no changes, so these tests should pass before AND + after the shim/transform fix. They are the regression guard against future + coordinate-space mistakes. +""" + +import math + +import pytest +import torch + +import tensorrt_llm._torch.auto_deploy # noqa: F401 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + + +def _sdpa_reference_with_sw( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + sliding_window: int, + cache_len_before_q: int, +) -> torch.Tensor: + """Reference SDPA with a per-query causal + sliding-window mask. + + Inputs are in [tokens, n_heads, head_dim] (flat). For each query token at + local position `q_pos = cache_len_before_q + i` (i ∈ [0, q_len)), the + valid keys are positions in [max(0, q_pos - W + 1), q_pos]. + """ + q_t = q.transpose(0, 1).unsqueeze(0) # [1, n_heads, q_len, head_dim] + k_t = k.transpose(0, 1).unsqueeze(0) # [1, n_kv_heads, kv_len, head_dim] + v_t = v.transpose(0, 1).unsqueeze(0) + + n_heads = q_t.shape[1] + n_kv_heads = k_t.shape[1] + if n_heads != n_kv_heads: + rep = n_heads // n_kv_heads + k_t = k_t.repeat_interleave(rep, dim=1) + v_t = v_t.repeat_interleave(rep, dim=1) + + q_len = q_t.shape[2] + kv_len = k_t.shape[2] + + head_dim = q_t.shape[-1] + scale = 1.0 / math.sqrt(head_dim) + scores = torch.matmul(q_t, k_t.transpose(-2, -1)) * scale # [1, h, q_len, kv_len] + + # mask in [q_len, kv_len]; q_pos = cache_len_before_q + i + q_positions = torch.arange(q_len, device=q.device) + cache_len_before_q + kv_positions = torch.arange(kv_len, device=q.device) + diff = q_positions.unsqueeze(1) - kv_positions.unsqueeze(0) + causal_ok = diff >= 0 + window_ok = diff < sliding_window + valid = causal_ok & window_ok + scores = scores.masked_fill(~valid.unsqueeze(0).unsqueeze(0), float("-inf")) + + weights = torch.softmax(scores, dim=-1) + out = torch.matmul(weights, v_t) # [1, h, q_len, head_dim] + return out.squeeze(0).transpose(0, 1) # [q_len, n_heads, head_dim] + + +def _write_kv_to_cache( + k: torch.Tensor, + v: torch.Tensor, + kv_cache: torch.Tensor, + page_table: torch.Tensor, + page_size: int, +): + """Write all of (k, v) sequentially into kv_cache at the listed pages. + + k, v are flat [tokens, n_kv_heads, head_dim]. page_table is a 1D int32 + tensor listing physical page ids in token-order. Assumes page_table + covers exactly enough pages for k.shape[0] tokens. + """ + from tensorrt_llm._torch.auto_deploy.custom_ops.attention.triton_paged_attention import ( + update_paged_kv_cache, + ) + + n_tokens = k.shape[0] + batch_indices = torch.zeros(n_tokens, dtype=torch.int32, device=k.device) + positions = torch.arange(n_tokens, dtype=torch.int32, device=k.device) + kv_indptr = torch.tensor([0, page_table.numel()], dtype=torch.int32, device=k.device) + update_paged_kv_cache(k, v, batch_indices, positions, kv_cache, page_table, kv_indptr) + + +class TestTritonPagedLongPrefillSWA: + """Prefill longer than the sliding window — no eviction yet. + + Verifies the kernel correctly applies the SW mask within a single long + prefill chunk. Covers BOTH the Phase-1 (full pages) and Phase-2 (boundary + pages) code paths since prefill_len > W * 2 forces both. + """ + + @pytest.mark.parametrize("page_size", [16, 64]) + @pytest.mark.parametrize("sliding_window", [128, 256]) + @pytest.mark.parametrize("prefill_len_mult", [1, 2, 4]) # 1*W, 2*W, 4*W + @pytest.mark.parametrize("n_heads,n_kv_heads", [(4, 4), (8, 2)]) + @pytest.mark.parametrize("head_dim", [64, 128]) + def test_long_prefill_sw_matches_sdpa( + self, + page_size: int, + sliding_window: int, + prefill_len_mult: int, + n_heads: int, + n_kv_heads: int, + head_dim: int, + ): + from tensorrt_llm._torch.auto_deploy.custom_ops.attention.triton_paged_attention import ( + triton_paged_context, + ) + + torch.manual_seed(0) + prefill_len = sliding_window * prefill_len_mult + # Round prefill_len up so it spans whole + boundary pages naturally. + num_pages = (prefill_len + page_size - 1) // page_size + num_blocks = num_pages + 4 + + q = torch.randn(prefill_len, n_heads, head_dim, dtype=torch.float16, device="cuda") + k = torch.randn(prefill_len, n_kv_heads, head_dim, dtype=torch.float16, device="cuda") + v = torch.randn(prefill_len, n_kv_heads, head_dim, dtype=torch.float16, device="cuda") + + kv_cache = torch.zeros( + num_blocks, 2, n_kv_heads, page_size, head_dim, dtype=torch.float16, device="cuda" + ) + page_table = torch.arange(num_pages, dtype=torch.int32, device="cuda") + _write_kv_to_cache(k, v, kv_cache, page_table, page_size) + + qo_indptr = torch.tensor([0, prefill_len], dtype=torch.int32, device="cuda") + kv_indptr = torch.tensor([0, num_pages], dtype=torch.int32, device="cuda") + kv_indices = page_table.clone() + last_token_in_page = prefill_len % page_size + kv_last_page_len = torch.tensor( + [last_token_in_page if last_token_in_page > 0 else page_size], + dtype=torch.int32, + device="cuda", + ) + seq_len_with_cache = torch.tensor([prefill_len], dtype=torch.int32, device="cuda") + + sm_scale = 1.0 / math.sqrt(head_dim) + output = triton_paged_context( + q, + kv_cache, + qo_indptr, + kv_indptr, + kv_indices, + kv_last_page_len, + seq_len_with_cache, + sm_scale, + sliding_window=sliding_window, + ) + + ref = _sdpa_reference_with_sw( + q.float(), + k.float(), + v.float(), + sliding_window=sliding_window, + cache_len_before_q=0, + ).to(output.dtype) + + torch.testing.assert_close(output.float(), ref.float(), rtol=1e-2, atol=1e-2) + + +class TestTritonPagedSWAFrontEviction: + """Front-eviction SWA — the shim hands the kernel a window-local view. + + The setup mirrors what `ad_executor.py` does after Phase 2: it slices + `cache_loc` to the live (post-eviction) pages and caps + `seq_len_with_cache` at the window. With Option B (everything local), the + kernel needs no changes — these tests are the regression guard for that + contract. + """ + + @pytest.mark.parametrize("front_removed", [0, 1, 4]) + @pytest.mark.parametrize("chunk", [1, 64, 128]) + @pytest.mark.parametrize("page_size", [16, 64]) + @pytest.mark.parametrize("with_sliding_window", [True, False]) + def test_front_evicted_context_matches_sdpa_on_live_window( + self, + front_removed: int, + chunk: int, + page_size: int, + with_sliding_window: bool, + ): + from tensorrt_llm._torch.auto_deploy.custom_ops.attention.triton_paged_attention import ( + triton_paged_context, + ) + + # Fixed simple geometry for clarity. Window W=256 → 16 pages at PAGE_SIZE=16, + # 4 pages at PAGE_SIZE=64. + torch.manual_seed(0) + sliding_window = 256 + n_heads = 4 + n_kv_heads = 4 + head_dim = 64 + + # Total cached tokens (post-fill, pre-chunk). Make this exactly equal to W + # so the "live" window contains exactly W tokens regardless of front_removed. + total_cache_tokens = sliding_window + # The historical page list must include the front-evicted pages too. + historical_pages = front_removed + (total_cache_tokens + page_size - 1) // page_size + num_blocks = historical_pages + 4 + (chunk + page_size - 1) // page_size + + # Build a historical KV history: front_removed*page_size junk tokens + # (will not be read) + total_cache_tokens real tokens, then `chunk` new + # tokens being processed. + live_cache_tokens = total_cache_tokens + new_chunk = chunk + + # Allocate cache and a "historical" page table (front_removed pages + # contain stale data; the rest contain live cached tokens). + kv_cache = torch.zeros( + num_blocks, 2, n_kv_heads, page_size, head_dim, dtype=torch.float16, device="cuda" + ) + + # Real (live + new) KV data we want the kernel to attend to. + live_k = torch.randn( + live_cache_tokens, n_kv_heads, head_dim, dtype=torch.float16, device="cuda" + ) + live_v = torch.randn( + live_cache_tokens, n_kv_heads, head_dim, dtype=torch.float16, device="cuda" + ) + new_k = torch.randn(new_chunk, n_kv_heads, head_dim, dtype=torch.float16, device="cuda") + new_v = torch.randn(new_chunk, n_kv_heads, head_dim, dtype=torch.float16, device="cuda") + # Stale data the kernel must NOT read (we'll keep these pages + # zero-filled — if the kernel reads them it'll skew the output). + + # Page table layout (historical, full list): + # [stale_0, stale_1, ..., stale_{front_removed-1}, live_0, ..., new_pages...] + live_pages_needed = (live_cache_tokens + page_size - 1) // page_size + new_pages_needed = (new_chunk + page_size - 1) // page_size + # Total physical pages used by this request, from page 0..end: + all_pages = torch.arange( + historical_pages + new_pages_needed, dtype=torch.int32, device="cuda" + ) + + # Write live KV to the live pages (positions [front_removed*page_size, + # front_removed*page_size + live_cache_tokens)) of the global cache. + live_page_table = all_pages[front_removed : front_removed + live_pages_needed] + _write_kv_to_cache(live_k, live_v, kv_cache, live_page_table, page_size) + # Write new chunk KV to the new pages (just after the live pages). + new_page_table = all_pages[ + front_removed + live_pages_needed : front_removed + live_pages_needed + new_pages_needed + ] + # Pad to page boundary for the write helper. + new_k_padded = torch.zeros( + new_pages_needed * page_size, n_kv_heads, head_dim, dtype=torch.float16, device="cuda" + ) + new_v_padded = torch.zeros_like(new_k_padded) + # Page-aligned local offset for the chunk's first token: it sits right + # after the live cache, in local coords that's live_cache_tokens. + # Since live_cache_tokens is a multiple of page_size (we set it = W and + # require page_size | W), the chunk starts at offset 0 within new_page_table. + assert live_cache_tokens % page_size == 0, ( + "test geometry assumes live cache is page-aligned" + ) + new_k_padded[:new_chunk] = new_k + new_v_padded[:new_chunk] = new_v + # We can't easily write a partial-page; only fill if new_chunk is multiple of page_size. + # For partial pages, fall back to per-token write. + from tensorrt_llm._torch.auto_deploy.custom_ops.attention.triton_paged_attention import ( + update_paged_kv_cache, + ) + + b_idx = torch.zeros(new_chunk, dtype=torch.int32, device="cuda") + positions = torch.arange(new_chunk, dtype=torch.int32, device="cuda") + kv_indptr_for_write = torch.tensor([0, new_pages_needed], dtype=torch.int32, device="cuda") + update_paged_kv_cache( + new_k, new_v, b_idx, positions, kv_cache, new_page_table, kv_indptr_for_write + ) + + # Shim's view: slice the historical page table to live + new only. + # This is exactly `all_indices[front_removed : front_removed + num_active]` + # in ad_executor.py after Phase 2. + live_view_pages = all_pages[ + front_removed : front_removed + live_pages_needed + new_pages_needed + ].contiguous() + num_kv_pages = live_view_pages.numel() + + # Q is the chunk's queries. + q = torch.randn(new_chunk, n_heads, head_dim, dtype=torch.float16, device="cuda") + + qo_indptr = torch.tensor([0, new_chunk], dtype=torch.int32, device="cuda") + kv_indptr = torch.tensor([0, num_kv_pages], dtype=torch.int32, device="cuda") + # Window-capped seq_len_with_cache: live_cache_tokens + new_chunk + # (NOT the global value which would include the evicted prefix). + seq_len_with_cache = torch.tensor( + [live_cache_tokens + new_chunk], dtype=torch.int32, device="cuda" + ) + last_token_in_page = (live_cache_tokens + new_chunk) % page_size + kv_last_page_len = torch.tensor( + [last_token_in_page if last_token_in_page > 0 else page_size], + dtype=torch.int32, + device="cuda", + ) + + sm_scale = 1.0 / math.sqrt(head_dim) + sw_arg = sliding_window if with_sliding_window else 0 + output = triton_paged_context( + q, + kv_cache, + qo_indptr, + kv_indptr, + live_view_pages, + kv_last_page_len, + seq_len_with_cache, + sm_scale, + sliding_window=sw_arg, + ) + + # SDPA reference computed on the LIVE window only: keys/values are + # [live_k; new_k], queries are `q`. Cache length before q is + # live_cache_tokens. Sliding window applied if requested (passing + # sw=0 corresponds to plain causal on the live window). + full_k = torch.cat([live_k, new_k], dim=0) + full_v = torch.cat([live_v, new_v], dim=0) + ref = _sdpa_reference_with_sw( + q.float(), + full_k.float(), + full_v.float(), + # sw=0 in the kernel ⇒ no SW pruning; reproduce that by feeding a + # window larger than the entire kv length. + sliding_window=sliding_window + if with_sliding_window + else (live_cache_tokens + new_chunk + 1), + cache_len_before_q=live_cache_tokens, + ).to(output.dtype) + + torch.testing.assert_close(output.float(), ref.float(), rtol=1e-2, atol=1e-2) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_trtllm_attention_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_trtllm_attention_op.py index d06ed725aff5..353938830fa1 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_trtllm_attention_op.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_trtllm_attention_op.py @@ -845,8 +845,8 @@ def test_extend_iteration_preserves_original_prompt_length(self, context_length: cu_seqlen=cu_seqlen, input_pos=input_pos, batch_info=batch_info, - cache_loc=torch.arange(num_pages), - cu_num_pages=torch.tensor([0, num_pages], dtype=torch.int), + cache_loc_per_pool=[torch.arange(num_pages)], + cu_num_pages_per_pool=[torch.tensor([0, num_pages], dtype=torch.int)], slot_idx=torch.arange(1), prompt_lens=[context_length], ) @@ -868,8 +868,8 @@ def test_metadata_handles_two_sequences_with_different_lengths(self): input_ids.flatten(), cu_seqlen=cu_seqlen, input_pos=0, - cache_loc=torch.arange(300 // info.tokens_per_block + 2), - cu_num_pages=torch.tensor([0, 4, 10], dtype=torch.int), + cache_loc_per_pool=[torch.arange(300 // info.tokens_per_block + 2)], + cu_num_pages_per_pool=[torch.tensor([0, 4, 10], dtype=torch.int)], slot_idx=torch.arange(2), ) _prepare_from_sequence_info(info) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_switch_to_generate_inplace.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_switch_to_generate_inplace.py index 6772e8cde779..16eba6cc79c5 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_switch_to_generate_inplace.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_switch_to_generate_inplace.py @@ -68,9 +68,9 @@ def _nest_prefill(si: SequenceInfo, input_ids, pages_per_seq, cache_loc, **kw): flat_ids, cu_seqlen, input_pos=0, - cache_loc=cache_loc, - cu_num_pages=cu_num_pages, - extra_page_per_seq=extra_page, + cache_loc_per_pool=[cache_loc], + cu_num_pages_per_pool=[cu_num_pages], + extra_page_per_seq_per_pool=[extra_page], **kw, ) return si @@ -158,9 +158,9 @@ def test_single_prefill_with_nonzero_input_pos(self): [1, 2, 3, 4, 5], cu_seqlen=[0, 5], input_pos=[10], - cache_loc=[0, 1, 2, 3], - cu_num_pages=[0, 4], - extra_page_per_seq=[-1], + cache_loc_per_pool=[[0, 1, 2, 3]], + cu_num_pages_per_pool=[[0, 4]], + extra_page_per_seq_per_pool=[[-1]], ) increment = torch.tensor([1], dtype=torch.int32, device=si.device) @@ -195,9 +195,9 @@ def _setup_decode_batch(self, si, positions, swc_vals, pages_per_seq, cache_loc) cu_seqlen=list(range(batch_size + 1)), input_pos=positions, batch_info=[0, 0, 0, 0, batch_size, batch_size], - cache_loc=cache_loc, - cu_num_pages=cu_num_pages, - extra_page_per_seq=[-1] * batch_size, + cache_loc_per_pool=[cache_loc], + cu_num_pages_per_pool=[cu_num_pages], + extra_page_per_seq_per_pool=[[-1] * batch_size], ) def test_decode_increment_by_one(self): @@ -282,9 +282,9 @@ def test_page_crossing_with_extra_page_inserts_into_cache_loc(self): cu_seqlen=[0, 1], input_pos=[7], batch_info=[0, 0, 0, 0, 1, 1], - cache_loc=[10, 11], - cu_num_pages=[0, 2], - extra_page_per_seq=[99], + cache_loc_per_pool=[[10, 11]], + cu_num_pages_per_pool=[[0, 2]], + extra_page_per_seq_per_pool=[[99]], ) increment = torch.tensor([1], dtype=torch.int32, device=si.device) @@ -304,9 +304,9 @@ def _setup_decode_at_page_boundary(si): cu_seqlen=[0, 1], input_pos=[7], batch_info=[0, 0, 0, 0, 1, 1], - cache_loc=[10, 11], - cu_num_pages=[0, 2], - extra_page_per_seq=[-1], + cache_loc_per_pool=[[10, 11]], + cu_num_pages_per_pool=[[0, 2]], + extra_page_per_seq_per_pool=[[-1]], ) @@ -342,9 +342,9 @@ def test_inactive_host_mirrors_not_synced_by_offset_pos_and_cache(self): cu_seqlen=[0, 1, 2], input_pos=[5, 10], batch_info=[0, 0, 0, 0, 2, 2], - cache_loc=[10, 11, 20, 21, 22], - cu_num_pages=[0, 2, 5], - extra_page_per_seq=[-1, -1], + cache_loc_per_pool=[[10, 11, 20, 21, 22]], + cu_num_pages_per_pool=[[0, 2, 5]], + extra_page_per_seq_per_pool=[[-1, -1]], ) assert not _has_active_host_arg(si) diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_ad_executor_swa_eviction.py b/tests/unittest/auto_deploy/singlegpu/shim/test_ad_executor_swa_eviction.py new file mode 100644 index 000000000000..95c873f8a160 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_ad_executor_swa_eviction.py @@ -0,0 +1,316 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Integration test for the shim's SWA front-eviction metadata path. + +Drives a real two-window CachedSequenceInterface + KVCacheManager (so the +historical `get_cache_indices` payload reflects real C++ page-table state), +then monkeypatches `get_num_front_blocks_removed` to simulate eviction at +controlled counts. The shim helper `_compute_window_local_view` must +produce a window-coherent local-coord view: live page slice, window-capped +seq_len_with_cache, derived last_page_len, and a coherent spillover slot. + +`get_num_front_blocks_removed` is monkeypatched (not driven by real decode) +because firing real SWA eviction from Python without a model requires +either spinning up a full ADEngine + model or adding a test-only C++ +binding to bump the counter directly — both are out of scope. The patch +covers exactly what the C++ side would otherwise do: bump a counter +without mutating mCacheBlockIds, so the historical page list still +contains the evicted-but-stale-entries-at-the-front pattern that the +shim is supposed to handle. +""" + +import pytest +import torch +from _model_test_utils import default_max_num_tokens + +from tensorrt_llm._torch.auto_deploy._compat import KvCacheConfig +from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import KVPagedResourceHandler +from tensorrt_llm._torch.auto_deploy.shim.ad_executor import _compute_window_local_view +from tensorrt_llm._torch.auto_deploy.shim.interface import CachedSequenceInterface + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + + +# Geometry shared across cases: SWA window of 64 tokens with tokens_per_block=32 +# gives 2 pages per window — small enough to exercise eviction boundaries +# without ballooning runtime. FULL_WINDOW doubles as the fixture's +# ``max_seq_len`` and therefore sets the SWA pool's per-sequence cap +# (``maxTokenNum = max(windowSize, maxSequenceLength)`` in kvCacheManager.cpp). +# It must exceed the longest scenario the tests construct (192-token prefill +# in ``test_helper_pre_eviction_long_prefill_passes_full_pages``) or the C++ +# side clamps allocation and the helper sees fewer live pages than expected. +SWA_WINDOW = 64 +FULL_WINDOW = 256 +TOKENS_PER_BLOCK = 32 +PAGES_PER_SWA_WINDOW = SWA_WINDOW // TOKENS_PER_BLOCK # 2 + + +@pytest.fixture +def two_window_interface(): + """A real CachedSequenceInterface hosting a dual-pool KVCacheManager. + + Layer 0 lives in an SWA window (W=64); layer 1 lives in a full window. + This is the Gemma-4 geometry — exactly the path Phase 2 fixes. + """ + interface = CachedSequenceInterface( + max_seq_len=FULL_WINDOW, + max_batch_size=2, + max_num_tokens=default_max_num_tokens(FULL_WINDOW, 2), + device="cuda", + kv_cache_config=KvCacheConfig( + tokens_per_block=TOKENS_PER_BLOCK, + max_tokens=1024, + free_gpu_memory_fraction=0.0, + ), + ) + interface.add_resource( + "kv_swa", KVPagedResourceHandler(4, 32, dtype=torch.float16, sliding_window=SWA_WINDOW) + ) + interface.add_resource("kv_full", KVPagedResourceHandler(4, 32, dtype=torch.float16)) + interface.initialize_resources() + return interface + + +def _add_request(manager, request_id: int, token_num: int): + """Add a dummy request that allocates pages for `token_num` tokens.""" + out = manager.add_dummy_requests([request_id], token_nums=[token_num]) + assert out is not None and len(out) == 1, "manager did not allocate the request" + return out[0] + + +def test_helper_no_eviction_passes_through(two_window_interface): + """Sanity guard: front_removed=0 reproduces the pre-Phase-2 view. + + Asserts `all_indices[:num_active]`, full window-capped + seq_len_with_cache, and the next slot as the spillover page. + """ + manager = two_window_interface.kv_cache_manager + req = _add_request(manager, request_id=42, token_num=SWA_WINDOW) + + all_indices = manager.get_cache_indices(req, window_size=SWA_WINDOW) + assert manager.get_num_front_blocks_removed(req.py_request_id, window_size=SWA_WINDOW) == 0 + + active_indices, extra_page, swc, lpl = _compute_window_local_view( + all_indices, + front_removed=0, + end_compute_i=SWA_WINDOW, + group_window=SWA_WINDOW, + tokens_per_block=TOKENS_PER_BLOCK, + ) + + assert active_indices == list(all_indices[:PAGES_PER_SWA_WINDOW]) + assert swc == SWA_WINDOW + assert lpl == TOKENS_PER_BLOCK + # No spillover slot when num_active fully consumes the historical list. + if len(all_indices) > PAGES_PER_SWA_WINDOW: + assert extra_page == all_indices[PAGES_PER_SWA_WINDOW] + else: + assert extra_page == -1 + + +@pytest.mark.parametrize("front_removed", [1, PAGES_PER_SWA_WINDOW]) +def test_helper_with_simulated_eviction_slices_correctly( + two_window_interface, monkeypatch, front_removed +): + """Eviction-aware slicing in window-local coordinates. + + With ``front_removed > 0`` the slice must start at ``front_removed`` and + ``seq_len_with_cache`` must equal the live (window-local) cache length: + ``end_compute_i - front_removed * tokens_per_block``. The C++ eviction + loop guarantees this stays within ``window + tokens_per_block``, so it + may exceed ``SWA_WINDOW`` by up to one spillover page — the helper does + not artificially clamp at ``SWA_WINDOW`` because doing so would discard + a live page's worth of data that the kernel still has to read. + + Simulates a request that has cumulatively seen + ``front_removed * tokens_per_block + SWA_WINDOW + 1`` tokens, i.e. + ``front_removed`` pages have been front-evicted and one spillover token + sits in the page after the window. + """ + manager = two_window_interface.kv_cache_manager + # Allocate enough pages to span (evicted + live + a spillover slot) so + # get_cache_indices returns a realistic-length list. + total_tokens = front_removed * TOKENS_PER_BLOCK + SWA_WINDOW + 1 + req = _add_request(manager, request_id=43, token_num=total_tokens) + + # Bump the eviction counter the way detachFrontBlock would (the C++ side + # bumps `mNumFrontBlocksRemovedPerWindow[ws]` without mutating + # mCacheBlockIds, so get_cache_indices STILL returns the full list and + # the shim must slice [front_removed:] off it). + real_query = manager.get_num_front_blocks_removed + + def fake_front_removed(request_id, window_size=None): + if request_id == req.py_request_id and window_size == SWA_WINDOW: + return front_removed + return real_query(request_id, window_size=window_size) + + monkeypatch.setattr(manager, "get_num_front_blocks_removed", fake_front_removed) + + all_indices = manager.get_cache_indices(req, window_size=SWA_WINDOW) + front = manager.get_num_front_blocks_removed(req.py_request_id, window_size=SWA_WINDOW) + assert front == front_removed + + end_compute_i = total_tokens + active_indices, extra_page, swc, lpl = _compute_window_local_view( + all_indices, + front_removed=front, + end_compute_i=end_compute_i, + group_window=SWA_WINDOW, + tokens_per_block=TOKENS_PER_BLOCK, + ) + + # Window-local cache length matches the C++ accounting. + expected_swc = end_compute_i - front_removed * TOKENS_PER_BLOCK + assert swc == expected_swc, ( + f"swc={swc} != expected {expected_swc} " + f"(end_compute_i={end_compute_i}, front_removed={front_removed})" + ) + # C++ eviction invariant: live cache ≤ window + one spillover page. + assert swc <= SWA_WINDOW + TOKENS_PER_BLOCK + + # Live slice starts at front_removed and covers every page needed to + # address swc tokens — including the spillover page when present. + expected_num_active = (expected_swc + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK + expected_slice = list(all_indices[front_removed : front_removed + expected_num_active]) + assert active_indices == expected_slice, ( + f"slice mismatch: front_removed={front_removed}, " + f"all_indices={list(all_indices)}, active={active_indices}" + ) + + # last_page_len matches the modular arithmetic on the local cache length. + expected_lpl = (swc - 1) % TOKENS_PER_BLOCK + 1 if swc > 0 else 0 + assert lpl == expected_lpl + + # Page table capacity covers the live cache length. + assert len(active_indices) * TOKENS_PER_BLOCK >= swc + + +def test_helper_caps_seq_len_with_cache_below_window(two_window_interface, monkeypatch): + """Below-window end_compute_i reports actual progress, not the window size. + + When end_compute_i is BELOW the window (early prefill / short request), + seq_len_with_cache reports the actual progress, not the window size — the + helper must not over-pad. + """ + manager = two_window_interface.kv_cache_manager + req = _add_request(manager, request_id=44, token_num=SWA_WINDOW) + + monkeypatch.setattr(manager, "get_num_front_blocks_removed", lambda req_id, window_size=None: 0) + all_indices = manager.get_cache_indices(req, window_size=SWA_WINDOW) + + short_end = TOKENS_PER_BLOCK + 5 # < SWA_WINDOW + _, _, swc, lpl = _compute_window_local_view( + all_indices, + front_removed=0, + end_compute_i=short_end, + group_window=SWA_WINDOW, + tokens_per_block=TOKENS_PER_BLOCK, + ) + assert swc == short_end + assert lpl == ((short_end - 1) % TOKENS_PER_BLOCK + 1) + + +def test_helper_pre_eviction_long_prefill_passes_full_pages(two_window_interface, monkeypatch): + """Long prefill that has not yet been front-evicted — hand FULL live pages. + + Captures the MMLU-on-Gemma3n regime: addSequenceBatch admits a prefill + longer than the SWA window and allocates blocks linearly for the full + prompt (kvCacheManager.cpp::addSequenceBatch comment "For SWA, blocks + are allocated linearly for the full prompt; out-of-window blocks are + only detached during generation in adjustBlocksIfNeeded"). At the + first forward pass ``front_removed == 0`` and the kernel needs every + allocated page plus ``seq_len_with_cache == end_compute_i`` so the + sliding-window mask is applied inside the kernel over the full prefill, + matching the contract validated by + ``test_long_prefill_sw_matches_sdpa``. + + Pre-fix the helper unconditionally clamped ``active_token_count`` to + ``group_window`` even when no eviction had fired, which truncated the + page list to the first ``window/page_size`` pages and silently + corrupted long prefills. + """ + manager = two_window_interface.kv_cache_manager + # End at 3x the window so the prefill clearly exceeds W on the SWA pool + # but front-eviction has not run yet. + prefill_len = SWA_WINDOW * 3 # 192 tokens + req = _add_request(manager, request_id=46, token_num=prefill_len) + monkeypatch.setattr(manager, "get_num_front_blocks_removed", lambda req_id, window_size=None: 0) + + all_indices = manager.get_cache_indices(req, window_size=SWA_WINDOW) + expected_pages = (prefill_len + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK + assert len(all_indices) >= expected_pages, ( + f"manager must over-allocate pages for SWA prefill: " + f"got {len(all_indices)}, expected >= {expected_pages}" + ) + + active_indices, extra_page, swc, lpl = _compute_window_local_view( + all_indices, + front_removed=0, + end_compute_i=prefill_len, + group_window=SWA_WINDOW, + tokens_per_block=TOKENS_PER_BLOCK, + ) + + # Hand the kernel every live page (not just the window's worth). + assert active_indices == list(all_indices[:expected_pages]), ( + f"expected full prefill page list, got {active_indices} vs all_indices={list(all_indices)}" + ) + # Window-local cache length equals the unclamped prefill length — the + # kernel will mask down to the window itself. + assert swc == prefill_len, f"swc={swc} should equal prefill_len={prefill_len}" + assert lpl == (prefill_len - 1) % TOKENS_PER_BLOCK + 1 + # Page table capacity covers the full prefill (the kernel writes 1 + # K/V token per query token). + assert len(active_indices) * TOKENS_PER_BLOCK >= swc + # extra_page only populated if the manager pre-allocated past the prefill. + if len(all_indices) > expected_pages: + assert extra_page == all_indices[expected_pages] + else: + assert extra_page == -1 + + +def test_helper_does_not_consume_evicted_extra_slot(two_window_interface, monkeypatch): + """extra_page must live AFTER the live region, never inside the evicted range. + + Pre-Phase-2 code that referenced ``all_indices[num_active]`` would land + inside the evicted region whenever front_removed > 0 — this test guards + against that. Under the live-page-derived view the live region already + covers every allocated page up to and including the spillover, so + ``extra_page == -1`` in the typical post-eviction case; we additionally + check that whenever extra_page IS populated, the index sits past the + evicted prefix. + """ + manager = two_window_interface.kv_cache_manager + front_removed = 2 + total_tokens = front_removed * TOKENS_PER_BLOCK + SWA_WINDOW + 1 + req = _add_request(manager, request_id=45, token_num=total_tokens) + monkeypatch.setattr( + manager, + "get_num_front_blocks_removed", + lambda req_id, window_size=None: front_removed, + ) + + all_indices = manager.get_cache_indices(req, window_size=SWA_WINDOW) + active_indices, extra_page, swc, _ = _compute_window_local_view( + all_indices, + front_removed=front_removed, + end_compute_i=total_tokens, + group_window=SWA_WINDOW, + tokens_per_block=TOKENS_PER_BLOCK, + ) + + # Live region exhausts every page the C++ side handed us, so the + # spillover slot beyond it is -1 (no allocated next page). + num_active = len(active_indices) + expected_extra_idx = front_removed + num_active + if len(all_indices) > expected_extra_idx: + assert extra_page == all_indices[expected_extra_idx] + # Crucially, the extra slot is NOT in the front-evicted range. + assert expected_extra_idx >= front_removed + else: + assert extra_page == -1 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py b/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py index 58bcc3985b06..91aca4a4f8a7 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py @@ -269,36 +269,10 @@ def test_initialize_resources_paged_only_creates_kv_cache_manager(paged_kv_cache assert not isinstance(interface.kv_cache_manager, MambaHybridCacheManager) -def test_initialize_resources_incompatible_kv_layers_can_be_unmanaged(paged_kv_cache_config): - """Compatible layers are managed; incompatible ones may fall back to unmanaged buffers.""" - interface = CachedSequenceInterface( - max_seq_len=128, - max_batch_size=4, - max_num_tokens=default_max_num_tokens(128, 4), - device="cuda", - kv_cache_config=paged_kv_cache_config, - requires_uniform_kv_caches=False, - ) - - managed_name = interface.add_resource( - "kv_cache_0", - KVPagedResourceHandler(8, 64, dtype=torch.float16, kv_layout="HND"), - ) - unmanaged_name = interface.add_resource( - "kv_cache_1", - KVPagedResourceHandler(8, 80, dtype=torch.float16, kv_layout="HND"), - ) - - interface.initialize_resources() - - assert managed_name not in interface._unmanaged_resources - assert unmanaged_name in interface._unmanaged_resources - - -def test_initialize_resources_incompatible_kv_layers_raise_when_uniform_managed_caches_required( +def test_initialize_resources_mixed_shape_pools_raise_when_uniform_managed_caches_required( paged_kv_cache_config, ): - """Strict interface configurations reject mixed managed and unmanaged KV layouts.""" + """Strict configurations reject more than one distinct window-pool.""" interface = CachedSequenceInterface( max_seq_len=128, max_batch_size=4, @@ -308,8 +282,10 @@ def test_initialize_resources_incompatible_kv_layers_raise_when_uniform_managed_ requires_uniform_kv_caches=True, ) - interface.add_resource("kv_cache_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) - interface.add_resource("kv_cache_1", KVPagedResourceHandler(8, 80, dtype=torch.float16)) + interface.add_resource("kv_cache_full", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource( + "kv_cache_swa", KVPagedResourceHandler(8, 80, dtype=torch.float16, sliding_window=32) + ) with pytest.raises(RuntimeError): interface.initialize_resources() @@ -875,6 +851,20 @@ def test_sequence_info_update_cache_information_resizes(): assert seq_info._input_buffer.get_capacity("cache_loc") >= expected_capacity +def test_sequence_info_update_cache_information_preserves_max_blocks(): + """Verify max_blocks_per_seq is always based on max_seq_len (not window).""" + seq_info = SequenceInfo( + max_seq_len=128, + max_batch_size=4, + max_num_tokens=default_max_num_tokens(128, 4), + tokens_per_block=32, + ) + + original_max_blocks = seq_info.max_blocks_per_seq # ceil(128 / 32) = 4 + seq_info.update_cache_information(num_blocks=100) + assert seq_info.max_blocks_per_seq == original_max_blocks + + def test_sequence_info_last_page_len_uses_tokens_per_block(): """Verify nest_sequences calculates last_page_len using tokens_per_block.""" seq_info = SequenceInfo( @@ -891,8 +881,8 @@ def test_sequence_info_last_page_len_uses_tokens_per_block(): input_ids, cu_seqlen=[0, 25], input_pos=[0], - cache_loc=[0, 1], # 2 pages - cu_num_pages=[0, 2], + cache_loc_per_pool=[[0, 1]], # 2 pages + cu_num_pages_per_pool=[[0, 2]], ) expected_last_page_len = (25 - 1) % 16 + 1 @@ -914,8 +904,8 @@ def test_sequence_info_page_assignments(): input_ids, cu_seqlen=[0, 10, 30], input_pos=[0, 0], - cache_loc=[0, 1, 2], # seq 0 has page 0, seq 1 has pages 1 and 2 - cu_num_pages=[0, 1, 3], + cache_loc_per_pool=[[0, 1, 2]], # seq 0 has page 0, seq 1 has pages 1 and 2 + cu_num_pages_per_pool=[[0, 1, 3]], ) cache_loc = seq_info.get_arg("cache_loc_host", truncate=True).tolist() @@ -1168,8 +1158,8 @@ def test_args_stored_to_input_buffer(): input_ids=[1, 2, 3], cu_seqlen=[0, 3], input_pos=[0], - cache_loc=[0], - cu_num_pages=[0, 1], + cache_loc_per_pool=[[0]], + cu_num_pages_per_pool=[[0, 1]], ) token_gather_indices = seq_info.get_arg("token_gather_indices", truncate=True) @@ -1183,8 +1173,8 @@ def test_args_stored_to_input_buffer(): input_ids=[1, 2, 3], cu_seqlen=[0, 3], input_pos=[0], - cache_loc=[0], - cu_num_pages=[0, 1], + cache_loc_per_pool=[[0]], + cu_num_pages_per_pool=[[0, 1]], gather_context_logits=True, ) @@ -1214,3 +1204,99 @@ def dummy_host_prepare(batch_info_host: torch.Tensor, cu_num_pages_host: torch.T assert "batch_info_host" in seq_info._active_host_prep_args assert "cu_num_pages_host" in seq_info._active_host_prep_args + + +# ============================================================================= +# Multi-Window (VSWA) KV Cache Tests — single unified KVCacheManager hosting +# multiple C++ pools via the per-window head_dim/dtype overrides. +# ============================================================================= + + +def test_identify_managed_kv_resources_single_window(): + """Single head_dim and no sliding window collapses into one window entry.""" + interface = CachedSequenceInterface( + max_seq_len=128, + max_batch_size=4, + max_num_tokens=default_max_num_tokens(128, 4), + device="cuda", + kv_cache_config=KvCacheConfig( + tokens_per_block=32, max_tokens=1024, free_gpu_memory_fraction=0.0 + ), + ) + interface.add_resource("kv_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource("kv_1", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + + kv_managed, pool_configurations = interface._identify_managed_kv_resources() + + assert len(kv_managed) == 2 + # No sliding_window → effective window = max_seq_len. + assert {pc.window_size: pc.head_dim for pc in pool_configurations} == {128: 64} + + +def test_identify_managed_kv_resources_dual_window_gemma4_pattern(): + """Gemma4-style mix: SWA layers with smaller head_dim + full-attn layer with larger head_dim.""" + interface = CachedSequenceInterface( + max_seq_len=128, + max_batch_size=4, + max_num_tokens=default_max_num_tokens(128, 4), + device="cuda", + kv_cache_config=KvCacheConfig( + tokens_per_block=32, max_tokens=1024, free_gpu_memory_fraction=0.0 + ), + ) + # SWA window: head_dim=64 + interface.add_resource( + "kv_0", KVPagedResourceHandler(8, 64, dtype=torch.float16, sliding_window=64) + ) + interface.add_resource( + "kv_1", KVPagedResourceHandler(8, 64, dtype=torch.float16, sliding_window=64) + ) + # Full-attention window: head_dim=128 + interface.add_resource("kv_2", KVPagedResourceHandler(4, 128, dtype=torch.float16)) + + kv_managed, pool_configurations = interface._identify_managed_kv_resources() + + assert len(kv_managed) == 3 + # Two windows, keyed by effective window size; full-attention falls back to max_seq_len. + assert {pc.window_size: pc.head_dim for pc in pool_configurations} == {64: 64, 128: 128} + + +def test_identify_managed_kv_resources_rejects_mixed_head_dim_in_same_window(): + """Layers sharing an effective window must agree on head_dim.""" + interface = CachedSequenceInterface( + max_seq_len=128, + max_batch_size=4, + max_num_tokens=default_max_num_tokens(128, 4), + device="cuda", + kv_cache_config=KvCacheConfig( + tokens_per_block=32, max_tokens=1024, free_gpu_memory_fraction=0.0 + ), + ) + # Both default to sliding_window=0 → effective window = max_seq_len. Different head_dims + # in the same window are not representable by one C++ pool. + interface.add_resource("kv_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource("kv_1", KVPagedResourceHandler(4, 128, dtype=torch.float16)) + + with pytest.raises(RuntimeError, match="head_dim"): + interface._identify_managed_kv_resources() + + +def test_single_window_creates_plain_kv_cache_manager(): + """One window creates a plain KVCacheManager with a single C++ pool.""" + interface = CachedSequenceInterface( + max_seq_len=128, + max_batch_size=4, + max_num_tokens=default_max_num_tokens(128, 4), + device="cuda", + kv_cache_config=KvCacheConfig( + tokens_per_block=32, max_tokens=1024, free_gpu_memory_fraction=0.0 + ), + ) + interface.add_resource("kv_0", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + interface.add_resource("kv_1", KVPagedResourceHandler(8, 64, dtype=torch.float16)) + + interface.initialize_resources() + + assert isinstance(interface.kv_cache_manager, KVCacheManager) + # Exactly one pool (max_seq_len, since sliding_window=0). + assert len(interface.kv_cache_manager.impl.pool_configurations) == 1 diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py b/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py index 2741e623dc60..a4ccba2610e2 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py @@ -175,7 +175,7 @@ def get_cache_indices(self, request): # Return many dummy page IDs; ADEngine will truncate as needed return list(range(1024)) - def get_batch_cache_indices(self, request_ids): + def get_batch_cache_indices(self, request_ids, layer_idx=None): return [list(range(1024)) for _ in request_ids] def get_num_kv_blocks(self, num_tokens: int) -> int: @@ -561,7 +561,7 @@ def get_cache_indices(self, request): return list(range(num_blocks)) return list(range(1024)) - def get_batch_cache_indices(self, request_ids): + def get_batch_cache_indices(self, request_ids, layer_idx=None): return [list(range(1024)) for _ in request_ids] def get_num_kv_blocks(self, num_tokens: int) -> int: diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache.py index 594e0f5a0d69..f2b3a0ac82ae 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache.py @@ -314,8 +314,8 @@ def _call_and_unnest(x, input_pos): x.flatten().tolist(), cu_seqlen=cu_seqlen, input_pos=input_pos, - cache_loc=list(range(bs)), - cu_num_pages=list(range(bs + 1)), + cache_loc_per_pool=[list(range(bs))], + cu_num_pages_per_pool=[list(range(bs + 1))], slot_idx=list(range(bs)), gather_context_logits=True, ) @@ -833,3 +833,379 @@ def test_insert_cached_mla_attention_preserves_non_flashinfer_backend(): backend = InsertCachedMLAAttention.resolve_backend_for_node("torch_mla", attn_node) assert backend == "torch_mla" + + +# ============================================================================= +# Sliding Window KV Cache Integration Tests +# ============================================================================= + + +class SlidingWindowGQA(GQAWithSdpaAndEmbedding): + """GQA model with a configurable per-layer sliding_window for testing.""" + + def __init__( + self, + num_attention_heads: int, + hidden_size: int, + num_key_value_heads: int, + vocab_size: int = 1000, + sliding_window: Optional[int] = None, + ): + super().__init__(num_attention_heads, hidden_size, num_key_value_heads, vocab_size) + self._sliding_window = sliding_window + + @torch.no_grad() + def forward( + self, input_ids: torch.Tensor, position_ids: Optional[torch.Tensor] = None + ) -> torch.Tensor: + x = self.embed_tokens(input_ids) + b, s, _ = x.shape + q = self.q_proj(x).view(b, s, self.num_heads, self.head_dim) + k = self.k_proj(x).view(b, s, self.num_kv_heads, self.head_dim) + v = self.v_proj(x).view(b, s, self.num_kv_heads, self.head_dim) + attn_output = torch.ops.auto_deploy.torch_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0, + is_causal=True, + scale=None, + sinks=None, + sliding_window=self._sliding_window, + logit_cap=None, + layout="bsnd", + ) + return self.o_proj(attn_output.reshape(b, s, -1)) + + +class VSWAModel(nn.Module): + """Model with two attention layers: one sliding-window, one full-attention (VSWA).""" + + def __init__( + self, + num_attention_heads: int, + hidden_size: int, + num_key_value_heads: int, + vocab_size: int = 1000, + sliding_window: int = 32, + ): + super().__init__() + self.num_heads = num_attention_heads + self.num_kv_heads = num_key_value_heads + self.head_dim = hidden_size // num_attention_heads + self.embed_tokens = nn.Embedding(vocab_size, hidden_size) + self.q_proj_0 = nn.Linear(hidden_size, hidden_size, bias=False) + self.k_proj_0 = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=False) + self.v_proj_0 = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=False) + self.o_proj_0 = nn.Linear(hidden_size, hidden_size, bias=False) + self.q_proj_1 = nn.Linear(hidden_size, hidden_size, bias=False) + self.k_proj_1 = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=False) + self.v_proj_1 = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=False) + self.o_proj_1 = nn.Linear(hidden_size, hidden_size, bias=False) + self._sliding_window = sliding_window + + @torch.no_grad() + def forward( + self, input_ids: torch.Tensor, position_ids: Optional[torch.Tensor] = None + ) -> torch.Tensor: + x = self.embed_tokens(input_ids) + b, s, _ = x.shape + + # Layer 0: sliding window attention + q0 = self.q_proj_0(x).view(b, s, self.num_heads, self.head_dim) + k0 = self.k_proj_0(x).view(b, s, self.num_kv_heads, self.head_dim) + v0 = self.v_proj_0(x).view(b, s, self.num_kv_heads, self.head_dim) + a0 = torch.ops.auto_deploy.torch_attention( + q0, + k0, + v0, + attn_mask=None, + dropout_p=0.0, + is_causal=True, + scale=None, + sinks=None, + sliding_window=self._sliding_window, + logit_cap=None, + layout="bsnd", + ) + x = x + self.o_proj_0(a0.reshape(b, s, -1)) + + # Layer 1: full attention (no sliding window) + q1 = self.q_proj_1(x).view(b, s, self.num_heads, self.head_dim) + k1 = self.k_proj_1(x).view(b, s, self.num_kv_heads, self.head_dim) + v1 = self.v_proj_1(x).view(b, s, self.num_kv_heads, self.head_dim) + a1 = torch.ops.auto_deploy.torch_attention( + q1, + k1, + v1, + attn_mask=None, + dropout_p=0.0, + is_causal=True, + scale=None, + sinks=None, + sliding_window=None, + logit_cap=None, + layout="bsnd", + ) + return x + self.o_proj_1(a1.reshape(b, s, -1)) + + +def _build_optimizer_with_backend(model, backend="triton_paged"): + """Helper to create an InferenceOptimizer for testing.""" + return InferenceOptimizer( + DummyFactory(model, cache_config_updates={}), + { + "build_model": { + "stage": "factory", + "run_per_gm": False, + "device": "cuda", + "run_graph_cleanup": False, + "requires_clean_graph": False, + }, + "export_to_gm": { + "stage": "export", + "strict": False, + "run_per_gm": False, + "clone_state_dict": True, + "run_graph_cleanup": False, + "requires_clean_graph": False, + }, + "cleanup_input_constraints": { + "stage": "post_export", + }, + "insert_cached_attention": { + "stage": "cache_init", + "backend": backend, + }, + }, + ) + + +@torch.inference_mode() +def test_insert_cached_attention_no_sliding_window_leaves_config_unchanged(): + """Verify insert_cached_attention does not set max_attention_window for non-SWA models.""" + max_seq_len = 128 + batch_size = 4 + + kv_cache_config = KvCacheConfig( + tokens_per_block=max_seq_len, + max_tokens=batch_size * max_seq_len, + free_gpu_memory_fraction=0.0, + ) + cm = CachedSequenceInterface( + max_seq_len=max_seq_len, + max_batch_size=batch_size, + max_num_tokens=default_max_num_tokens(max_seq_len, batch_size), + device="cuda", + kv_cache_config=kv_cache_config, + ) + + model = GQAWithSdpaAndEmbedding( + num_attention_heads=8, + hidden_size=512, + num_key_value_heads=8, + ).to(dtype=torch.float16, device="cuda") + + optimizer = _build_optimizer_with_backend(model) + optimizer(cm) + + assert cm.kv_cache_config.max_attention_window is None + + +@torch.inference_mode() +def test_insert_cached_attention_respects_user_override(): + """Verify insert_cached_attention does not overwrite user-set max_attention_window.""" + max_seq_len = 128 + batch_size = 4 + user_window = [64] + + kv_cache_config = KvCacheConfig( + tokens_per_block=max_seq_len, + max_tokens=batch_size * max_seq_len, + free_gpu_memory_fraction=0.0, + max_attention_window=user_window, + ) + cm = CachedSequenceInterface( + max_seq_len=max_seq_len, + max_batch_size=batch_size, + max_num_tokens=default_max_num_tokens(max_seq_len, batch_size), + device="cuda", + kv_cache_config=kv_cache_config, + ) + + model = SlidingWindowGQA( + num_attention_heads=8, + hidden_size=512, + num_key_value_heads=8, + sliding_window=32, + ).to(dtype=torch.float16, device="cuda") + + optimizer = _build_optimizer_with_backend(model) + optimizer(cm) + + # User-provided value must be preserved + assert cm.kv_cache_config.max_attention_window == user_window + + +@torch.inference_mode() +def test_insert_cached_attention_vswa_preserves_per_layer_windows(): + """Verify VSWA model preserves per-layer window sizes for proportional allocation.""" + sliding_window = 32 + max_seq_len = 128 + batch_size = 4 + + kv_cache_config = KvCacheConfig( + tokens_per_block=32, + max_tokens=batch_size * max_seq_len, + free_gpu_memory_fraction=0.0, + ) + cm = CachedSequenceInterface( + max_seq_len=max_seq_len, + max_batch_size=batch_size, + max_num_tokens=default_max_num_tokens(max_seq_len, batch_size), + device="cuda", + kv_cache_config=kv_cache_config, + ) + + model = VSWAModel( + num_attention_heads=8, + hidden_size=512, + num_key_value_heads=8, + sliding_window=sliding_window, + ).to(dtype=torch.float16, device="cuda") + + optimizer = _build_optimizer_with_backend(model) + optimizer(cm) + + # VSWA preserves per-layer windows: [32, 128] (not collapsed) + assert cm.kv_cache_config.max_attention_window is not None + assert len(cm.kv_cache_config.max_attention_window) == 2 + assert cm.kv_cache_config.max_attention_window == [sliding_window, max_seq_len] + + # Window groups should be registered on SequenceInfo + assert cm.info.num_window_groups == 2 + assert cm.info.window_groups == [sliding_window, max_seq_len] + assert cm.info.window_group_map == {sliding_window: 0, max_seq_len: 1} + + +@torch.inference_mode() +def test_kv_cache_manager_initialized_with_sliding_window(): + """Verify KVCacheManager receives max_attention_window_vec from SWA model. + + Runs the full insert_cached_attention + initialize_cache pipeline. + """ + import math + + sliding_window = 32 + max_seq_len = 128 + batch_size = 4 + tokens_per_block = 16 + + kv_cache_config = KvCacheConfig( + tokens_per_block=tokens_per_block, + max_tokens=batch_size * max_seq_len, + free_gpu_memory_fraction=0.0, + ) + cm = CachedSequenceInterface( + max_seq_len=max_seq_len, + max_batch_size=batch_size, + max_num_tokens=default_max_num_tokens(max_seq_len, batch_size), + device="cuda", + kv_cache_config=kv_cache_config, + ) + + model = SlidingWindowGQA( + num_attention_heads=8, + hidden_size=512, + num_key_value_heads=8, + sliding_window=sliding_window, + ).to(dtype=torch.float16, device="cuda") + + # Run insert_cached_attention + initialize_cache + optimizer = InferenceOptimizer( + DummyFactory(model, cache_config_updates={}), + { + "build_model": { + "stage": "factory", + "run_per_gm": False, + "device": "cuda", + "run_graph_cleanup": False, + "requires_clean_graph": False, + }, + "export_to_gm": { + "stage": "export", + "strict": False, + "run_per_gm": False, + "clone_state_dict": True, + "run_graph_cleanup": False, + "requires_clean_graph": False, + }, + "cleanup_input_constraints": { + "stage": "post_export", + }, + "insert_cached_attention": { + "stage": "cache_init", + "backend": "triton_paged", + }, + "initialize_cache": { + "stage": "cache_init", + "run_per_gm": False, + }, + }, + ) + optimizer(cm) + + # KVCacheManager should exist and carry the window vector + mgr = cm.kv_cache_manager + assert mgr.max_attention_window_vec == [sliding_window] + + # max_blocks_per_seq is based on max_seq_len (not window) because sequences + # temporarily need full blocks during prefill; SWA eviction frees them during decode. + expected_max_blocks = math.ceil(max_seq_len / tokens_per_block) + assert cm.info.max_blocks_per_seq == expected_max_blocks + + +# ============================================================================= +# VSWA SequenceInfo and Graph Wiring Tests +# ============================================================================= + + +@torch.inference_mode() +def test_sequence_info_register_window_groups(): + """Verify register_window_groups creates per-group tensors in InputBuffer.""" + from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import SequenceInfo + + max_seq_len = 128 + batch_size = 4 + tokens_per_block = 16 + + seq_info = SequenceInfo( + max_seq_len=max_seq_len, + max_batch_size=batch_size, + max_num_tokens=default_max_num_tokens(max_seq_len, batch_size), + tokens_per_block=tokens_per_block, + ) + + # Before registration: no window groups + assert seq_info.num_window_groups == 0 + assert seq_info.window_groups == [] + + # Register two groups: [32, 128] + seq_info.register_window_groups([32, 128]) + + assert seq_info.num_window_groups == 2 + assert seq_info.window_groups == [32, 128] + assert seq_info.window_group_map == {32: 0, 128: 1} + + # Group 0 reuses existing tensors (cache_loc, cu_num_pages, etc.) + # Group 1 gets new tensors + assert "cache_loc_g1" in seq_info.available_args + assert "cu_num_pages_g1" in seq_info.available_args + assert "cu_num_pages_g1_host" in seq_info.available_args + assert "last_page_len_g1" in seq_info.available_args + assert "last_page_len_g1_host" in seq_info.available_args + assert "extra_page_per_seq_g1" in seq_info.available_args + + # Group 0 names should NOT appear with _g0 suffix + assert "cache_loc_g0" not in seq_info.available_args diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kvcache_vswa_metadata.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kvcache_vswa_metadata.py new file mode 100644 index 000000000000..6e9043bc96b7 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kvcache_vswa_metadata.py @@ -0,0 +1,213 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Graph-shape tests for the VSWA per-group metadata wiring in the kvcache transform. + +The kvcache transform builds a per-window page table layout. Phase-1 of the +Gemma-4 enablement (PR #13745) only routed *standard* metadata (cache_loc, +cu_num_pages, last_page_len) per group; the *extra*-metadata op (the host-side +prepare that produces e.g. update_paged_kv_cache write positions) was built +once from group-0 inputs and reused for every layer. With seq_len_with_cache +added to the swappable set (Phase 2), the extra-metadata op MUST also run per +group, otherwise non-zero-group layers silently consume group-0's +seq_len_with_cache (window-uncapped) → wrong write positions under eviction. + +These tests are graph-shape only — no GPU, no kernel execution. +""" + +import pytest +import torch + +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.shim.interface import CachedSequenceInterface +from tensorrt_llm._torch.auto_deploy.transform.interface import SharedConfig, Stages +from tensorrt_llm._torch.auto_deploy.transform.library.kvcache import ( + InsertCachedAttention, + InsertCachedAttentionConfig, +) + + +class _TwoWindowModule(torch.nn.Module): + """A tiny two-layer model with one SWA layer and one full-attention layer. + + Layer 0 is SWA (sliding_window=256), layer 1 is full attention. This + forces the transform to create two window groups. + """ + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + qkv = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], 2, 4) + swa_layer = torch.ops.auto_deploy.torch_attention( + qkv, + qkv, + qkv, + None, # attn_mask + 0.0, # dropout_p + True, # is_causal + 1.0, # scale + None, # sinks + 256, # sliding_window <-- SWA group + None, # logit_cap + "bsnd", + 0, # layer_idx + ) + full_layer = torch.ops.auto_deploy.torch_attention( + qkv, + qkv, + qkv, + None, + 0.0, + True, + 1.0, + None, + None, # sliding_window=None <-- full-attention group + None, + "bsnd", + 1, + ) + return swa_layer + full_layer + + +def _run_transform(backend: str): + module = _TwoWindowModule().eval() + gm = torch_export_to_gm(module, (torch.randn(1, 4, 8),)) + cm = CachedSequenceInterface( + max_seq_len=512, + max_batch_size=2, + max_num_tokens=512, + device="cpu", + ) + transform = InsertCachedAttention( + InsertCachedAttentionConfig(stage=Stages.CACHE_INIT, backend=backend) + ) + gm, info = transform._apply(gm, cm, factory=None, shared_config=SharedConfig()) + return gm, info, cm + + +def _placeholder_names(gm) -> list: + return [node.target for node in gm.graph.nodes if node.op == "placeholder"] + + +def test_vswa_registers_per_group_seq_len_with_cache_placeholders(): + """Group 1 must have its own seq_len_with_cache_g1{_host} placeholders. + + This is required so the SWA group's window-capped value reaches the + kernel and the prepare-extra-metadata op. + """ + gm, info, cm = _run_transform(backend="triton_paged") + assert info.num_matches == 2, "expected 2 cached-attention insertions" + + names = _placeholder_names(gm) + # Group 0 keeps the unsuffixed name; group 1 must have its own _g1 variant. + assert "seq_len_with_cache_host" in names + assert "seq_len_with_cache_g1_host" in names, ( + "seq_len_with_cache must be in vswa_swappable_bases — without per-group " + "wiring, SWA layers silently consume group-0's window-uncapped value." + ) + + +def test_vswa_extra_metadata_is_per_group(): + """Prepare-extra-metadata must be invoked once per window group. + + The Triton backend's prepare-extra op consumes seq_len_with_cache. With + seq_len_with_cache routed per-group, the transform MUST invoke the + extra-op once per group; otherwise non-zero-group layers route through + group-0's prepare-extra output and the swappable seq_len_with_cache_g1 + placeholder is dead. + + This is the regression guard against the pre-fix transform behavior + described in the autodeploy-kvcache-vswa-meta-nodes-extra backlog. + """ + gm, info, cm = _run_transform(backend="triton_paged") + + prep_meta_op = torch.ops.auto_deploy.triton_paged_prepare_metadata.default + prep_calls = [ + node + for node in gm.graph.nodes + if node.op == "call_function" and node.target == prep_meta_op + ] + assert len(prep_calls) == 2, ( + f"expected one prepare-extra call per group (2 groups), got {len(prep_calls)}" + ) + + # The two calls must consume DIFFERENT seq_len_with_cache placeholders: + # one from group 0 (unsuffixed), one from group 1 (suffixed). + swc_inputs_by_call = [] + for call in prep_calls: + swc_input = None + for arg in call.args: + if ( + isinstance(arg, torch.fx.Node) + and arg.op == "placeholder" + and "seq_len_with_cache" in str(arg.target) + ): + swc_input = str(arg.target) + break + assert swc_input is not None, ( + f"prepare-extra call {call.format_node()} did not consume any " + f"seq_len_with_cache placeholder" + ) + swc_inputs_by_call.append(swc_input) + + assert len(set(swc_inputs_by_call)) == 2, ( + f"both prepare-extra calls consume the same seq_len_with_cache " + f"input ({swc_inputs_by_call}); per-group routing failed" + ) + + +def test_vswa_each_layer_routes_to_its_groups_extra_metadata(): + """Each layer's cached-attn must wire to its own group's prepare-extra outputs. + + The cached-attn call for group 1 (the SWA layer) must consume the + group-1 prepare-extra outputs, not group 0's. + """ + gm, info, cm = _run_transform(backend="triton_paged") + + cached_op = torch.ops.auto_deploy.triton_paged_mha_with_cache.default + cached_calls = [ + node for node in gm.graph.nodes if node.op == "call_function" and node.target == cached_op + ] + assert len(cached_calls) == 2 + + # For each cached-attn call, find which seq_len_with_cache placeholder its + # extra-metadata args ultimately depend on. We do a simple recursive walk + # through args (bounded by graph size) since the prepare-extra outputs are + # getitem nodes of a call_function whose args include the placeholder. + def find_swc_dep(node, visited=None): + if visited is None: + visited = set() + if id(node) in visited: + return None + visited.add(id(node)) + if isinstance(node, torch.fx.Node): + if node.op == "placeholder" and "seq_len_with_cache" in str(node.target): + return str(node.target) + if node.op in ("call_function", "call_method"): + for arg in list(node.args) + list(node.kwargs.values()): + result = find_swc_dep(arg, visited) + if result is not None: + return result + return None + + deps_per_call = [] + for call in cached_calls: + # The first three args are q/k/v; standard metadata follows. Walk all + # args until we find a seq_len_with_cache placeholder dependency. + dep = None + for arg in call.args: + dep = find_swc_dep(arg) + if dep is not None: + break + assert dep is not None, "no seq_len_with_cache dependency found" + deps_per_call.append(dep) + + # The two cached-attn calls (one per layer = one per group) must each route + # to a distinct seq_len_with_cache placeholder. + assert len(set(deps_per_call)) == 2, ( + f"both cached-attn calls depend on the same seq_len_with_cache " + f"placeholder ({deps_per_call}); per-group wiring failed" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])