diff --git a/benchmarks/cpp/disaggServerBenchmark.cpp b/benchmarks/cpp/disaggServerBenchmark.cpp index 38f74314d099..057e7898afeb 100644 --- a/benchmarks/cpp/disaggServerBenchmark.cpp +++ b/benchmarks/cpp/disaggServerBenchmark.cpp @@ -610,7 +610,7 @@ class DisaggExecutorServer texec::KvCacheConfig kvCacheConfig(benchmarkParams.enableBlockReuse, benchmarkParams.maxTokensInPagedKvCache, benchmarkParams.maxAttentionWindowVec, benchmarkParams.sinkTokenLength, benchmarkParams.freeGpuMemoryFractions.at(in), - benchmarkParams.kvHostCacheSize, benchmarkParams.kvOnboardBlocks); + benchmarkParams.kvHostCacheSize); texec::ExtendedRuntimePerfKnobConfig extendedRuntimePerfKnobConfig(benchmarkParams.multiBlockMode, benchmarkParams.enableContextFMHAFP32Acc, benchmarkParams.cudaGraphMode, benchmarkParams.cudaGraphCacheSize); @@ -1213,8 +1213,6 @@ int main(int argc, char* argv[]) options.add_options()("kv_host_cache_bytes", "Size of secondary memory pool used for offloading kv cache blocks (in bytes).", cxxopts::value()->default_value("0")); - options.add_options()("kv_onboard_blocks", "If offloaded blocks should be onboarded to primary memory before reuse", - cxxopts::value()->default_value("true")); options.add_options()( "max_prompt_len", "Truncate all prompts from dataset to the length specified.", cxxopts::value()); options.add_options()("gpu_weights_percent", @@ -1482,10 +1480,6 @@ int main(int argc, char* argv[]) TLLM_CHECK_WITH_INFO( benchmarkParams.kvHostCacheSize == false, "Currently disaggServer don't support kv_host_cache!"); - // Argument: If offloaded blocks should be onboarded to primary memory before they are reused. - benchmarkParams.kvOnboardBlocks = result["kv_onboard_blocks"].as(); - TLLM_CHECK_WITH_INFO( - benchmarkParams.kvOnboardBlocks == true, "Currently disaggServer don't support kv_onboard_blocks =false!"); // Argument: Medusa choices for the Medusa speculative decoding. if (result.count("medusa_choices")) { diff --git a/benchmarks/cpp/gptManagerBenchmark.cpp b/benchmarks/cpp/gptManagerBenchmark.cpp index 1dbbc518012e..20a1f1c343da 100644 --- a/benchmarks/cpp/gptManagerBenchmark.cpp +++ b/benchmarks/cpp/gptManagerBenchmark.cpp @@ -630,7 +630,7 @@ class ExecutorServer texec::KvCacheConfig kvCacheConfig(benchmarkParams.enableBlockReuse, benchmarkParams.maxTokensInPagedKvCache, benchmarkParams.maxAttentionWindowVec, benchmarkParams.sinkTokenLength, - benchmarkParams.freeGpuMemoryFraction, benchmarkParams.kvHostCacheSize, benchmarkParams.kvOnboardBlocks, + benchmarkParams.freeGpuMemoryFraction, benchmarkParams.kvHostCacheSize, benchmarkParams.crossKvCacheFraction); texec::PeftCacheConfig peftCacheConfig(0, benchmarkParams.loraDeviceNumModLayers, 8, 64, 4, 4, 4, 24, 8, std::nullopt, benchmarkParams.loraHostCacheSize); @@ -1355,9 +1355,6 @@ int main(int argc, char* argv[]) // Argument: How many KV cache blocks (as fraction of number of GPU kv cache blocks). benchmarkParams.kvHostCacheSize = result["kv_host_cache_bytes"].as(); - // Argument: If offloaded blocks should be onboarded to primary memory before they are reused. - benchmarkParams.kvOnboardBlocks = result["kv_onboard_blocks"].as(); - // Argument: Medusa choices for the Medusa speculative decoding. if (result.count("medusa_choices")) { diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index c8400d481f32..4d6735ae54ff 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -584,7 +584,7 @@ class WindowBlockManager std::vector const& managedLayers, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, bool isSWA, SizeType32 blocksInPrimaryPool, SizeType32 blocksInSecondaryPool, SizeType32 maxNumSequences, std::shared_ptr stream, - bool onboardBlocks, CacheType cacheType, std::optional secondaryOffloadMinPriority, + CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, std::shared_ptr loopbackAgent = nullptr); @@ -899,8 +899,6 @@ class WindowBlockManager // getPoolLayerIdx std::unordered_map mLayerToIndexWithinPool; - // Whether offloaded blocks should be onboarded before reuse. - bool mOnboardBlocks; // Buffer manager runtime::BufferManager mBufferManager; @@ -976,7 +974,7 @@ class BlockManager CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, std::optional const& tempAttentionWindowInputs, nvinfer1::DataType dtype, - SizeType32 sinkBubbleLength, bool onboardBlocks, CacheType cacheType = CacheType::kSELF, + SizeType32 sinkBubbleLength, CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, bool copyOnPartialReuse = true, @@ -1584,7 +1582,7 @@ class KVCacheManager : public BaseKVCacheManager std::vector const& maxAttentionWindowVec, std::optional const& tempAttentionWindowInputs, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, SizeType32 maxSequenceLength, bool enableBlockReuse = false, - bool onboardBlocks = true, CacheType cacheType = CacheType::kSELF, + CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, bool copyOnpartialReuse = true, @@ -1595,7 +1593,7 @@ class KVCacheManager : public BaseKVCacheManager std::vector const& maxAttentionWindowVec, std::optional const& tempAttentionWindowInputs, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, SizeType32 maxSequenceLength, bool enableBlockReuse = false, - bool onboardBlocks = true, CacheType cacheType = CacheType::kSELF, + CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, bool copyOnpartialReuse = true, @@ -1606,7 +1604,7 @@ class KVCacheManager : public BaseKVCacheManager std::vector const& maxAttentionWindowVec, std::optional const& tempAttentionWindowInputs, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, SizeType32 maxSequenceLength, bool enableBlockReuse = true, - bool onboardBlocks = true, CacheType cacheType = CacheType::kSELF, + CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, bool copyOnpartialReuse = true, @@ -1617,8 +1615,7 @@ class KVCacheManager : public BaseKVCacheManager std::vector const& maxAttentionWindowVec, std::optional const& tempAttentionWindowInputs, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, SizeType32 maxSequenceLength, bool enableBlockReuse = false, - bool onboardBlocks = true, CacheType cacheType = CacheType::kSELF, bool enablePartialReuse = true, - bool copyOnpartialReuse = true); + CacheType cacheType = CacheType::kSELF, bool enablePartialReuse = true, bool copyOnpartialReuse = true); ~KVCacheManager() override = default; diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index b5769177b289..7df24ebf4b14 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -1003,7 +1003,7 @@ class KvCacheConfig std::optional> const& maxAttentionWindowVec = std::nullopt, std::optional const& sinkTokenLength = std::nullopt, std::optional const& freeGpuMemoryFraction = std::nullopt, - std::optional const& hostCacheSize = std::nullopt, bool onboardBlocks = true, + std::optional const& hostCacheSize = std::nullopt, std::optional const& crossKvCacheFraction = std::nullopt, std::optional secondaryOffloadMinPriority = std::nullopt, size_t eventBufferMaxSize = 0, bool enablePartialReuse = true, bool copyOnPartialReuse = true, bool useUvm = false, @@ -1020,7 +1020,6 @@ class KvCacheConfig [[nodiscard]] std::optional getFreeGpuMemoryFraction() const; [[nodiscard]] std::optional getCrossKvCacheFraction() const; [[nodiscard]] std::optional getHostCacheSize() const; - [[nodiscard]] bool getOnboardBlocks() const; [[nodiscard]] std::optional getSecondaryOffloadMinPriority() const; [[nodiscard]] size_t getEventBufferMaxSize() const; [[nodiscard]] bool getUseUvm() const; @@ -1036,7 +1035,6 @@ class KvCacheConfig void setFreeGpuMemoryFraction(FloatType freeGpuMemoryFraction); void setCrossKvCacheFraction(FloatType crossKvCacheFraction); void setHostCacheSize(size_t hostCacheSize); - void setOnboardBlocks(bool onboardBlocks); void setSecondaryOffloadMinPriority(std::optional secondaryOffloadMinPriority); void setEventBufferMaxSize(size_t eventBufferMaxSize); void setUseUvm(bool useUvm); @@ -1080,9 +1078,6 @@ class KvCacheConfig /// Having a secondary memory pool increases KV cache block reuse potential. std::optional mHostCacheSize; - /// @brief Controls whether offloaded blocks should be onboarded back into primary memory before being reused. - bool mOnboardBlocks; - /// @brief Only blocks with priority > mSecondaryOfflineMinPriority can be offloaded to secondary memory. std::optional mSecondaryOffloadMinPriority; diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 0f8ee70ea349..9273e9362c50 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -540,7 +540,7 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si std::shared_ptr stream, SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, std::optional const& tempAttentionWindowInputs, nvinfer1::DataType dtype, - SizeType32 sinkBubbleLength, bool onboardBlocks, CacheType cacheType, + SizeType32 sinkBubbleLength, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, @@ -583,8 +583,8 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si TLLM_CHECK(allottedPrimaryBlocks > 0); // You can't have a model with negative primary blocks... mWindowBlockManagers.try_emplace(windowSize, dtype, windowSize, layersWithWindowSize, numKvHeadsPerLayer, sizePerHead, tokensPerBlock, /*isSWA=*/windowSize < maxSequenceLength, allottedPrimaryBlocks, - allottedSecondaryBlocks, maxNumSequences, stream, onboardBlocks, cacheType, secondaryOffloadMinPriority, - mEventManager, enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, mLoopbackAgent); + allottedSecondaryBlocks, maxNumSequences, stream, cacheType, secondaryOffloadMinPriority, mEventManager, + enablePartialReuse, copyOnPartialReuse, kvCacheConnectorManager, mLoopbackAgent); } auto const numAllPools = getNumPools(); @@ -638,7 +638,7 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind std::vector const& managedLayers, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, bool isSWA, SizeType32 blocksInPrimaryPool, SizeType32 blocksInSecondaryPool, SizeType32 maxNumSequences, std::shared_ptr stream, - bool onboardBlocks, CacheType cacheType, std::optional secondaryOffloadMinPriority, + CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, std::shared_ptr loopbackAgent) @@ -646,7 +646,6 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind , mWindowSize{windowSize} , mNumPrimaryBlocks{blocksInPrimaryPool} , mNumSecondaryBlocks{blocksInSecondaryPool} - , mOnboardBlocks(onboardBlocks) , mBufferManager{std::move(stream)} , mSchedulingNumFreeBlocks{0} , mTokensPerBlock{tokensPerBlock} @@ -937,9 +936,7 @@ BlockPtr WindowBlockManager::getFreeBlock(GenerationRequest& sequence, executor: // 1. Block contains state (evidenced by presence of tokens) // 2. Eviction policy indicated block can be offloaded // 3. At least one free block in secondary memory - // 4. Onboarding is enabled (allowing block to be brought back into primary) - if (!block->getUniqueTokens().empty() && canOffload && mEvictionPolicy->getNumFreeBlocks(kSecondaryLevel) > 0 - && mOnboardBlocks) + if (!block->getUniqueTokens().empty() && canOffload && mEvictionPolicy->getNumFreeBlocks(kSecondaryLevel) > 0) { // Offload block in primary memory before repurposing auto offloadBlock = std::get<0>(mEvictionPolicy->getFreeBlock(kSecondaryLevel)); @@ -1034,7 +1031,7 @@ void BlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr const& off void WindowBlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr const& offloadBlock, executor::KvCacheTransferMode mode, std::string const& directory) { - if (mOnboardBlocks && !offloadBlock->isPrimary()) + if (!offloadBlock->isPrimary()) { auto block = getFreeBlock( sequence, executor::KvCacheRetentionConfig::kDefaultRetentionPriority, std::nullopt, mode, directory); @@ -1068,7 +1065,7 @@ void WindowBlockManager::offloadBlock( // block is useful or not and may just lead to more traffic instead. // The ideal way of this is to dedicate the offloading of the block // to the eviction policy. - if (mOnboardBlocks && block->isPrimary()) + if (block->isPrimary()) { // Offload block in primary memory before repurposing auto offloadBlock = std::get<0>(mEvictionPolicy->getFreeBlock(kSecondaryLevel)); @@ -1898,11 +1895,11 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, std::optional const& tempAttentionWindowInputs, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, runtime::SizeType32 maxSequenceLength, bool enableBlockReuse, - bool onboardBlocks, CacheType cacheType, bool enablePartialReuse, bool copyOnPartialReuse) + CacheType cacheType, bool enablePartialReuse, bool copyOnPartialReuse) : KVCacheManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, tempAttentionWindowInputs, dtype, sinkTokenLength, std::make_shared(reinterpret_cast(stream)), maxSequenceLength, - enableBlockReuse, onboardBlocks, cacheType, std::nullopt, nullptr, enablePartialReuse, copyOnPartialReuse) + enableBlockReuse, cacheType, std::nullopt, nullptr, enablePartialReuse, copyOnPartialReuse) { } @@ -1911,14 +1908,14 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, std::optional const& tempAttentionWindowInputs, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, runtime::SizeType32 maxSequenceLength, bool enableBlockReuse, - bool onboardBlocks, CacheType cacheType, std::optional secondaryOffloadMinPriority, + CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager) : KVCacheManager(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, tempAttentionWindowInputs, dtype, sinkTokenLength, std::make_shared(reinterpret_cast(stream)), maxSequenceLength, - enableBlockReuse, onboardBlocks, cacheType, secondaryOffloadMinPriority, eventManager, enablePartialReuse, - copyOnPartialReuse, kvCacheConnectorManager) + enableBlockReuse, cacheType, secondaryOffloadMinPriority, eventManager, enablePartialReuse, copyOnPartialReuse, + kvCacheConnectorManager) { } @@ -1927,7 +1924,7 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, std::optional const& tempAttentionWindowInputs, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, bool enableBlockReuse, - bool onboardBlocks, CacheType cacheType, std::optional secondaryOffloadMinPriority, + CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager) : mMaxBeamWidth(maxBeamWidth) @@ -1938,8 +1935,8 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer , mSinkBlockTokenLength(mSinkBubbleLength + sinkTokenLength) , mBlockManager(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, std::move(stream), maxSequenceLength, maxBeamWidth, maxAttentionWindowVec, tempAttentionWindowInputs, dtype, - mSinkBubbleLength, onboardBlocks, cacheType, secondaryOffloadMinPriority, std::move(eventManager), - enablePartialReuse, copyOnPartialReuse, std::move(kvCacheConnectorManager)) + mSinkBubbleLength, cacheType, secondaryOffloadMinPriority, std::move(eventManager), enablePartialReuse, + copyOnPartialReuse, std::move(kvCacheConnectorManager)) // disable block reuse for sink bubble since chopVectorIntoBlocks does not match KV cache blocks in this case , mEnableBlockReuse{mSinkBubbleLength > 0 ? false : enableBlockReuse} { @@ -1961,12 +1958,12 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, std::optional const& tempAttentionWindowInputs, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, bool enableBlockReuse, - bool onboardBlocks, CacheType cacheType, std::optional secondaryOffloadMinPriority, + CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager) : KVCacheManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, tempAttentionWindowInputs, dtype, sinkTokenLength, - std::move(stream), maxSequenceLength, enableBlockReuse, onboardBlocks, cacheType, secondaryOffloadMinPriority, + std::move(stream), maxSequenceLength, enableBlockReuse, cacheType, secondaryOffloadMinPriority, std::move(eventManager), enablePartialReuse, copyOnPartialReuse, std::move(kvCacheConnectorManager)) { } @@ -2618,9 +2615,7 @@ BlocksPerWindow BaseKVCacheManager::calculateMaxNumBlocks(executor::KvCacheConfi = static_cast(allottedSecondaryMemBytes * windowSizeShare / cacheSizeBytesPerToken); SizeType32 const blocksInSecondaryPool = std::max(0, maxTokensSecondary / tokensPerBlock); TLLM_LOG_DEBUG( - "Number of blocks in KV cache secondary pool for windowSize %d: %d, onboard blocks to primary memory " - "before reuse: %s", - windowSize, blocksInSecondaryPool, config.getOnboardBlocks() ? "true" : "false"); + "Number of blocks in KV cache secondary pool for windowSize %d: %d", windowSize, blocksInSecondaryPool); return blocksInSecondaryPool; }; diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index c54e02642caf..39d7e7ebe29f 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -218,12 +218,6 @@ TrtGptModelInflightBatching::TrtGptModelInflightBatching(std::shared_ptr TrtGptModelInflightBatching::c blocksPerWindow, getMaxNumSequences(), getMaxBeamWidth(), maxAttentionWindowVec, tempAttentionWindowInputs, kvDtype, getSinkTokenLen(), mRuntime->getStreamPtr(), kvCacheType == KvCacheType::kCROSS ? mModelConfig.getMaxEncoderLen() : getMaxSequenceLen(), enableBlockReuse, - kvCacheConfig.getOnboardBlocks(), kvCacheType, kvCacheConfig.getSecondaryOffloadMinPriority(), + kvCacheType, kvCacheConfig.getSecondaryOffloadMinPriority(), kvCacheConfig.getEventBufferMaxSize() > 0 ? std::make_unique(kvCacheConfig.getEventBufferMaxSize()) : nullptr, diff --git a/cpp/tensorrt_llm/executor/kvCacheConfig.cpp b/cpp/tensorrt_llm/executor/kvCacheConfig.cpp index 1e83ba4b3a64..6a281f939ef4 100644 --- a/cpp/tensorrt_llm/executor/kvCacheConfig.cpp +++ b/cpp/tensorrt_llm/executor/kvCacheConfig.cpp @@ -24,14 +24,12 @@ namespace tensorrt_llm::executor KvCacheConfig::KvCacheConfig(bool enableBlockReuse, std::optional const& maxTokens, std::optional> const& maxAttentionWindowVec, std::optional const& sinkTokenLength, std::optional const& freeGpuMemoryFraction, - std::optional const& hostCacheSize, bool onboardBlocks, - std::optional const& crossKvCacheFraction, std::optional secondaryOffloadMinPriority, - size_t eventBufferMaxSize, bool enablePartialReuse, bool copyOnPartialReuse, bool useUvm, - SizeType32 attentionDpEventsGatherPeriodMs, + std::optional const& hostCacheSize, std::optional const& crossKvCacheFraction, + std::optional secondaryOffloadMinPriority, size_t eventBufferMaxSize, bool enablePartialReuse, + bool copyOnPartialReuse, bool useUvm, SizeType32 attentionDpEventsGatherPeriodMs, std::optional const& runtimeDefaults, uint64_t const& maxGpuTotalBytes) : mEnableBlockReuse(enableBlockReuse) , mHostCacheSize(hostCacheSize) - , mOnboardBlocks(onboardBlocks) , mSecondaryOffloadMinPriority(secondaryOffloadMinPriority) , mEventBufferMaxSize{eventBufferMaxSize} , mEnablePartialReuse{enablePartialReuse} @@ -117,11 +115,6 @@ std::optional KvCacheConfig::getHostCacheSize() const return mHostCacheSize; } -bool KvCacheConfig::getOnboardBlocks() const -{ - return mOnboardBlocks; -} - std::optional KvCacheConfig::getSecondaryOffloadMinPriority() const { return mSecondaryOffloadMinPriority; @@ -206,11 +199,6 @@ void KvCacheConfig::setHostCacheSize(size_t hostCacheSize) mHostCacheSize = hostCacheSize; } -void KvCacheConfig::setOnboardBlocks(bool onboardBlocks) -{ - mOnboardBlocks = onboardBlocks; -} - void KvCacheConfig::setSecondaryOffloadMinPriority(std::optional secondaryOffloadMinPriority) { mSecondaryOffloadMinPriority = secondaryOffloadMinPriority; diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index 1786a43bdbec..b84930a6a0b1 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -1172,7 +1172,6 @@ KvCacheConfig Serialization::deserializeKvCacheConfig(std::istream& is) auto sinkTokenLength = su::deserialize>(is); auto freeGpuMemoryFraction = su::deserialize>(is); auto hostCacheSize = su::deserialize>(is); - auto onboardBlocks = su::deserialize(is); auto crossKvCacheFraction = su::deserialize>(is); auto secondaryOffloadMinPriority = su::deserialize>(is); auto eventBufferMaxSize = su::deserialize(is); @@ -1180,8 +1179,8 @@ KvCacheConfig Serialization::deserializeKvCacheConfig(std::istream& is) auto attentionDpEventsGatherPeriodMs = su::deserialize(is); return KvCacheConfig{enableBlockReuse, maxTokens, maxAttentionWindowVec, sinkTokenLength, freeGpuMemoryFraction, - hostCacheSize, onboardBlocks, crossKvCacheFraction, secondaryOffloadMinPriority, eventBufferMaxSize, - enablePartialReuse, copyOnPartialReuse, useUvm, attentionDpEventsGatherPeriodMs}; + hostCacheSize, crossKvCacheFraction, secondaryOffloadMinPriority, eventBufferMaxSize, enablePartialReuse, + copyOnPartialReuse, useUvm, attentionDpEventsGatherPeriodMs}; } void Serialization::serialize(KvCacheConfig const& kvCacheConfig, std::ostream& os) @@ -1194,7 +1193,6 @@ void Serialization::serialize(KvCacheConfig const& kvCacheConfig, std::ostream& su::serialize(kvCacheConfig.getSinkTokenLength(), os); su::serialize(kvCacheConfig.getFreeGpuMemoryFraction(), os); su::serialize(kvCacheConfig.getHostCacheSize(), os); - su::serialize(kvCacheConfig.getOnboardBlocks(), os); su::serialize(kvCacheConfig.getCrossKvCacheFraction(), os); su::serialize(kvCacheConfig.getSecondaryOffloadMinPriority(), os); su::serialize(kvCacheConfig.getEventBufferMaxSize(), os); @@ -1214,7 +1212,6 @@ size_t Serialization::serializedSize(KvCacheConfig const& kvCacheConfig) totalSize += su::serializedSize(kvCacheConfig.getSinkTokenLength()); totalSize += su::serializedSize(kvCacheConfig.getFreeGpuMemoryFraction()); totalSize += su::serializedSize(kvCacheConfig.getHostCacheSize()); - totalSize += su::serializedSize(kvCacheConfig.getOnboardBlocks()); totalSize += su::serializedSize(kvCacheConfig.getCrossKvCacheFraction()); totalSize += su::serializedSize(kvCacheConfig.getSecondaryOffloadMinPriority()); totalSize += su::serializedSize(kvCacheConfig.getEventBufferMaxSize()); diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 589bf6e2bf5f..95e471fd0f69 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -490,18 +490,17 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) .def(nb::init const&, SizeType32, SizeType32, std::map> const&, SizeType32, SizeType32, std::vector const&, std::optional const&, - nvinfer1::DataType, SizeType32, int64_t, runtime::SizeType32, bool, bool, tbk::CacheType, + nvinfer1::DataType, SizeType32, int64_t, runtime::SizeType32, bool, tbk::CacheType, std::optional, std::shared_ptr, bool, bool, std::shared_ptr>(), 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("temp_attention_window_inputs").none(), nb::arg("dtype"), nb::arg("sink_token_length"), nb::arg("stream"), nb::arg("max_sequence_length").none(), - nb::arg("enable_block_reuse") = false, nb::arg("onboard_blocks") = true, - nb::arg("cache_type") = tbk::CacheType::kSELF, nb::arg("secondary_offload_min_priority") = std::nullopt, - nb::arg("event_manager") = nullptr, nb::arg("enable_partial_reuse") = true, - nb::arg("copy_on_partial_reuse") = true, nb::arg("kv_connector_manager") = nullptr, - nb::call_guard()); + nb::arg("enable_block_reuse") = false, nb::arg("cache_type") = tbk::CacheType::kSELF, + nb::arg("secondary_offload_min_priority") = std::nullopt, nb::arg("event_manager") = nullptr, + nb::arg("enable_partial_reuse") = true, nb::arg("copy_on_partial_reuse") = true, + nb::arg("kv_connector_manager") = nullptr, nb::call_guard()); } void tb::BasePeftCacheManagerBindings::initBindings(nb::module_& m) diff --git a/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp index 0334eb14f6a5..232f072ff8b6 100644 --- a/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp @@ -110,35 +110,34 @@ void initConfigBindings(nb::module_& m) { return nb::make_tuple(self.getEnableBlockReuse(), self.getMaxTokens(), self.getMaxAttentionWindowVec(), self.getSinkTokenLength(), self.getFreeGpuMemoryFraction(), self.getHostCacheSize(), - self.getOnboardBlocks(), self.getCrossKvCacheFraction(), self.getSecondaryOffloadMinPriority(), - self.getEventBufferMaxSize(), self.getEnablePartialReuse(), self.getCopyOnPartialReuse(), self.getUseUvm(), + self.getCrossKvCacheFraction(), self.getSecondaryOffloadMinPriority(), self.getEventBufferMaxSize(), + self.getEnablePartialReuse(), self.getCopyOnPartialReuse(), self.getUseUvm(), self.getAttentionDpEventsGatherPeriodMs(), self.getMaxGpuTotalBytes()); }; auto kvCacheConfigSetstate = [](tle::KvCacheConfig& self, nb::tuple const& state) { - if (state.size() != 15) + if (state.size() != 14) { throw std::runtime_error("Invalid state!"); } new (&self) tle::KvCacheConfig(nb::cast(state[0]), nb::cast>(state[1]), nb::cast>>(state[2]), nb::cast>(state[3]), nb::cast>(state[4]), nb::cast>(state[5]), - nb::cast(state[6]), nb::cast>(state[7]), - nb::cast>(state[8]), nb::cast(state[9]), - nb::cast(state[10]), nb::cast(state[11]), nb::cast(state[12]), - nb::cast(state[13]), std::nullopt, nb::cast(state[14])); + nb::cast>(state[6]), nb::cast>(state[7]), + nb::cast(state[8]), nb::cast(state[9]), nb::cast(state[10]), nb::cast(state[11]), + nb::cast(state[12]), std::nullopt, nb::cast(state[13])); }; nb::class_(m, "KvCacheConfig") .def(nb::init const&, std::optional> const&, - std::optional const&, std::optional const&, std::optional const&, bool, + std::optional const&, std::optional const&, std::optional const&, std::optional const&, std::optional, size_t const&, bool, bool, bool, SizeType32, std::optional const&, uint64_t const&>(), nb::arg("enable_block_reuse") = true, nb::arg("max_tokens") = nb::none(), nb::arg("max_attention_window") = nb::none(), nb::arg("sink_token_length") = nb::none(), nb::arg("free_gpu_memory_fraction") = nb::none(), nb::arg("host_cache_size") = nb::none(), - nb::arg("onboard_blocks") = true, nb::arg("cross_kv_cache_fraction") = nb::none(), - nb::arg("secondary_offload_min_priority") = nb::none(), nb::arg("event_buffer_max_size") = 0, nb::kw_only(), - nb::arg("enable_partial_reuse") = true, nb::arg("copy_on_partial_reuse") = true, nb::arg("use_uvm") = false, + nb::arg("cross_kv_cache_fraction") = nb::none(), nb::arg("secondary_offload_min_priority") = nb::none(), + nb::arg("event_buffer_max_size") = 0, nb::kw_only(), nb::arg("enable_partial_reuse") = true, + nb::arg("copy_on_partial_reuse") = true, nb::arg("use_uvm") = false, nb::arg("attention_dp_events_gather_period_ms") = 5, nb::arg("runtime_defaults") = nb::none(), nb::arg("max_gpu_total_bytes") = 0) .def_prop_rw( @@ -151,7 +150,6 @@ void initConfigBindings(nb::module_& m) .def_prop_rw("free_gpu_memory_fraction", &tle::KvCacheConfig::getFreeGpuMemoryFraction, &tle::KvCacheConfig::setFreeGpuMemoryFraction) .def_prop_rw("host_cache_size", &tle::KvCacheConfig::getHostCacheSize, &tle::KvCacheConfig::setHostCacheSize) - .def_prop_rw("onboard_blocks", &tle::KvCacheConfig::getOnboardBlocks, &tle::KvCacheConfig::setOnboardBlocks) .def_prop_rw("cross_kv_cache_fraction", &tle::KvCacheConfig::getCrossKvCacheFraction, &tle::KvCacheConfig::setCrossKvCacheFraction) .def_prop_rw("secondary_offload_min_priority", &tle::KvCacheConfig::getSecondaryOffloadMinPriority, diff --git a/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp index 2f40cd9de9dc..387c42e81ab6 100644 --- a/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/pybind/batch_manager/kvCacheManager.cpp @@ -483,14 +483,14 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(py::module_& m) .def(py::init const&, SizeType32, SizeType32, std::map> const&, SizeType32, SizeType32, std::vector const&, std::optional const&, - nvinfer1::DataType, SizeType32, bool, int64_t, bool, bool, tbk::CacheType, + nvinfer1::DataType, SizeType32, bool, int64_t, bool, tbk::CacheType, std::optional, std::shared_ptr, bool, bool, std::shared_ptr>(), py::arg("num_kv_heads_per_layer"), py::arg("size_per_head"), py::arg("tokens_per_block"), py::arg("blocks_per_window"), py::arg("max_num_sequences"), py::arg("max_beam_width"), py::arg("max_attention_window_vec"), py::arg("temp_attention_window_inputs"), py::arg("dtype"), py::arg("sink_token_length"), py::arg("stream"), py::arg("max_sequence_length"), - py::arg("enable_block_reuse") = false, py::arg("onboard_blocks") = true, + py::arg("enable_block_reuse") = false, py::arg_v("cache_type", tbk::CacheType::kSELF, "bindings.internal.batch_manager.CacheType.SELF"), py::arg("secondary_offload_min_priority") = std::nullopt, py::arg("event_manager") = nullptr, py::arg("enable_partial_reuse") = true, py::arg("copy_on_partial_reuse") = true, diff --git a/cpp/tensorrt_llm/pybind/executor/executorConfig.cpp b/cpp/tensorrt_llm/pybind/executor/executorConfig.cpp index 74e2fe56c168..b652789958b4 100644 --- a/cpp/tensorrt_llm/pybind/executor/executorConfig.cpp +++ b/cpp/tensorrt_llm/pybind/executor/executorConfig.cpp @@ -102,34 +102,34 @@ void initConfigBindings(pybind11::module_& m) { return py::make_tuple(self.getEnableBlockReuse(), self.getMaxTokens(), self.getMaxAttentionWindowVec(), self.getSinkTokenLength(), self.getFreeGpuMemoryFraction(), self.getHostCacheSize(), - self.getOnboardBlocks(), self.getCrossKvCacheFraction(), self.getSecondaryOffloadMinPriority(), - self.getEventBufferMaxSize(), self.getEnablePartialReuse(), self.getCopyOnPartialReuse(), self.getUseUvm(), + self.getCrossKvCacheFraction(), self.getSecondaryOffloadMinPriority(), self.getEventBufferMaxSize(), + self.getEnablePartialReuse(), self.getCopyOnPartialReuse(), self.getUseUvm(), self.getAttentionDpEventsGatherPeriodMs(), self.getMaxGpuTotalBytes()); }; auto kvCacheConfigSetstate = [](py::tuple const& state) { - if (state.size() != 15) + if (state.size() != 14) { throw std::runtime_error("Invalid state!"); } return tle::KvCacheConfig(state[0].cast(), state[1].cast>(), state[2].cast>>(), state[3].cast>(), - state[4].cast>(), state[5].cast>(), state[6].cast(), - state[7].cast>(), state[8].cast>(), - state[9].cast(), state[10].cast(), state[11].cast(), state[12].cast(), - state[13].cast(), std::nullopt, state[14].cast()); + state[4].cast>(), state[5].cast>(), + state[6].cast>(), state[7].cast>(), + state[8].cast(), state[9].cast(), state[10].cast(), state[11].cast(), + state[12].cast(), std::nullopt, state[13].cast()); }; py::class_(m, "KvCacheConfig") .def(py::init const&, std::optional> const&, - std::optional const&, std::optional const&, std::optional const&, bool, + std::optional const&, std::optional const&, std::optional const&, std::optional const&, std::optional, size_t const&, bool, bool, bool, SizeType32, std::optional const&, uint64_t const&>(), py::arg("enable_block_reuse") = true, py::arg("max_tokens") = py::none(), py::arg("max_attention_window") = py::none(), py::arg("sink_token_length") = py::none(), py::arg("free_gpu_memory_fraction") = py::none(), py::arg("host_cache_size") = py::none(), - py::arg("onboard_blocks") = true, py::arg("cross_kv_cache_fraction") = py::none(), - py::arg("secondary_offload_min_priority") = py::none(), py::arg("event_buffer_max_size") = 0, py::kw_only(), - py::arg("enable_partial_reuse") = true, py::arg("copy_on_partial_reuse") = true, py::arg("use_uvm") = false, + py::arg("cross_kv_cache_fraction") = py::none(), py::arg("secondary_offload_min_priority") = py::none(), + py::arg("event_buffer_max_size") = 0, py::kw_only(), py::arg("enable_partial_reuse") = true, + py::arg("copy_on_partial_reuse") = true, py::arg("use_uvm") = false, py::arg("attention_dp_events_gather_period_ms") = 5, py::arg("runtime_defaults") = py::none(), py::arg("max_gpu_total_bytes") = 0) .def_property( @@ -144,7 +144,6 @@ void initConfigBindings(pybind11::module_& m) .def_property( "max_gpu_total_bytes", &tle::KvCacheConfig::getMaxGpuTotalBytes, &tle::KvCacheConfig::setMaxGpuTotalBytes) .def_property("host_cache_size", &tle::KvCacheConfig::getHostCacheSize, &tle::KvCacheConfig::setHostCacheSize) - .def_property("onboard_blocks", &tle::KvCacheConfig::getOnboardBlocks, &tle::KvCacheConfig::setOnboardBlocks) .def_property("cross_kv_cache_fraction", &tle::KvCacheConfig::getCrossKvCacheFraction, &tle::KvCacheConfig::setCrossKvCacheFraction) .def_property("secondary_offload_min_priority", &tle::KvCacheConfig::getSecondaryOffloadMinPriority, diff --git a/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp b/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp index 42ba14d84dcb..def08a756388 100644 --- a/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp @@ -50,7 +50,6 @@ class CacheTransBufferTest : public ::testing::Test auto constexpr blocksInSecondaryPool = 0; auto constexpr enableBlockReuse = true; - auto constexpr onboardBlocks = true; auto constexpr dataType = nvinfer1::DataType::kFLOAT; using BlocksPerWindow = std::map>; @@ -58,8 +57,8 @@ class CacheTransBufferTest : public ::testing::Test mCacheManager = std::make_unique(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, - std::nullopt, dataType, sinkTokenLength, stream, kvMaxNumTokens, enableBlockReuse, onboardBlocks, cacheType, - std::nullopt, nullptr, true); + std::nullopt, dataType, sinkTokenLength, stream, kvMaxNumTokens, enableBlockReuse, cacheType, std::nullopt, + nullptr, true); mCacheManager->allocatePools(false); diff --git a/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp b/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp index 58f1f7d4fa04..0b151b26d0b2 100644 --- a/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp @@ -126,7 +126,6 @@ class CapacitySchedulerTest : public ::testing::Test // NOLINT(cppcoreguidelines auto constexpr sizePerHead = 1; auto const maxNumBlocks = tc::divUp(maxNumTokens, tokensPerBlock); auto const kvDtype = nvinfer1::DataType::kHALF; - bool onboardBlocks = true; CudaStreamPtr streamPtr = std::make_shared(); using BlocksPerWindow = std::map>; @@ -135,7 +134,7 @@ class CapacitySchedulerTest : public ::testing::Test // NOLINT(cppcoreguidelines // init KV cache block manager return std::make_shared(numLayers, nbKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumRequests, 1, std::vector{maxNumTokensPerSeq}, std::nullopt, kvDtype, - sinkTokenLength, streamPtr, maxNumTokensPerSeq, enableReuse, onboardBlocks, cacheType); + sinkTokenLength, streamPtr, maxNumTokensPerSeq, enableReuse, cacheType); } static std::shared_ptr getPeftCacheManager() diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 2e4fa943ab0b..39c2b11e39b8 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -119,7 +119,6 @@ TEST_F(KVCacheManagerTest, BlockManagerTest) auto constexpr blocksInSecondaryPool = 0; auto constexpr maxNumSequences = 8; auto const stream = std::make_shared(); - auto constexpr onboardBlocks = true; auto constexpr beamWidth = 8; auto constexpr numBlocksPerBeam = blocksInPrimaryPool / beamWidth; @@ -130,8 +129,7 @@ TEST_F(KVCacheManagerTest, BlockManagerTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0, - onboardBlocks); + std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -229,7 +227,6 @@ void runPartialCopyTest() auto constexpr blocksInSecondaryPool = 4; auto constexpr maxNumSequences = 8; auto const stream = std::make_shared(); - auto constexpr onboardBlocks = true; auto constexpr batchSize = 1; auto constexpr maxBlocksPerSeq = 10; @@ -259,7 +256,7 @@ void runPartialCopyTest() BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, std::nullopt, type, 0, onboardBlocks); + std::vector{maxAttentionWindow}, std::nullopt, type, 0); blockManager.allocatePools(false); auto oneLayerBlockSize = blockManager.getBlockSize(0); @@ -609,7 +606,6 @@ TEST_F(KVCacheManagerTest, FindBlocksInReuseTreeByBlockKeysTest) auto constexpr blocksInSecondaryPool = 4; auto constexpr maxNumSequences = 8; auto const stream = std::make_shared(); - auto constexpr onboardBlocks = true; auto constexpr batchSize = 1; auto constexpr maxBlocksPerSeq = 10; @@ -628,7 +624,7 @@ TEST_F(KVCacheManagerTest, FindBlocksInReuseTreeByBlockKeysTest) auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - false, stream, true, onboardBlocks); + false, stream, true); // Add sequence [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16] (17 tokens, three blocks) auto inputTokens = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); @@ -669,7 +665,6 @@ TEST_F(KVCacheManagerTest, FP4BlockScaleManagementTest) auto constexpr blocksInSecondaryPool = 16; auto constexpr numFp4EltsPerContainer = 2; auto constexpr vectorSize = 16; - auto constexpr onboardBlocks = true; auto const stream = std::make_shared(); auto constexpr beamWidth = 1; @@ -679,7 +674,7 @@ TEST_F(KVCacheManagerTest, FP4BlockScaleManagementTest) KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kFP4, - false, stream, maxAttentionWindow, true, onboardBlocks); + false, stream, maxAttentionWindow, true); kvCacheManager.allocatePools(/*useUvm=*/false); @@ -707,7 +702,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) auto constexpr blocksInSecondaryPool = 0; auto constexpr maxNumSequences = 8; auto const stream = std::make_shared(); - auto constexpr onboardBlocks = true; auto constexpr maxAttentionWindow = tokensPerBlock * maxBlocksPerSeq; auto constexpr beamWidth = 1; @@ -716,8 +710,7 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0, - onboardBlocks); + std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -1006,7 +999,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) auto constexpr blocksInSecondaryPool = 0; auto constexpr maxNumSequences = 8; auto const stream = std::make_shared(); - auto constexpr onboardBlocks = true; auto constexpr numReturnSequences = 1; auto constexpr maxAttentionWindow = tokensPerBlock * maxBlocksPerSeq; @@ -1016,8 +1008,7 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0, - onboardBlocks); + std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -1217,7 +1208,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest) auto constexpr blocksInSecondaryPool = 0; auto constexpr maxNumSequences = 8; auto const stream = std::make_shared(); - auto constexpr onboardBlocks = true; auto constexpr numReturnSequences = 1; auto constexpr maxAttentionWindow = tokensPerBlock * maxBlocksPerSeq; auto constexpr beamWidth = 1; @@ -1226,8 +1216,7 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0, - onboardBlocks); + std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -1402,7 +1391,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) auto constexpr blocksInSecondaryPool = 0; auto constexpr maxNumSequences = 8; auto const stream = std::make_shared(); - auto constexpr onboardBlocks = true; auto constexpr maxAttentionWindow = tokensPerBlock * maxBlocksPerSeq; auto constexpr beamWidth = 1; @@ -1411,8 +1399,7 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0, - onboardBlocks); + std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -1660,7 +1647,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) auto constexpr blocksInSecondaryPool = 0; auto constexpr maxNumSequences = 8; auto const stream = std::make_shared(); - auto constexpr onboardBlocks = true; auto constexpr maxAttentionWindow = tokensPerBlock * maxBlocksPerSeq; auto constexpr beamWidth = 1; @@ -1669,8 +1655,7 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0, - onboardBlocks); + std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -1901,7 +1886,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) auto constexpr blocksInSecondaryPool = 0; auto constexpr maxNumSequences = 8; auto const stream = std::make_shared(); - auto constexpr onboardBlocks = true; auto constexpr numReturnSequences = 1; auto constexpr maxAttentionWindow = tokensPerBlock * maxBlocksPerSeq; auto constexpr beamWidth = 1; @@ -1910,8 +1894,7 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0, - onboardBlocks); + std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -2107,7 +2090,6 @@ TEST_F(KVCacheManagerTest, KVCacheManagerPerRequestStatsTest) auto constexpr maxNumSequences = 8; auto constexpr blocksInPrimaryPool = 16; auto constexpr blocksInSecondaryPool = 0; - auto constexpr onboardBlocks = true; auto const stream = std::make_shared(); auto constexpr beamWidth = 1; @@ -2121,7 +2103,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerPerRequestStatsTest) KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - 0, stream, maxSequenceLength, true, onboardBlocks); + 0, stream, maxSequenceLength, true); kvCacheManager.allocatePools(false); auto inputTokens = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8}); @@ -2167,7 +2149,6 @@ TEST_F(KVCacheManagerTest, BlockManagerBlockPriorityTest) auto constexpr blocksInSecondaryPool = 0; auto constexpr maxNumSequences = 4; auto const stream = std::make_shared(); - auto constexpr onboardBlocks = true; auto constexpr maxAttentionWindow = tokensPerBlock * maxBlocksPerSeq; auto constexpr beamWidth = 1; @@ -2176,8 +2157,7 @@ TEST_F(KVCacheManagerTest, BlockManagerBlockPriorityTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxAttentionWindow, beamWidth, - std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0, - onboardBlocks); + std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); @@ -2278,7 +2258,6 @@ TEST_F(KVCacheManagerTest, KVCacheManagerDecodeBlockPriorityTest) auto constexpr maxNumSequences = 8; auto constexpr blocksInPrimaryPool = 8; auto constexpr blocksInSecondaryPool = 0; - auto constexpr onboardBlocks = true; auto const stream = std::make_shared(); auto constexpr beamWidth = 1; @@ -2292,7 +2271,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerDecodeBlockPriorityTest) KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - 0, stream, maxSequenceLength, true, onboardBlocks); + 0, stream, maxSequenceLength, true); kvCacheManager.allocatePools(false); auto const& blockManager = kvCacheManager.getBlockManager(); @@ -2382,7 +2361,6 @@ TEST_F(KVCacheManagerTest, KVCacheManagerTimedEvictionTest) auto constexpr maxNumSequences = 8; auto constexpr blocksInPrimaryPool = 8; auto constexpr blocksInSecondaryPool = 0; - auto constexpr onboardBlocks = true; auto const stream = std::make_shared(); auto constexpr beamWidth = 1; @@ -2396,7 +2374,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerTimedEvictionTest) KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - 0, stream, maxSequenceLength, true, onboardBlocks); + 0, stream, maxSequenceLength, true); kvCacheManager.allocatePools(false); auto inputTokens0 = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}); @@ -2451,7 +2429,6 @@ TEST_F(KVCacheManagerTest, KVCacheManagerDecodeTimedEvictionTest) auto constexpr maxNumSequences = 8; auto constexpr blocksInPrimaryPool = 9; auto constexpr blocksInSecondaryPool = 0; - auto constexpr onboardBlocks = true; auto const stream = std::make_shared(); auto constexpr beamWidth = 1; @@ -2465,7 +2442,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerDecodeTimedEvictionTest) KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - 0, stream, maxSequenceLength, true, onboardBlocks); + 0, stream, maxSequenceLength, true); kvCacheManager.allocatePools(false); { // 12 tokens, occupy 3 blocks 0, 1, 2. @@ -2539,7 +2516,6 @@ TEST_F(KVCacheManagerTest, KVCacheManagerSecondaryBlockPrimaryChildTest) auto constexpr maxSequenceLength = tokensPerBlock * maxBlocksPerSeq; auto constexpr blocksInPrimaryPool = 4; auto constexpr blocksInSecondaryPool = 4; - auto constexpr onboardBlocks = true; auto const stream = std::make_shared(); auto constexpr beamWidth = 1; @@ -2553,7 +2529,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerSecondaryBlockPrimaryChildTest) KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - 0, stream, maxSequenceLength, true, onboardBlocks); + 0, stream, maxSequenceLength, true); kvCacheManager.allocatePools(false); auto inputTokens0 = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}); @@ -2632,7 +2608,6 @@ TEST_F(KVCacheManagerTest, KVCacheManagerLeafBlockTest) auto constexpr maxNumSequences = 8; auto constexpr blocksInPrimaryPool = 4; auto constexpr blocksInSecondaryPool = 0; - auto constexpr onboardBlocks = true; auto const stream = std::make_shared(); auto constexpr beamWidth = 1; @@ -2645,7 +2620,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerLeafBlockTest) KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - 0, stream, maxSequenceLength, true, onboardBlocks); + 0, stream, maxSequenceLength, true); kvCacheManager.allocatePools(false); auto inputTokens0 = std::make_shared(VecTokens{0, 1, 2, 3}); @@ -2710,7 +2685,6 @@ TEST_F(KVCacheManagerTest, KVCacheManagerLeafBlockWithDependentTest) auto constexpr maxNumSequences = 8; auto constexpr blocksInPrimaryPool = 4; auto constexpr blocksInSecondaryPool = 1; - auto constexpr onboardBlocks = true; auto const stream = std::make_shared(); auto constexpr beamWidth = 1; @@ -2725,7 +2699,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerLeafBlockWithDependentTest) KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - 0, stream, maxSequenceLength, true, onboardBlocks); + 0, stream, maxSequenceLength, true); kvCacheManager.allocatePools(false); // Create sequence with one block worth of context tokens @@ -2838,7 +2812,6 @@ TEST_P(KVCacheManagerTest, DISABLED_KVCacheManagerAllocationTest) auto constexpr enableBlockReuse = false; auto constexpr useUvm = false; - auto constexpr onboardBlocks = true; auto const homogeneousLayers = GetParam(); auto const granularity = tensorrt_llm::common::getAllocationGranularity(); @@ -2848,11 +2821,10 @@ TEST_P(KVCacheManagerTest, DISABLED_KVCacheManagerAllocationTest) KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, - nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, onboardBlocks) + nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse) : KVCacheManager(std::vector(numLayers, numHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, - std::nullopt, nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, - onboardBlocks); + std::nullopt, nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse); auto const& blockManager = kvCacheManager.getBlockManager(); auto const& bufferManager = blockManager.getBufferManager(theOnlyWindowSize(kvCacheManager)); @@ -2900,7 +2872,6 @@ TEST_P(KVCacheManagerTest, KVCacheManagerTest) auto constexpr blocksInSecondaryPool = 0; auto constexpr enableBlockReuse = false; - auto constexpr onboardBlocks = true; auto const homogeneousLayers = GetParam(); auto const expectedNumPools = homogeneousLayers ? 1 : static_cast(expectedHeadsPerPool.size()); @@ -2909,10 +2880,10 @@ TEST_P(KVCacheManagerTest, KVCacheManagerTest) KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, - nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, onboardBlocks) + nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse) : KVCacheManager(numHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, onboardBlocks); + sinkTokenLength, stream, maxSequenceLength, enableBlockReuse); kvCacheManager.allocatePools(false); EXPECT_EQ(kvCacheManager.getOffsetTableDimensions().maxBlocksPerSeq, maxBlocksPerSeq); @@ -3049,7 +3020,6 @@ TEST_P(KVCacheManagerTest, KVCacheManagerRewindTokensTest) auto constexpr blocksInSecondaryPool = 0; auto constexpr enableBlockReuse = false; - auto constexpr onboardBlocks = true; auto const homogeneousLayers = GetParam(); auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {totalNumBlocks, blocksInSecondaryPool}}}; @@ -3057,11 +3027,10 @@ TEST_P(KVCacheManagerTest, KVCacheManagerRewindTokensTest) KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, - nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, onboardBlocks) + nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse) : KVCacheManager(std::vector(numLayers, numHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, - std::nullopt, nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, - onboardBlocks); + std::nullopt, nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse); kvCacheManager.allocatePools(false); EXPECT_EQ(kvCacheManager.getTokensPerBlock(), tokensPerBlock); @@ -3135,7 +3104,6 @@ TEST_P(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowTest) auto constexpr blocksInSecondaryPool = 0; auto constexpr enableBlockReuse = false; - auto constexpr onboardBlocks = true; auto const homogeneousLayers = GetParam(); auto const expectedNumPools = homogeneousLayers ? 1 : static_cast(expectedHeadsPerPool.size()); @@ -3144,10 +3112,10 @@ TEST_P(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowTest) KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, - nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, onboardBlocks) + nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse) : KVCacheManager(numHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, onboardBlocks); + sinkTokenLength, stream, maxSequenceLength, enableBlockReuse); kvCacheManager.allocatePools(false); EXPECT_EQ(kvCacheManager.getOffsetTableDimensions().maxBlocksPerSeq, maxBlocksPerSeq); @@ -3257,12 +3225,11 @@ TEST_F(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowSmallerThanBlockSizeT auto constexpr blocksInSecondaryPool = 0; auto constexpr enableBlockReuse = false; - auto constexpr onboardBlocks = true; auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, - nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, onboardBlocks); + nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse); kvCacheManager.allocatePools(false); auto const& blockManager = kvCacheManager.getBlockManager(); @@ -3353,7 +3320,6 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStream) auto constexpr maxNumSequences = 8; auto constexpr blocksInPrimaryPool = 8; auto constexpr blocksInSecondaryPool = 2; - auto constexpr onboardBlocks = true; auto constexpr dtype = nvinfer1::DataType::kHALF; auto const stream = std::make_shared(); @@ -3369,8 +3335,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStream) KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, dtype, 0, stream, - maxSequenceLength, true, onboardBlocks, CacheType::kSELF, std::nullopt, - std::make_unique(1024)); + maxSequenceLength, true, CacheType::kSELF, std::nullopt, std::make_unique(1024)); kvCacheManager.allocatePools(false); auto events = getEvents(kvCacheManager); @@ -3527,13 +3492,12 @@ TEST_F(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowWithReuseTest) auto constexpr blocksInSecondaryPool = 16; auto constexpr enableBlockReuse = true; - auto constexpr onboardBlocks = true; auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, blocksInSecondaryPool}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, - nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, onboardBlocks); + nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse); kvCacheManager.allocatePools(false); auto const& blockManager = kvCacheManager.getBlockManager(); @@ -3664,13 +3628,12 @@ TEST_F(KVCacheManagerTest, KVCacheManagerSWAInvalidateReuseTest) auto constexpr blocksInSecondaryPoolPerWindow = 0; auto constexpr enableBlockReuse = true; - auto constexpr onboardBlocks = true; auto const blocksPerWindow = BlocksPerWindow{{attentionWindow, {blocksInPrimaryPoolPerWindow, blocksInSecondaryPoolPerWindow}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, std::nullopt, dtype, sinkTokenLength, stream, maxSequenceLength, - enableBlockReuse, onboardBlocks); + enableBlockReuse); kvCacheManager.allocatePools(false); auto const& blockManager = kvCacheManager.getBlockManager(); @@ -3743,14 +3706,13 @@ TEST_F(KVCacheManagerTest, KVCacheManagerVariableWindowAttentionWithReuseTest) auto constexpr blocksInSecondaryPoolPerWindow = 16; auto constexpr enableBlockReuse = true; - auto constexpr onboardBlocks = true; auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPoolPerWindow, blocksInSecondaryPoolPerWindow}}, {minAttentionWindow, {blocksInPrimaryPoolPerWindow, blocksInSecondaryPoolPerWindow}}}; KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, maxAttentionWindowVec, std::nullopt, dtype, sinkTokenLength, stream, maxSequenceLength, - enableBlockReuse, onboardBlocks); + enableBlockReuse); kvCacheManager.allocatePools(false); auto const& blockManager = kvCacheManager.getBlockManager(); @@ -3845,7 +3807,6 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamOverflow) auto constexpr maxNumSequences = 8; auto constexpr blocksInPrimaryPool = 8; auto constexpr blocksInSecondaryPool = 2; - auto constexpr onboardBlocks = true; auto constexpr dtype = nvinfer1::DataType::kHALF; auto const stream = std::make_shared(); @@ -3861,8 +3822,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamOverflow) KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, dtype, 0, stream, - maxSequenceLength, true, onboardBlocks, CacheType::kSELF, std::nullopt, - std::make_unique(1)); + maxSequenceLength, true, CacheType::kSELF, std::nullopt, std::make_unique(1)); kvCacheManager.allocatePools(false); auto inputTokens0 = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); @@ -3904,7 +3864,6 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamPriority) auto constexpr maxNumSequences = 8; auto constexpr blocksInPrimaryPool = 8; auto constexpr blocksInSecondaryPool = 2; - auto constexpr onboardBlocks = true; auto constexpr dtype = nvinfer1::DataType::kHALF; auto const stream = std::make_shared(); @@ -3920,8 +3879,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamPriority) KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, dtype, 0, stream, - maxSequenceLength, true, onboardBlocks, CacheType::kSELF, std::nullopt, - std::make_unique(1024)); + maxSequenceLength, true, CacheType::kSELF, std::nullopt, std::make_unique(1024)); kvCacheManager.allocatePools(false); auto inputTokens0 = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7}); @@ -4015,7 +3973,6 @@ TEST_F(KVCacheManagerTest, PinAndUnpinBlocksById) auto constexpr blocksInSecondaryPool = 0; auto constexpr maxNumSequences = 8; auto const stream = std::make_shared(); - auto constexpr onboardBlocks = true; auto constexpr beamWidth = 1; auto const maxAttentionWindow = tokensPerBlock * blocksInPrimaryPool; @@ -4023,7 +3980,7 @@ TEST_F(KVCacheManagerTest, PinAndUnpinBlocksById) KVCacheManager kvCacheManager(numLayers, numKvHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - 0, stream, maxAttentionWindow, true, onboardBlocks); + 0, stream, maxAttentionWindow, true); kvCacheManager.allocatePools(false); LlmRequest::RequestIdType requestId{0}; @@ -4059,7 +4016,6 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamBlocking) auto constexpr maxNumSequences = 8; auto constexpr blocksInPrimaryPool = 8; auto constexpr blocksInSecondaryPool = 2; - auto constexpr onboardBlocks = true; auto constexpr dtype = nvinfer1::DataType::kHALF; auto const stream = std::make_shared(); @@ -4075,13 +4031,13 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamBlocking) KVCacheManager kvCacheManagerTest(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, dtype, 0, - stream, maxSequenceLength, true, onboardBlocks, CacheType::kSELF, std::nullopt); + stream, maxSequenceLength, true, CacheType::kSELF, std::nullopt); EXPECT_EQ(getEvents(kvCacheManagerTest).size(), 0); KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - 0, stream, maxSequenceLength, true, onboardBlocks, CacheType::kSELF, std::nullopt, + 0, stream, maxSequenceLength, true, CacheType::kSELF, std::nullopt, std::make_unique(1024)); kvCacheManager.allocatePools(false); @@ -4113,7 +4069,6 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamWindowSize) auto constexpr maxNumSequences = 8; auto blocksInPool = std::vector{8, 2}; auto blocksInSlidingWindowPool = std::vector{4, 2}; - auto constexpr onboardBlocks = true; auto constexpr dtype = nvinfer1::DataType::kHALF; auto const stream = std::make_shared(); @@ -4131,7 +4086,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamWindowSize) KVCacheManager kvCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, beamWidth, std::vector{maxAttentionWindow, slidingWindow}, std::nullopt, dtype, 0, - stream, maxSequenceLength, true, onboardBlocks, CacheType::kSELF, std::nullopt, + stream, maxSequenceLength, true, CacheType::kSELF, std::nullopt, std::make_unique(1024)); kvCacheManager.allocatePools(false); @@ -4232,7 +4187,6 @@ TEST_P(KVCacheManagerTest, DISABLED_KVCacheManagerSinkTokenLengthTest) auto constexpr blocksInSecondaryPool = 0; auto constexpr enableBlockReuse = false; - auto constexpr onboardBlocks = true; auto const homogeneousLayers = GetParam(); auto const expectedNumPools = homogeneousLayers ? 1 : static_cast(expectedHeadsPerPool.size()); @@ -4242,10 +4196,10 @@ TEST_P(KVCacheManagerTest, DISABLED_KVCacheManagerSinkTokenLengthTest) KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, - nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, onboardBlocks) + nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse) : KVCacheManager(numHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, onboardBlocks); + sinkTokenLength, stream, maxSequenceLength, enableBlockReuse); kvCacheManager.allocatePools(false); EXPECT_EQ(kvCacheManager.getOffsetTableDimensions().maxBlocksPerSeq, maxBlocksPerSeq); @@ -4373,7 +4327,6 @@ TEST_P(KVCacheManagerTest, KVCacheManagerBatchTest) auto constexpr blocksInSecondaryPool = 0; auto constexpr enableBlockReuse = false; - auto constexpr onboardBlocks = true; auto const homogeneousLayers = GetParam(); auto const expectedNumPools = homogeneousLayers ? 1 : static_cast(expectedHeadsPerPool.size()); @@ -4382,10 +4335,10 @@ TEST_P(KVCacheManagerTest, KVCacheManagerBatchTest) KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, - nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, onboardBlocks) + nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, enableBlockReuse) : KVCacheManager(numHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, - sinkTokenLength, stream, maxSequenceLength, enableBlockReuse, onboardBlocks); + sinkTokenLength, stream, maxSequenceLength, enableBlockReuse); kvCacheManager.allocatePools(false); EXPECT_EQ(kvCacheManager.getOffsetTableDimensions().maxBlocksPerSeq, maxBlocksPerSeq); @@ -4509,7 +4462,6 @@ void testNeededBlocksOneStep(bool kv_cache_block_reuse, int beamWidth, int draft { auto constexpr maxAttentionWindow = 46; auto constexpr blocksInSecondaryPool = 0; - auto constexpr onboardBlocks = true; auto constexpr maxSequenceLength = 256; auto constexpr maxBlocksPerSeq = tc::ceilDiv(maxSequenceLength, tokensPerBlock); auto constexpr totalNumBlocks = maxNumSequences * maxBlocksPerSeq; @@ -4518,12 +4470,10 @@ void testNeededBlocksOneStep(bool kv_cache_block_reuse, int beamWidth, int draft KVCacheManager kvCacheManager = homogeneousLayers ? KVCacheManager(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, - nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, kv_cache_block_reuse, - onboardBlocks) + nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, kv_cache_block_reuse) : KVCacheManager(numHeadsPerLayer, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, - nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, kv_cache_block_reuse, - onboardBlocks); + nvinfer1::DataType::kHALF, sinkTokenLength, stream, maxSequenceLength, kv_cache_block_reuse); kvCacheManager.allocatePools(false); EXPECT_EQ(kvCacheManager.getOffsetTableDimensions().maxBlocksPerSeq, maxBlocksPerSeq); @@ -4765,7 +4715,7 @@ std::shared_ptr createKvCacheManager( numBlocksInPrimaryPool, kvCacheInstantiationParameters.maxBeamWidth, std::vector{kvCacheInstantiationParameters.maxAttentionWindow}, temporaryKvCacheInputs, kvCacheInstantiationParameters.dtype, kvCacheInstantiationParameters.sinkTokenLength, stream, - maxSequenceLength, kvCacheInstantiationParameters.kvCacheBlockReuse, true, CacheType::kSELF); + maxSequenceLength, kvCacheInstantiationParameters.kvCacheBlockReuse, CacheType::kSELF); } if (std::holds_alternative>(kvCacheInstantiationParameters.numHeadsPerLayer)) { @@ -4776,7 +4726,7 @@ std::shared_ptr createKvCacheManager( numBlocksInPrimaryPool, kvCacheInstantiationParameters.maxBeamWidth, std::vector{kvCacheInstantiationParameters.maxAttentionWindow}, temporaryKvCacheInputs, kvCacheInstantiationParameters.dtype, kvCacheInstantiationParameters.sinkTokenLength, stream, - maxSequenceLength, kvCacheInstantiationParameters.kvCacheBlockReuse, true, CacheType::kSELF); + maxSequenceLength, kvCacheInstantiationParameters.kvCacheBlockReuse, CacheType::kSELF); } TLLM_THROW("Unhandled type of num heads per layer provided."); } diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheUtilsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheUtilsTest.cpp index c4de5b6a8c68..216506b647b3 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheUtilsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheUtilsTest.cpp @@ -82,7 +82,6 @@ TEST_F(BlockIteratorTest, CacheManagerTest) auto constexpr maxAttentionWindow = tokensPerBlock * maxBlocksPerSeq; auto const stream = std::make_shared(); - auto constexpr onboardBlocks = true; // TODO: Support and add coverage for beamWidth > 1 auto constexpr beamWidth = 1; @@ -96,7 +95,7 @@ TEST_F(BlockIteratorTest, CacheManagerTest) BlockManager blockManager(std::vector(numLayers, numKvHeads), sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, stream, maxSequenceLength, beamWidth, maxAttentionWindowVec, std::nullopt, - dataType, 0, onboardBlocks); + dataType, 0); blockManager.allocatePools(false); EXPECT_EQ(blockManager.getTokensPerBlock(), tokensPerBlock); diff --git a/cpp/tests/unit_tests/executor/agentCommTest.cpp b/cpp/tests/unit_tests/executor/agentCommTest.cpp index fd7a7a23de73..bbf48ebce379 100644 --- a/cpp/tests/unit_tests/executor/agentCommTest.cpp +++ b/cpp/tests/unit_tests/executor/agentCommTest.cpp @@ -74,7 +74,6 @@ class AgentCommTest : public ::testing::Test auto constexpr blocksInSecondaryPool = 0; auto constexpr enableBlockReuse = true; - auto constexpr onboardBlocks = true; auto constexpr dataType = nvinfer1::DataType::kFLOAT; using BlocksPerWindow = std::map>; @@ -83,8 +82,8 @@ class AgentCommTest : public ::testing::Test mCacheManager = std::make_unique(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, maxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, - std::nullopt, dataType, sinkTokenLength, stream, kvMaxNumTokens, enableBlockReuse, onboardBlocks, cacheType, - std::nullopt, nullptr, true); + std::nullopt, dataType, sinkTokenLength, stream, kvMaxNumTokens, enableBlockReuse, cacheType, std::nullopt, + nullptr, true); mCacheManager->allocatePools(false); diff --git a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp index 71a4c3732d62..af5aafce3366 100644 --- a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp +++ b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp @@ -476,7 +476,7 @@ TEST(SerializeUtilsTest, VectorResponses) TEST(SerializeUtilsTest, KvCacheConfig) { texec::KvCacheConfig kvCacheConfig( - true, 10, std::vector(1, 100), 2, 0.1, 10000, false, 0.5, 50, 1024, false, false, true, 77); + true, 10, std::vector(1, 100), 2, 0.1, 10000, 0.5, 50, 1024, false, false, true, 77); auto kvCacheConfig2 = serializeDeserialize(kvCacheConfig); EXPECT_EQ(kvCacheConfig.getEnableBlockReuse(), kvCacheConfig2.getEnableBlockReuse()); @@ -487,7 +487,6 @@ TEST(SerializeUtilsTest, KvCacheConfig) EXPECT_EQ(kvCacheConfig.getSinkTokenLength(), kvCacheConfig2.getSinkTokenLength()); EXPECT_EQ(kvCacheConfig.getFreeGpuMemoryFraction(), kvCacheConfig2.getFreeGpuMemoryFraction()); EXPECT_EQ(kvCacheConfig.getHostCacheSize(), kvCacheConfig2.getHostCacheSize()); - EXPECT_EQ(kvCacheConfig.getOnboardBlocks(), kvCacheConfig2.getOnboardBlocks()); EXPECT_EQ(kvCacheConfig.getCrossKvCacheFraction(), kvCacheConfig2.getCrossKvCacheFraction()); EXPECT_EQ(kvCacheConfig.getSecondaryOffloadMinPriority(), kvCacheConfig2.getSecondaryOffloadMinPriority()); EXPECT_EQ(kvCacheConfig.getEventBufferMaxSize(), kvCacheConfig2.getEventBufferMaxSize()); diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index af9c1b096bb3..da8990502762 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -209,7 +209,6 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- auto constexpr blocksInSecondaryPool = 0; auto constexpr enableBlockReuse = false; - auto constexpr onboardBlocks = true; auto constexpr dataType = nvinfer1::DataType::kFLOAT; using BlocksPerWindow = std::map>; @@ -217,8 +216,8 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- mManager = std::make_unique(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, mMaxNumSequences, maxBeamWidth, std::vector{maxAttentionWindow}, std::nullopt, - dataType, sinkTokenLength, stream, maxNumTokens, enableBlockReuse, onboardBlocks, CacheType::kSELF, - std::nullopt, nullptr, true); + dataType, sinkTokenLength, stream, maxNumTokens, enableBlockReuse, CacheType::kSELF, std::nullopt, nullptr, + true); auto attentionLayerNumPerPP = std::vector{numLayers}; mCacheState = std::make_unique( numLayers, numHeads, sizePerHead, tokensPerBlock, 1, 1, 1, attentionLayerNumPerPP, dataType); @@ -578,7 +577,6 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam(layerNumthisRank, numHeadsPerRank, sizePerHead, tokensPerBlock, blocksPerWindow, mMaxNumSequences, maxBeamWidth, maxAttentionWindowVec, std::nullopt, dataType, - sinkTokenLength, stream, maxNumTokens, enableBlockReuse, onboardBlocks, cacheType, std::nullopt, nullptr, - true); + sinkTokenLength, stream, maxNumTokens, enableBlockReuse, cacheType, std::nullopt, nullptr, true); texec::kv_cache::CacheState::AttentionType attentionType = isMLA ? texec::kv_cache::CacheState::AttentionType::kMLA : texec::kv_cache::CacheState::AttentionType::kDEFAULT; diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 82249e6748a2..b0ed0c41ab66 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -332,7 +332,6 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], 'stream': self._stream.cuda_stream, 'max_sequence_length': max_seq_len, 'enable_block_reuse': kv_cache_config.enable_block_reuse, - 'onboard_blocks': kv_cache_config.onboard_blocks, 'cache_type': kv_cache_type, 'enable_partial_reuse': kv_cache_config.enable_partial_reuse, 'copy_on_partial_reuse': kv_cache_config.copy_on_partial_reuse, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 31d38db73e11..261de87630fc 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1228,8 +1228,6 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): description= "Size of the host cache in bytes. If both `max_tokens` and `host_cache_size` are specified, memory corresponding to the minimum will be used." ) - onboard_blocks: bool = Field( - default=True, description="Controls if blocks are onboarded.") cross_kv_cache_fraction: Optional[float] = Field( default=None, description= @@ -1290,7 +1288,6 @@ def _to_pybind(self): sink_token_length=self.sink_token_length, free_gpu_memory_fraction=self.free_gpu_memory_fraction, host_cache_size=self.host_cache_size, - onboard_blocks=self.onboard_blocks, cross_kv_cache_fraction=self.cross_kv_cache_fraction, secondary_offload_min_priority=self.secondary_offload_min_priority, event_buffer_max_size=self.event_buffer_max_size, diff --git a/tests/unittest/bindings/test_bindings_ut.py b/tests/unittest/bindings/test_bindings_ut.py index f049a4437cbf..4a63abd172c1 100644 --- a/tests/unittest/bindings/test_bindings_ut.py +++ b/tests/unittest/bindings/test_bindings_ut.py @@ -475,8 +475,6 @@ def test_KvCache_events_binding(): max_sequence_length, 'enable_block_reuse': True, - 'onboard_blocks': - False, 'cache_type': _tb.internal.batch_manager.CacheType.SELF, 'event_manager': diff --git a/tests/unittest/bindings/test_executor_bindings.py b/tests/unittest/bindings/test_executor_bindings.py index c91353129867..ae01772c121a 100644 --- a/tests/unittest/bindings/test_executor_bindings.py +++ b/tests/unittest/bindings/test_executor_bindings.py @@ -1308,7 +1308,6 @@ def test_kv_cache_config(): assert config.free_gpu_memory_fraction is None assert config.cross_kv_cache_fraction is None assert config.host_cache_size is None - assert config.onboard_blocks == True assert config.secondary_offload_min_priority is None assert config.event_buffer_max_size == 0 assert config.enable_partial_reuse == True @@ -1323,7 +1322,6 @@ def test_kv_cache_config(): config.free_gpu_memory_fraction = 0.5 config.cross_kv_cache_fraction = 0.5 config.host_cache_size = 4 - config.onboard_blocks = False config.secondary_offload_min_priority = 50 config.event_buffer_max_size = 1024 config.enable_partial_reuse = False @@ -1337,7 +1335,6 @@ def test_kv_cache_config(): assert config.free_gpu_memory_fraction == 0.5 assert config.cross_kv_cache_fraction == 0.5 assert config.host_cache_size == 4 - assert config.onboard_blocks == False assert config.secondary_offload_min_priority == 50 assert config.event_buffer_max_size == 1024 assert config.enable_partial_reuse == False @@ -1353,7 +1350,6 @@ def test_kv_cache_config(): "free_gpu_memory_fraction": 0.5, "cross_kv_cache_fraction": 0.5, "host_cache_size": 1024, - "onboard_blocks": False, "event_buffer_max_size": 2048, "enable_partial_reuse": True, "copy_on_partial_reuse": False, @@ -2400,7 +2396,6 @@ def test_kv_cache_config_pickle(): config.free_gpu_memory_fraction = 0.3 config.cross_kv_cache_fraction = 0.5 config.host_cache_size = 4 - config.onboard_blocks = False config.secondary_offload_min_priority = 50 config.event_buffer_max_size = 1024 config.enable_partial_reuse = False @@ -2414,7 +2409,6 @@ def test_kv_cache_config_pickle(): assert config.free_gpu_memory_fraction == config_copy.free_gpu_memory_fraction assert config.cross_kv_cache_fraction == config_copy.cross_kv_cache_fraction assert config.host_cache_size == config_copy.host_cache_size - assert config.onboard_blocks == config_copy.onboard_blocks assert config.secondary_offload_min_priority == config_copy.secondary_offload_min_priority assert config.event_buffer_max_size == config_copy.event_buffer_max_size assert config.enable_partial_reuse == config_copy.enable_partial_reuse diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 910c372a4b97..4752a019c3c3 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -157,7 +157,6 @@ def test_KvCacheConfig_declaration(): sink_token_length=32, free_gpu_memory_fraction=0.5, host_cache_size=1024, - onboard_blocks=True, cross_kv_cache_fraction=0.5, secondary_offload_min_priority=1, event_buffer_max_size=0, @@ -172,7 +171,6 @@ def test_KvCacheConfig_declaration(): assert pybind_config.sink_token_length == 32 assert pybind_config.free_gpu_memory_fraction == 0.5 assert pybind_config.host_cache_size == 1024 - assert pybind_config.onboard_blocks == True assert pybind_config.cross_kv_cache_fraction == 0.5 assert pybind_config.secondary_offload_min_priority == 1 assert pybind_config.event_buffer_max_size == 0 diff --git a/tests/unittest/llmapi/test_llm_kv_cache_events.py b/tests/unittest/llmapi/test_llm_kv_cache_events.py index db90a34413e6..590091d29f6f 100644 --- a/tests/unittest/llmapi/test_llm_kv_cache_events.py +++ b/tests/unittest/llmapi/test_llm_kv_cache_events.py @@ -21,7 +21,6 @@ global_kvcache_config = KvCacheConfig(free_gpu_memory_fraction=0.4, event_buffer_max_size=1024, enable_block_reuse=True, - onboard_blocks=True, max_tokens=256) diff --git a/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/1/model.py b/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/1/model.py index ab165323c1f3..9beb46d3bfdc 100755 --- a/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/1/model.py +++ b/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/1/model.py @@ -914,8 +914,6 @@ def get_kv_cache_config(self, model_config): get_parameter(model_config, "cross_kv_cache_fraction", float), "host_cache_size": get_parameter(model_config, "kv_cache_host_memory_bytes", int), - "onboard_blocks": - get_parameter(model_config, "kv_cache_onboard_blocks", bool), } max_attention_window_size = get_parameter(model_config, "max_attention_window_size") diff --git a/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt b/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt index f5f6cb41a4b8..6b62c00a8a8d 100644 --- a/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt +++ b/triton_backend/all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt @@ -679,13 +679,6 @@ parameters: { string_value: "${kv_cache_host_memory_bytes}" } } -# kv_cache_onboard_blocks is for internal implementation. -parameters: { - key: "kv_cache_onboard_blocks" - value: { - string_value: "${kv_cache_onboard_blocks}" - } -} # enable_trt_overlap is deprecated and doesn't have any effect on the runtime # parameters: { # key: "enable_trt_overlap" diff --git a/triton_backend/all_models/tests/test_python_backend.py b/triton_backend/all_models/tests/test_python_backend.py index 8e17f2b09fbd..c17e08bdc50a 100644 --- a/triton_backend/all_models/tests/test_python_backend.py +++ b/triton_backend/all_models/tests/test_python_backend.py @@ -784,7 +784,6 @@ def model_config() -> Dict: "kv_cache_free_gpu_mem_fraction": "0.5", "cross_kv_cache_fraction": "0.5", "kv_cache_host_memory_bytes": "4", - "kv_cache_onboard_blocks": "false", "gpu_device_ids": "0,1,2,3", "executor_worker_path": str(os.path.abspath(__file__)), "lora_cache_optimal_adapter_size": "1", @@ -814,7 +813,6 @@ def test_get_executor_config(model_config: Dict): assert config.kv_cache_config.free_gpu_memory_fraction == 0.5 assert config.kv_cache_config.cross_kv_cache_fraction == 0.5 assert config.kv_cache_config.host_cache_size == 4 - assert config.kv_cache_config.onboard_blocks == False assert config.parallel_config.device_ids == [0, 1, 2, 3] assert config.parallel_config.orchestrator_config is None assert config.peft_cache_config.optimal_adapter_size == 1 @@ -859,7 +857,6 @@ def test_get_executor_config_minimal(): assert config.kv_cache_config.free_gpu_memory_fraction is None assert config.kv_cache_config.cross_kv_cache_fraction is None assert config.kv_cache_config.host_cache_size is None - assert config.kv_cache_config.onboard_blocks == True assert config.parallel_config is None assert config.peft_cache_config.optimal_adapter_size == 8 assert config.peft_cache_config.max_adapter_size == 64 diff --git a/triton_backend/inflight_batcher_llm/src/model_instance_state.cc b/triton_backend/inflight_batcher_llm/src/model_instance_state.cc index 82ee70bc992b..7177ea2cce5c 100644 --- a/triton_backend/inflight_batcher_llm/src/model_instance_state.cc +++ b/triton_backend/inflight_batcher_llm/src/model_instance_state.cc @@ -127,17 +127,6 @@ executor::KvCacheConfig ModelInstanceState::getKvCacheConfigFromParams() TLLM_LOG_WARNING("kv_cache_host_memory_bytes not set, defaulting to 0"); } - bool kvCacheOnboardBlocks = true; - try - { - kvCacheOnboardBlocks = model_state_->GetParameter("kv_cache_onboard_blocks"); - } - catch (std::exception const& e) - { - // If parameter is not specified, just ignore - TLLM_LOG_WARNING("kv_cache_onboard_blocks not set, defaulting to true"); - } - std::optional> maxAttentionWindow = std::nullopt; try { @@ -183,7 +172,7 @@ executor::KvCacheConfig ModelInstanceState::getKvCacheConfigFromParams() } return executor::KvCacheConfig(enableKVCacheReuse, maxTokensInPagedKvCache, maxAttentionWindowVec, sinkTokenLength, - kvCacheFreeGpuMemFraction, kvCacheHostCacheSize, kvCacheOnboardBlocks, crossKvCacheFraction); + kvCacheFreeGpuMemFraction, kvCacheHostCacheSize, crossKvCacheFraction); } executor::ExtendedRuntimePerfKnobConfig ModelInstanceState::getExtendedRuntimePerfKnobConfigFromParams()